generated from coulomb/repo-seed
Bernd Worsch
f76a58d6e9
- Remove redundant async_execute_prompt overrides from OpenAI/Gemini/OpenRouter adapters (identical to base class default — asyncio import also removed) - Cache prompt.split() result in MockLLMAdapter to avoid double evaluation - Promote deferred LLMBudgetExceededError imports to module level in models.py and adapter.py (no circular dependency) - Auto-populate context dict in LLMBudgetExceededError.__init__ so callers need not pass redundant context= kwarg Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
132 lines
4.4 KiB
Python
132 lines
4.4 KiB
Python
"""
|
|
Shared data models for LLM execution.
|
|
|
|
These classes are the canonical definitions; they are re-exported by
|
|
markitect.prompts.execution.models for backward compatibility.
|
|
"""
|
|
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, Any, Optional
|
|
|
|
from llm_connect.exceptions import LLMBudgetExceededError
|
|
|
|
|
|
class BudgetTracker:
|
|
"""Shared token budget for a call or delegation chain.
|
|
|
|
Thread-safe. Tracks cumulative token spend across multiple adapter
|
|
calls. Raises ``LLMBudgetExceededError`` when the cap is exceeded.
|
|
|
|
Example::
|
|
|
|
tracker = BudgetTracker(total=4000)
|
|
config = RunConfig(budget_tracker=tracker)
|
|
# All adapter calls sharing this config will consume from the same cap.
|
|
"""
|
|
|
|
def __init__(self, total: int) -> None:
|
|
if total <= 0:
|
|
raise ValueError(f"BudgetTracker total must be positive, got {total}")
|
|
self.total = total
|
|
self.spent = 0
|
|
self._lock = threading.Lock()
|
|
|
|
def remaining(self) -> int:
|
|
"""Return tokens remaining in the budget."""
|
|
return max(0, self.total - self.spent)
|
|
|
|
def consume(self, tokens: int) -> None:
|
|
"""Record *tokens* as spent. Raises ``LLMBudgetExceededError`` if cap exceeded."""
|
|
with self._lock:
|
|
new_spent = self.spent + tokens
|
|
if new_spent > self.total:
|
|
raise LLMBudgetExceededError(
|
|
f"Token budget exceeded: {new_spent} tokens used, cap is {self.total}",
|
|
total=self.total,
|
|
spent=self.spent,
|
|
requested=tokens,
|
|
)
|
|
self.spent = new_spent
|
|
|
|
def __repr__(self) -> str:
|
|
return f"BudgetTracker(total={self.total}, spent={self.spent}, remaining={self.remaining()})"
|
|
|
|
|
|
@dataclass
|
|
class RunConfig:
|
|
"""
|
|
Configuration for prompt execution.
|
|
|
|
Attributes:
|
|
model_name: LLM model to use
|
|
temperature: Model temperature (0.0-1.0)
|
|
max_tokens: Maximum tokens to generate
|
|
model_params: Additional model parameters
|
|
max_depth: Maximum generation depth for nested runs
|
|
skip_if_exists: Skip if identical InputBundleHash exists
|
|
timeout_seconds: Execution timeout
|
|
"""
|
|
model_name: str = "gpt-4"
|
|
temperature: float = 0.7
|
|
max_tokens: int = 2000
|
|
model_params: Dict[str, Any] = field(default_factory=dict)
|
|
max_depth: int = 3
|
|
skip_if_exists: bool = True
|
|
timeout_seconds: int = 300
|
|
budget_tracker: Optional["BudgetTracker"] = field(default=None, repr=False)
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert to dictionary. ``budget_tracker`` is excluded (runtime object)."""
|
|
return {
|
|
"model_name": self.model_name,
|
|
"temperature": self.temperature,
|
|
"max_tokens": self.max_tokens,
|
|
"model_params": self.model_params,
|
|
"max_depth": self.max_depth,
|
|
"skip_if_exists": self.skip_if_exists,
|
|
"timeout_seconds": self.timeout_seconds,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Dict[str, Any]) -> "RunConfig":
|
|
"""Create from dictionary."""
|
|
return cls(
|
|
model_name=data.get("model_name", "gpt-4"),
|
|
temperature=data.get("temperature", 0.7),
|
|
max_tokens=data.get("max_tokens", 2000),
|
|
model_params=data.get("model_params", {}),
|
|
max_depth=data.get("max_depth", 3),
|
|
skip_if_exists=data.get("skip_if_exists", True),
|
|
timeout_seconds=data.get("timeout_seconds", 300),
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class LLMResponse:
|
|
"""
|
|
Response from LLM execution.
|
|
|
|
Attributes:
|
|
content: Generated content
|
|
model: Model used
|
|
usage: Token usage statistics
|
|
finish_reason: Why generation stopped
|
|
metadata: Additional response metadata
|
|
"""
|
|
content: str
|
|
model: str
|
|
usage: Dict[str, int] = field(default_factory=dict)
|
|
finish_reason: str = "stop"
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert to dictionary."""
|
|
return {
|
|
"content": self.content,
|
|
"model": self.model,
|
|
"usage": self.usage,
|
|
"finish_reason": self.finish_reason,
|
|
"metadata": self.metadata,
|
|
}
|