Add RAG tool logging and update to use chunks toggle, ragK value

This commit is contained in:
akhisud3195 2025-05-09 10:26:16 +05:30
parent 770a080232
commit 889b7ee166
2 changed files with 16 additions and 14 deletions

View file

@ -153,7 +153,7 @@ def get_rag_tool(config: dict, complete_request: dict) -> FunctionTool:
"""
project_id = complete_request.get("projectId", "")
if config.get("ragDataSources", None):
print("rag_search")
print(f"Creating rag_search tool with params:\n-Data Sources: {config.get('ragDataSources', [])}\n-Return Type: {config.get('ragReturnType', 'chunks')}\n-K: {config.get('ragK', 3)}")
params = {
"type": "object",
"properties": {
@ -171,7 +171,7 @@ def get_rag_tool(config: dict, complete_request: dict) -> FunctionTool:
name="rag_search",
description="Get information about an article",
params_json_schema=params,
on_invoke_tool=lambda ctx, args: call_rag_tool(project_id, json.loads(args)['query'], config.get("ragDataSources", []), "chunks", 3)
on_invoke_tool=lambda ctx, args: call_rag_tool(project_id, json.loads(args)['query'], config.get("ragDataSources", []), config.get("ragReturnType", "chunks"), config.get("ragK", 3))
)
return tool
else:

View file

@ -75,18 +75,19 @@ async def call_rag_tool(
"active": True
}).to_list(length=None)
print(sources)
print(f"Sources: {sources}")
# Filter sources to those in source_ids
valid_source_ids = [
str(s["_id"]) for s in sources if str(s["_id"]) in source_ids
]
print(valid_source_ids)
print(f"Valid source ids: {valid_source_ids}")
# If no valid sources are found, return empty results
if not valid_source_ids:
return ''
# Perform Qdrant vector search
print(f"Calling Qdrant search with limit {k}")
qdrant_results = qdrant_client.search(
collection_name="embeddings",
query_vector=embed_result["embedding"],
@ -112,11 +113,13 @@ async def call_rag_tool(
for point in qdrant_results
]
print(return_type)
print(results)
print(f"Return type: {return_type}")
print(f"Results: {results}")
# If return_type is 'chunks', return the results directly
if return_type == "chunks":
return json.dumps({"Information": results}, indent=2)
chunks = json.dumps({"Information": results}, indent=2)
print(f"Returning chunks: {chunks}")
return chunks
# Otherwise, fetch the full document contents from MongoDB
doc_ids = [ObjectId(r["docId"]) for r in results]
@ -132,10 +135,9 @@ async def call_rag_tool(
]
# Convert results to a JSON string
formatted_string = json.dumps({"Information": results}, indent=2)
print(formatted_string)
return formatted_string
docs = json.dumps({"Information": results}, indent=2)
print(f"Returning docs: {docs}")
return docs
if __name__ == "__main__":
asyncio.run(call_rag_tool(