83 lines
3.0 KiB
Python

"""Structured model output -- the hinge between single-turn eval and agents.
Every ModelAdapter returns ModelOutput, never a bare string:
- single-turn recipes read .text
- agent loops read .tool_calls and feed observations back
- accounting/monitoring reads .usage
"""
import time
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field
class ToolCall(BaseModel):
"""One function call the model wants executed (OpenAI tool_calls shape)."""
id: str = ''
name: str
arguments: str = '' # JSON-encoded args string
arguments_dict: Dict[str, Any] = Field(default_factory=dict) # parsed convenience
def to_openai(self) -> Dict[str, Any]:
return {'id': self.id or f'call_{self.name}', 'type': 'function',
'function': {'name': self.name, 'arguments': self.arguments or '{}'}}
class Usage(BaseModel):
input_tokens: int = 0
output_tokens: int = 0
total_tokens: int = 0
cost: float = 0.0
latency_s: float = 0.0
finish_reason: str = ''
def __add__(self, other: 'Usage') -> 'Usage':
return Usage(
input_tokens=self.input_tokens + other.input_tokens,
output_tokens=self.output_tokens + other.output_tokens,
total_tokens=self.total_tokens + other.total_tokens,
cost=round(self.cost + other.cost, 6),
latency_s=round(self.latency_s + other.latency_s, 3),
finish_reason=self.finish_reason or other.finish_reason,
)
class ChoiceLogprob(BaseModel):
"""Continuation scoring for one choice (paper-faithful MCQ evaluation)."""
text: str = ''
logprob: float = 0.0 # total logprob of the choice tokens
token_logprobs: List[float] = Field(default_factory=list)
num_tokens: int = 0
@property
def logprob_per_char(self) -> float:
"""lm-eval-harness acc_norm uses per-BYTE normalization; chars are the
practical proxy (identical ranking for ascii-dominant choices)."""
return self.logprob / max(len(self.text), 1)
class ModelOutput(BaseModel):
"""What every adapter returns. text may be '' when the model only calls tools."""
text: str = ''
tool_calls: List[ToolCall] = Field(default_factory=list)
usage: Usage = Field(default_factory=Usage)
raw: Optional[Dict[str, Any]] = None # provider response (audit/retry)
model: str = ''
created_at: str = Field(default_factory=lambda: time.strftime('%Y-%m-%d %H:%M:%S'))
choice_logprobs: Optional[List[ChoiceLogprob]] = None # P1 continuation scoring
@property
def is_tool_call(self) -> bool:
return bool(self.tool_calls)
def best_choice(self, normalized: bool = False) -> int:
"""argmax over choice_logprobs (acc / acc_norm)."""
if not self.choice_logprobs:
return -1
key = (lambda c: c.logprob_per_char if normalized else c.logprob)
return max(range(len(self.choice_logprobs)), key=lambda i: key(self.choice_logprobs[i]))