diff --git a/agent_platform/runtime/provider.py b/agent_platform/runtime/provider.py index 07db2df..bf8c449 100644 --- a/agent_platform/runtime/provider.py +++ b/agent_platform/runtime/provider.py @@ -30,6 +30,96 @@ def retry_delay_seconds(response: httpx.Response | None, attempt: int) -> float: return min(MAX_RETRY_DELAY_SECONDS, float(2 ** (attempt - 1))) +async def decode_complete_response(response: httpx.Response) -> dict[str, Any]: + content_type = response.headers.get("Content-Type", "").lower() + if "text/event-stream" not in content_type: + await response.aread() + data = response.json() + if not data.get("choices"): + raise RuntimeError("Model provider returned no choices") + return data + + envelope: dict[str, Any] = {"object": "chat.completion"} + message: dict[str, Any] = {"role": "assistant", "content": None} + content_parts: list[str] = [] + reasoning_parts: list[str] = [] + tool_calls: dict[int, dict[str, Any]] = {} + finish_reason: str | None = None + usage: dict[str, Any] | None = None + saw_choice = False + + async for line in response.aiter_lines(): + if not line.startswith("data:"): + continue + raw = line[5:].strip() + if not raw or raw == "[DONE]": + continue + chunk = httpx.Response(200, content=raw).json() + if chunk.get("error"): + raise RuntimeError(f"Model provider stream error: {chunk['error']}") + for key in ("id", "created", "model", "system_fingerprint"): + if chunk.get(key) is not None: + envelope[key] = chunk[key] + if chunk.get("usage"): + usage = chunk["usage"] + choices = chunk.get("choices") or [] + if not choices: + continue + saw_choice = True + choice = choices[0] + delta = choice.get("delta") or {} + if delta.get("role"): + message["role"] = delta["role"] + if isinstance(delta.get("content"), str): + content_parts.append(delta["content"]) + reasoning = delta.get("reasoning_content") + if not isinstance(reasoning, str): + reasoning = delta.get("reasoning") + if isinstance(reasoning, str): + reasoning_parts.append(reasoning) + for position, fragment in enumerate(delta.get("tool_calls") or []): + index = int(fragment.get("index", position)) + merged = tool_calls.setdefault( + index, + { + "id": "", + "type": "function", + "function": {"name": "", "arguments": ""}, + }, + ) + if fragment.get("id"): + merged["id"] = fragment["id"] + if fragment.get("type"): + merged["type"] = fragment["type"] + function = fragment.get("function") or {} + if function.get("name"): + merged["function"]["name"] += function["name"] + if function.get("arguments"): + merged["function"]["arguments"] += function["arguments"] + if choice.get("finish_reason") is not None: + finish_reason = choice["finish_reason"] + + if not saw_choice: + raise RuntimeError("Model provider stream returned no choices") + if content_parts: + message["content"] = "".join(content_parts) + if reasoning_parts: + message["reasoning_content"] = "".join(reasoning_parts) + if tool_calls: + message["tool_calls"] = [tool_calls[index] for index in sorted(tool_calls)] + + envelope["choices"] = [ + { + "index": 0, + "message": message, + "finish_reason": finish_reason, + } + ] + if usage is not None: + envelope["usage"] = usage + return envelope + + class ModelProvider: def __init__(self, settings: Settings, client: httpx.AsyncClient | None = None) -> None: self.settings = settings @@ -61,7 +151,7 @@ class ModelProvider: payload: dict[str, Any] = { "model": model, "messages": messages, - "stream": False, + "stream": True, } if tools: payload["tools"] = tools @@ -69,29 +159,27 @@ class ModelProvider: if temperature is not None: payload["temperature"] = temperature for attempt in range(1, COMPLETE_MAX_ATTEMPTS + 1): + response: httpx.Response | None = None try: - response = await self.client.post( + async with self.client.stream( + "POST", completions_url(self.settings.model_api_base_url), headers=self.headers, json=payload, - ) + ) as response: + if response.status_code in TRANSIENT_COMPLETE_STATUS_CODES and attempt < COMPLETE_MAX_ATTEMPTS: + await response.aread() + await asyncio.sleep(retry_delay_seconds(response, attempt)) + continue + + response.raise_for_status() + return await decode_complete_response(response) except httpx.TransportError: if attempt >= COMPLETE_MAX_ATTEMPTS: raise - await asyncio.sleep(retry_delay_seconds(None, attempt)) - continue - - if response.status_code in TRANSIENT_COMPLETE_STATUS_CODES and attempt < COMPLETE_MAX_ATTEMPTS: - await response.aread() await asyncio.sleep(retry_delay_seconds(response, attempt)) continue - response.raise_for_status() - data = response.json() - if not data.get("choices"): - raise RuntimeError("Model provider returned no choices") - return data - raise RuntimeError("Model provider retry loop exited unexpectedly") async def forward(self, payload: dict[str, Any]) -> httpx.Response: diff --git a/tests/test_provider.py b/tests/test_provider.py index 2eb1f3a..28ed69b 100644 --- a/tests/test_provider.py +++ b/tests/test_provider.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json from collections.abc import Awaitable, Callable import httpx @@ -78,3 +79,80 @@ async def test_complete_does_not_retry_client_error(settings: Settings) -> None: await _complete_with_handler(settings, handler) assert calls == 1 + + +async def test_complete_aggregates_streamed_tool_call(settings: Settings) -> None: + chunks = [ + { + "id": "completion-1", + "object": "chat.completion.chunk", + "created": 123, + "model": "test-model", + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "id": "call-1", + "type": "function", + "function": {"name": "ping", "arguments": '{"val'}, + } + ], + }, + "finish_reason": None, + } + ], + }, + { + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "function": {"arguments": 'ue":"ok"}'}, + } + ] + }, + "finish_reason": None, + } + ] + }, + { + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7}, + }, + ] + body = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n" + + def handler(request: httpx.Request) -> httpx.Response: + assert json.loads(request.content)["stream"] is True + return httpx.Response( + 200, + request=request, + headers={"Content-Type": "text/event-stream"}, + content=body, + ) + + result = await _complete_with_handler(settings, handler) + + choice = result["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"] == [ + { + "id": "call-1", + "type": "function", + "function": {"name": "ping", "arguments": '{"value":"ok"}'}, + } + ] + assert result["usage"]["total_tokens"] == 7