63 lines
2.5 KiB
Python
63 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from app.modules.rag.explain.models import ExplainPack
|
|
|
|
|
|
class PromptBudgeter:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
max_paths: int = 3,
|
|
max_symbols: int = 25,
|
|
max_excerpts: int = 40,
|
|
max_chars: int = 30000,
|
|
) -> None:
|
|
self._max_paths = max_paths
|
|
self._max_symbols = max_symbols
|
|
self._max_excerpts = max_excerpts
|
|
self._max_chars = max_chars
|
|
|
|
def build_prompt_input(self, question: str, pack: ExplainPack) -> str:
|
|
symbol_ids: list[str] = []
|
|
for path in pack.trace_paths[: self._max_paths]:
|
|
for symbol_id in path.symbol_ids:
|
|
if symbol_id and symbol_id not in symbol_ids and len(symbol_ids) < self._max_symbols:
|
|
symbol_ids.append(symbol_id)
|
|
excerpts = []
|
|
total_chars = 0
|
|
for excerpt in pack.code_excerpts:
|
|
if symbol_ids and excerpt.symbol_id and excerpt.symbol_id not in symbol_ids:
|
|
continue
|
|
body = excerpt.content.strip()
|
|
remaining = self._max_chars - total_chars
|
|
if remaining <= 0 or len(excerpts) >= self._max_excerpts:
|
|
break
|
|
if len(body) > remaining:
|
|
body = body[:remaining].rstrip() + "...[truncated]"
|
|
excerpts.append(
|
|
{
|
|
"evidence_id": excerpt.evidence_id,
|
|
"title": excerpt.title,
|
|
"path": excerpt.path,
|
|
"start_line": excerpt.start_line,
|
|
"end_line": excerpt.end_line,
|
|
"focus": excerpt.focus,
|
|
"content": body,
|
|
}
|
|
)
|
|
total_chars += len(body)
|
|
payload = {
|
|
"question": question,
|
|
"intent": pack.intent.model_dump(mode="json"),
|
|
"selected_entrypoints": [item.model_dump(mode="json") for item in pack.selected_entrypoints[:5]],
|
|
"seed_symbols": [item.model_dump(mode="json") for item in pack.seed_symbols[: self._max_symbols]],
|
|
"trace_paths": [path.model_dump(mode="json") for path in pack.trace_paths[: self._max_paths]],
|
|
"evidence_index": {key: value.model_dump(mode="json") for key, value in pack.evidence_index.items()},
|
|
"code_excerpts": excerpts,
|
|
"missing": pack.missing,
|
|
"conflicts": pack.conflicts,
|
|
}
|
|
return json.dumps(payload, ensure_ascii=False, indent=2)
|