Repository navigation
Expand file tree
/
Copy pathcodebase_query.py
More file actions
94 lines (83 loc) · 3.12 KB
/
Copy pathcodebase_query.py
File metadata and controls
94 lines (83 loc) · 3.12 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
from __future__ import annotations
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 .. import exceptions as ex
from .. import logs as ls
from ..constants import (
QUERY_NOT_AVAILABLE,
QUERY_RESULTS_PANEL_TITLE,
QUERY_SUMMARY_DB_ERROR,
QUERY_SUMMARY_SUCCESS,
QUERY_SUMMARY_TRANSLATION_FAILED,
)
from ..schemas import QueryGraphData
from ..services import QueryProtocol
from ..services.llm import CypherGenerator
from . import tool_descriptions as td
def create_query_tool(
ingestor: QueryProtocol,
cypher_gen: CypherGenerator,
console: Console | None = None,
) -> Tool:
if console is None:
console = Console(width=None, force_terminal=True)
async def query_codebase_knowledge_graph(
natural_language_query: str,
) -> QueryGraphData:
logger.info(ls.TOOL_QUERY_RECEIVED.format(query=natural_language_query))
cypher_query = QUERY_NOT_AVAILABLE
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):
renderable_values.append("✓" if value else "✗")
elif isinstance(value, int | float):
renderable_values.append(str(value))
else:
renderable_values.append(str(value))
table.add_row(*renderable_values)
console.print(
Panel(
table,
title=QUERY_RESULTS_PANEL_TITLE,
expand=False,
)
)
summary = QUERY_SUMMARY_SUCCESS.format(count=len(results))
return QueryGraphData(
query_used=cypher_query, results=results, summary=summary
)
except ex.LLMGenerationError as e:
return QueryGraphData(
query_used=QUERY_NOT_AVAILABLE,
results=[],
summary=QUERY_SUMMARY_TRANSLATION_FAILED.format(error=e),
)
except Exception as e:
logger.exception(ls.TOOL_QUERY_ERROR.format(error=e))
return QueryGraphData(
query_used=cypher_query,
results=[],
summary=QUERY_SUMMARY_DB_ERROR.format(error=e),
)
return Tool(
function=query_codebase_knowledge_graph,
name=td.AgenticToolName.QUERY_GRAPH,
description=td.CODEBASE_QUERY,
)