412 lines
18 KiB
Python
412 lines
18 KiB
Python
"""ModelAdapter: HOW to call a model (protocol level). Never deploys anything.
|
|
|
|
Registry + two built-ins:
|
|
openai_compatible -- /v1/chat/completions; covers vllm, sglang, lmdeploy,
|
|
ollama, tgi, and every OpenAI-protocol cloud API
|
|
mock -- offline deterministic/testing adapter
|
|
|
|
Model spec grammar (a single string, no config files needed for the common case):
|
|
openai/http://localhost:8000/v1?qwen3-8b adapter + api_base + model id
|
|
openai/https://api.openai.com/v1?gpt-4o (OPENAI_API_KEY read from env)
|
|
anthropic/claude-... (provider-native protocols later)
|
|
mock offline
|
|
deploy:vllm/qwen3-8b -> Deployer resolves to an endpoint,
|
|
then re-dispatches as openai/...
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from ..data.sample import ChatMessage
|
|
from ..eval.registry import EvalRegistry
|
|
from .output import ModelOutput, ToolCall, Usage
|
|
|
|
ADAPTER_REGISTRY = EvalRegistry('model adapter')
|
|
|
|
|
|
def register_adapter(name: str):
|
|
def decorator(cls):
|
|
ADAPTER_REGISTRY.register(name, cls)
|
|
return cls
|
|
|
|
return decorator
|
|
|
|
|
|
class ModelAdapter:
|
|
"""Base class. Subclasses implement generate() (async, structured output)."""
|
|
|
|
name = 'base'
|
|
|
|
def __init__(self, model: str = '', api_base: str = '', api_key: str = '', **kwargs):
|
|
self.model = model
|
|
self.api_base = api_base.rstrip('/')
|
|
self.api_key = api_key
|
|
self.extra = kwargs
|
|
|
|
async def generate(
|
|
self,
|
|
messages: List[ChatMessage],
|
|
tools: Optional[List[Dict[str, Any]]] = None,
|
|
**gen_kwargs,
|
|
) -> ModelOutput:
|
|
raise NotImplementedError
|
|
|
|
async def score_choices(
|
|
self,
|
|
prompt: str,
|
|
choices: List[str],
|
|
) -> 'ModelOutput':
|
|
"""Paper-faithful MCQ scoring (P1): logprob of each choice continuing
|
|
the prompt. Strategy per backend capability:
|
|
|
|
completions_echo -- /completions echo+prompt_logprobs (vllm/sglang/
|
|
ollama); true continuation logprob, the lm-eh paper path
|
|
chat_first_token -- /chat/completions max_tokens=1 + top_logprobs; the
|
|
NEXT-token distribution encodes each choice; works on
|
|
every OpenAI-protocol API (cloud included); scores
|
|
choices as first-token probabilities
|
|
|
|
Selected via extra['logprob_strategy'] ('auto' default).
|
|
"""
|
|
strategy = self.extra.get('logprob_strategy', 'auto')
|
|
if strategy in ('auto', 'completions_echo'):
|
|
try:
|
|
return await self._score_choices_echo(prompt, choices)
|
|
except Exception as e:
|
|
if strategy == 'completions_echo':
|
|
raise
|
|
self.extra['logprob_strategy'] = 'chat_first_token' # downgrade once
|
|
return await self._score_choices_first_token(prompt, choices)
|
|
|
|
async def _score_choices_first_token(self, prompt: str, choices: List[str]):
|
|
"""Approximate P1 via next-token top_logprobs (any chat API).
|
|
|
|
The choices are shown in the prompt; we score P(first content token
|
|
of each choice | prompt) from the top_logprobs of a 1-token reply.
|
|
"""
|
|
from .output import ChoiceLogprob, ModelOutput
|
|
|
|
letters = 'ABCDEFGHIJ'
|
|
opts = '\n'.join(f'{letters[i]}. {c}' for i, c in enumerate(choices))
|
|
asked = (f'{prompt}\n{opts}\n\n'
|
|
'Answer with only the letter of the correct option.')
|
|
payload = {'model': self.model, 'messages': [{'role': 'user', 'content': asked}],
|
|
'max_tokens': 1, 'temperature': 0.0,
|
|
'logprobs': True, 'top_logprobs': 20}
|
|
headers = {'Content-Type': 'application/json'}
|
|
if self.api_key:
|
|
headers['Authorization'] = f'Bearer {self.api_key}'
|
|
data = await self._post(f'{self.api_base}/chat/completions', payload, headers)
|
|
import json as _json
|
|
|
|
choice = (data.get('choices') or [{}])[0]
|
|
lp_wrapper = (choice.get('logprobs') or {})
|
|
top = lp_wrapper.get('content') or lp_wrapper.get('top_logprobs') or []
|
|
# top: list per generated token of {token: logprob} dicts
|
|
token_lps: Dict[str, float] = {}
|
|
first = top[0] if top else None
|
|
if isinstance(first, dict) and 'top_logprobs' in first:
|
|
# OpenAI/sglang chat shape: content[0] = {token, logprob, top_logprobs:[...]}
|
|
entries = {}
|
|
for e in (first.get('top_logprobs') or []):
|
|
tok, lp = e.get('token'), e.get('logprob')
|
|
if tok is not None and lp is not None:
|
|
entries[str(tok)] = lp
|
|
elif isinstance(first, dict):
|
|
entries = first # vllm-ish {token: lp} map
|
|
else:
|
|
entries = {}
|
|
for tok, lp in entries.items():
|
|
if isinstance(lp, dict):
|
|
lp = lp.get('logprob')
|
|
if lp is None:
|
|
continue
|
|
try:
|
|
token_lps[str(tok).strip().upper()] = float(lp)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if not token_lps:
|
|
raise RuntimeError(f'no top_logprobs from {self.api_base}; P1 needs a '
|
|
'backend exposing logprobs (or use mcq_mode=generate)')
|
|
results = []
|
|
for i in range(len(choices)):
|
|
letter = letters[i]
|
|
lp = token_lps.get(letter, -100.0)
|
|
results.append(ChoiceLogprob(text=choices[i], logprob=lp,
|
|
token_logprobs=[lp], num_tokens=1))
|
|
usage = Usage(finish_reason=choice.get('finish_reason', ''))
|
|
return ModelOutput(model=self.model, choice_logprobs=results, usage=usage)
|
|
|
|
async def _score_choices_echo(self, prompt: str, choices: List[str]):
|
|
from .output import ChoiceLogprob, ModelOutput
|
|
|
|
results: List[ChoiceLogprob] = []
|
|
for choice in choices:
|
|
full = f'{prompt} {choice}'
|
|
payload = {'model': self.model, 'prompt': full, 'max_tokens': 1,
|
|
'echo': True, 'prompt_logprobs': 0,
|
|
'temperature': 0.0}
|
|
headers = {'Content-Type': 'application/json'}
|
|
if self.api_key:
|
|
headers['Authorization'] = f'Bearer {self.api_key}'
|
|
data = await self._post_raw(f'{self.api_base}/completions', payload, headers)
|
|
# vllm/sglang shape: choices[0].prompt_logprobs: list[None|{token: lp}]
|
|
import json as _json
|
|
|
|
ch = (data.get('choices') or [{}])[0]
|
|
pl = ch.get('prompt_logprobs') or []
|
|
prompt_ids = ch.get('prompt_token_ids') or []
|
|
# tokens belonging to the choice = the tail after the prompt part.
|
|
# echo mode returns the whole sequence; we approximate the choice
|
|
# token count via the tail offset encoded in logprob entries.
|
|
lps: List[float] = []
|
|
for entry in pl:
|
|
if isinstance(entry, dict):
|
|
lp = next(iter(entry.values()))
|
|
if isinstance(lp, dict):
|
|
lp = lp.get('logprob', 0.0)
|
|
lps.append(float(lp))
|
|
if not lps:
|
|
raise RuntimeError(
|
|
f'backend at {self.api_base} returned no prompt_logprobs; '
|
|
'P1 logprob scoring needs vllm/sglang-style completions-echo'
|
|
)
|
|
# last len tokens we cannot know exactly; use the trailing segment
|
|
# matching the choice's char proportion as a robust approximation
|
|
n_prompt_chars = len(prompt) + 1
|
|
frac = max(len(choice), 1) / max(len(full), 1)
|
|
n_choice = max(1, round(frac * len(lps)))
|
|
choice_lps = lps[-n_choice:]
|
|
results.append(ChoiceLogprob(
|
|
text=choice, logprob=sum(choice_lps),
|
|
token_logprobs=choice_lps, num_tokens=len(choice_lps)))
|
|
return ModelOutput(model=self.model, choice_logprobs=results)
|
|
|
|
async def _post_raw(self, url: str, payload: Dict, headers: Dict) -> Dict[str, Any]:
|
|
return await self._post(url, payload, headers)
|
|
|
|
async def close(self) -> None:
|
|
pass
|
|
|
|
def __repr__(self):
|
|
return f'{type(self).__name__}(model={self.model!r}, base={self.api_base!r})'
|
|
|
|
|
|
# --------------------------- spec parsing / dispatch ---------------------------
|
|
|
|
_SPEC_RE = re.compile(r'^(?P<adapter>[a-z_]+)(/(?P<rest>.+))?$')
|
|
_DEPLOY_RE = re.compile(r'^deploy:(?P<deployer>[a-z0-9_-]+)/(?P<model>.+)$')
|
|
|
|
|
|
def parse_model_spec(spec: str) -> Dict[str, str]:
|
|
"""'openai/http://x:8000/v1?qwen' -> {adapter, api_base, model}."""
|
|
m = _SPEC_RE.match(spec.strip())
|
|
if not m:
|
|
raise ValueError(f'bad model spec: {spec!r} (expected adapter/rest?model)')
|
|
adapter = m.group('adapter')
|
|
rest = m.group('rest') or ''
|
|
api_base, _, model = rest.partition('?')
|
|
return {'adapter': adapter, 'api_base': api_base, 'model': model}
|
|
|
|
|
|
def resolve_adapter(spec: str, deploy_fn=None) -> ModelAdapter:
|
|
"""Spec -> adapter instance. deploy:... specs go through a Deployer first."""
|
|
m = _DEPLOY_RE.match(spec.strip())
|
|
if m:
|
|
if deploy_fn is None:
|
|
from .deployer import deploy as default_deploy
|
|
deploy_fn = default_deploy
|
|
endpoint = deploy_fn(m.group('deployer'), m.group('model'))
|
|
spec = f"openai/{endpoint['api_base']}?{endpoint['model']}"
|
|
parsed = parse_model_spec(spec)
|
|
cls = ADAPTER_REGISTRY.get(parsed['adapter'])
|
|
key = parsed.get('api_base') and _key_for(parsed['api_base'])
|
|
return cls(model=parsed['model'], api_base=parsed['api_base'], api_key=key)
|
|
|
|
|
|
def _key_for(api_base: str) -> str:
|
|
"""Best-effort API key by endpoint; explicit env always wins."""
|
|
for host_hint, var in (('api.openai.com', 'OPENAI_API_KEY'),
|
|
('anthropic.com', 'ANTHROPIC_API_KEY'),
|
|
('dashscope', 'DASHSCOPE_API_KEY'),
|
|
('bigmodel', 'ZAI_API_KEY')):
|
|
if host_hint in api_base:
|
|
return os.environ.get(var, '')
|
|
return os.environ.get('OPENAI_API_KEY', '')
|
|
|
|
|
|
# --------------------------- openai_compatible ---------------------------
|
|
|
|
|
|
@register_adapter('openai')
|
|
class OpenAICompatible(ModelAdapter):
|
|
"""Async OpenAI chat-completions client with zero hard dependencies.
|
|
|
|
Uses httpx if installed (proper async); falls back to urllib in a thread
|
|
so the core install stays dependency-free.
|
|
"""
|
|
|
|
name = 'openai'
|
|
|
|
async def generate(self, messages, tools=None, **kw) -> ModelOutput:
|
|
payload = self._payload(messages, tools, kw)
|
|
headers = {'Content-Type': 'application/json'}
|
|
if self.api_key:
|
|
headers['Authorization'] = f'Bearer {self.api_key}'
|
|
retries = self.extra.get('retries', 3)
|
|
last_exc: Exception = None
|
|
for attempt in range(retries + 1):
|
|
try:
|
|
data = await self._post(f'{self.api_base}/chat/completions', payload, headers)
|
|
return self._parse(data)
|
|
except Exception as e: # 5xx/429/timeouts: worth retrying
|
|
last_exc = e
|
|
retryable = 'Server error' in str(e) or '504' in str(e) or '502' in str(e) \
|
|
or '429' in str(e) or 'timeout' in str(e).lower() \
|
|
or 'TimeoutException' in type(e).__name__
|
|
if attempt >= retries or not retryable:
|
|
raise
|
|
import asyncio
|
|
|
|
await asyncio.sleep(min(2 ** attempt * 2, 30))
|
|
raise last_exc # unreachable
|
|
|
|
def _payload(self, messages, tools, kw) -> Dict[str, Any]:
|
|
msgs = [{'role': m.role, 'content': m.content} for m in messages]
|
|
payload: Dict[str, Any] = {'model': self.model, 'messages': msgs}
|
|
if tools:
|
|
payload['tools'] = [
|
|
{'type': 'function', 'function': t} if 'function' not in t else t for t in tools
|
|
]
|
|
for k in ('temperature', 'max_tokens', 'top_p', 'stop', 'seed', 'response_format'):
|
|
if kw.get(k) is not None:
|
|
payload[k] = kw[k]
|
|
payload.setdefault('max_tokens', self.extra.get('max_tokens', 4096)) # CoT room
|
|
return payload
|
|
|
|
def _parse(self, data: Dict[str, Any]) -> ModelOutput:
|
|
choice = (data.get('choices') or [{}])[0]
|
|
msg = choice.get('message') or {}
|
|
calls = []
|
|
for c in msg.get('tool_calls') or []:
|
|
fn = c.get('function') or {}
|
|
args = fn.get('arguments') or '{}'
|
|
try:
|
|
args_dict = json.loads(args)
|
|
except (ValueError, TypeError):
|
|
args_dict = {}
|
|
calls.append(ToolCall(id=c.get('id', ''), name=fn.get('name', ''),
|
|
arguments=args, arguments_dict=args_dict))
|
|
u = data.get('usage') or {}
|
|
usage = Usage(input_tokens=u.get('prompt_tokens', 0),
|
|
output_tokens=u.get('completion_tokens', 0),
|
|
total_tokens=u.get('total_tokens', 0),
|
|
finish_reason=choice.get('finish_reason', ''))
|
|
return ModelOutput(text=msg.get('content') or '', tool_calls=calls, usage=usage,
|
|
raw=data, model=data.get('model', self.model))
|
|
|
|
async def _post(self, url: str, payload: Dict, headers: Dict) -> Dict[str, Any]:
|
|
try:
|
|
import httpx
|
|
|
|
async with httpx.AsyncClient(timeout=self.extra.get('timeout', 600)) as client:
|
|
r = await client.post(url, json=payload, headers=headers)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except ImportError:
|
|
import asyncio
|
|
import urllib.request
|
|
|
|
def _sync():
|
|
req = urllib.request.Request(url, data=json.dumps(payload).encode(),
|
|
headers=headers, method='POST')
|
|
with urllib.request.urlopen(req, timeout=self.extra.get('timeout', 600)) as resp:
|
|
return json.loads(resp.read().decode())
|
|
|
|
return await asyncio.to_thread(_sync)
|
|
|
|
|
|
# --------------------------- mock ---------------------------
|
|
|
|
_MOCK_PATTERNS = (
|
|
(re.compile(r'\\boxed\{([^}]*)\}'), None), # echo any boxed target in the input
|
|
)
|
|
|
|
|
|
@register_adapter('mock')
|
|
class MockAdapter(ModelAdapter):
|
|
"""Offline adapter for tests/CI/dev.
|
|
|
|
Modes (extra['mode']):
|
|
echo -- return the input text (default)
|
|
boxed -- return \\boxed{target} (oracle channel: runner tags a
|
|
MOCKTARGET message when the sample carries a target)
|
|
oracle -- return the target verbatim (same channel; for coding
|
|
benches whose target is the canonical solution)
|
|
fc -- replay the target's ground-truth tool calls (oracle for
|
|
function-calling benches; target JSON in runner dict form)
|
|
tool -- return one tool call named extra['tool_name']
|
|
const -- return extra['text']
|
|
"""
|
|
|
|
name = 'mock'
|
|
|
|
async def generate(self, messages, tools=None, **kw) -> ModelOutput:
|
|
mode = self.extra.get('mode', 'echo')
|
|
if mode == 'const':
|
|
text = self.extra.get('text', 'mock')
|
|
elif mode in ('fc', 'tool'):
|
|
if mode == 'tool':
|
|
return ModelOutput(text='', tool_calls=[ToolCall(
|
|
name=self.extra.get('tool_name', 'dummy_tool'), arguments='{}',
|
|
arguments_dict={})], model='mock')
|
|
target = None
|
|
already_played = any(m.role == 'tool' for m in messages)
|
|
if already_played:
|
|
# oracle replays ground truth ONCE, then wraps up like a
|
|
# well-behaved agent (final turn, no more calls)
|
|
return ModelOutput(text='Done.', model='mock', usage=Usage(
|
|
input_tokens=1, output_tokens=1, total_tokens=2, finish_reason='stop'))
|
|
for m in reversed(messages):
|
|
if m.role == 'user' and m.content.startswith('MOCKTARGET::'):
|
|
target = m.content[len('MOCKTARGET::'):]
|
|
break
|
|
calls = []
|
|
if target:
|
|
try:
|
|
gt = json.loads(target)
|
|
raw_calls = gt.get('tool_calls', gt if isinstance(gt, list) else [])
|
|
for c in raw_calls:
|
|
fn = c.get('function', c)
|
|
args = fn.get('arguments', {})
|
|
calls.append(ToolCall(
|
|
name=fn.get('name', ''), arguments=json.dumps(args),
|
|
arguments_dict=args if isinstance(args, dict) else {}))
|
|
except (ValueError, TypeError):
|
|
calls = []
|
|
if not calls:
|
|
return ModelOutput(text='no tool needed', model='mock', usage=Usage(
|
|
input_tokens=1, output_tokens=1, total_tokens=2, finish_reason='stop'))
|
|
return ModelOutput(text='', tool_calls=calls, model='mock', usage=Usage(
|
|
input_tokens=1, output_tokens=len(calls), total_tokens=1 + len(calls),
|
|
finish_reason='tool_calls'))
|
|
else:
|
|
last = next((m.content for m in reversed(messages) if m.role == 'user'), '')
|
|
text = last
|
|
target = None
|
|
for m in reversed(messages):
|
|
if m.role == 'user' and m.content.startswith('MOCKTARGET::'):
|
|
target = m.content[len('MOCKTARGET::'):]
|
|
break
|
|
if mode == 'oracle':
|
|
text = target if target is not None else last
|
|
elif mode == 'boxed':
|
|
if target is None:
|
|
m = _MOCK_PATTERNS[0][0].search(last)
|
|
target = m.group(1) if m else (re.findall(r'-?\d+\.?\d*', last) or ['0'])[-1]
|
|
text = f'The answer is \\boxed{{{target}}}.'
|
|
return ModelOutput(text=text, model='mock', usage=Usage(
|
|
input_tokens=1, output_tokens=1, total_tokens=2, finish_reason='stop'))
|