Repository navigation
Expand file tree
/
Copy pathcodebase_query.py
More file actions
51 lines (48 loc) · 2.08 KB
/
Copy pathcodebase_query.py
File metadata and controls
51 lines (48 loc) · 2.08 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
# codebase_rag/tools/codebase_query.py
from pydantic_ai import Tool, RunContext
from ..schemas import GraphData
from ..services.graph_db import MemgraphService, GraphQueryError
from ..services.llm import CypherGenerator, LLMGenerationError
from loguru import logger
def create_query_tool(
db_service: MemgraphService,
cypher_gen: CypherGenerator
) -> Tool:
"""
Factory function that creates the knowledge graph query tool,
injecting its dependencies.
"""
async def query_codebase_knowledge_graph(
ctx: RunContext,
natural_language_query: str
) -> GraphData:
"""
Queries the codebase knowledge graph. Translates a natural language question
into a Cypher query, executes it against the Memgraph database, and returns
the structured results.
"""
logger.info(f"[Tool:QueryGraph] Received NL query: '{natural_language_query}'")
try:
cypher_query = await cypher_gen.generate(natural_language_query)
results = db_service.execute_query(cypher_query)
summary = f"Successfully retrieved {len(results)} item(s) from the graph."
return GraphData(query_used=cypher_query, results=results, summary=summary)
except LLMGenerationError as e:
return GraphData(
query_used="N/A", results=[],
summary=f"I couldn't translate your request into a database query. Error: {e}"
)
except GraphQueryError as e:
return GraphData(
query_used=cypher_query, results=[],
summary=f"There was an error querying the database: {e}"
)
except Exception as e:
return GraphData(
query_used="N/A", results=[],
summary=f"An unexpected error occurred in the tool: {e}"
)
return Tool(
function=query_codebase_knowledge_graph,
description="Use this tool to query the codebase knowledge graph for specific information like classes, functions, methods, dependencies, or code structure."
)