60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
from app.modules.rag.explain.models import CodeExcerpt, LayeredRetrievalItem
|
|
|
|
|
|
class ExcerptPlanner:
|
|
_FOCUS_TOKENS = ("raise", "except", "db", "select", "insert", "update", "delete", "http", "publish", "emit")
|
|
|
|
def plan(self, chunk: LayeredRetrievalItem, *, evidence_id: str, symbol_id: str | None) -> list[CodeExcerpt]:
|
|
location = chunk.location
|
|
if location is None:
|
|
return []
|
|
excerpts = [
|
|
CodeExcerpt(
|
|
evidence_id=evidence_id,
|
|
symbol_id=symbol_id,
|
|
title=chunk.title,
|
|
path=location.path,
|
|
start_line=location.start_line,
|
|
end_line=location.end_line,
|
|
content=chunk.content.strip(),
|
|
focus="overview",
|
|
)
|
|
]
|
|
focus = self._focus_excerpt(chunk, evidence_id=evidence_id, symbol_id=symbol_id)
|
|
if focus is not None:
|
|
excerpts.append(focus)
|
|
return excerpts
|
|
|
|
def _focus_excerpt(
|
|
self,
|
|
chunk: LayeredRetrievalItem,
|
|
*,
|
|
evidence_id: str,
|
|
symbol_id: str | None,
|
|
) -> CodeExcerpt | None:
|
|
location = chunk.location
|
|
if location is None:
|
|
return None
|
|
lines = chunk.content.splitlines()
|
|
for index, line in enumerate(lines):
|
|
lowered = line.lower()
|
|
if not any(token in lowered for token in self._FOCUS_TOKENS):
|
|
continue
|
|
start = max(0, index - 2)
|
|
end = min(len(lines), index + 3)
|
|
if end - start >= len(lines):
|
|
return None
|
|
return CodeExcerpt(
|
|
evidence_id=evidence_id,
|
|
symbol_id=symbol_id,
|
|
title=f"{chunk.title}:focus",
|
|
path=location.path,
|
|
start_line=(location.start_line or 1) + start,
|
|
end_line=(location.start_line or 1) + end - 1,
|
|
content="\n".join(lines[start:end]).strip(),
|
|
focus="focus",
|
|
)
|
|
return None
|