ccs-llmconnector 1.0.6__py3-none-any.whl → 1.1.1__py3-none-any.whl

llmconnector/grok_client.py

@@ -6,13 +6,15 @@ import base64
 import mimetypes
 from pathlib import Path
 import logging
-from typing import Optional, Sequence, Union
+from typing import Optional, Sequence
 
 from xai_sdk import Client
 from xai_sdk.chat import image as chat_image
 from xai_sdk.chat import user
 
-ImageInput = Union[str, Path]
+from .types import ImageInput, MessageSequence, normalize_messages
+from .utils import clamp_retries, run_sync_in_thread, run_with_retries
+
 logger = logging.getLogger(__name__)
 
 
@@ -23,11 +25,16 @@ class GrokClient:
         self,
         *,
         api_key: str,
-        prompt: str,
+        prompt: Optional[str] = None,
         model: str,
         max_tokens: int = 32000,
         reasoning_effort: Optional[str] = None,
         images: Optional[Sequence[ImageInput]] = None,
+        messages: Optional[MessageSequence] = None,
+        request_id: Optional[str] = None,
+        timeout_s: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        retry_backoff_s: float = 0.5,
     ) -> str:
         """Generate a response from the specified Grok model.
 
@@ -38,6 +45,11 @@ class GrokClient:
             max_tokens: Cap for tokens in the generated response, defaults to 32000.
             reasoning_effort: Optional hint for reasoning-focused models (``"low"`` or ``"high"``).
             images: Optional collection of image references (local paths, URLs, or data URLs).
+            messages: Optional list of chat-style messages (role/content).
+            request_id: Optional request identifier for tracing/logging.
+            timeout_s: Optional request timeout in seconds (best-effort).
+            max_retries: Optional retry count for transient failures.
+            retry_backoff_s: Base delay (seconds) for exponential backoff between retries.
 
         Returns:
             The text output produced by the model.
@@ -48,14 +60,17 @@ class GrokClient:
         """
         if not api_key:
             raise ValueError("api_key must be provided.")
-        if not prompt and not images:
-            raise ValueError("At least one of prompt or images must be provided.")
+        if not prompt and not messages and not images:
+            raise ValueError("At least one of prompt, messages, or images must be provided.")
         if not model:
             raise ValueError("model must be provided.")
 
-        message_parts = []
-        if prompt:
-            message_parts.append(prompt)
+        normalized_messages = normalize_messages(prompt=prompt, messages=messages)
+        message_parts: list[object] = []
+        if normalized_messages:
+            for message in normalized_messages:
+                if message["content"]:
+                    message_parts.append(f"{message['role']}: {message['content']}")
 
         if images:
             for image in images:
@@ -64,53 +79,104 @@ class GrokClient:
         if not message_parts:
             raise ValueError("No content provided for response generation.")
 
-        grok_client = Client(api_key=api_key)
-
-        create_kwargs = {
-            "model": model,
-            "max_tokens": max_tokens,
-            "messages": [user(*message_parts)],
-        }
-
-        normalized_effort = (reasoning_effort or "").strip().lower()
-        if normalized_effort in {"low", "high"}:
-            create_kwargs["reasoning_effort"] = normalized_effort
-
-        try:
-            chat = grok_client.chat.create(**create_kwargs)
-            response = chat.sample()
-        except Exception as exc:
-            logger.exception("xAI Grok chat request failed: %s", exc)
-            raise
-
-        content = getattr(response, "content", None)
-        if content:
+        retry_count = clamp_retries(max_retries)
+
+        def _run_request() -> str:
+            grok_client = Client(api_key=api_key)
+
+            create_kwargs = {
+                "model": model,
+                "max_tokens": max_tokens,
+                "messages": [user(*message_parts)],
+            }
+
+            normalized_effort = (reasoning_effort or "").strip().lower()
+            if normalized_effort in {"low", "high"}:
+                create_kwargs["reasoning_effort"] = normalized_effort
+
+            if timeout_s is not None:
+                create_kwargs["timeout"] = timeout_s
+
+            try:
+                chat = grok_client.chat.create(**create_kwargs)
+                response = chat.sample()
+            except Exception as exc:
+                logger.exception(
+                    "xAI Grok chat request failed: %s request_id=%s",
+                    exc,
+                    request_id,
+                )
+                raise
+
+            content = getattr(response, "content", None)
+            if content:
+                logger.info(
+                    "xAI chat succeeded: model=%s images=%d text_len=%d request_id=%s",
+                    model,
+                    len(images or []),
+                    len(content or ""),
+                    request_id,
+                )
+                return content
+
+            reasoning_content = getattr(response, "reasoning_content", None)
+            if reasoning_content:
+                logger.info(
+                    "xAI chat succeeded (reasoning): model=%s images=%d text_len=%d request_id=%s",
+                    model,
+                    len(images or []),
+                    len(reasoning_content or ""),
+                    request_id,
+                )
+                return reasoning_content
+
+            # Treat successful calls without textual content as a successful, empty response
+            # rather than raising. This aligns with callers that handle empty outputs gracefully.
             logger.info(
-                "xAI chat succeeded: model=%s images=%d text_len=%d",
+                "xAI chat succeeded with no text: model=%s images=%d request_id=%s",
                 model,
                 len(images or []),
-                len(content or ""),
+                request_id,
             )
-            return content
+            return ""
 
-        reasoning_content = getattr(response, "reasoning_content", None)
-        if reasoning_content:
-            logger.info(
-                "xAI chat succeeded (reasoning): model=%s images=%d text_len=%d",
-                model,
-                len(images or []),
-                len(reasoning_content or ""),
+        return run_with_retries(
+            func=_run_request,
+            max_retries=retry_count,
+            retry_backoff_s=retry_backoff_s,
+            request_id=request_id,
+        )
+
+    async def async_generate_response(
+        self,
+        *,
+        api_key: str,
+        prompt: Optional[str] = None,
+        model: str,
+        max_tokens: int = 32000,
+        reasoning_effort: Optional[str] = None,
+        images: Optional[Sequence[ImageInput]] = None,
+        messages: Optional[MessageSequence] = None,
+        request_id: Optional[str] = None,
+        timeout_s: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        retry_backoff_s: float = 0.5,
+    ) -> str:
+        return await run_sync_in_thread(
+            lambda: self.generate_response(
+                api_key=api_key,
+                prompt=prompt,
+                model=model,
+                max_tokens=max_tokens,
+                reasoning_effort=reasoning_effort,
+                images=images,
+                messages=messages,
+                request_id=request_id,
+                timeout_s=timeout_s,
+                max_retries=max_retries,
+                retry_backoff_s=retry_backoff_s,
             )
-            return reasoning_content
-
-        # Treat successful calls without textual content as a successful, empty response
-        # rather than raising. This aligns with callers that handle empty outputs gracefully.
-        logger.info(
-            "xAI chat succeeded with no text: model=%s images=%d",
-            model,
-            len(images or []),
         )
-        return ""
 
     def generate_image(
         self,
@@ -128,42 +194,106 @@ class GrokClient:
         """
         raise NotImplementedError("Image generation is not implemented for Grok.")
 
-    def list_models(self, *, api_key: str) -> list[dict[str, Optional[str]]]:
+    async def async_generate_image(
+        self,
+        *,
+        api_key: str,
+        prompt: str,
+        model: str,
+        image_size: str = "2K",
+        image: Optional[ImageInput] = None,
+    ) -> bytes:
+        return await run_sync_in_thread(
+            lambda: self.generate_image(
+                api_key=api_key,
+                prompt=prompt,
+                model=model,
+                image_size=image_size,
+                image=image,
+            )
+        )
+
+    def list_models(
+        self,
+        *,
+        api_key: str,
+        request_id: Optional[str] = None,
+        timeout_s: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        retry_backoff_s: float = 0.5,
+    ) -> list[dict[str, Optional[str]]]:
         """Return the Grok language models available to the authenticated account."""
         if not api_key:
             raise ValueError("api_key must be provided.")
 
-        grok_client = Client(api_key=api_key)
-        models: list[dict[str, Optional[str]]] = []
-
-        try:
-            iterator = grok_client.models.list_language_models()
-        except Exception as exc:
-            logger.exception("xAI list language models failed: %s", exc)
-            raise
-
-        for model in iterator:
-            model_id = getattr(model, "name", None)
-            if not model_id:
-                continue
-
-            aliases = getattr(model, "aliases", None)
-            display_name: Optional[str] = None
-            if aliases:
-                try:
-                    display_name = next(iter(aliases)) or None
-                except (StopIteration, TypeError):
-                    display_name = None
-
-            models.append(
-                {
-                    "id": model_id,
-                    "display_name": display_name,
-                }
+        retry_count = clamp_retries(max_retries)
+
+        def _run_request() -> list[dict[str, Optional[str]]]:
+            grok_client = Client(api_key=api_key)
+            models: list[dict[str, Optional[str]]] = []
+
+            try:
+                iterator = grok_client.models.list_language_models()
+            except Exception as exc:
+                logger.exception(
+                    "xAI list language models failed: %s request_id=%s",
+                    exc,
+                    request_id,
+                )
+                raise
+
+            for model in iterator:
+                model_id = getattr(model, "name", None)
+                if not model_id:
+                    continue
+
+                aliases = getattr(model, "aliases", None)
+                display_name: Optional[str] = None
+                if aliases:
+                    try:
+                        display_name = next(iter(aliases)) or None
+                    except (StopIteration, TypeError):
+                        display_name = None
+
+                models.append(
+                    {
+                        "id": model_id,
+                        "display_name": display_name,
+                    }
+                )
+
+            logger.info(
+                "xAI list_language_models succeeded: count=%d request_id=%s",
+                len(models),
+                request_id,
             )
+            return models
+
+        return run_with_retries(
+            func=_run_request,
+            max_retries=retry_count,
+            retry_backoff_s=retry_backoff_s,
+            request_id=request_id,
+        )
 
-        logger.info("xAI list_language_models succeeded: count=%d", len(models))
-        return models
+    async def async_list_models(
+        self,
+        *,
+        api_key: str,
+        request_id: Optional[str] = None,
+        timeout_s: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        retry_backoff_s: float = 0.5,
+    ) -> list[dict[str, Optional[str]]]:
+        return await run_sync_in_thread(
+            lambda: self.list_models(
+                api_key=api_key,
+                request_id=request_id,
+                timeout_s=timeout_s,
+                max_retries=max_retries,
+                retry_backoff_s=retry_backoff_s,
+            )
+        )
 
     @staticmethod
     def _to_image_url(image: ImageInput) -> str:

llmconnector/openai_client.py

@@ -6,16 +6,13 @@ import base64
 import mimetypes
 from pathlib import Path
 import logging
-from typing import Optional, Sequence, Union
+from typing import Optional, Sequence
 
 from openai import OpenAI
 
-try:
-    from openai import APIError as OpenAIError  # type: ignore
-except ImportError:  # pragma: no cover - fallback for older SDKs
-    from openai.error import OpenAIError  # type: ignore
+from .types import ImageInput, MessageSequence, normalize_messages
+from .utils import clamp_retries, run_sync_in_thread, run_with_retries
 
-ImageInput = Union[str, Path]
 logger = logging.getLogger(__name__)
 
 
@@ -26,11 +23,16 @@ class OpenAIResponsesClient:
         self,
         *,
         api_key: str,
-        prompt: str,
+        prompt: Optional[str] = None,
         model: str,
         max_tokens: int = 32000,
         reasoning_effort: Optional[str] = None,
         images: Optional[Sequence[ImageInput]] = None,
+        messages: Optional[MessageSequence] = None,
+        request_id: Optional[str] = None,
+        timeout_s: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        retry_backoff_s: float = 0.5,
     ) -> str:
         """Generate a response from the specified model.
 
@@ -41,6 +43,11 @@ class OpenAIResponsesClient:
             max_tokens: Cap for tokens across the entire exchange, defaults to 32000.
             reasoning_effort: Optional reasoning effort hint (``"low"``, ``"medium"``, or ``"high"``).
             images: Optional collection of image references (local paths or URLs).
+            messages: Optional list of chat-style messages (role/content).
+            request_id: Optional request identifier for tracing/logging.
+            timeout_s: Optional request timeout in seconds.
+            max_retries: Optional retry count for transient failures.
+            retry_backoff_s: Base delay (seconds) for exponential backoff between retries.
 
         Returns:
             The text output produced by the model.
@@ -51,52 +58,110 @@ class OpenAIResponsesClient:
         """
         if not api_key:
             raise ValueError("api_key must be provided.")
-        if not prompt and not images:
-            raise ValueError("At least one of prompt or images must be provided.")
+        if not prompt and not messages and not images:
+            raise ValueError("At least one of prompt, messages, or images must be provided.")
         if not model:
             raise ValueError("model must be provided.")
 
-        client = OpenAI(api_key=api_key)
-
-        content_blocks = []
-        if prompt:
-            content_blocks.append({"type": "input_text", "text": prompt})
+        normalized_messages = normalize_messages(prompt=prompt, messages=messages)
+        content_messages: list[dict] = []
+        for message in normalized_messages:
+            content_blocks = []
+            if message["content"]:
+                content_blocks.append({"type": "input_text", "text": message["content"]})
+            content_messages.append({"role": message["role"], "content": content_blocks})
 
         if images:
-            for image in images:
-                content_blocks.append(self._to_image_block(image))
-
-        if not content_blocks:
+            image_blocks = [self._to_image_block(image) for image in images]
+            target_index = next(
+                (
+                    index
+                    for index in range(len(content_messages) - 1, -1, -1)
+                    if content_messages[index]["role"] == "user"
+                ),
+                None,
+            )
+            if target_index is None:
+                content_messages.append({"role": "user", "content": image_blocks})
+            else:
+                content_messages[target_index]["content"].extend(image_blocks)
+
+        if not any(message["content"] for message in content_messages):
             raise ValueError("No content provided for response generation.")
 
         request_payload = {
             "model": model,
-            "input": [
-                {
-                    "role": "user",
-                    "content": content_blocks,
-                }
-            ],
+            "input": content_messages,
             "max_output_tokens": max_tokens,
         }
 
         if reasoning_effort:
             request_payload["reasoning"] = {"effort": reasoning_effort}
 
-        try:
-            response = client.responses.create(**request_payload)
-        except Exception as exc:  # Log and re-raise to preserve default behavior
-            logger.exception("OpenAI Responses API request failed: %s", exc)
-            raise
-
-        output_text = response.output_text
-        logger.info(
-            "OpenAI generate_response succeeded: model=%s images=%d text_len=%d",
-            model,
-            len(images or []),
-            len(output_text or ""),
+        retry_count = clamp_retries(max_retries)
+
+        def _run_request() -> str:
+            client_kwargs = {"api_key": api_key}
+            if timeout_s is not None:
+                client_kwargs["timeout"] = timeout_s
+            client = OpenAI(**client_kwargs)
+            try:
+                response = client.responses.create(**request_payload)
+            except Exception as exc:  # Log and re-raise to preserve default behavior
+                logger.exception(
+                    "OpenAI Responses API request failed: %s request_id=%s",
+                    exc,
+                    request_id,
+                )
+                raise
+
+            output_text = response.output_text
+            logger.info(
+                "OpenAI generate_response succeeded: model=%s images=%d text_len=%d request_id=%s",
+                model,
+                len(images or []),
+                len(output_text or ""),
+                request_id,
+            )
+            return output_text
+
+        return run_with_retries(
+            func=_run_request,
+            max_retries=retry_count,
+            retry_backoff_s=retry_backoff_s,
+            request_id=request_id,
+        )
+
+    async def async_generate_response(
+        self,
+        *,
+        api_key: str,
+        prompt: Optional[str] = None,
+        model: str,
+        max_tokens: int = 32000,
+        reasoning_effort: Optional[str] = None,
+        images: Optional[Sequence[ImageInput]] = None,
+        messages: Optional[MessageSequence] = None,
+        request_id: Optional[str] = None,
+        timeout_s: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        retry_backoff_s: float = 0.5,
+    ) -> str:
+        return await run_sync_in_thread(
+            lambda: self.generate_response(
+                api_key=api_key,
+                prompt=prompt,
+                model=model,
+                max_tokens=max_tokens,
+                reasoning_effort=reasoning_effort,
+                images=images,
+                messages=messages,
+                request_id=request_id,
+                timeout_s=timeout_s,
+                max_retries=max_retries,
+                retry_backoff_s=retry_backoff_s,
+            )
         )
-        return output_text
 
     @staticmethod
     def _to_image_block(image: ImageInput) -> dict:
@@ -135,35 +200,102 @@ class OpenAIResponsesClient:
         """
         raise NotImplementedError("Image generation is not implemented for OpenAI.")
 
-    def list_models(self, *, api_key: str) -> list[dict[str, Optional[str]]]:
+    async def async_generate_image(
+        self,
+        *,
+        api_key: str,
+        prompt: str,
+        model: str,
+        image_size: Optional[str] = None,
+        aspect_ratio: Optional[str] = None,
+        image: Optional[ImageInput] = None,
+    ) -> bytes:
+        return await run_sync_in_thread(
+            lambda: self.generate_image(
+                api_key=api_key,
+                prompt=prompt,
+                model=model,
+                image_size=image_size,
+                aspect_ratio=aspect_ratio,
+                image=image,
+            )
+        )
+
+    def list_models(
+        self,
+        *,
+        api_key: str,
+        request_id: Optional[str] = None,
+        timeout_s: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        retry_backoff_s: float = 0.5,
+    ) -> list[dict[str, Optional[str]]]:
         """Return the models available to the authenticated OpenAI account."""
         if not api_key:
             raise ValueError("api_key must be provided.")
 
-        client = OpenAI(api_key=api_key)
-        try:
-            response = client.models.list()
-        except Exception as exc:
-            logger.exception("OpenAI list models failed: %s", exc)
-            raise
-        data = getattr(response, "data", []) or []
-
-        models: list[dict[str, Optional[str]]] = []
-        for model in data:
-            model_id = getattr(model, "id", None)
-            if model_id is None and isinstance(model, dict):
-                model_id = model.get("id")
-            if not model_id:
-                continue
-
-            display_name = getattr(model, "display_name", None)
-            if display_name is None and isinstance(model, dict):
-                display_name = model.get("display_name")
-
-            models.append({"id": model_id, "display_name": display_name})
-
-        logger.info("OpenAI list_models succeeded: count=%d", len(models))
-        return models
+        retry_count = clamp_retries(max_retries)
+
+        def _run_request() -> list[dict[str, Optional[str]]]:
+            client_kwargs = {"api_key": api_key}
+            if timeout_s is not None:
+                client_kwargs["timeout"] = timeout_s
+            client = OpenAI(**client_kwargs)
+            try:
+                response = client.models.list()
+            except Exception as exc:
+                logger.exception(
+                    "OpenAI list models failed: %s request_id=%s", exc, request_id
+                )
+                raise
+            data = getattr(response, "data", []) or []
+
+            models: list[dict[str, Optional[str]]] = []
+            for model in data:
+                model_id = getattr(model, "id", None)
+                if model_id is None and isinstance(model, dict):
+                    model_id = model.get("id")
+                if not model_id:
+                    continue
+
+                display_name = getattr(model, "display_name", None)
+                if display_name is None and isinstance(model, dict):
+                    display_name = model.get("display_name")
+
+                models.append({"id": model_id, "display_name": display_name})
+
+            logger.info(
+                "OpenAI list_models succeeded: count=%d request_id=%s",
+                len(models),
+                request_id,
+            )
+            return models
+
+        return run_with_retries(
+            func=_run_request,
+            max_retries=retry_count,
+            retry_backoff_s=retry_backoff_s,
+            request_id=request_id,
+        )
+
+    async def async_list_models(
+        self,
+        *,
+        api_key: str,
+        request_id: Optional[str] = None,
+        timeout_s: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        retry_backoff_s: float = 0.5,
+    ) -> list[dict[str, Optional[str]]]:
+        return await run_sync_in_thread(
+            lambda: self.list_models(
+                api_key=api_key,
+                request_id=request_id,
+                timeout_s=timeout_s,
+                max_retries=max_retries,
+                retry_backoff_s=retry_backoff_s,
+            )
+        )
 
 
 def _encode_image_path(path: Path) -> str: