trulens-providers-google 2.2.3__tar.gz

trulens_providers_google-2.2.3/PKG-INFO

@@ -0,0 +1,28 @@
+Metadata-Version: 2.1
+Name: trulens-providers-google
+Version: 2.2.3
+Summary: Library to systematically track and evaluate LLM based applications.
+Home-page: https://trulens.org/
+License: MIT
+Author: Snowflake Inc.
+Author-email: ml-observability-wg-dl@snowflake.com
+Requires-Python: >=3.9,<4.0
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: google-auth (>=2.20.0)
+Requires-Dist: google-genai (>=1.27.0)
+Requires-Dist: trulens-core (>=2.0.0,<3.0.0)
+Requires-Dist: trulens-feedback (>=2.0.0,<3.0.0)
+Project-URL: Documentation, https://trulens.org/getting_started/
+Project-URL: Repository, https://github.com/truera/trulens
+Description-Content-Type: text/markdown
+
+# trulens-providers-google
+

trulens_providers_google-2.2.3/README.md

@@ -0,0 +1 @@
+# trulens-providers-google

trulens_providers_google-2.2.3/pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = [
+  "poetry-core",
+]
+
+[tool.poetry]
+name = "trulens-providers-google"
+version = "2.2.3"
+description = "Library to systematically track and evaluate LLM based applications."
+authors = [
+  "Snowflake Inc. <ml-observability-wg-dl@snowflake.com>",
+]
+license = "MIT"
+readme = "README.md"
+packages = [
+  { include = "trulens" },
+]
+homepage = "https://trulens.org/"
+documentation = "https://trulens.org/getting_started/"
+repository = "https://github.com/truera/trulens"
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "Operating System :: OS Independent",
+  "Development Status :: 3 - Alpha",
+  "License :: OSI Approved :: MIT License",
+]
+
+[tool.poetry.dependencies]
+python = "^3.9"
+trulens-core = { version = "^2.0.0" }
+trulens-feedback = { version = "^2.0.0" }
+google-genai = ">=1.27.0"
+google-auth = ">=2.20.0"
+
+[tool.poetry.group.dev.dependencies]
+trulens-core = { path = "../../core" }
+trulens-feedback = { path = "../../feedback" }

trulens_providers_google-2.2.3/trulens/providers/google/__init__.py

@@ -0,0 +1,22 @@
+"""
+!!! note "Additional Dependency Required"
+
+    To use this module, you must have the `trulens-providers-google` package installed.
+
+    ```bash
+    pip install trulens-providers-google
+    ```
+"""
+
+# WARNING: This file does not follow the no-init aliases import standard.
+
+from importlib.metadata import version
+
+from trulens.core.utils import imports as import_utils
+from trulens.providers.google.provider import Google
+
+__version__ = version(
+    import_utils.safe_importlib_package_name(__package__ or __name__)
+)
+
+__all__ = ["Google"]

trulens_providers_google-2.2.3/trulens/providers/google/endpoint.py

@@ -0,0 +1,178 @@
+import inspect
+import logging
+import os
+import pprint
+from typing import (
+    Any,
+    Callable,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+)
+
+from trulens.core.feedback import endpoint as core_endpoint
+from trulens.otel.semconv.trace import SpanAttributes
+
+from google import genai
+from google.auth.credentials import Credentials
+from google.genai import Client
+from google.genai.types import GenerateContentResponse
+
+logger = logging.getLogger(__name__)
+
+pp = pprint.PrettyPrinter()
+
+
+def _get_env_api_key() -> Optional[str]:
+    """Gets the API key from environment variables, prioritizing GOOGLE_API_KEY.
+    Returns:
+        The API key string if found, otherwise None. Empty string is considered
+        invalid.
+    """
+    env_google_api_key = os.environ.get("GOOGLE_API_KEY", None)
+    env_gemini_api_key = os.environ.get("GEMINI_API_KEY", None)
+    if env_google_api_key and env_gemini_api_key:
+        logger.warning(
+            "Both GOOGLE_API_KEY and GEMINI_API_KEY are set. Using GOOGLE_API_KEY."
+        )
+
+    return env_google_api_key or env_gemini_api_key or None
+
+
+class GoogleCostComputer:
+    @staticmethod
+    def handle_response(response: Any) -> Dict[str, Any]:
+        usage = response.usage_metadata
+        return {
+            SpanAttributes.COST.NUM_TOKENS: usage.total_token_count or 0,
+            SpanAttributes.COST.NUM_PROMPT_TOKENS: usage.prompt_token_count
+            or 0,
+            SpanAttributes.COST.NUM_COMPLETION_TOKENS: usage.candidates_token_count
+            or 0,
+            SpanAttributes.COST.NUM_REASONING_TOKENS: usage.thoughts_token_count
+            or 0,
+            # TODO: Check the cost computation functionality
+            # SpanAttributes.COST.COST: completion_cost(response),
+            SpanAttributes.COST.CURRENCY: "USD",
+            SpanAttributes.COST.MODEL: response.model_version,
+        }
+
+
+class GoogleCallback(core_endpoint.EndpointCallback):
+    _FIELDS_MAP: ClassVar[List[Tuple[str, str]]] = [
+        ("cost", "total_cost"),
+        ("n_tokens", "total_tokens"),
+        ("n_successful_requests", "successful_requests"),
+        ("n_prompt_tokens", "prompt_tokens"),
+        ("n_completion_tokens", "completion_tokens"),
+    ]
+
+    def handle_generation(self, response: Any):
+        """Get the usage information from GoogleGenAI LLM function response's usage_metadata field."""
+        response_dict = response
+        if isinstance(response, GenerateContentResponse):
+            response_dict = response.to_json_dict()
+
+        usage = response_dict.get("usage_metadata")
+        super().handle_generation(response_dict)
+        self.cost.n_successful_requests += 1
+
+        for cost_field, google_field in [
+            ("n_tokens", "total_token_count"),
+            ("n_prompt_tokens", "prompt_token_count"),
+            ("n_completion_tokens", "candidates_token_count"),
+            ("n_reasoning_tokens", "thoughts_token_count"),
+        ]:
+            setattr(
+                self.cost,
+                cost_field,
+                getattr(self.cost, cost_field, 0) + usage.get(google_field, 0),
+            )
+
+        # TODO: missing code for cost calculation
+
+
+class GoogleEndpoint(core_endpoint.Endpoint):
+    client: Optional["Client"] = None
+    vertexai: Optional[bool] = None
+    api_key: Optional[str] = None
+    credentials: Optional["Credentials"] = None
+    project: Optional[str] = None
+    location: Optional[str] = None
+
+    def __init__(
+        self,
+        client: Optional["Client"] = None,
+        vertexai: Optional[bool] = None,
+        api_key: Optional[str] = None,
+        credentials: Optional["Credentials"] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        **kwargs: dict,
+    ):
+        actual_client_instance = client
+        if not actual_client_instance:
+            if vertexai:
+                actual_client_instance = genai.Client(
+                    vertexai=vertexai,
+                    credentials=credentials,
+                    project=project,
+                    location=location,
+                )
+            else:
+                actual_client_instance = genai.Client(
+                    api_key=api_key or _get_env_api_key()
+                )
+
+        kwargs_for_super = {
+            # These are the fields declared in GoogleEndpoint
+            "client": actual_client_instance,
+            "vertexai": vertexai,
+            "api_key": api_key,
+            "credentials": credentials,
+            "project": project,
+            "location": location,
+            # This is a kwarg your Endpoint base class expects,
+            # and it's also set for the parent.
+            "callback_class": GoogleCallback,
+            **kwargs,  # Pass through any other arbitrary kwargs
+        }
+
+        super().__init__(**kwargs_for_super)
+
+    def handle_wrapped_call(
+        self,
+        func: Callable[..., Any],
+        bindings: inspect.BoundArguments,
+        response: Any,
+        callback: Optional[core_endpoint.EndpointCallback],
+    ) -> Any:
+        try:
+            if isinstance(response, GenerateContentResponse):
+                response_dict = response.to_json_dict()
+        except Exception as e:
+            logger.error(f"Error occurred while parsing response: {e}")
+            raise e
+
+        if (
+            isinstance(response_dict, dict)
+            and "usage_metadata" in response_dict
+        ):
+            candidate = response_dict["candidates"][0]
+            if (
+                "finish_reason" in candidate
+                and candidate["finish_reason"] == "STOP"
+            ):
+                self.global_callback.handle_generation(response=response_dict)
+
+                if callback is not None:
+                    callback.handle_generation(response=response_dict)
+        else:
+            logger.warning(
+                "Unrecognized Google content response format. It did not have usage information:\n%s",
+                pp.pformat(response_dict),
+            )
+
+        return response

trulens_providers_google-2.2.3/trulens/providers/google/provider.py

@@ -0,0 +1,174 @@
+import logging
+from typing import ClassVar, Dict, Optional, Sequence, Type, Union
+
+import pydantic
+from trulens.feedback import llm_provider
+from trulens.providers.google import endpoint as google_endpoint
+
+from google.auth.credentials import Credentials
+from google.genai import Client
+from google.genai.types import GenerateContentConfig
+
+logger = logging.getLogger(__name__)
+
+
+class Google(llm_provider.LLMProvider):
+    """Google provides access to Google's generative models via the Gemini Developer API
+    or Vertex AI, depending on the configuration.
+
+    For more details, see the official Gemini documentation.
+
+    Examples:
+
+    === "Connecting with a Gemini Developer API client"
+        ```python
+        from google import genai
+        from trulens.providers.google import Google
+
+        google_client = genai.Client(api_key="GOOGLE_API_KEY")
+        provider = Google(client=google_client)
+        ```
+
+    === "Connecting with a Vertex AI client"
+        ```python
+        from google import genai
+        from trulens.providers.google import Google
+
+        PROJECT_ID = "your_project_id"
+        LOCATION = "us-central1"
+
+        vertex_client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)
+        provider = Google(client=vertex_client)
+        ```
+
+    === "Using only an API key (Gemini Developer API)"
+        ```python
+        from trulens.providers.google import Google
+
+        provider = Google(api_key="GOOGLE_API_KEY")
+        ```
+
+    === "Using Vertex AI configuration directly"
+        ```python
+        from trulens.providers.google import Google
+
+        PROJECT_ID = "your_project_id"
+        LOCATION = "us-central1"
+
+        provider = Google(vertexai=True, project=PROJECT_ID, location=LOCATION)
+        ```
+
+    Args:
+        model_engine: Model engine to use. Defaults to `"gemini-2.5-flash"`.
+        api_key: API key for authenticating with the Gemini Developer API. If not provided,
+        the key will be read from the environment variable `GOOGLE_API_KEY` or `GEMINI_API_KEY`, if available.
+        vertexai: Whether to use Vertex AI endpoints. Set to `True` to use Vertex AI instead of the Gemini Developer API. Defaults to `False`.
+        credentials: Credentials to authenticate with Vertex AI. If not provided, default application credentials are used.
+        project: Google Cloud project ID used for billing and quota when using Vertex AI. Can be set via environment variables.
+        location: Region to send Vertex AI API requests to (e.g., `"us-central1"`). Can also be set via environment variables.
+    """
+
+    DEFAULT_MODEL_ENGINE: ClassVar[str] = "gemini-2.5-flash"
+
+    def __init__(
+        self,
+        endpoint=None,
+        client: Optional["Client"] = None,
+        vertexai: Optional[bool] = None,
+        api_key: Optional[str] = None,
+        credentials: Optional["Credentials"] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        model_engine: Optional[str] = None,
+        **kwargs: Dict,
+    ):
+        if model_engine is None:
+            model_engine = self.DEFAULT_MODEL_ENGINE
+
+        self_kwargs = dict()
+        self_kwargs.update(**kwargs)
+        self_kwargs["model_engine"] = model_engine
+
+        endpoint = google_endpoint.GoogleEndpoint(
+            client=client,
+            vertexai=vertexai,
+            api_key=api_key,
+            credentials=credentials,
+            project=project,
+            location=location,
+            **kwargs,
+        )
+        self_kwargs["endpoint"] = endpoint
+        super().__init__(**self_kwargs)
+
+    def _create_chat_completion(
+        self,
+        prompt: Optional[str] = None,
+        messages: Optional[Sequence[Dict]] = None,
+        response_format: Optional[Type[pydantic.BaseModel]] = None,
+        **kwargs,
+    ) -> Optional[Union[str, pydantic.BaseModel]]:
+        contents = []
+        system_instruction = ""
+        if messages is not None:
+            for message in messages:
+                if message["role"] == "system":
+                    system_instruction = message["content"]
+                elif message["role"] == "user":
+                    # TODO: Add multi-modal (text + image) handling here for Google models
+                    contents.append({
+                        "parts": [{"text": message["content"]}],
+                        "role": "user",
+                    })
+                else:
+                    logger.warning(
+                        f"Ignoring role '{message['role']}' — only 'system' and 'user' are supported."
+                    )
+        elif prompt is not None:
+            contents.append({
+                "parts": [{"text": prompt}],
+                "role": "user",
+            })
+        else:
+            raise ValueError("`prompt` or `messages` must be specified.")
+        config_kwargs = dict(**kwargs)
+
+        # Ensure seed is set if response_format is not used
+        if response_format is None and "seed" not in config_kwargs:
+            config_kwargs["seed"] = 123
+
+        # Add optional fields based on conditions
+        if response_format is not None and self._structured_output_supported():
+            config_kwargs.update({
+                "response_mime_type": "application/json",
+                "response_schema": response_format,
+            })
+
+        if system_instruction:
+            config_kwargs["system_instruction"] = system_instruction
+
+        response = self.endpoint.client.models.generate_content(
+            model=self.model_engine,
+            contents=contents,
+            config=GenerateContentConfig(**config_kwargs),
+        )
+
+        if response_format:
+            return response.parsed
+        return response.text
+
+    def _structured_output_supported(self) -> bool:
+        """Whether the provider supports structured output.
+        For more details: https://ai.google.dev/gemini-api/docs/models
+        """
+        # Models with only output audio do not support structured output
+        # generation (very logical)
+        audio_only_output_models = [
+            "gemini-2.5-pro-preview-tts",
+            "gemini-2.5-flash-preview-tts",
+            "gemini-2.5-flash-preview-native-audio-dialog",
+            "gemini-2.5-flash-exp-native-audio-thinking-dialog",
+        ]
+        if self.model_engine in audio_only_output_models:
+            return False
+        return True