Repository navigation
Expand file tree
/
Copy pathcodebase_query.py
More file actions
120 lines (104 loc) · 4.71 KB
/
Copy pathcodebase_query.py
File metadata and controls
120 lines (104 loc) · 4.71 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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
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 ..schemas import GraphData
from ..services import QueryProtocol
from ..services.cache import get_query_cache
from ..services.llm import CypherGenerator, LLMGenerationError
class GraphQueryError(Exception):
"""Custom exception for graph query failures."""
pass
def create_query_tool(
ingestor: QueryProtocol,
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)
# Get query cache for result caching
query_cache = get_query_cache()
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)
# Check result cache first
cached_results = query_cache.get_results(cypher_query)
if cached_results is not None:
logger.info(
f"[Tool:QueryGraph] Result cache HIT ({len(cached_results)} items)"
)
results = cached_results
else:
results = ingestor.fetch_all(cypher_query)
# Cache the results
query_cache.set_results(cypher_query, results)
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'.",
)