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

llmconnector/grok_client.py

@@ -2,54 +2,54 @@
 
 from __future__ import annotations
 
-import base64
-import mimetypes
-from pathlib import Path
-import logging
-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
-
-from .types import ImageInput, MessageSequence, normalize_messages
-from .utils import clamp_retries, run_sync_in_thread, run_with_retries
-
-logger = logging.getLogger(__name__)
+import base64
+import mimetypes
+from pathlib import Path
+import logging
+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
+
+from .types import ImageInput, MessageSequence, normalize_messages
+from .utils import clamp_retries, run_sync_in_thread, run_with_retries
+
+logger = logging.getLogger(__name__)
 
 
 class GrokClient:
     """Convenience wrapper around the xAI Grok chat API."""
 
-    def 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:
+    def 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:
         """Generate a response from the specified Grok model.
 
         Args:
             api_key: API key used to authenticate with xAI.
-            prompt: Natural-language instruction or query for the model.
-            model: Identifier of the Grok model to target (for example, ``"grok-3"``).
-            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.
+            prompt: Natural-language instruction or query for the model.
+            model: Identifier of the Grok model to target (for example, ``"grok-3"``).
+            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.
@@ -57,243 +57,243 @@ class GrokClient:
         Raises:
             ValueError: If required arguments are missing or the request payload is empty.
             RuntimeError: If the Grok response does not contain any textual content.
-        """
-        if not api_key:
-            raise ValueError("api_key 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.")
-
-        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:
-                message_parts.append(chat_image(self._to_image_url(image)))
-
-        if not message_parts:
-            raise ValueError("No content provided for response generation.")
-
-        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 with no text: model=%s images=%d request_id=%s",
-                model,
-                len(images or []),
-                request_id,
-            )
-            return ""
-
-        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,
-            )
-        )
-
-    def generate_image(
-        self,
-        *,
-        api_key: str,
-        prompt: str,
+        """
+        if not api_key:
+            raise ValueError("api_key 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.")
+
+        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:
+                message_parts.append(chat_image(self._to_image_url(image)))
+
+        if not message_parts:
+            raise ValueError("No content provided for response generation.")
+
+        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 with no text: model=%s images=%d request_id=%s",
+                model,
+                len(images or []),
+                request_id,
+            )
+            return ""
+
+        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,
+            )
+        )
+
+    def generate_image(
+        self,
+        *,
+        api_key: str,
+        prompt: str,
         model: str,
         image_size: str = "2K",
         image: Optional[ImageInput] = None,
     ) -> bytes:
         """Generate an image using the Grok API.
 
-        Raises:
-            NotImplementedError: This method is not yet implemented for Grok.
-        """
-        raise NotImplementedError("Image generation is not implemented for Grok.")
-
-    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.")
-
-        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,
-        )
-
-    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,
-            )
-        )
+        Raises:
+            NotImplementedError: This method is not yet implemented for Grok.
+        """
+        raise NotImplementedError("Image generation is not implemented for Grok.")
+
+    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.")
+
+        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,
+        )
+
+    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: