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", "") project_id = complete_request.get("projectId", "")
if config.get("ragDataSources", None): 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 = { params = {
"type": "object", "type": "object",
"properties": { "properties": {
@ -171,7 +171,7 @@ def get_rag_tool(config: dict, complete_request: dict) -> FunctionTool:
name="rag_search", name="rag_search",
description="Get information about an article", description="Get information about an article",
params_json_schema=params, 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 return tool
else: else:

View file

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