1"""Cloudflare Workers AI client for the Lexigram LLM routing system.
2
3Implements the :class:`~lexigram.contracts.ai.protocols.LLMClientProtocol` protocol
4against the Cloudflare Workers AI REST API.
5
6Authentication uses a Bearer token (CF API token) and an ``account_id`` in
7the URL path. Both are sourced from ``ClientConfig.extra`` because
8``ClientConfig.api_key`` represents a generic key and Cloudflare requires an
9additional account-scoped identifier.
10
11Notes:
12 Cloudflare Workers AI does not consistently return token usage metadata.
13 ``Completion.usage`` may contain zeroes for token counts when omitted.
14"""
15
16from __future__ import annotations
17
18from typing import TYPE_CHECKING, Any, cast
19
20from lexigram.ai.llm.clients._message_utils import serialize_text_only
21from lexigram.ai.llm.clients._tools_utils import (
22 serialize_openai_tool_calls,
23 tool_to_openai_format,
24)
25from lexigram.ai.llm.clients.base import AbstractLLMClient
26from lexigram.ai.llm.exceptions import (
27 LLMAuthenticationError,
28 LLMError,
29 LLMModelNotFoundError,
30 LLMRateLimitError,
31)
32from lexigram.ai.llm.http.client import ResilientHTTPClient
33from lexigram.ai.llm.types import (
34 AIError,
35 ChatMessage,
36 Completion,
37 FunctionCall,
38 StreamChunk,
39 TokenUsage,
40 ToolCall,
41)
42from lexigram.contracts.core.health import HealthCheckResult, HealthStatus
43from lexigram.contracts.web.http_models import HttpStatusError
44from lexigram.logging import (
45 get_logger,
46)
47from lexigram.result import Err, Ok, Result
48from lexigram.serialization import dumps_str
49from lexigram.serialization import loads as _loads
50
51logger = get_logger(__name__)
52
53_CF_BASE_URL = "https://api.cloudflare.com"
54
55if TYPE_CHECKING:
56 from collections.abc import AsyncIterator
57
58 from lexigram.ai.llm.config import ClientConfig
59
60__all__ = ["CloudflareWorkersClient"]
61
62
63class CloudflareWorkersClient(AbstractLLMClient):
64 """Client for the Cloudflare Workers AI REST API.
65
66 Routes requests to
67 ``/client/v4/accounts/{account_id}/ai/run/{model}``
68 and normalises the response to a standard :class:`~lexigram.ai.llm.types.Completion`.
69
70 Required ``config.extra`` keys:
71
72 * ``cf_account_id`` — Cloudflare account identifier.
73 * ``cf_api_token`` — Cloudflare API token with Workers AI permission.
74
75 Example:
76 >>> config = ClientConfig(
77 ... provider="cloudflare",
78 ... model="@cf/meta/llama-3.2-11b-vision-instruct",
79 ... extra={
80 ... "cf_account_id": "abc123",
81 ... "cf_api_token": "…",
82 ... },
83 ... timeout=45.0,
84 ... )
85 >>> client = CloudflareWorkersClient(config)
86 >>> completion = await client.complete(
87 ... messages=[{"role": "user", "content": "Hello!"}],
88 ... )
89 """
90
91 def __init__(self, config: ClientConfig) -> None:
92 """Initialise the Cloudflare Workers AI client.
93
94 Args:
95 config: LLM configuration. Required extra fields:
96 ``cf_account_id`` (str) and ``cf_api_token`` (str).
97
98 Raises:
99 LLMAuthenticationError: When ``cf_account_id`` or ``cf_api_token``
100 are absent from ``config.extra``.
101 """
102 super().__init__(config=config)
103 self._http: ResilientHTTPClient | None = None
104
105 account_id: str = config.extra.get("cf_account_id", "")
106 api_token: str = config.extra.get("cf_api_token", "")
107 if not account_id or not api_token:
108 msg = (
109 "CloudflareWorkersClient requires config.extra['cf_account_id'] "
110 "and config.extra['cf_api_token']"
111 )
112 raise LLMAuthenticationError(msg)
113 self._account_id = account_id
114 self._api_token = api_token
115
116 @property
117 def _base_url(self) -> str:
118 """Return the API base URL, using config override when provided."""
119 return self.config.api_base or _CF_BASE_URL
120
121 def _get_http(self) -> ResilientHTTPClient:
122 """Return a lazily-created HTTP client instance."""
123 if self._http is None:
124 self._http = ResilientHTTPClient(
125 base_url=self._base_url,
126 headers={
127 "Authorization": f"Bearer {self._api_token}",
128 "Content-Type": "application/json",
129 },
130 timeout=self.config.timeout,
131 name="cloudflare-workers-client",
132 )
133 return self._http
134
135 # ──────────────────────────────────────────────────────────────────
136 # LLMClientProtocol protocol implementation
137 # ──────────────────────────────────────────────────────────────────
138
139 async def _do_complete(
140 self,
141 messages: list[ChatMessage],
142 *,
143 model: str | None = None,
144 temperature: float = 0.2,
145 max_tokens: int | None = None,
146 **kwargs: Any,
147 ) -> Result[Completion, LLMError]:
148 """Generate a completion from Cloudflare Workers AI.
149
150 Args:
151 messages: Chat messages.
152 model: Model override. When ``None``, uses ``config.model``.
153 temperature: Sampling temperature.
154 max_tokens: Maximum output tokens.
155 **kwargs: Ignored for protocol compatibility.
156
157 Returns:
158 ``Ok(Completion)`` on success. ``Err(LLMError)`` for recoverable
159 failures (rate limit, quota, content filter, model not found).
160
161 Raises:
162 LLMLLMAuthenticationError: When credentials are invalid (HTTP 401/403).
163 AIError: For unexpected infrastructure failures.
164 """
165 active_model = model or self.config.model
166 cf_messages = _convert_messages_for_cloudflare(messages)
167 payload: dict[str, Any] = {
168 "messages": cf_messages,
169 "temperature": temperature,
170 "stream": False,
171 }
172 if max_tokens is not None:
173 payload["max_tokens"] = max_tokens
174 tools = kwargs.pop("tools", None)
175 if tools:
176 payload["tools"] = [
177 converted
178 for tool in tools
179 if (converted := tool_to_openai_format(tool)) is not None
180 ]
181
182 path = f"/client/v4/accounts/{self._account_id}/ai/run/{active_model}"
183
184 try:
185 http = self._get_http()
186 response = await http.post(path, json=payload)
187 response.raise_for_status()
188 except (
189 HttpStatusError,
190 OSError,
191 ConnectionError,
192 TimeoutError,
193 RuntimeError,
194 ) as exc:
195 return self._handle_error_as_result(exc)
196
197 data: dict[str, Any] = response.json
198 if not data.get("success", False):
199 errors = data.get("errors", [])
200 return Err(LLMError(f"Cloudflare Workers AI request failed: {errors}"))
201
202 return Ok(_parse_cf_response_with_tools(data, active_model))
203
204 async def _do_stream_chat(
205 self,
206 messages: list[ChatMessage],
207 *,
208 model: str | None = None,
209 temperature: float = 0.2,
210 max_tokens: int | None = None,
211 **kwargs: Any,
212 ) -> Result[AsyncIterator[StreamChunk], LLMError]:
213 """Stream completion tokens from Cloudflare Workers AI.
214
215 Sends ``stream: true`` to the Cloudflare REST endpoint, which responds
216 with an SSE body. Each ``data:`` line carries a JSON object with a
217 ``response`` field containing the incremental text delta.
218
219 Args:
220 messages: Chat messages.
221 model: Model override.
222 temperature: Sampling temperature.
223 max_tokens: Maximum output tokens.
224 **kwargs: Ignored for protocol compatibility.
225
226 Returns:
227 ``Ok(AsyncIterator[StreamChunk])`` on successful connection.
228 ``Err(LLMError)`` for recoverable failures.
229
230 Raises:
231 LLMLLMAuthenticationError: When credentials are invalid.
232 AIError: For unexpected infrastructure failures.
233 """
234 active_model = model or self.config.model
235 cf_messages = _convert_messages_for_cloudflare(messages)
236 payload: dict[str, Any] = {
237 "messages": cf_messages,
238 "temperature": temperature,
239 "stream": True,
240 }
241 if max_tokens is not None:
242 payload["max_tokens"] = max_tokens
243
244 path = f"/client/v4/accounts/{self._account_id}/ai/run/{active_model}"
245 try:
246 http = self._get_http()
247 response = await http.post(path, json=payload)
248 response.raise_for_status()
249 except (
250 HttpStatusError,
251 OSError,
252 ConnectionError,
253 TimeoutError,
254 RuntimeError,
255 ) as exc:
256 return self._handle_error_as_result(exc)
257
258 return Ok(_parse_cf_sse_body(response.text or "", active_model))
259
260 async def _do_chat(
261 self,
262 messages: list[ChatMessage],
263 tools: list[Any] | None = None,
264 *,
265 model: str | None = None,
266 temperature: float = 0.2,
267 max_tokens: int | None = None,
268 **kwargs: Any,
269 ) -> Result[Completion, LLMError]:
270 """Generate completion with optional Cloudflare tool/function calling.
271
272 Tool calling is handled by :meth:`_do_complete` via
273 ``complete(..., tools=...)``; this method forwards the tool
274 descriptors to keep the ``chat`` code path consistent.
275
276 Args:
277 messages: Chat messages.
278 tools: Optional tool descriptors.
279 model: Model override.
280 temperature: Sampling temperature.
281 max_tokens: Maximum output tokens.
282 **kwargs: Ignored for protocol compatibility.
283
284 Returns:
285 ``Ok(Completion)`` on success. ``Err(LLMError)`` for recoverable
286 failures.
287 """
288 return await self._do_complete(
289 messages,
290 model=model,
291 temperature=temperature,
292 max_tokens=max_tokens,
293 tools=tools,
294 **kwargs,
295 )
296
297 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult:
298 """Perform a lightweight health check against Cloudflare Workers AI.
299
300 Lists the first model as a zero-inference probe to verify credentials
301 and service reachability.
302
303 Args:
304 timeout: Seconds to wait (informational; uses client-level timeout).
305
306 Returns:
307 :class:`~lexigram.contracts.core.health.HealthCheckResult`.
308 """
309 try:
310 http = self._get_http()
311 path = f"/client/v4/accounts/{self._account_id}/ai/models/search?per_page=1"
312 response = await http.get(path)
313 response.raise_for_status()
314 return HealthCheckResult(
315 component="cloudflare", status=HealthStatus.HEALTHY
316 )
317 except (OSError, ConnectionError, TimeoutError, RuntimeError) as exc:
318 return HealthCheckResult(
319 component="cloudflare",
320 status=HealthStatus.UNHEALTHY,
321 error=str(exc),
322 )
323
324 async def close(self) -> None:
325 """Close the underlying HTTP client and release connections."""
326 if self._http is not None:
327 await self._http.close()
328 self._http = None
329 await super().close()
330
331 def _handle_error_as_result(self, error: Exception) -> Result[Any, LLMError]:
332 """Map a caught exception to ``Err`` (recoverable) or re-raise (infra)."""
333 status: int | None = None
334 if isinstance(error, HttpStatusError):
335 status = error.status
336
337 if status in (401, 403):
338 raise LLMAuthenticationError(
339 f"cloudflare: authentication failed ({status}): {error}"
340 ) from error
341 if status == 429:
342 return Err(LLMRateLimitError(f"cloudflare: rate limit exceeded: {error}"))
343 if status == 400:
344 err_str = str(error).lower()
345 if "not found" in err_str or "does not exist" in err_str:
346 return Err(
347 LLMModelNotFoundError(f"cloudflare: model not found: {error}")
348 )
349 return Err(LLMError(f"cloudflare: invalid request: {error}"))
350 if status == 404:
351 return Err(LLMModelNotFoundError(f"cloudflare: model not found: {error}"))
352 raise AIError(f"cloudflare: infrastructure error: {error}") from error
353
354
355# ──────────────────────────────────────────────────────────────────────
356# Private helpers
357# ──────────────────────────────────────────────────────────────────────
358
359
360def _convert_messages_for_cloudflare(
361 messages: list[ChatMessage] | list[dict[str, Any]],
362) -> list[dict[str, Any]]:
363 """Convert ChatMessage objects to Cloudflare API format.
364
365 Cloudflare Workers AI accepts ``role`` and ``content`` (string), plus
366 the OpenAI ``tool_calls`` / ``tool_call_id`` fields for tool round
367 trips. Uses ``serialize_text_only`` to extract text from multimodal
368 content, warning when image parts are encountered.
369
370 Args:
371 messages: List of ChatMessage objects or dicts (for backward compatibility).
372
373 Returns:
374 Cleaned message list compatible with the Cloudflare API.
375 """
376 result: list[dict[str, Any]] = []
377 for msg in messages:
378 if isinstance(msg, dict):
379 # Backward compatibility: handle dict messages
380 role: str = msg.get("role", "user")
381 content: Any = msg.get("content", "")
382 tool_call_id: str | None = msg.get("tool_call_id")
383 tool_calls: Any = msg.get("tool_calls")
384 else:
385 # ChatMessage object
386 role = msg.role.value
387 content = msg.content
388 tool_call_id = msg.tool_call_id
389 tool_calls = msg.tool_calls
390
391 # Serialize content using text-only for non-vision clients
392 text_content = serialize_text_only(
393 content,
394 logger=logger,
395 client_name="cloudflare",
396 )
397 entry: dict[str, Any] = {"role": role, "content": text_content}
398 if tool_call_id:
399 entry["tool_call_id"] = tool_call_id
400 serialized_calls = serialize_openai_tool_calls(
401 cast("list[ToolCall] | None", tool_calls)
402 )
403 if serialized_calls:
404 entry["tool_calls"] = serialized_calls
405 result.append(entry)
406 return result
407
408
409def _parse_cf_response_with_tools(data: dict[str, Any], model: str) -> Completion:
410 """Parse a Cloudflare Workers AI response that may include tool calls.
411
412 Args:
413 data: Parsed JSON response from the Cloudflare API.
414 model: Model identifier.
415
416 Returns:
417 Normalised :class:`~lexigram.ai.llm.types.Completion` with optional
418 :attr:`~Completion.tool_calls` populated.
419 """
420 result_obj = data.get("result", {})
421 text: str = result_obj.get("response", "")
422 usage_data = result_obj.get("usage", {})
423 usage = TokenUsage(
424 prompt_tokens=usage_data.get("prompt_tokens", 0),
425 completion_tokens=usage_data.get("completion_tokens", 0),
426 total_tokens=usage_data.get("total_tokens", 0),
427 )
428
429 raw_calls: list[dict[str, Any]] = result_obj.get("tool_calls", [])
430 tool_calls: list[ToolCall] | None = None
431 if raw_calls:
432 tool_calls = []
433 for tc in raw_calls:
434 func = tc.get("function", {})
435 args = func.get("arguments", {})
436 if isinstance(args, dict):
437 args = dumps_str(args)
438 tool_calls.append(
439 ToolCall(
440 id=tc.get("id", func.get("name", "")),
441 type="function",
442 function=FunctionCall(
443 name=func.get("name", ""),
444 arguments=args,
445 ),
446 )
447 )
448
449 return Completion(content=text, model=model, usage=usage, tool_calls=tool_calls)
450
451
452async def _parse_cf_sse_body(body: str, model: str) -> AsyncIterator[StreamChunk]:
453 """Yield :class:`StreamChunk` objects from a Cloudflare Workers AI SSE body.
454
455 Each ``data:`` line carries a JSON object with a ``response`` field
456 containing the incremental text delta.
457
458 Args:
459 body: Full SSE response body text.
460 model: Model label embedded in each :class:`StreamChunk`.
461
462 Yields:
463 :class:`StreamChunk` with incremental text deltas.
464 """
465 for line in body.splitlines():
466 stripped = line.strip()
467 if not stripped.startswith("data:"):
468 continue
469 raw = stripped[len("data:") :].strip()
470 if not raw or raw == "[DONE]":
471 continue
472 try:
473 chunk: dict[str, Any] = _loads(raw)
474 except (ValueError, TypeError):
475 continue
476 delta = chunk.get("response", "")
477 yield StreamChunk(
478 delta=delta,
479 model=model,
480 finish_reason=None,
481 index=0,
482 )