opik 1.6.4__py3-none-any.whl → 1.9.71__py3-none-any.whl

opik/guardrails/guards/topic.py

@@ -0,0 +1,76 @@
+from typing import List, Dict, Any, Optional
+from . import guard
+from .. import schemas
+import functools
+
+
+class Topic(guard.Guard):
+    """
+    Guard that validates text against a list of topics.
+
+    Can be configured to allow or restrict specific topics.
+    """
+
+    def __init__(
+        self,
+        allowed_topics: Optional[List[str]] = None,
+        restricted_topics: Optional[List[str]] = None,
+        threshold: float = 0.5,
+    ) -> None:
+        """
+        Initialize a Topic guard.
+
+        Args:
+            allowed_topics: List of topics that are allowed (text must match at least one)
+            restricted_topics: List of topics that are restricted (text must not match any)
+            threshold: Confidence threshold for topic detection (default: 0.5)
+
+        Raises:
+            ValueError: If both allowed_topics and restricted_topics are None
+
+        Note:
+            At least one of allowed_topics or restricted_topics must be provided.
+        """
+        if allowed_topics is None and restricted_topics is None:
+            raise ValueError("Must specify either allowed_topics or restricted_topics")
+
+        self._allowed_topics = allowed_topics
+        self._restricted_topics = restricted_topics
+        self._threshold = threshold
+
+    @functools.lru_cache()
+    def get_validation_configs(self) -> List[Dict[str, Any]]:
+        """
+        Get the validation configuration for topic matching.
+
+        Returns:
+            List containing the topic validation configuration
+        """
+
+        result = []
+
+        if self._allowed_topics:
+            result.append(
+                {
+                    "type": schemas.ValidationType.TOPIC,
+                    "config": {
+                        "topics": self._allowed_topics,
+                        "threshold": self._threshold,
+                        "mode": "allow",
+                    },
+                }
+            )
+
+        if self._restricted_topics:
+            result.append(
+                {
+                    "type": schemas.ValidationType.TOPIC,
+                    "config": {
+                        "topics": self._restricted_topics,
+                        "threshold": self._threshold,
+                        "mode": "restrict",
+                    },
+                }
+            )
+
+        return result

opik/guardrails/rest_api_client.py

@@ -0,0 +1,34 @@
+from typing import Any, Dict, List
+
+import httpx
+
+import opik.guardrails.schemas as schemas
+from opik.rest_api import core as rest_api_core
+
+
+class GuardrailsApiClient:  # TODO: in future this class can be autogenerated with Fern
+    def __init__(
+        self,
+        httpx_client: httpx.Client,
+        host_url: str,
+    ) -> None:
+        self._api_url = host_url.rstrip("/")
+        self._httpx_client = httpx_client
+
+    def __del__(self) -> None:
+        self._httpx_client.close()
+
+    def validate(
+        self, text: str, validations: List[Dict[str, Any]]
+    ) -> schemas.ValidationResponse:
+        payload = {"text": text, "validations": validations}
+
+        response = self._httpx_client.post(
+            f"{self._api_url}/api/v1/guardrails/validations", json=payload
+        )
+        if response.status_code != 200:
+            raise rest_api_core.ApiError(
+                status_code=response.status_code, body=response.json()
+            )
+
+        return schemas.ValidationResponse(**response.json())

opik/guardrails/schemas.py

@@ -0,0 +1,24 @@
+import pydantic
+import enum
+from typing import Dict, Any, List, Optional
+
+from opik.rest_api.types.check_public_result import CheckPublicResult
+
+
+class ValidationType(str, enum.Enum):
+    PII = "PII"
+    TOPIC = "TOPIC"
+
+
+class ValidationResult(pydantic.BaseModel):
+    validation_passed: bool
+    type: ValidationType
+    validation_config: Dict[str, Any]
+    validation_details: Dict[str, Any]
+
+
+class ValidationResponse(pydantic.BaseModel):
+    validation_passed: bool
+    validations: List[ValidationResult]
+    # This is client side injected
+    guardrail_result: Optional[CheckPublicResult] = None

opik/guardrails/tracing.py

@@ -0,0 +1,61 @@
+from typing import (
+    List,
+    Optional,
+    Callable,
+    AsyncGenerator,
+    Generator,
+    Union,
+    Any,
+    Dict,
+    Tuple,
+)
+
+
+from opik.decorator import (
+    arguments_helpers,
+)
+from opik.decorator import base_track_decorator
+from opik.api_objects import span
+
+
+class GuardrailsTrackDecorator(base_track_decorator.BaseTrackDecorator):
+    """
+    An implementation of BaseTrackDecorator designed specifically for guardrails span.
+    """
+
+    def _start_span_inputs_preprocessor(
+        self,
+        func: Callable,
+        track_options: arguments_helpers.TrackOptions,
+        args: Optional[Tuple],
+        kwargs: Optional[Dict[str, Any]],
+    ) -> arguments_helpers.StartSpanParameters:
+        assert isinstance(kwargs, dict)
+
+        result = arguments_helpers.StartSpanParameters(
+            name="Guardrail",
+            input={"generation": kwargs["generation"]},
+            type="guardrail",
+        )
+
+        return result
+
+    def _end_span_inputs_preprocessor(
+        self,
+        output: Any,
+        capture_output: bool,
+        current_span_data: span.SpanData,
+    ) -> arguments_helpers.EndSpanParameters:
+        result = arguments_helpers.EndSpanParameters(
+            output=output,
+        )
+
+        return result
+
+    def _streams_handler(
+        self,
+        output: Any,
+        capture_output: bool,
+        generations_aggregator: Optional[Callable[[List[Any]], str]],
+    ) -> Optional[Union[Generator, AsyncGenerator]]:
+        return super()._streams_handler(output, capture_output, generations_aggregator)

opik/healthcheck/__init__.py

@@ -1,5 +1,6 @@
 import opik
-from opik import config, environment
+import opik.config as config
+import opik.environment as environment
 from opik.healthcheck import checks, rich_representation
 
 

opik/healthcheck/checks.py

@@ -3,7 +3,7 @@ from typing import Optional, Tuple
 import httpx
 from rich.console import Console
 
-from opik import Opik
+import opik
 
 console = Console()
 
@@ -12,7 +12,7 @@ def get_backend_workspace_availability() -> Tuple[bool, Optional[str]]:
     is_available = False
     err_msg = None
 
-    opik_obj = Opik(_show_misconfiguration_message=False)
+    opik_obj = opik.Opik(_show_misconfiguration_message=False)
     console.print(
         f"--> Checking backend workspace availability at: {opik_obj.config.url_override}"
     )

opik/healthcheck/rich_representation.py

@@ -5,7 +5,7 @@ from rich.console import Console
 from rich.table import Table
 from rich.text import Text
 
-from opik import config
+import opik.config as config
 from opik.config import OpikConfig
 
 DEFAULT_KEY_COLOR = "green"

opik/hooks/__init__.py

@@ -0,0 +1,23 @@
+from .httpx_client_hook import (
+    HttpxClientHook,
+    add_httpx_client_hook,
+    register_httpx_client_hook,
+)
+from .anonymizer_hook import (
+    has_anonymizers,
+    add_anonymizer,
+    apply_anonymizers,
+    get_anonymizers,
+    clear_anonymizers,
+)
+
+__all__ = (
+    "HttpxClientHook",
+    "add_httpx_client_hook",
+    "register_httpx_client_hook",
+    "add_anonymizer",
+    "apply_anonymizers",
+    "clear_anonymizers",
+    "get_anonymizers",
+    "has_anonymizers",
+)

opik/hooks/anonymizer_hook.py

@@ -0,0 +1,36 @@
+from typing import List
+
+from opik.anonymizer import anonymizer
+
+
+# holder for a global list of anonymizers
+_anonymizers: List[anonymizer.Anonymizer] = []
+
+
+def add_anonymizer(anonymizer_hook: anonymizer.Anonymizer) -> None:
+    """Register a new anonymizer to be applied to all sensitive data logged by Opik."""
+    _anonymizers.append(anonymizer_hook)
+
+
+def clear_anonymizers() -> None:
+    """Clear all registered anonymizers."""
+    _anonymizers.clear()
+
+
+def has_anonymizers() -> bool:
+    """Check if any anonymizers have been registered."""
+    return len(_anonymizers) > 0
+
+
+def get_anonymizers() -> List[anonymizer.Anonymizer]:
+    """Get a list of all registered anonymizers."""
+    return _anonymizers
+
+
+def apply_anonymizers(
+    data: anonymizer.AnonymizerDataType,
+) -> anonymizer.AnonymizerDataType:
+    """Apply all registered anonymizers to the given data."""
+    for anonymizer_ in _anonymizers:
+        data = anonymizer_.anonymize(data)
+    return data

opik/hooks/httpx_client_hook.py

@@ -0,0 +1,112 @@
+import logging
+
+import httpx
+from typing import Any, Callable, List, Optional, Dict
+
+
+_deprecated_httpx_client_hooks: List[Callable[[httpx.Client], httpx.Client]] = []
+
+# holder for the global httpx client hook
+_httpx_client_hooks: List["HttpxClientHook"] = []
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+class HttpxClientHook:
+    def __init__(
+        self,
+        client_modifier: Optional[Callable[[httpx.Client], None]],
+        client_init_arguments: Optional[Dict[str, Any]],
+    ) -> None:
+        """Provides a means to customize an `httpx.Client` instance used by Opik.
+
+        This class allows users to attach a callable hook to modify or interact
+        with an `httpx.Client` instance and pass initialization arguments to
+        create customized client configurations. The primary usage is to allow
+        pre-processing or setup of HTTP clients used in a broader application.
+
+        Args:
+            client_modifier: Optional callable that accepts an `httpx.Client` instance and
+                returns a modified httpx.Client instance.
+            client_init_arguments: Dictionary containing additional `httpx.Client`
+                initialization arguments to be passed to the default `httpx.Client`.
+        """
+        self._hook = client_modifier
+        self._httpx_client_arguments = client_init_arguments
+
+    def update_init_arguments(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        if self._httpx_client_arguments is not None:
+            kwargs.update(self._httpx_client_arguments)
+
+        return kwargs
+
+    def __call__(self, client: httpx.Client) -> None:
+        if self._hook is not None:
+            self._hook(client)
+
+
+def add_httpx_client_hook(hook: HttpxClientHook) -> None:
+    """
+    Adds an HttpxClientHook to the list of hooks to be used with HTTPX clients.
+
+    Injects a new hook into the global list of HTTPX client hooks, allowing for
+    custom behavior or additional functionality when making requests with
+    an HTTPX client.
+
+    Args:
+        hook (HttpxClientHook): A callable to be added as an HTTPX client hook.
+
+    """
+    global _httpx_client_hooks
+    _httpx_client_hooks.append(hook)
+
+
+def build_init_arguments(default_kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Modifies and returns initialization arguments by applying pre-defined hooks.
+
+    This function iterates through a collection of hooks and applies their logic to
+    update the initialization arguments provided.
+
+    Args:
+        default_kwargs: A dictionary containing default initialization
+            arguments.
+
+    Returns:
+        Dict[str, Any]: The modified dictionary of initialization arguments.
+    """
+    for hook in _httpx_client_hooks:
+        default_kwargs = hook.update_init_arguments(default_kwargs)
+
+    return default_kwargs
+
+
+def apply_httpx_client_hooks(client: httpx.Client) -> None:
+    """Applies registered httpx client hooks."""
+    for hook in _httpx_client_hooks:
+        hook(client)
+
+    # apply deprecated hooks if any
+    for deprecated_hook in _deprecated_httpx_client_hooks:
+        deprecated_hook(client)
+
+
+def register_httpx_client_hook(hook: Callable[[httpx.Client], httpx.Client]) -> None:
+    """
+    Deprecated: This method is deprecated and will be removed in a future release. Please use `add_httpx_client_hook` instead.
+
+    Registers a hook for the customization of `httpx.Client` instances. The provided
+    hook function will be invoked with an `httpx.Client` instance and is expected
+    to return a customized `httpx.Client`.
+
+
+    Args:
+        hook: A callable that takes an `httpx.Client` instance and returns a
+              customized `httpx.Client`.
+    """
+    _deprecated_httpx_client_hooks.append(hook)
+
+    LOGGER.warning(
+        "register_httpx_client_hook is deprecated and will be removed in a future release. Please use add_httpx_client_hook instead."
+    )

opik/httpx_client.py

@@ -1,6 +1,9 @@
-from typing import Optional, Dict, Any, Union
+import gzip
+from typing import Optional, Dict, Any, Union, Iterable, AsyncIterable, Mapping
 import httpx
 import os
+import json as jsonlib
+
 from . import hooks, package_version
 import platform
 
@@ -15,7 +18,10 @@ POOL_TIMEOUT_SECONDS = 20
 
 
 def get(
-    workspace: Optional[str], api_key: Optional[str], check_tls_certificate: bool
+    workspace: Optional[str],
+    api_key: Optional[str],
+    check_tls_certificate: bool,
+    compress_json_requests: bool,
 ) -> httpx.Client:
     limits = httpx.Limits(keepalive_expiry=KEEPALIVE_EXPIRY_SECONDS)
 
@@ -24,6 +30,8 @@ def get(
         if check_tls_certificate is True and "SSL_CERT_FILE" in os.environ
         else check_tls_certificate
     )
+    # we need this to enable proxy server to analyze the request/response session during debugging
+    proxy = os.environ.get("_OPIK_HTTP_PROXY")
 
     timeout = httpx.Timeout(
         connect=CONNECT_TIMEOUT_SECONDS,
@@ -32,12 +40,22 @@ def get(
         pool=POOL_TIMEOUT_SECONDS,
     )
 
-    client = httpx.Client(limits=limits, verify=verify, timeout=timeout)
+    # build HTTPX client arguments
+    kwargs = {
+        "limits": limits,
+        "verify": verify,
+        "timeout": timeout,
+        "follow_redirects": True,
+        "proxy": proxy,
+    }
+    kwargs = hooks.httpx_client_hook.build_init_arguments(kwargs)
+
+    client = OpikHttpxClient(compress_json_requests=compress_json_requests, **kwargs)
 
     headers = _prepare_headers(workspace=workspace, api_key=api_key)
     client.headers.update(headers)
 
-    hooks.run_httpx_client_hooks(client)
+    hooks.httpx_client_hook.apply_httpx_client_hooks(client)
 
     return client
 
@@ -48,6 +66,7 @@ def _prepare_headers(
     result = {
         "X-OPIK-DEBUG-SDK-VERSION": package_version.VERSION,
         "X-OPIK-DEBUG-PY-VERSION": platform.python_version(),
+        "Accept-Encoding": "gzip",
     }
 
     if workspace is not None:
@@ -57,3 +76,55 @@ def _prepare_headers(
         result["Authorization"] = api_key
 
     return result
+
+
+class OpikHttpxClient(httpx.Client):
+    def __init__(self, compress_json_requests: bool = True, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.compress_json_requests = compress_json_requests
+
+    def build_request(
+        self,
+        method: str,
+        url: Union[httpx.URL, str],
+        *,
+        content: Optional[
+            Union[str, bytes, Iterable[bytes], AsyncIterable[bytes]]
+        ] = None,
+        data: Optional[Mapping[str, Any]] = None,
+        files: Any = None,
+        json: Any = None,
+        params: Any = None,
+        headers: Any = None,
+        cookies: Any = None,
+        timeout: Any = httpx.USE_CLIENT_DEFAULT,
+        extensions: Any = None,
+    ) -> httpx.Request:
+        # we override this method to allow compression of JSON requests that is handled
+        # by httpx.Client.request() as well as by httpx.Client.stream() (both used in the OPIK)
+        if self.compress_json_requests:
+            if method in ("POST", "PUT", "PATCH") and json is not None:
+                json_data = jsonlib.dumps(json).encode("utf-8")
+                content = gzip.compress(json_data)
+                json = None
+                if headers is None:
+                    headers = {}
+                headers["Content-Length"] = str(len(content))
+                headers["Content-Encoding"] = "gzip"
+                if "content-type" not in headers:
+                    # to avoid having it in headers two times with different cases in keys (e.g., streaming operations)
+                    headers["Content-Type"] = "application/json;charset=utf-8"
+
+        return super().build_request(
+            method=method,
+            url=url,
+            content=content,
+            data=data,
+            files=files,
+            json=json,
+            params=params,
+            headers=headers,
+            cookies=cookies,
+            timeout=timeout,
+            extensions=extensions,
+        )

opik/id_helpers.py

@@ -1,5 +1,7 @@
 from datetime import datetime
 from typing import Optional
+import random
+import string
 import uuid
 import uuid6
 
@@ -12,6 +14,22 @@ def generate_id(timestamp: Optional[datetime] = None) -> str:
     return str(uuid6.uuid7())
 
 
+def generate_random_alphanumeric_string(length: int) -> str:
+    """Generate a random alphanumeric string of the specified length.
+
+    Args:
+        length: The length of the string to generate.
+
+    Returns:
+        A random string containing only alphanumeric characters (a-z, A-Z, 0-9).
+    """
+    if length < 0:
+        raise ValueError("Length must be non-negative")
+
+    characters = string.ascii_letters + string.digits
+    return "".join(random.choice(characters) for _ in range(length))
+
+
 def uuid4_to_uuid7(user_datetime: datetime, user_uuid: str) -> uuid.UUID:
     """Convert a UUID v4 into a UUID v7 following RFC draft specification."""
     # Get Unix timestamp in milliseconds

opik/integrations/adk/__init__.py

@@ -0,0 +1,14 @@
+import google.adk
+from opik import semantic_version
+
+if semantic_version.SemanticVersion.parse(google.adk.__version__) < "1.3.0":  # type: ignore
+    from .legacy_opik_tracer import LegacyOpikTracer as OpikTracer
+else:
+    # Keep this import second to avoid static analyzers confusion
+    from .opik_tracer import OpikTracer  # type: ignore
+
+from .recursive_callback_injector import track_adk_agent_recursive
+from .graph.mermaid_graph_builder import build_mermaid_graph_definition
+
+
+__all__ = ["OpikTracer", "track_adk_agent_recursive", "build_mermaid_graph_definition"]

opik/integrations/adk/callback_context_info_extractors.py

@@ -0,0 +1,32 @@
+import logging
+from typing import Optional, Tuple, Dict, Any
+import google.adk.agents
+from google.adk.agents import callback_context
+
+LOGGER = logging.getLogger(__name__)
+
+
+def try_get_session_info(
+    callback_context: callback_context.CallbackContext,
+) -> Tuple[Optional[str], Dict[str, Any]]:
+    try:
+        session = callback_context._invocation_context.session
+        return session.id, {"user_id": session.user_id, "app_name": session.app_name}
+    except Exception:
+        LOGGER.error(
+            "Failed to get session information from ADK callback context", exc_info=True
+        )
+        return None, {}
+
+
+def try_get_current_agent_instance(
+    callback_context: callback_context.CallbackContext,
+) -> Optional[google.adk.agents.BaseAgent]:
+    try:
+        invocation_context = callback_context._invocation_context
+        return invocation_context.agent
+    except Exception:
+        LOGGER.error(
+            "Failed to get agent information from ADK callback context", exc_info=True
+        )
+        return None