49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from app.modules.rag.contracts.enums import RagLayer
|
|
from app.modules.agent.intent_router_v2 import ConversationState, IntentRouterV2, RepoContext
|
|
|
|
|
|
def repo_context() -> RepoContext:
|
|
return RepoContext(
|
|
languages=["python"],
|
|
available_domains=["CODE", "DOCS"],
|
|
available_layers=[
|
|
RagLayer.CODE_ENTRYPOINTS,
|
|
RagLayer.CODE_SYMBOL_CATALOG,
|
|
RagLayer.CODE_DEPENDENCY_GRAPH,
|
|
RagLayer.CODE_SEMANTIC_ROLES,
|
|
RagLayer.CODE_SOURCE_CHUNKS,
|
|
RagLayer.DOCS_MODULE_CATALOG,
|
|
RagLayer.DOCS_FACT_INDEX,
|
|
RagLayer.DOCS_SECTION_INDEX,
|
|
RagLayer.DOCS_POLICY_INDEX,
|
|
],
|
|
)
|
|
|
|
|
|
def run_sequence(queries: list[str], *, router: IntentRouterV2 | None = None, trace_label: str = "intent-router") -> list:
|
|
active_router = router or IntentRouterV2()
|
|
state = ConversationState()
|
|
results = []
|
|
for index, query in enumerate(queries, start=1):
|
|
result = active_router.route(query, state, repo_context())
|
|
print_trace(index, query, result, label=trace_label)
|
|
results.append(result)
|
|
state = state.advance(result)
|
|
return results
|
|
|
|
|
|
def run_single(query: str, *, router: IntentRouterV2 | None = None, trace_label: str = "intent-router"):
|
|
result = run_sequence([query], router=router, trace_label=trace_label)[0]
|
|
return result
|
|
|
|
|
|
def print_trace(index: int, query: str, result, *, label: str = "intent-router") -> None:
|
|
print(f"[{label}][turn {index}] input: {query}")
|
|
print()
|
|
print(f"[{label}][turn {index}] output: {json.dumps(result.model_dump(), ensure_ascii=False)}")
|
|
print("=" * 50)
|