FastAPI app, upstream Confluence/model adapters, authoritative history, rootless container lifecycle, artifact storage and downloads, fake peers under backend/dev, tests under tests/backend. Root pytest.ini deselects the live marker by default; requirements gain the backend dependencies.
359 lines
14 KiB
Python
359 lines
14 KiB
Python
"""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()
|