File size: 12,093 Bytes
73ecef8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 | """
LLM Backend — Swappable inference layer.
Supports: HuggingFace Inference Providers, OpenAI, Anthropic, local models,
or any custom backend. Swap by changing one constructor call.
Design: Abstract base class with structured output support.
Inspired by smolagents Model interface + HF Inference Providers API.
"""
from __future__ import annotations
import json
import logging
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Message types (OpenAI-compatible chat format)
# ---------------------------------------------------------------------------
@dataclass
class ChatMessage:
role: str # "system", "user", "assistant"
content: str
# ---------------------------------------------------------------------------
# Abstract LLM Backend
# ---------------------------------------------------------------------------
class LLMBackend(ABC):
"""
Abstract LLM backend. All modules call this — swap the implementation
to change the underlying model without touching any other code.
Subclasses must implement `generate()` which takes messages and returns
a string. Optionally implement `generate_structured()` for JSON-schema
constrained generation (used by the Purpose Function for reliable scoring).
"""
@abstractmethod
def generate(
self,
messages: list[ChatMessage],
temperature: float = 0.7,
max_tokens: int = 2048,
stop: list[str] | None = None,
) -> str:
"""Generate a text completion from chat messages."""
...
def generate_structured(
self,
messages: list[ChatMessage],
schema: dict[str, Any],
temperature: float = 0.3,
max_tokens: int = 1024,
) -> dict[str, Any]:
"""
Generate with JSON schema constraint.
Default implementation: append schema instruction to last message
and parse JSON from response. Override for native structured output.
"""
schema_instruction = (
f"\n\nYou MUST respond with valid JSON matching this schema:\n"
f"```json\n{json.dumps(schema, indent=2)}\n```\n"
f"Respond ONLY with the JSON object, no other text."
)
augmented = list(messages)
last = augmented[-1]
augmented[-1] = ChatMessage(
role=last.role, content=last.content + schema_instruction
)
raw = self.generate(augmented, temperature=temperature, max_tokens=max_tokens)
# Extract JSON from response (handle markdown code blocks)
text = raw.strip()
if text.startswith("```"):
lines = text.split("\n")
# Remove first and last ``` lines
json_lines = []
inside = False
for line in lines:
if line.strip().startswith("```") and not inside:
inside = True
continue
elif line.strip() == "```" and inside:
break
elif inside:
json_lines.append(line)
text = "\n".join(json_lines)
return json.loads(text)
# ---------------------------------------------------------------------------
# HuggingFace Inference Provider Backend
# ---------------------------------------------------------------------------
class HFInferenceBackend(LLMBackend):
"""
Uses huggingface_hub InferenceClient for HF Inference Providers.
Supports: Cerebras, Novita, Fireworks, Together, SambaNova, etc.
Models: Qwen, Llama, Mistral, DeepSeek — anything on HF Hub.
Example:
backend = HFInferenceBackend(
model_id="Qwen/Qwen3-32B",
provider="cerebras",
)
"""
def __init__(
self,
model_id: str = "Qwen/Qwen3-32B",
provider: str = "auto",
api_key: str | None = None,
):
from huggingface_hub import InferenceClient
self.model_id = model_id
self.provider = provider
self.client = InferenceClient(
provider=provider,
api_key=api_key or os.environ.get("HF_TOKEN"),
)
def generate(
self,
messages: list[ChatMessage],
temperature: float = 0.7,
max_tokens: int = 2048,
stop: list[str] | None = None,
) -> str:
msg_dicts = [{"role": m.role, "content": m.content} for m in messages]
response = self.client.chat_completion(
model=self.model_id,
messages=msg_dicts,
temperature=temperature,
max_tokens=max_tokens,
stop=stop or [],
)
return response.choices[0].message.content
def generate_structured(
self,
messages: list[ChatMessage],
schema: dict[str, Any],
temperature: float = 0.3,
max_tokens: int = 1024,
) -> dict[str, Any]:
msg_dicts = [{"role": m.role, "content": m.content} for m in messages]
response = self.client.chat_completion(
model=self.model_id,
messages=msg_dicts,
temperature=temperature,
max_tokens=max_tokens,
response_format={
"type": "json_schema",
"json_schema": {"schema": schema},
},
)
return json.loads(response.choices[0].message.content)
# ---------------------------------------------------------------------------
# OpenAI-Compatible Backend (OpenAI, Azure, vLLM, Ollama, LiteLLM)
# ---------------------------------------------------------------------------
class OpenAICompatibleBackend(LLMBackend):
"""
Works with any OpenAI-compatible API endpoint.
Examples:
# OpenAI
backend = OpenAICompatibleBackend(model="gpt-4o")
# Local Ollama
backend = OpenAICompatibleBackend(
model="llama3.2",
base_url="http://localhost:11434/v1",
api_key="ollama",
)
# vLLM server
backend = OpenAICompatibleBackend(
model="meta-llama/Llama-3.2-3B-Instruct",
base_url="http://localhost:8000/v1",
api_key="token-placeholder",
)
# HF Inference via OpenAI SDK (for structured output with .parse())
backend = OpenAICompatibleBackend(
model="Qwen/Qwen3-32B",
base_url="https://router.huggingface.co/cerebras/v1",
api_key=os.environ["HF_TOKEN"],
)
"""
def __init__(
self,
model: str = "gpt-4o",
base_url: str | None = None,
api_key: str | None = None,
):
from openai import OpenAI
self.model = model
self.client = OpenAI(
base_url=base_url,
api_key=api_key or os.environ.get("OPENAI_API_KEY"),
)
def generate(
self,
messages: list[ChatMessage],
temperature: float = 0.7,
max_tokens: int = 2048,
stop: list[str] | None = None,
) -> str:
msg_dicts = [{"role": m.role, "content": m.content} for m in messages]
response = self.client.chat.completions.create(
model=self.model,
messages=msg_dicts,
temperature=temperature,
max_tokens=max_tokens,
stop=stop,
)
return response.choices[0].message.content
def generate_structured(
self,
messages: list[ChatMessage],
schema: dict[str, Any],
temperature: float = 0.3,
max_tokens: int = 1024,
) -> dict[str, Any]:
msg_dicts = [{"role": m.role, "content": m.content} for m in messages]
response = self.client.chat.completions.create(
model=self.model,
messages=msg_dicts,
temperature=temperature,
max_tokens=max_tokens,
response_format={
"type": "json_schema",
"json_schema": {"name": "purpose_score", "schema": schema},
},
)
return json.loads(response.choices[0].message.content)
# ---------------------------------------------------------------------------
# Mock Backend (for testing without API calls)
# ---------------------------------------------------------------------------
class MockLLMBackend(LLMBackend):
"""
Deterministic mock backend for testing the framework without LLM calls.
Returns canned responses based on keywords in the prompt, or a default.
You can register custom response handlers.
"""
def __init__(self):
self._handlers: list[tuple[str, str | callable]] = []
self._structured_default: dict[str, Any] = {}
self._call_log: list[dict] = []
def register_handler(
self, keyword: str, response: str | callable
) -> "MockLLMBackend":
"""Add a keyword-matched response handler. Checked in order."""
self._handlers.append((keyword, response))
return self
def set_structured_default(self, default: dict[str, Any]) -> "MockLLMBackend":
"""Set the default response for structured generation."""
self._structured_default = default
return self
@property
def call_log(self) -> list[dict]:
return self._call_log
def generate(
self,
messages: list[ChatMessage],
temperature: float = 0.7,
max_tokens: int = 2048,
stop: list[str] | None = None,
) -> str:
full_text = " ".join(m.content for m in messages)
self._call_log.append({
"method": "generate",
"messages": [{"role": m.role, "content": m.content[:200]} for m in messages],
})
for keyword, response in self._handlers:
if keyword.lower() in full_text.lower():
if callable(response):
return response(messages)
return response
# Default: echo the last user message with a generic response
last_user = next(
(m.content for m in reversed(messages) if m.role == "user"),
"no input",
)
return f"[MockLLM] Acknowledged: {last_user[:100]}"
def generate_structured(
self,
messages: list[ChatMessage],
schema: dict[str, Any],
temperature: float = 0.3,
max_tokens: int = 1024,
) -> dict[str, Any]:
self._call_log.append({
"method": "generate_structured",
"schema_keys": list(schema.get("properties", {}).keys()),
})
# Try keyword handlers first — they may return JSON strings or dicts
full_text = " ".join(m.content for m in messages)
for keyword, response in self._handlers:
if keyword.lower() in full_text.lower():
if callable(response):
result = response(messages)
else:
result = response
# If handler returned a string, try to parse as JSON
if isinstance(result, str):
try:
return json.loads(result)
except (json.JSONDecodeError, TypeError):
pass
elif isinstance(result, dict):
return result
# Fall back to structured default
if self._structured_default:
return self._structured_default
# Build a minimal valid response from the schema
props = schema.get("properties", {})
result = {}
for key, prop in props.items():
ptype = prop.get("type", "string")
if ptype == "number":
result[key] = 5.0
elif ptype == "integer":
result[key] = 5
elif ptype == "boolean":
result[key] = True
else:
result[key] = f"mock_{key}"
return result
|