Repository navigation
Expand file tree
/
Copy pathcodebase_query.py
More file actions
106 lines (91 loc) · 4.17 KB
/
Copy pathcodebase_query.py
File metadata and controls
106 lines (91 loc) · 4.17 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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
from loguru import logger
from pydantic_ai import Tool
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from ..graph_updater import MemgraphIngestor
from ..schemas import GraphData
from ..services.llm import CypherGenerator, LLMGenerationError
class GraphQueryError(Exception):
"""Custom exception for graph query failures."""
pass
def create_query_tool(
ingestor: MemgraphIngestor,
cypher_gen: CypherGenerator,
console: Console | None = None,
) -> Tool:
"""
Factory function that creates the knowledge graph query tool,
injecting its dependencies.
"""
# Use provided console or create a default one
if console is None:
console = Console(width=None, force_terminal=True)
async def query_codebase_knowledge_graph(natural_language_query: str) -> GraphData:
"""
Queries the codebase knowledge graph using natural language.
Provide your question in plain English about the codebase structure,
functions, classes, dependencies, or relationships. The tool will
automatically translate your natural language question into the
appropriate database query and return the results.
Examples:
- "Find all functions that call each other"
- "What classes are in the user authentication module"
- "Show me functions with the longest call chains"
- "Which files contain functions related to database operations"
"""
logger.info(f"[Tool:QueryGraph] Received NL query: '{natural_language_query}'")
cypher_query = "N/A"
try:
cypher_query = await cypher_gen.generate(natural_language_query)
results = ingestor.fetch_all(cypher_query)
if results:
table = Table(
show_header=True,
header_style="bold magenta",
)
headers = results[0].keys()
for header in headers:
table.add_column(header)
for row in results:
renderable_values = []
for value in row.values():
if value is None:
renderable_values.append("")
elif isinstance(value, bool):
# Check bool first since bool is a subclass of int in Python
renderable_values.append("✓" if value else "✗")
elif isinstance(value, int | float):
# Let Rich handle number formatting by converting to string
renderable_values.append(str(value))
else:
renderable_values.append(str(value))
table.add_row(*renderable_values)
console.print(
Panel(
table,
title="[bold blue]Cypher Query Results[/bold blue]",
expand=False,
)
)
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 Exception as e:
logger.error(
f"[Tool:QueryGraph] Error during query execution: {e}", exc_info=True
)
return GraphData(
query_used=cypher_query,
results=[],
summary=f"There was an error querying the database: {e}",
)
return Tool(
function=query_codebase_knowledge_graph,
description="Query the codebase knowledge graph using natural language questions. Ask in plain English about classes, functions, methods, dependencies, or code structure. Examples: 'Find all functions that call each other', 'What classes are in the user module', 'Show me functions with the longest call chains'.",
)