252 lines
9.5 KiB
Python
252 lines
9.5 KiB
Python
"""A bounded tool message with one copy of evidence and no retrieval diagnostics."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import json
|
|
import os
|
|
import re
|
|
|
|
from tools.rag.build_chunks import TokenCounter
|
|
from tools.retrieval.text_budget import split_to_fit
|
|
|
|
SOURCE_FIELDS = (
|
|
"source_id",
|
|
"title",
|
|
"author",
|
|
"document_path",
|
|
"source_url",
|
|
"published",
|
|
"categories",
|
|
"tags",
|
|
"section",
|
|
)
|
|
FIELDS = (
|
|
"status",
|
|
"reason",
|
|
"query",
|
|
"engine",
|
|
"strategies",
|
|
"instructions",
|
|
"answer_format",
|
|
"complete",
|
|
"truncated",
|
|
"truncation_reasons",
|
|
"index_version",
|
|
"source_count",
|
|
"clarification",
|
|
"named_subjects",
|
|
"work_scopes",
|
|
"question_plan",
|
|
"resolved_query",
|
|
"name_resolution",
|
|
)
|
|
|
|
|
|
def compact_rag_payload(
|
|
data: dict, *, max_tokens: int | None = None, include_diagnostics: bool = False
|
|
) -> dict:
|
|
from tools.rag.rag_utils import build_context_text
|
|
|
|
limit = max_tokens if max_tokens is not None else int(os.getenv("RAG_LLM_MAX_TOKENS", "12000"))
|
|
if limit < 256:
|
|
raise ValueError("RAG_LLM_MAX_TOKENS musí byť aspoň 256")
|
|
encoding = os.getenv("RAG_LLM_TOKEN_ENCODING", "cl100k_base")
|
|
counter = TokenCounter(encoding)
|
|
if counter.encoding is None:
|
|
raise RuntimeError("Kontextový rozpočet vyžaduje platný tiktoken tokenizér")
|
|
base = {key: copy.deepcopy(data[key]) for key in FIELDS if key in data}
|
|
base.setdefault("status", "ok")
|
|
base.setdefault("complete", None)
|
|
base.setdefault("truncated", False)
|
|
base.setdefault("truncation_reasons", [])
|
|
input_sources = [
|
|
{**source, "source_id": source.get("source_id") or f"S{number}"}
|
|
for number, source in enumerate(data.get("sources", []), 1)
|
|
]
|
|
references = [
|
|
{key: source[key] for key in SOURCE_FIELDS if key in source} for source in input_sources
|
|
]
|
|
|
|
def stamp(payload):
|
|
payload["context_budget"] = {
|
|
"max_tokens": limit,
|
|
"used_tokens": 0,
|
|
"tokenizer": counter.name,
|
|
"tokens_are_estimate": True,
|
|
}
|
|
for _ in range(8):
|
|
used = counter.count(json.dumps(payload, ensure_ascii=False, separators=(",", ":")))
|
|
if payload["context_budget"]["used_tokens"] == used:
|
|
break
|
|
payload["context_budget"]["used_tokens"] = used
|
|
return payload
|
|
|
|
def render(sources, *, truncated=False, raw_context=None):
|
|
payload = copy.deepcopy(base)
|
|
if truncated:
|
|
payload["complete"] = False
|
|
payload["truncated"] = True
|
|
payload["truncation_reasons"] = list(
|
|
dict.fromkeys([*payload["truncation_reasons"], "llm_context_budget"])
|
|
)
|
|
payload["sources"] = [
|
|
{key: source[key] for key in SOURCE_FIELDS if key in source} for source in sources
|
|
]
|
|
payload["source_count"] = len(sources)
|
|
coverage = {
|
|
key: payload.get(key)
|
|
for key in ["complete", "truncated", "truncation_reasons", "index_version"]
|
|
}
|
|
context = raw_context if raw_context is not None else build_context_text(sources)
|
|
payload["context"] = (
|
|
"METADÁTA VYHĽADANIA\n" + json.dumps(coverage, ensure_ascii=False) + "\n\n" + context
|
|
)
|
|
return stamp(payload)
|
|
|
|
def fits(payload):
|
|
return payload["context_budget"]["used_tokens"] <= limit
|
|
|
|
def whole_prefix(body, end):
|
|
boundaries = [
|
|
m.end()
|
|
for m in re.finditer(
|
|
r"\n\s*\n|(?=^[ \t]*(?:\d+[.)]|[-+*])\s+|^#{1,6}\s|^Sekcia:)", body, re.M
|
|
)
|
|
if m.end() <= end
|
|
]
|
|
return body[: max(boundaries, default=0)]
|
|
|
|
def balanced_sources(sources):
|
|
"""Reserve evidence space for every requested person before expanding."""
|
|
suffix = "\n[OBSAH ZDROJA OREZANÝ]"
|
|
selected = [{**s, "text": suffix if s.get("text") else ""} for s in sources]
|
|
if not fits(render(selected, truncated=True)):
|
|
raise ValueError("Rozpočet nestačí na metadáta všetkých porovnávaných osôb")
|
|
families = {
|
|
s.get("document_path", "").removesuffix("README.md"): s["title"]
|
|
for s in sources
|
|
if s.get("title") in base["named_subjects"] and s.get("document_path")
|
|
}
|
|
groups = {}
|
|
for i, s in enumerate(sources):
|
|
owner = next(
|
|
(p for path, p in families.items() if s.get("document_path", "").startswith(path)),
|
|
f"other:{i}",
|
|
)
|
|
groups.setdefault(owner, []).append(i)
|
|
|
|
def used():
|
|
return render(selected, truncated=True)["context_budget"]["used_tokens"]
|
|
|
|
def expand(index, budget):
|
|
body = sources[index].get("text") or ""
|
|
if not body:
|
|
return
|
|
ceiling = min(limit, used() + max(0, budget))
|
|
|
|
def candidate(text):
|
|
replacement = {
|
|
**selected[index],
|
|
"text": text + (suffix if len(text) < len(body) else ""),
|
|
}
|
|
return render(
|
|
[*selected[:index], replacement, *selected[index + 1 :]], truncated=True
|
|
)
|
|
|
|
if candidate(body)["context_budget"]["used_tokens"] <= ceiling:
|
|
text = body
|
|
else:
|
|
low, high = 0, len(body) - 1
|
|
while low < high:
|
|
mid = (low + high + 1) // 2
|
|
if candidate(body[:mid])["context_budget"]["used_tokens"] <= ceiling:
|
|
low = mid
|
|
else:
|
|
high = mid - 1
|
|
text = body[:low]
|
|
# Keep a whole paragraph/list item, including wrapped caveats.
|
|
text = whole_prefix(body, low)
|
|
selected[index] = {
|
|
**selected[index],
|
|
"text": text + (suffix if len(text) < len(body) else ""),
|
|
}
|
|
|
|
remaining = limit - used()
|
|
quota = remaining // max(1, len(groups))
|
|
for indexes in groups.values():
|
|
start = used()
|
|
for pos, index in enumerate(indexes):
|
|
expand(index, max(0, quota - (used() - start)) // (len(indexes) - pos))
|
|
# Reuse the space left by short profiles, still sharing it by person.
|
|
remaining = limit - used()
|
|
quota = remaining // max(1, len(groups))
|
|
for indexes in groups.values():
|
|
start = used()
|
|
for pos, index in enumerate(indexes):
|
|
expand(index, max(0, quota - (used() - start)) // (len(indexes) - pos))
|
|
return render(selected, truncated=True)
|
|
|
|
# An already projected server response needs no reconstruction by its client.
|
|
if data.get("context_budget"):
|
|
result = stamp({**base, "sources": references, "context": data.get("context", "")})
|
|
if not fits(result):
|
|
parts = split_to_fit(
|
|
data.get("context", ""),
|
|
lambda text: fits(render(references, truncated=True, raw_context=text)),
|
|
)
|
|
result = render(references, truncated=True, raw_context=parts[0])
|
|
elif any("text" in source for source in input_sources):
|
|
sources = input_sources
|
|
result = render(sources)
|
|
if not fits(result) and len(base.get("named_subjects", [])) > 1:
|
|
result = balanced_sources(sources)
|
|
elif not fits(result):
|
|
selected = []
|
|
if not fits(render([], truncated=True)):
|
|
raise ValueError("Tokenový rozpočet je menší ako povinné pokyny a metadáta")
|
|
for source in sources:
|
|
candidate = render([*selected, source], truncated=True)
|
|
if fits(candidate):
|
|
selected.append(source)
|
|
continue
|
|
body = source.get("text") or ""
|
|
suffix = "\n[OBSAH ZDROJA OREZANÝ]"
|
|
|
|
def body_fits(text):
|
|
return fits(
|
|
render([*selected, {**source, "text": text + suffix}], truncated=True)
|
|
)
|
|
|
|
if body and body_fits(body[:1]):
|
|
# Only the first fitting prefix is delivered; the complete
|
|
# original remains available through the diagnostics option.
|
|
low, high = 1, len(body)
|
|
while low < high:
|
|
middle = (low + high + 1) // 2
|
|
if body_fits(body[:middle]):
|
|
low = middle
|
|
else:
|
|
high = middle - 1
|
|
selected.append({**source, "text": whole_prefix(body, low) + suffix})
|
|
break
|
|
result = render(selected, truncated=True)
|
|
if not selected:
|
|
result["status"] = "no_answer"
|
|
result["reason"] = "llm_context_budget_exhausted"
|
|
result = stamp(result)
|
|
else:
|
|
# Compatibility with older tool servers that expose only context + refs.
|
|
context = data.get("context", "")
|
|
result = render(references, raw_context=context)
|
|
if not fits(result):
|
|
parts = split_to_fit(
|
|
context, lambda text: fits(render(references, truncated=True, raw_context=text))
|
|
)
|
|
result = render(references, truncated=True, raw_context=parts[0])
|
|
if not fits(result):
|
|
raise ValueError("Povinné metadáta prekračujú kontextový rozpočet")
|
|
if include_diagnostics:
|
|
result["diagnostics"] = copy.deepcopy(data)
|
|
return result
|