context_admission.py

"""Pre-generation context admission. Only generated distractors may be removed."""
import asyncio,copy,json,re
from urllib.parse import urlsplit,urlunsplit
import httpx


class ContextAdmissionError(RuntimeError):pass


_ARCHIVED=re.compile(r'Archived record \d{6}: queue-\d+ sample=\d+ status=nominal\.\n')


async def tokenize(provider,model,providers,messages,tools):
    from advisor import _private_url
    if provider.get('kind')!='openai' or not _private_url(provider.get('base_url','')):return None
    u=urlsplit(provider['base_url']); path=u.path.rstrip('/'); path=path[:-3] if path.endswith('/v1') else path
    url=urlunsplit((u.scheme,u.netloc,path+'/tokenize','',''))
    key=providers.fernet.decrypt(provider['secret']).decode() if provider.get('secret') else ''
    payload={'model':model['model_id'],'messages':messages,'add_generation_prompt':True,'return_token_strs':False}
    if tools:payload['tools']=tools
    try:
        async with asyncio.timeout(15):
            async with httpx.AsyncClient(timeout=15,follow_redirects=False,trust_env=False) as c:
                async with c.stream('POST',url,headers={'Authorization':'Bearer '+key} if key else {},json=payload) as r:
                    if r.status_code!=200:return None
                    chunks=[];size=0
                    async for chunk in r.aiter_bytes():
                        size+=len(chunk)
                        if size>32*1024*1024:return None
                        chunks.append(chunk)
        data=json.loads(b''.join(chunks)); n=data.get('count')
        return {'count':n,'max_model_len':data.get('max_model_len')} if isinstance(n,int) and not isinstance(n,bool) and n>=0 else None
    except (httpx.HTTPError,TimeoutError,ValueError):return None


def _remove_distractors(messages,keep_ratio):
    removed=0
    for msg in messages:
        if msg.get('role')!='user' or not isinstance(msg.get('content'),str):continue
        text=msg['content'];matches=list(_ARCHIVED.finditer(text)); count=len(matches)
        if not count:continue
        keep=max(0,min(count-1,int(count*keep_ratio)))
        chosen={int(i*count/keep) for i in range(keep)} if keep else set()
        output=[];offset=0
        for i,match in enumerate(matches):
            output.append(text[offset:match.start()])
            if i in chosen:output.append(match.group())
            else:removed+=1
            offset=match.end()
        output.append(text[offset:]);msg['content']=''.join(output)
    return removed


async def prepare_messages(provider,model,providers,messages,max_tokens,target_context,tools=None,measure=None):
    measured=measure or tokenize; result=copy.deepcopy(messages); removed=0
    registry=model['context_limit']; reserve=max_tokens+512
    info={'registered_context_limit':registry,'requested_context_size':target_context,'removed_distractor_records':0}
    for attempt in range(5):
        count=await measured(provider,model,providers,result,tools)
        if count is not None:
            engine=count.get('max_model_len'); engine=engine if isinstance(engine,int) and not isinstance(engine,bool) and engine>0 else registry
            budget=min(target_context,registry-reserve,engine-reserve)
            info.update(count_source='serving_tokenizer',engine_context_limit=engine,admitted_prompt_budget=budget,actual_prompt_tokens=count['count'])
            n=count['count']
        else:
            # UTF-8 bytes plus a generous template allowance bound ordinary text tokenization;
            # no claim of exact target-band control without the serving tokenizer.
            n=len(json.dumps({'messages':result,'tools':tools},ensure_ascii=False).encode())+4096
            budget=registry-reserve
            info.update(count_source='conservative_byte_admission',engine_context_limit=None,admitted_prompt_budget=budget,actual_prompt_tokens=None,estimated_prompt_upper_bound=n,note='Tokenizer unavailable; target band remains approximate.')
        if budget<=0:raise ContextAdmissionError('Context admission: output reservation leaves no prompt room')
        if n<=budget:
            info['removed_distractor_records']=removed
            return result,info
        ratio=max(0,min(.95,budget/max(n,1)*.90))
        changed=_remove_distractors(result,ratio);removed+=changed
        if not changed:raise ContextAdmissionError('Context admission: required evidence or tool conversation exceeds the prompt budget; no evidence was truncated')
    raise ContextAdmissionError('Context admission: unable to fit after bounded distractor reduction')