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

llmconnector/client_cli.py

@@ -82,35 +82,35 @@ def _build_parser() -> argparse.ArgumentParser:
         default=32000,
         help="Maximum output tokens (provider-specific meaning)",
     )
-    p_respond.add_argument(
-        "--reasoning-effort",
-        choices=["low", "medium", "high"],
-        default=None,
-        help="Optional reasoning effort hint if supported",
-    )
-    p_respond.add_argument(
-        "--request-id",
-        default=None,
-        help="Optional request identifier for tracing/logging",
-    )
-    p_respond.add_argument(
-        "--timeout-s",
-        type=float,
-        default=None,
-        help="Optional timeout in seconds",
-    )
-    p_respond.add_argument(
-        "--max-retries",
-        type=int,
-        default=0,
-        help="Number of retries for transient failures",
-    )
-    p_respond.add_argument(
-        "--retry-backoff-s",
-        type=float,
-        default=0.5,
-        help="Base delay in seconds for exponential backoff",
-    )
+    p_respond.add_argument(
+        "--reasoning-effort",
+        choices=["low", "medium", "high"],
+        default=None,
+        help="Optional reasoning effort hint if supported",
+    )
+    p_respond.add_argument(
+        "--request-id",
+        default=None,
+        help="Optional request identifier for tracing/logging",
+    )
+    p_respond.add_argument(
+        "--timeout-s",
+        type=float,
+        default=None,
+        help="Optional timeout in seconds",
+    )
+    p_respond.add_argument(
+        "--max-retries",
+        type=int,
+        default=0,
+        help="Number of retries for transient failures",
+    )
+    p_respond.add_argument(
+        "--retry-backoff-s",
+        type=float,
+        default=0.5,
+        help="Base delay in seconds for exponential backoff",
+    )
 
     # models: list available models
     p_models = subparsers.add_parser(
@@ -172,19 +172,19 @@ def _cmd_respond(args: argparse.Namespace) -> int:
             print("Error: provide a prompt or at least one image.", file=sys.stderr)
             return 2
     try:
-        output = client.generate_response(
-            provider=provider,
-            api_key=api_key,
-            prompt=prompt,
-            model=model,
-            max_tokens=args.max_tokens,
-            reasoning_effort=args.reasoning_effort,
-            images=images,
-            request_id=args.request_id,
-            timeout_s=args.timeout_s,
-            max_retries=args.max_retries,
-            retry_backoff_s=args.retry_backoff_s,
-        )
+        output = client.generate_response(
+            provider=provider,
+            api_key=api_key,
+            prompt=prompt,
+            model=model,
+            max_tokens=args.max_tokens,
+            reasoning_effort=args.reasoning_effort,
+            images=images,
+            request_id=args.request_id,
+            timeout_s=args.timeout_s,
+            max_retries=args.max_retries,
+            retry_backoff_s=args.retry_backoff_s,
+        )
     except Exception as exc:  # pragma: no cover - CLI surface
         print(f"Error: {exc}", file=sys.stderr)
         return 2

llmconnector/gemini_client.py

@@ -6,38 +6,45 @@ import base64
 import mimetypes
 from pathlib import Path
 import logging
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Union
 from urllib.request import urlopen
 
 from google import genai
 from google.genai import types
 
-from .types import ImageInput, MessageSequence, normalize_messages
+from .types import (
+    EmbeddingVector,
+    ImageInput,
+    LLMResponse,
+    MessageSequence,
+    TokenUsage,
+    normalize_messages,
+)
 from .utils import clamp_retries, run_sync_in_thread, run_with_retries
 
-logger = logging.getLogger(__name__)
-
-
-_GEMINI_MIN_TIMEOUT_S = 10.0
-_GEMINI_MIN_TIMEOUT_MS = int(_GEMINI_MIN_TIMEOUT_S * 1000)
-
-
-def _normalize_gemini_timeout_ms(timeout_s: float) -> int:
-    """Convert a seconds timeout into the millisecond value expected by google-genai HttpOptions."""
-    # google-genai HttpOptions expects milliseconds, but our public API uses seconds.
-    effective_timeout_s = max(_GEMINI_MIN_TIMEOUT_S, timeout_s)
-    if effective_timeout_s != timeout_s:
-        logger.warning(
-            "Gemini timeout %ss is too short, clamping to %ss.",
-            timeout_s,
-            effective_timeout_s,
-        )
-    timeout_ms = int(effective_timeout_s * 1000)
-    return max(_GEMINI_MIN_TIMEOUT_MS, timeout_ms)
-
-
+logger = logging.getLogger(__name__)
+
+
+_GEMINI_MIN_TIMEOUT_S = 10.0
+_GEMINI_MIN_TIMEOUT_MS = int(_GEMINI_MIN_TIMEOUT_S * 1000)
+
+
+def _normalize_gemini_timeout_ms(timeout_s: float) -> int:
+    """Convert a seconds timeout into the millisecond value expected by google-genai HttpOptions."""
+    # google-genai HttpOptions expects milliseconds, but our public API uses seconds.
+    effective_timeout_s = max(_GEMINI_MIN_TIMEOUT_S, timeout_s)
+    if effective_timeout_s != timeout_s:
+        logger.warning(
+            "Gemini timeout %ss is too short, clamping to %ss.",
+            timeout_s,
+            effective_timeout_s,
+        )
+    timeout_ms = int(effective_timeout_s * 1000)
+    return max(_GEMINI_MIN_TIMEOUT_MS, timeout_ms)
+
+
 class GeminiClient:
-    """Convenience wrapper around the Google Gemini SDK."""
+    """Convenience wrapper around the Google Gemini SDK."""
 
     def generate_response(
         self,
@@ -120,13 +127,13 @@ class GeminiClient:
 
         retry_count = clamp_retries(max_retries)
 
-        def _build_client() -> genai.Client:
-            client_kwargs: dict[str, object] = {"api_key": api_key}
-            if timeout_s is not None:
-                client_kwargs["http_options"] = types.HttpOptions(
-                    timeout=_normalize_gemini_timeout_ms(timeout_s)
-                )
-            return genai.Client(**client_kwargs)
+        def _build_client() -> genai.Client:
+            client_kwargs: dict[str, object] = {"api_key": api_key}
+            if timeout_s is not None:
+                client_kwargs["http_options"] = types.HttpOptions(
+                    timeout=_normalize_gemini_timeout_ms(timeout_s)
+                )
+            return genai.Client(**client_kwargs)
 
         def _run_request() -> str:
             client = _build_client()
@@ -194,17 +201,153 @@ class GeminiClient:
             )
             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,
+        return run_with_retries(
+            func=_run_request,
+            max_retries=retry_count,
+            retry_backoff_s=retry_backoff_s,
+            request_id=request_id,
+        )
+
+    def generate_response_with_usage(
+        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,
+    ) -> LLMResponse:
+        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)
+        contents: list[types.Content] = []
+        for message in normalized_messages:
+            parts: list[types.Part] = []
+            if message["content"]:
+                parts.append(types.Part.from_text(text=message["content"]))
+            contents.append(types.Content(role=message["role"], parts=parts))
+
+        if images:
+            image_parts = [self._to_image_part(image) for image in images]
+            target_index = next(
+                (
+                    index
+                    for index in range(len(contents) - 1, -1, -1)
+                    if contents[index].role == "user"
+                ),
+                None,
+            )
+            if target_index is None:
+                contents.append(types.Content(role="user", parts=image_parts))
+            else:
+                existing_parts = list(contents[target_index].parts or [])
+                existing_parts.extend(image_parts)
+                contents[target_index] = types.Content(role="user", parts=existing_parts)
+
+        if not contents or not any(content.parts for content in contents):
+            raise ValueError("No content provided for response generation.")
+
+        config = types.GenerateContentConfig(max_output_tokens=max_tokens)
+        _ = reasoning_effort  # accepted for API parity; not currently applied by the Gemini SDK.
+
+        retry_count = clamp_retries(max_retries)
+
+        def _build_client() -> genai.Client:
+            client_kwargs: dict[str, object] = {"api_key": api_key}
+            if timeout_s is not None:
+                client_kwargs["http_options"] = types.HttpOptions(
+                    timeout=_normalize_gemini_timeout_ms(timeout_s)
+                )
+            return genai.Client(**client_kwargs)
+
+        def _run_request() -> LLMResponse:
+            client = _build_client()
+            try:
+                try:
+                    response = client.models.generate_content(
+                        model=model,
+                        contents=contents,
+                        config=config,
+                    )
+                except Exception as exc:
+                    logger.exception(
+                        "Gemini generate_content failed: %s request_id=%s",
+                        exc,
+                        request_id,
+                    )
+                    raise
+            finally:
+                closer = getattr(client, "close", None)
+                if callable(closer):
+                    try:
+                        closer()
+                    except Exception:
+                        pass
+
+            usage = _extract_gemini_usage(response)
+
+            if response.text:
+                result_text = response.text
+                logger.info(
+                    "Gemini generate_content succeeded: model=%s images=%d text_len=%d request_id=%s",
+                    model,
+                    len(images or []),
+                    len(result_text or ""),
+                    request_id,
+                )
+                return LLMResponse(text=result_text, usage=usage, provider="gemini", model=model)
+
+            candidate_texts: list[str] = []
+            for candidate in getattr(response, "candidates", []) or []:
+                content_obj = getattr(candidate, "content", None)
+                if not content_obj:
+                    continue
+                for part in getattr(content_obj, "parts", []) or []:
+                    text = getattr(part, "text", None)
+                    if text:
+                        candidate_texts.append(text)
+
+            if candidate_texts:
+                result_text = "\n".join(candidate_texts)
+                logger.info(
+                    "Gemini generate_content succeeded (candidates): model=%s images=%d text_len=%d request_id=%s",
+                    model,
+                    len(images or []),
+                    len(result_text or ""),
+                    request_id,
+                )
+                return LLMResponse(text=result_text, usage=usage, provider="gemini", model=model)
+
+            logger.info(
+                "Gemini generate_content succeeded with no text: model=%s images=%d request_id=%s",
+                model,
+                len(images or []),
+                request_id,
+            )
+            return LLMResponse(text="", usage=usage, provider="gemini", model=model)
+
+        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,
@@ -216,21 +359,52 @@ class GeminiClient:
         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 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,
+            )
+        )
+
+    async def async_generate_response_with_usage(
+        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,
+    ) -> LLMResponse:
+        return await run_sync_in_thread(
+            lambda: self.generate_response_with_usage(
+                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,
@@ -284,13 +458,13 @@ class GeminiClient:
 
         retry_count = clamp_retries(max_retries)
 
-        def _build_client() -> genai.Client:
-            client_kwargs: dict[str, object] = {"api_key": api_key}
-            if timeout_s is not None:
-                client_kwargs["http_options"] = types.HttpOptions(
-                    timeout=_normalize_gemini_timeout_ms(timeout_s)
-                )
-            return genai.Client(**client_kwargs)
+        def _build_client() -> genai.Client:
+            client_kwargs: dict[str, object] = {"api_key": api_key}
+            if timeout_s is not None:
+                client_kwargs["http_options"] = types.HttpOptions(
+                    timeout=_normalize_gemini_timeout_ms(timeout_s)
+                )
+            return genai.Client(**client_kwargs)
 
         def _run_request() -> bytes:
             client = _build_client()
@@ -376,13 +550,13 @@ class GeminiClient:
 
         retry_count = clamp_retries(max_retries)
 
-        def _build_client() -> genai.Client:
-            client_kwargs: dict[str, object] = {"api_key": api_key}
-            if timeout_s is not None:
-                client_kwargs["http_options"] = types.HttpOptions(
-                    timeout=_normalize_gemini_timeout_ms(timeout_s)
-                )
-            return genai.Client(**client_kwargs)
+        def _build_client() -> genai.Client:
+            client_kwargs: dict[str, object] = {"api_key": api_key}
+            if timeout_s is not None:
+                client_kwargs["http_options"] = types.HttpOptions(
+                    timeout=_normalize_gemini_timeout_ms(timeout_s)
+                )
+            return genai.Client(**client_kwargs)
 
         def _run_request() -> list[dict[str, Optional[str]]]:
             models: list[dict[str, Optional[str]]] = []
@@ -435,7 +609,7 @@ class GeminiClient:
             request_id=request_id,
         )
 
-    async def async_list_models(
+    async def async_list_models(
         self,
         *,
         api_key: str,
@@ -444,15 +618,126 @@ class GeminiClient:
         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,
-            )
-        )
+        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 embed_content(
+        self,
+        *,
+        api_key: str,
+        model: str,
+        contents: Union[str, Sequence[str]],
+        task_type: Optional[str] = None,
+        output_dimensionality: Optional[int] = None,
+        request_id: Optional[str] = None,
+        timeout_s: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        retry_backoff_s: float = 0.5,
+    ) -> list[EmbeddingVector]:
+        if not api_key:
+            raise ValueError("api_key must be provided.")
+        if not model:
+            raise ValueError("model must be provided.")
+
+        if isinstance(contents, str):
+            payload: Union[str, list[str]] = contents
+        else:
+            payload = list(contents)
+            if not payload:
+                raise ValueError("contents must not be empty.")
+
+        retry_count = clamp_retries(max_retries)
+
+        def _build_client() -> genai.Client:
+            client_kwargs: dict[str, object] = {"api_key": api_key}
+            if timeout_s is not None:
+                http_options = getattr(types, "HttpOptions", None)
+                if http_options is not None:
+                    try:
+                        client_kwargs["http_options"] = http_options(timeout=timeout_s)
+                    except Exception:
+                        logger.debug("Gemini HttpOptions timeout not applied.", exc_info=True)
+            return genai.Client(**client_kwargs)
+
+        config_kwargs: dict[str, object] = {}
+        if task_type is not None:
+            config_kwargs["task_type"] = task_type
+        if output_dimensionality is not None:
+            config_kwargs["output_dimensionality"] = output_dimensionality
+        config = types.EmbedContentConfig(**config_kwargs) if config_kwargs else None
+
+        def _run_request() -> list[EmbeddingVector]:
+            client = _build_client()
+            try:
+                result = client.models.embed_content(
+                    model=model,
+                    contents=payload,
+                    config=config,
+                )
+                embeddings = getattr(result, "embeddings", None)
+                if embeddings is None:
+                    raise ValueError("Gemini embeddings response missing embeddings field.")
+                vectors: list[EmbeddingVector] = []
+                for embedding in embeddings:
+                    values = getattr(embedding, "values", None)
+                    if values is None:
+                        raise ValueError("Gemini embedding missing values field.")
+                    vectors.append(list(values))
+                return vectors
+            finally:
+                closer = getattr(client, "close", None)
+                if callable(closer):
+                    try:
+                        closer()
+                    except Exception:
+                        pass
+
+        vectors = run_with_retries(
+            func=_run_request,
+            max_retries=retry_count,
+            retry_backoff_s=retry_backoff_s,
+            request_id=request_id,
+        )
+        logger.info(
+            "Gemini embed_content succeeded: count=%d request_id=%s",
+            len(vectors),
+            request_id,
+        )
+        return vectors
+
+    async def async_embed_content(
+        self,
+        *,
+        api_key: str,
+        model: str,
+        contents: Union[str, Sequence[str]],
+        task_type: Optional[str] = None,
+        output_dimensionality: Optional[int] = None,
+        request_id: Optional[str] = None,
+        timeout_s: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        retry_backoff_s: float = 0.5,
+    ) -> list[EmbeddingVector]:
+        return await run_sync_in_thread(
+            lambda: self.embed_content(
+                api_key=api_key,
+                model=model,
+                contents=contents,
+                task_type=task_type,
+                output_dimensionality=output_dimensionality,
+                request_id=request_id,
+                timeout_s=timeout_s,
+                max_retries=max_retries,
+                retry_backoff_s=retry_backoff_s,
+            )
+        )
 
     @staticmethod
     def _to_image_part(image: ImageInput) -> types.Part:
@@ -469,18 +754,18 @@ class GeminiClient:
         return _part_from_path(Path(image))
 
 
-def _part_from_path(path: Path) -> types.Part:
-    """Create an image part from a local filesystem path."""
-    # Ensure common audio types are recognized across platforms (used for transcription as well).
-    mimetypes.add_type("audio/mp4", ".m4a")
-    mimetypes.add_type("audio/mpeg", ".mp3")
-    mimetypes.add_type("audio/wav", ".wav")
-    mimetypes.add_type("audio/aac", ".aac")
-
-    expanded = path.expanduser()
-    data = expanded.read_bytes()
-    mime_type = mimetypes.guess_type(expanded.name)[0] or "application/octet-stream"
-    return types.Part.from_bytes(data=data, mime_type=mime_type)
+def _part_from_path(path: Path) -> types.Part:
+    """Create an image part from a local filesystem path."""
+    # Ensure common audio types are recognized across platforms (used for transcription as well).
+    mimetypes.add_type("audio/mp4", ".m4a")
+    mimetypes.add_type("audio/mpeg", ".mp3")
+    mimetypes.add_type("audio/wav", ".wav")
+    mimetypes.add_type("audio/aac", ".aac")
+
+    expanded = path.expanduser()
+    data = expanded.read_bytes()
+    mime_type = mimetypes.guess_type(expanded.name)[0] or "application/octet-stream"
+    return types.Part.from_bytes(data=data, mime_type=mime_type)
 
 
 def _part_from_url(url: str) -> types.Part:
@@ -495,7 +780,7 @@ def _part_from_url(url: str) -> types.Part:
     return types.Part.from_bytes(data=data, mime_type=mime_type)
 
 
-def _part_from_data_url(data_url: str) -> types.Part:
+def _part_from_data_url(data_url: str) -> types.Part:
     """Create an image part from a data URL."""
     header, encoded = data_url.split(",", 1)
     metadata = header[len("data:") :]
@@ -510,5 +795,35 @@ def _part_from_data_url(data_url: str) -> types.Part:
         data = base64.b64decode(encoded)
     else:
         data = encoded.encode("utf-8")
-
-    return types.Part.from_bytes(data=data, mime_type=mime_type or "application/octet-stream")
+
+    return types.Part.from_bytes(data=data, mime_type=mime_type or "application/octet-stream")
+
+
+def _extract_gemini_usage(response: object) -> TokenUsage | None:
+    usage_obj = getattr(response, "usage_metadata", None)
+    if usage_obj is None:
+        usage_obj = getattr(response, "usage", None)
+    if usage_obj is None:
+        return None
+
+    input_tokens = getattr(usage_obj, "prompt_token_count", None)
+    output_tokens = getattr(usage_obj, "candidates_token_count", None)
+    total_tokens = getattr(usage_obj, "total_token_count", None)
+
+    if input_tokens is None:
+        input_tokens = getattr(usage_obj, "input_tokens", None)
+    if output_tokens is None:
+        output_tokens = getattr(usage_obj, "output_tokens", None)
+    if total_tokens is None:
+        total_tokens = getattr(usage_obj, "total_tokens", None)
+
+    if isinstance(usage_obj, dict):
+        input_tokens = usage_obj.get("prompt_token_count", usage_obj.get("input_tokens"))
+        output_tokens = usage_obj.get("candidates_token_count", usage_obj.get("output_tokens"))
+        total_tokens = usage_obj.get("total_token_count", usage_obj.get("total_tokens"))
+
+    return TokenUsage(
+        input_tokens=int(input_tokens) if isinstance(input_tokens, int) else None,
+        output_tokens=int(output_tokens) if isinstance(output_tokens, int) else None,
+        total_tokens=int(total_tokens) if isinstance(total_tokens, int) else None,
+    )