add new agent components

src/openai_compat.py

@@ -1,10 +1,17 @@
 from __future__ import annotations
 
 import json
-from typing import Any
+from typing import Any, Iterator
 from urllib import error, request
 
-from .agent_types import AssistantTurn, ModelConfig, ToolCall
+from .agent_types import (
+    AssistantTurn,
+    ModelConfig,
+    OutputSchemaConfig,
+    StreamEvent,
+    ToolCall,
+    UsageStats,
+)
 
 
 class OpenAICompatError(RuntimeError):
@@ -66,6 +73,64 @@ def _parse_tool_arguments(raw_arguments: Any) -> dict[str, Any]:
     )
 
 
+def _optional_int(value: Any) -> int:
+    if isinstance(value, bool):
+        return 0
+    if isinstance(value, int):
+        return value
+    if isinstance(value, float):
+        return int(value)
+    if isinstance(value, str):
+        try:
+            return int(value)
+        except ValueError:
+            return 0
+    return 0
+
+
+def _parse_usage(payload: Any) -> UsageStats:
+    if not isinstance(payload, dict):
+        return UsageStats()
+    completion_details = payload.get('completion_tokens_details')
+    if not isinstance(completion_details, dict):
+        completion_details = {}
+    return UsageStats(
+        input_tokens=(
+            _optional_int(payload.get('input_tokens'))
+            or _optional_int(payload.get('prompt_tokens'))
+            or _optional_int(payload.get('prompt_eval_count'))
+        ),
+        output_tokens=(
+            _optional_int(payload.get('output_tokens'))
+            or _optional_int(payload.get('completion_tokens'))
+            or _optional_int(payload.get('eval_count'))
+        ),
+        cache_creation_input_tokens=_optional_int(
+            payload.get('cache_creation_input_tokens')
+        ),
+        cache_read_input_tokens=_optional_int(payload.get('cache_read_input_tokens')),
+        reasoning_tokens=(
+            _optional_int(payload.get('reasoning_tokens'))
+            or _optional_int(completion_details.get('reasoning_tokens'))
+        ),
+    )
+
+
+def _build_response_format(
+    schema: OutputSchemaConfig | None,
+) -> dict[str, Any] | None:
+    if schema is None:
+        return None
+    return {
+        'type': 'json_schema',
+        'json_schema': {
+            'name': schema.name,
+            'schema': schema.schema,
+            'strict': schema.strict,
+        },
+    }
+
+
 class OpenAICompatClient:
     """Minimal OpenAI-compatible chat client for local model servers."""
 
@@ -76,15 +141,81 @@ class OpenAICompatClient:
         self,
         messages: list[dict[str, Any]],
         tools: list[dict[str, Any]],
+        *,
+        output_schema: OutputSchemaConfig | None = None,
     ) -> AssistantTurn:
-        payload = {
-            'model': self.config.model,
-            'messages': messages,
-            'tools': tools,
-            'tool_choice': 'auto',
-            'temperature': self.config.temperature,
-            'stream': False,
-        }
+        payload = self._request_json(
+            self._build_payload(
+                messages=messages,
+                tools=tools,
+                stream=False,
+                output_schema=output_schema,
+            )
+        )
+        choices = payload.get('choices')
+        if not isinstance(choices, list) or not choices:
+            raise OpenAICompatError('Local model backend returned no choices')
+        first_choice = choices[0]
+        if not isinstance(first_choice, dict):
+            raise OpenAICompatError('Local model backend returned malformed choice data')
+
+        message = first_choice.get('message')
+        if not isinstance(message, dict):
+            raise OpenAICompatError('Local model backend returned no assistant message')
+
+        content = _normalize_content(message.get('content'))
+        tool_calls = self._parse_tool_calls_from_message(message)
+
+        finish_reason = first_choice.get('finish_reason')
+        if finish_reason is not None and not isinstance(finish_reason, str):
+            finish_reason = str(finish_reason)
+
+        return AssistantTurn(
+            content=content,
+            tool_calls=tuple(tool_calls),
+            finish_reason=finish_reason,
+            raw_message=message,
+            usage=_parse_usage(payload.get('usage')),
+        )
+
+    def stream(
+        self,
+        messages: list[dict[str, Any]],
+        tools: list[dict[str, Any]],
+        *,
+        output_schema: OutputSchemaConfig | None = None,
+    ) -> Iterator[StreamEvent]:
+        payload = self._build_payload(
+            messages=messages,
+            tools=tools,
+            stream=True,
+            output_schema=output_schema,
+        )
+        req = request.Request(
+            _join_url(self.config.base_url, '/chat/completions'),
+            data=json.dumps(payload).encode('utf-8'),
+            headers={
+                'Authorization': f'Bearer {self.config.api_key}',
+                'Content-Type': 'application/json',
+            },
+            method='POST',
+        )
+        try:
+            with request.urlopen(req, timeout=self.config.timeout_seconds) as response:
+                yield StreamEvent(type='message_start')
+                for event_payload in self._iter_sse_payloads(response):
+                    yield from self._parse_stream_payload(event_payload)
+        except error.HTTPError as exc:
+            detail = exc.read().decode('utf-8', errors='replace')
+            raise OpenAICompatError(
+                f'HTTP {exc.code} from local model backend: {detail}'
+            ) from exc
+        except error.URLError as exc:
+            raise OpenAICompatError(
+                f'Unable to reach local model backend at {self.config.base_url}: {exc.reason}'
+            ) from exc
+
+    def _request_json(self, payload: dict[str, Any]) -> dict[str, Any]:
         body = json.dumps(payload).encode('utf-8')
         req = request.Request(
             _join_url(self.config.base_url, '/chat/completions'),
@@ -112,19 +243,34 @@ class OpenAICompatClient:
             payload = json.loads(raw.decode('utf-8'))
         except json.JSONDecodeError as exc:
             raise OpenAICompatError('Local model backend returned invalid JSON') from exc
+        if not isinstance(payload, dict):
+            raise OpenAICompatError('Local model backend returned malformed JSON payload')
+        return payload
 
-        choices = payload.get('choices')
-        if not isinstance(choices, list) or not choices:
-            raise OpenAICompatError('Local model backend returned no choices')
-        first_choice = choices[0]
-        if not isinstance(first_choice, dict):
-            raise OpenAICompatError('Local model backend returned malformed choice data')
+    def _build_payload(
+        self,
+        *,
+        messages: list[dict[str, Any]],
+        tools: list[dict[str, Any]],
+        stream: bool,
+        output_schema: OutputSchemaConfig | None,
+    ) -> dict[str, Any]:
+        payload: dict[str, Any] = {
+            'model': self.config.model,
+            'messages': messages,
+            'tools': tools,
+            'tool_choice': 'auto',
+            'temperature': self.config.temperature,
+            'stream': stream,
+        }
+        if stream:
+            payload['stream_options'] = {'include_usage': True}
+        response_format = _build_response_format(output_schema)
+        if response_format is not None:
+            payload['response_format'] = response_format
+        return payload
 
-        message = first_choice.get('message')
-        if not isinstance(message, dict):
-            raise OpenAICompatError('Local model backend returned no assistant message')
-
-        content = _normalize_content(message.get('content'))
+    def _parse_tool_calls_from_message(self, message: dict[str, Any]) -> list[ToolCall]:
         tool_calls: list[ToolCall] = []
         raw_tool_calls = message.get('tool_calls')
         if isinstance(raw_tool_calls, list):
@@ -149,14 +295,119 @@ class OpenAICompatClient:
                 raise OpenAICompatError('Function call missing name')
             arguments = _parse_tool_arguments(function_call.get('arguments'))
             tool_calls.append(ToolCall(id='call_0', name=name, arguments=arguments))
+        return tool_calls
 
-        finish_reason = first_choice.get('finish_reason')
-        if finish_reason is not None and not isinstance(finish_reason, str):
-            finish_reason = str(finish_reason)
+    def _iter_sse_payloads(self, response: Any) -> Iterator[dict[str, Any]]:
+        buffer: list[str] = []
+        while True:
+            line = response.readline()
+            if not line:
+                break
+            if isinstance(line, bytes):
+                text = line.decode('utf-8', errors='replace')
+            else:
+                text = str(line)
+            stripped = text.strip()
+            if not stripped:
+                if not buffer:
+                    continue
+                joined = '\n'.join(buffer)
+                buffer.clear()
+                if joined == '[DONE]':
+                    break
+                try:
+                    payload = json.loads(joined)
+                except json.JSONDecodeError as exc:
+                    raise OpenAICompatError(
+                        f'Invalid JSON in streaming response: {joined!r}'
+                    ) from exc
+                if not isinstance(payload, dict):
+                    raise OpenAICompatError('Malformed SSE payload from model backend')
+                yield payload
+                continue
+            if stripped.startswith('data:'):
+                buffer.append(stripped[5:].strip())
 
-        return AssistantTurn(
-            content=content,
-            tool_calls=tuple(tool_calls),
-            finish_reason=finish_reason,
-            raw_message=message,
-        )
+        if buffer:
+            joined = '\n'.join(buffer)
+            if joined != '[DONE]':
+                try:
+                    payload = json.loads(joined)
+                except json.JSONDecodeError as exc:
+                    raise OpenAICompatError(
+                        f'Invalid trailing JSON in streaming response: {joined!r}'
+                    ) from exc
+                if not isinstance(payload, dict):
+                    raise OpenAICompatError('Malformed trailing SSE payload from model backend')
+                yield payload
+
+    def _parse_stream_payload(
+        self,
+        payload: dict[str, Any],
+    ) -> Iterator[StreamEvent]:
+        usage = _parse_usage(payload.get('usage'))
+        if usage.total_tokens:
+            yield StreamEvent(
+                type='usage',
+                usage=usage,
+                raw_event=payload,
+            )
+
+        choices = payload.get('choices')
+        if not isinstance(choices, list):
+            return
+
+        for choice in choices:
+            if not isinstance(choice, dict):
+                continue
+            delta = choice.get('delta')
+            if not isinstance(delta, dict):
+                delta = {}
+            content = delta.get('content')
+            if isinstance(content, str) and content:
+                yield StreamEvent(
+                    type='content_delta',
+                    delta=content,
+                    raw_event=choice,
+                )
+            tool_calls = delta.get('tool_calls')
+            if isinstance(tool_calls, list):
+                for raw_tool_call in tool_calls:
+                    if not isinstance(raw_tool_call, dict):
+                        continue
+                    function_block = raw_tool_call.get('function')
+                    if not isinstance(function_block, dict):
+                        function_block = {}
+                    yield StreamEvent(
+                        type='tool_call_delta',
+                        tool_call_index=(
+                            raw_tool_call.get('index')
+                            if isinstance(raw_tool_call.get('index'), int)
+                            else 0
+                        ),
+                        tool_call_id=(
+                            raw_tool_call.get('id')
+                            if isinstance(raw_tool_call.get('id'), str)
+                            else None
+                        ),
+                        tool_name=(
+                            function_block.get('name')
+                            if isinstance(function_block.get('name'), str)
+                            else None
+                        ),
+                        arguments_delta=(
+                            function_block.get('arguments')
+                            if isinstance(function_block.get('arguments'), str)
+                            else ''
+                        ),
+                        raw_event=raw_tool_call,
+                    )
+            finish_reason = choice.get('finish_reason')
+            if finish_reason is not None:
+                if not isinstance(finish_reason, str):
+                    finish_reason = str(finish_reason)
+                yield StreamEvent(
+                    type='message_stop',
+                    finish_reason=finish_reason,
+                    raw_event=choice,
+                )