Coverage for astrocyte/providers/openai.py: 71%
95 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""OpenAI LLMProvider adapter.
3Wraps the openai Python SDK to implement the Astrocyte LLMProvider protocol.
5Usage via config YAML:
6 llm_provider: openai
7 llm_provider_config:
8 api_key: ${OPENAI_API_KEY}
9 model: gpt-4o-mini
10 embedding_model: text-embedding-3-small
12Usage programmatically:
13 from astrocyte.providers.openai import OpenAIProvider
15 llm = OpenAIProvider(api_key="sk-...", model="gpt-4o-mini")
16"""
18from __future__ import annotations
20import asyncio
21import json
22import os
23from typing import ClassVar
25from astrocyte.types import (
26 Completion,
27 LLMCapabilities,
28 Message,
29 TokenUsage,
30 ToolCall,
31 ToolDefinition,
32)
35class OpenAIProvider:
36 """LLMProvider backed by the OpenAI API."""
38 SPI_VERSION: ClassVar[int] = 1
40 def __init__(
41 self,
42 *,
43 api_key: str | None = None,
44 model: str = "gpt-4o-mini",
45 embedding_model: str = "text-embedding-3-small",
46 base_url: str | None = None,
47 read_timeout: float | None = None,
48 max_concurrency: int | None = None,
49 ) -> None:
50 try:
51 import openai
52 except ImportError as e:
53 raise ImportError(
54 "The 'openai' package is required for OpenAIProvider. Install it with: pip install 'astrocyte[openai]'"
55 ) from e
57 resolved_key = api_key or os.environ.get("OPENAI_API_KEY")
58 if not resolved_key:
59 raise ValueError("OpenAI API key is required. Pass api_key= or set OPENAI_API_KEY.")
61 # Use HTTP/2 multiplexing when ``h2`` is installed. Without it, the
62 # OpenAI SDK's default httpx client falls back to HTTP/1.1, which can
63 # only carry one in-flight request per TCP connection — a hard
64 # serialisation bottleneck for concurrent retain/recall workloads.
65 # With HTTP/2, a single connection multiplexes dozens of requests
66 # (typical 10-30x throughput improvement on embedding/completion
67 # heavy phases). On benchmark hardware retaining 272 LoCoMo sessions
68 # drops from ~6h on HTTP/1.1 to under a minute with HTTP/2.
69 try:
70 import h2 # noqa: F401 — presence enables httpx HTTP/2
71 import httpx
73 http_client: httpx.AsyncClient | None = httpx.AsyncClient(
74 http2=True,
75 # Match OpenAI's defaults but with explicit pool settings so
76 # we don't suddenly serialise on bursty workloads.
77 limits=httpx.Limits(
78 max_connections=100,
79 max_keepalive_connections=20,
80 keepalive_expiry=30.0,
81 ),
82 # Read timeout was 600s historically — far too long when an
83 # HTTP/2 stream stalls (no RST/FIN, just no bytes). Such
84 # stalls block for the full timeout × ``max_retries`` (≈30min
85 # per stuck call) and silently murder benchmark wall-clock.
86 # OpenAI completions normally take 5-30s; 90s is generous
87 # while still failing fast on stuck streams. ``max_retries=3``
88 # below handles transient stalls.
89 # `read_timeout` exists for local backends: a hosted API answers
90 # in seconds, but a local model on consumer hardware — especially
91 # a reasoning model that thinks before emitting — routinely needs
92 # longer than 90s. Too short and calls fail as APITimeoutError
93 # mid-run; where the caller is fire-and-forget (observation
94 # consolidation) that degrades memory silently rather than erroring.
95 timeout=httpx.Timeout(
96 connect=5.0,
97 read=read_timeout or 90.0,
98 write=read_timeout or 90.0,
99 pool=read_timeout or 90.0,
100 ),
101 )
102 except ImportError:
103 # h2 not installed — let the OpenAI SDK use its default HTTP/1.1
104 # client. Throughput will be lower for concurrent workloads but
105 # functionality is unchanged.
106 http_client = None
108 client_kwargs: dict[str, object] = {
109 "api_key": resolved_key,
110 "base_url": base_url,
111 "max_retries": 3, # Retry on 429 / 5xx with exponential backoff
112 }
113 if http_client is not None:
114 client_kwargs["http_client"] = http_client
116 self._client = openai.AsyncOpenAI(**client_kwargs)
118 # Ceiling on in-flight API calls. Without one this provider is an
119 # unbounded fan-out: Astrocyte's retain path spawns background
120 # consolidation per call, and while that backlog is now capped
121 # (``max_pending_consolidations``), the cap bounds *pending tasks*, not
122 # concurrent HTTP requests — every admitted task can be in flight at
123 # once, alongside foreground traffic. Against a provider that bills and
124 # rate-limits per request, that turns a load spike into 429s during the
125 # exact sustained-ingest window an evaluator applies.
126 #
127 # Deliberately opt-in via config rather than a low default: the SDK
128 # already multiplexes over HTTP/2 and a tight default would throttle
129 # healthy deployments. ``None``/0 preserves today's unbounded
130 # behaviour.
131 conc = max_concurrency if max_concurrency is not None else int(
132 os.environ.get("ASTROCYTE_OPENAI_MAX_CONCURRENCY", "0")
133 )
134 self._sem: asyncio.Semaphore | None = (
135 asyncio.Semaphore(conc) if conc and conc > 0 else None
136 )
137 self._model = model
138 self._embedding_model = embedding_model
140 def capabilities(self) -> LLMCapabilities:
141 return LLMCapabilities(
142 supports_multimodal_completion=True,
143 modalities_supported=("text", "image_url", "image_base64"),
144 supports_multimodal_embedding=False,
145 supports_batch_embed=True,
146 )
148 async def complete(
149 self,
150 messages: list[Message],
151 model: str | None = None,
152 max_tokens: int = 1024,
153 temperature: float = 0.0,
154 tools: list[ToolDefinition] | None = None,
155 tool_choice: str | None = None,
156 response_format: dict | None = None,
157 ) -> Completion:
158 oai_messages = [_to_oai_message(m) for m in messages]
159 use_model = model or self._model
161 # Native function-calling pass-through (Hindsight parity).
162 # When ``tools`` is set, translate to OpenAI's wire format and
163 # forward; otherwise the request is exactly as before so the
164 # legacy text-only callers see no change.
165 kwargs: dict = {
166 "model": use_model,
167 "messages": oai_messages,
168 "max_tokens": max_tokens,
169 "temperature": temperature,
170 }
171 # Structured outputs (Phase 2 of Hindsight cost-control port).
172 # When set, forwards ``response_format`` straight to the OpenAI
173 # ``chat.completions.create`` call. Eliminates malformed-JSON
174 # failures by constraining the decoder to the supplied schema.
175 # Caller is responsible for the schema shape; we don't validate
176 # it here so any future ``response_format`` variant (json_object,
177 # json_schema, custom) works without an SDK bump.
178 if response_format is not None:
179 kwargs["response_format"] = response_format
180 if tools:
181 kwargs["tools"] = [
182 {
183 "type": "function",
184 "function": {
185 "name": t.name,
186 "description": t.description,
187 "parameters": t.parameters,
188 },
189 }
190 for t in tools
191 ]
192 # OpenAI accepts "auto" / "required" / "none" / specific function ref.
193 if tool_choice in ("auto", "required", "none"):
194 kwargs["tool_choice"] = tool_choice
195 elif tool_choice:
196 kwargs["tool_choice"] = {"type": "function", "function": {"name": tool_choice}}
197 else:
198 kwargs["tool_choice"] = "auto"
200 if self._sem is not None:
201 async with self._sem:
202 response = await self._client.chat.completions.create(**kwargs)
203 else:
204 response = await self._client.chat.completions.create(**kwargs)
206 choice = response.choices[0]
207 usage = None
208 if response.usage:
209 usage = TokenUsage(
210 input_tokens=response.usage.prompt_tokens,
211 output_tokens=response.usage.completion_tokens,
212 )
214 # Parse tool calls when the model emitted them. The OpenAI SDK
215 # returns ``message.tool_calls`` as a list (or ``None``); each
216 # item has ``id``, ``function.name``, ``function.arguments`` (a
217 # JSON string we parse to a dict).
218 tool_calls: list[ToolCall] | None = None
219 raw_tool_calls = getattr(choice.message, "tool_calls", None)
220 if raw_tool_calls:
221 tool_calls = []
222 for tc in raw_tool_calls:
223 fn = getattr(tc, "function", None)
224 if fn is None:
225 continue
226 try:
227 args = json.loads(fn.arguments) if fn.arguments else {}
228 except json.JSONDecodeError:
229 args = {}
230 tool_calls.append(
231 ToolCall(
232 id=tc.id,
233 name=fn.name,
234 arguments=args if isinstance(args, dict) else {},
235 )
236 )
238 return Completion(
239 text=choice.message.content or "",
240 model=response.model,
241 usage=usage,
242 tool_calls=tool_calls,
243 )
245 async def embed(
246 self,
247 texts: list[str],
248 model: str | None = None,
249 ) -> list[list[float]]:
250 use_model = model or self._embedding_model
252 # Truncate texts that would exceed the model's token limit.
253 # text-embedding-3-small has an 8192-token limit (~30K chars).
254 # Truncating at 28K chars provides a safe margin.
255 max_chars = 28_000
256 safe_texts = [_sanitize_text(t)[:max_chars] for t in texts]
258 if self._sem is not None:
259 async with self._sem:
260 response = await self._client.embeddings.create(
261 model=use_model,
262 input=safe_texts,
263 )
264 else:
265 response = await self._client.embeddings.create(
266 model=use_model,
267 input=safe_texts,
268 )
270 # Sort by index to guarantee order matches input
271 sorted_data = sorted(response.data, key=lambda d: d.index)
272 return [d.embedding for d in sorted_data]
275def _sanitize_text(text: str) -> str:
276 """Remove control characters that break the OpenAI API's JSON parser."""
277 # Remove null bytes and other C0 control chars except \t \n \r
278 return "".join(ch for ch in text if ch in ("\t", "\n", "\r") or (ord(ch) >= 32))
281def _to_oai_message(msg: Message) -> dict:
282 """Convert an Astrocyte Message to an OpenAI API message dict.
284 Tool-call round-trip support (Hindsight-parity agentic reflect):
285 - ``role="tool"`` messages carry ``tool_call_id`` (required by
286 OpenAI) and the tool's serialized output as ``content``.
287 - ``role="assistant"`` messages with ``tool_calls`` are translated
288 back into the OpenAI ``tool_calls`` array so the model sees its
289 own prior calls when the loop continues.
290 """
291 if msg.role == "tool":
292 return {
293 "role": "tool",
294 "tool_call_id": msg.tool_call_id or "",
295 "content": _sanitize_text(msg.content) if isinstance(msg.content, str) else "",
296 }
297 if msg.role == "assistant" and msg.tool_calls:
298 return {
299 "role": "assistant",
300 "content": _sanitize_text(msg.content) if isinstance(msg.content, str) else None,
301 "tool_calls": [
302 {
303 "id": tc.id,
304 "type": "function",
305 "function": {
306 "name": tc.name,
307 "arguments": json.dumps(tc.arguments),
308 },
309 }
310 for tc in msg.tool_calls
311 ],
312 }
313 if isinstance(msg.content, str):
314 return {"role": msg.role, "content": _sanitize_text(msg.content)}
316 # Multimodal: list of ContentPart
317 parts = []
318 for part in msg.content:
319 if part.type == "text" and part.text:
320 parts.append({"type": "text", "text": _sanitize_text(part.text)})
321 elif part.type == "image_url" and part.image_url:
322 parts.append(
323 {
324 "type": "image_url",
325 "image_url": {"url": part.image_url},
326 }
327 )
328 elif part.type == "image_base64" and part.image_base64:
329 parts.append(
330 {
331 "type": "image_url",
332 "image_url": {"url": f"data:image/png;base64,{part.image_base64}"},
333 }
334 )
336 return {"role": msg.role, "content": parts}