"""Neutral model contract adapter and provider implementations.""" from __future__ import annotations import abc import asyncio import json import uuid from typing import Any, Dict, List, Optional, Tuple import httpx from backend.errors import AppError, InvalidInputError, ModelContextExceededError, ModelOutputLimitError, UpstreamFailedError from backend.errors import UpstreamResponseTooLargeError from backend.upstream import read_json from backend.validation import validate_model_request MAX_MODEL_RESPONSE_BYTES = 128 * 1024 * 1024 # 128 MiB class ModelAdapter(abc.ABC): """Abstract model provider adapter.""" @abc.abstractmethod async def complete( self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], system_instruction: str, provider_state_meta: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Execute a model completion turn. Returns ModelResponse dict: { "content": List[Content], "stop_reason": "stop" | "tool_calls" | "length", "usage": {"input_tokens": int, "output_tokens": int}, "provider_state": Optional[str] } """ raise NotImplementedError async def close(self) -> None: pass class FakeModelAdapter(ModelAdapter): """Deterministic model adapter for testing and dev mode.""" def __init__(self, script: Optional[List[Dict[str, Any]]] = None): self.script = list(script) if script else [] self.call_count = 0 self.received_requests: List[Dict[str, Any]] = [] def queue_response(self, response: Dict[str, Any]) -> None: self.script.append(response) async def complete( self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], system_instruction: str, provider_state_meta: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: self.call_count += 1 self.received_requests.append({ "messages": messages, "tools": tools, "system_instruction": system_instruction, "provider_state_meta": provider_state_meta, }) if self.script: return self.script.pop(0) # Default fallback response: simple text stop return { "content": [{"type": "text", "text": "I have completed the research."}], "stop_reason": "stop", "usage": {"input_tokens": 10, "output_tokens": 10}, } class OpenAIModelAdapter(ModelAdapter): """OpenAI-compatible Chat Completions API adapter.""" def __init__( self, api_key: Optional[str], model_name: str = "gpt-4o", endpoint: Optional[str] = None, context_window_tokens: int = 128_000, max_output_tokens: int = 4096, timeout: float = 60.0, transport: Optional[httpx.AsyncBaseTransport] = None, ): if not api_key or not api_key.strip(): raise ValueError("api_key is required for OpenAIModelAdapter") self.api_key = api_key.strip() self.model_name = model_name self.endpoint = endpoint or "https://api.openai.com/v1/chat/completions" self.context_window_tokens = context_window_tokens self.max_output_tokens = max_output_tokens self._client = httpx.AsyncClient( timeout=httpx.Timeout(timeout), headers={ "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", "Accept-Encoding": "gzip, deflate", }, transport=transport, ) async def close(self) -> None: await self._client.aclose() async def complete( self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], system_instruction: str, provider_state_meta: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: # Translate neutral messages to OpenAI messages openai_messages: List[Dict[str, Any]] = [ {"role": "system", "content": system_instruction} ] for msg in messages: role = msg.get("role") if role == "user": content_chunks = msg.get("content", []) text = "".join(c.get("text", "") for c in content_chunks if c.get("type") == "text") openai_messages.append({"role": "user", "content": text}) elif role == "assistant": content_chunks = msg.get("content", []) text_parts = [c.get("text", "") for c in content_chunks if c.get("type") == "text"] tool_calls = [] for c in content_chunks: if c.get("type") == "tool_call": tool_calls.append({ "id": c.get("id"), "type": "function", "function": { "name": c.get("name"), "arguments": json.dumps(c.get("arguments", {})), }, }) entry: Dict[str, Any] = {"role": "assistant"} if text_parts: entry["content"] = "".join(text_parts) if tool_calls: entry["tool_calls"] = tool_calls openai_messages.append(entry) elif role == "tool": openai_messages.append({ "role": "tool", "tool_call_id": msg.get("tool_call_id"), "content": str(msg.get("content", "")), }) else: raise InvalidInputError(f"Unsupported message role: {role}") # Translate neutral tools to OpenAI tools openai_tools = [] for t in tools: openai_tools.append({ "type": "function", "function": { "name": t.get("name"), "description": t.get("description", ""), "parameters": t.get("input_schema", {}), }, }) req_body: Dict[str, Any] = { "model": self.model_name, "messages": openai_messages, "max_tokens": self.max_output_tokens, } if openai_tools: req_body["tools"] = openai_tools try: req = self._client.build_request("POST", self.endpoint, json=req_body) response = await self._client.send(req, stream=True) except Exception as e: raise UpstreamFailedError(f"Model request failed: {type(e).__name__}") from e if response.status_code != 200: status = response.status_code try: error_data = await read_json(response, 4096) except AppError: error_data = {} provider_error = error_data.get("error", {}) if isinstance(error_data, dict) else {} if status == 400 and isinstance(provider_error, dict) and provider_error.get("code") == "context_length_exceeded": raise ModelContextExceededError() raise UpstreamFailedError(f"Model provider returned HTTP {status}") data = await read_json(response, MAX_MODEL_RESPONSE_BYTES) if not isinstance(data, dict): raise UpstreamFailedError("Invalid model response structure") choices = data.get("choices", []) if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict): raise UpstreamFailedError("Model returned empty choices list") choice = choices[0] msg_resp = choice.get("message", {}) finish_reason = choice.get("finish_reason") if not isinstance(msg_resp, dict) or finish_reason not in ("stop", "tool_calls", "length", "content_filter"): raise UpstreamFailedError("Invalid model choice structure") content: List[Dict[str, Any]] = [] text = msg_resp.get("content") if text is not None and not isinstance(text, str): raise UpstreamFailedError("Invalid model text content") if text: content.append({"type": "text", "text": text}) tool_calls = msg_resp.get("tool_calls", []) if not isinstance(tool_calls, list): raise UpstreamFailedError("Invalid model tool calls") seen_tool_ids = set() for tc in tool_calls: if not isinstance(tc, dict) or not isinstance(tc.get("function"), dict): raise UpstreamFailedError("Invalid model tool call") fn = tc.get("function", {}) call_id = tc.get("id") name = fn.get("name", "") raw_args = fn.get("arguments", "{}") try: args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args except Exception as e: raise UpstreamFailedError(f"Model returned invalid JSON arguments for tool '{name}'") from e if not isinstance(args, dict) or not isinstance(call_id, str) or not call_id or not isinstance(name, str) or not name: raise UpstreamFailedError("Invalid model tool arguments or identity") if call_id in seen_tool_ids: raise UpstreamFailedError("Model returned duplicate tool call IDs") seen_tool_ids.add(call_id) content.append({ "type": "tool_call", "id": call_id, "name": name, "arguments": args, }) if finish_reason == "tool_calls": stop_reason = "tool_calls" elif finish_reason == "length": stop_reason = "length" else: stop_reason = "stop" usage_data = data.get("usage", {}) if not isinstance(usage_data, dict): raise UpstreamFailedError("Invalid model usage") usage = { "input_tokens": usage_data.get("prompt_tokens", 0), "output_tokens": usage_data.get("completion_tokens", 0), } if any(type(v) is not int or v < 0 for v in usage.values()): raise UpstreamFailedError("Invalid model token usage") if finish_reason == "content_filter": raise UpstreamFailedError("Model provider could not complete response") # Check token limits if usage["input_tokens"] > self.context_window_tokens: raise ModelContextExceededError( f"Input tokens ({usage['input_tokens']}) exceeded context window ({self.context_window_tokens})" ) if usage["output_tokens"] > self.max_output_tokens: raise ModelOutputLimitError( f"Output tokens ({usage['output_tokens']}) exceeded limit ({self.max_output_tokens})" ) return { "content": content, "stop_reason": stop_reason, "usage": usage, } class ModelDispatcher: """Dispatches ModelRequest messages with call budgeting and state tracking.""" def __init__(self, adapter: ModelAdapter, max_calls: int = 50): self.adapter = adapter self.max_calls = max_calls self.call_count = 0 self._valid_handles: Dict[str, Dict[str, Any]] = {} def issue_handle(self, meta: Dict[str, Any]) -> str: handle = uuid.uuid4().hex self._valid_handles[handle] = meta return handle def validate_and_get_handle(self, handle: Optional[str]) -> Optional[Dict[str, Any]]: if not handle: return None if handle not in self._valid_handles: raise InvalidInputError("Forged or invalid provider_state handle") return self._valid_handles[handle] async def dispatch( self, request_payload: Dict[str, Any], system_instruction: str, ) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: """Dispatch a model request. Returns (response_dict, error_dict).""" try: validate_model_request(request_payload) except AppError as exc: return None, exc.to_error_dict() if self.call_count >= self.max_calls: err = UpstreamFailedError(f"Model call limit ({self.max_calls}) reached") return None, err.to_error_dict() self.call_count += 1 messages = request_payload["messages"] tools = request_payload["tools"] # Validate provider_state handles if present on assistant messages last_meta = None for m in messages: if m.get("role") == "assistant" and "provider_state" in m: state_handle = m.get("provider_state") try: last_meta = self.validate_and_get_handle(state_handle) if last_meta and "content" in last_meta and last_meta["content"] != m["content"]: raise InvalidInputError("provider_state is not associated with this assistant message") except AppError as e: return None, e.to_error_dict() try: resp = await self.adapter.complete( messages=messages, tools=tools, system_instruction=system_instruction, provider_state_meta=last_meta, ) if await asyncio.to_thread(lambda: len(json.dumps(resp).encode("utf-8"))) > MAX_MODEL_RESPONSE_BYTES: raise UpstreamResponseTooLargeError("Serialized model result exceeds 128 MiB") # If provider_state returned or needed, store handle if resp.get("provider_state"): handle = self.issue_handle({"state": resp["provider_state"], "content": resp["content"]}) resp["provider_state"] = handle return resp, None except AppError as e: return None, e.to_error_dict() except Exception as e: err = UpstreamFailedError(f"Model completion failed: {type(e).__name__}") return None, err.to_error_dict()