172 lines
8.6 KiB
Python
172 lines
8.6 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')
|
|
|
|
|
|
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
|