115 lines
5.7 KiB
Python
115 lines
5.7 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
|
|
|
|
from tools.rag.build_chunks import TokenCounter
|
|
from tools.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')
|
|
|
|
|
|
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
|
|
|
|
# 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):
|
|
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': 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
|