aeri-python 4.0.0__py3-none-any.whl

aeri/langchain/__init__.py

@@ -0,0 +1,5 @@
+from .CallbackHandler import (
+    LangchainCallbackHandler as CallbackHandler,
+)  # For backward compatibility
+
+__all__ = ["CallbackHandler"]

aeri/langchain/utils.py

@@ -0,0 +1,212 @@
+"""@private"""
+
+import re
+from typing import Any, Dict, List, Literal, Optional, cast
+
+# NOTE ON DEPENDENCIES:
+# - since Jan 2024, there is https://pypi.org/project/langchain-openai/ which is a separate package and imports openai models.
+#   Decided to not make this a dependency of aeri as few people will have this. Need to match these models manually
+# - langchain_community is loaded as a dependency of langchain, but extremely unreliable. Decided to not depend on it.
+
+
+def _extract_model_name(
+    serialized: Optional[Dict[str, Any]],
+    **kwargs: Any,
+) -> Optional[str]:
+    """Extracts the model name from the serialized or kwargs object. This is used to get the model names for Aeri."""
+    # In this function we return on the first match, so the order of operations is important
+
+    # First, extract known models where we know the model name aka id
+    # Extract the model name from the provided path (aray) in the serialized or kwargs object
+    models_by_id = [
+        ("ChatGoogleGenerativeAI", ["kwargs", "model"], "serialized"),
+        ("ChatMistralAI", ["kwargs", "model"], "serialized"),
+        ("ChatVertexAi", ["kwargs", "model_name"], "serialized"),
+        ("ChatVertexAI", ["kwargs", "model_name"], "serialized"),
+        ("OpenAI", ["invocation_params", "model_name"], "kwargs"),
+        ("ChatOpenAI", ["invocation_params", "model_name"], "kwargs"),
+        ("AzureChatOpenAI", ["invocation_params", "model"], "kwargs"),
+        ("AzureChatOpenAI", ["invocation_params", "model_name"], "kwargs"),
+        ("AzureChatOpenAI", ["invocation_params", "azure_deployment"], "kwargs"),
+        ("HuggingFacePipeline", ["invocation_params", "model_id"], "kwargs"),
+        ("BedrockChat", ["kwargs", "model_id"], "serialized"),
+        ("Bedrock", ["kwargs", "model_id"], "serialized"),
+        ("BedrockLLM", ["kwargs", "model_id"], "serialized"),
+        ("ChatBedrock", ["kwargs", "model_id"], "serialized"),
+        ("LlamaCpp", ["invocation_params", "model_path"], "kwargs"),
+        ("WatsonxLLM", ["invocation_params", "model_id"], "kwargs"),
+    ]
+
+    for model_name, keys, select_from in models_by_id:
+        model = _extract_model_by_path_for_id(
+            model_name,
+            serialized,
+            kwargs,
+            keys,
+            cast(Literal["serialized", "kwargs"], select_from),
+        )
+        if model:
+            return model
+
+    # Second, we match AzureOpenAI as we need to extract the model name, fdeployment version and deployment name
+    if serialized:
+        serialized_id = serialized.get("id")
+        if (
+            serialized_id
+            and isinstance(serialized_id, list)
+            and len(serialized_id) > 0
+            and serialized_id[-1] == "AzureOpenAI"
+        ):
+            invocation_params = kwargs.get("invocation_params")
+            if invocation_params and isinstance(invocation_params, dict):
+                if invocation_params.get("model"):
+                    return str(invocation_params.get("model"))
+
+                if invocation_params.get("model_name"):
+                    return str(invocation_params.get("model_name"))
+
+            deployment_name = None
+            deployment_version = None
+
+            serialized_kwargs = serialized.get("kwargs")
+            if serialized_kwargs and isinstance(serialized_kwargs, dict):
+                if serialized_kwargs.get("openai_api_version"):
+                    deployment_version = serialized_kwargs.get("deployment_version")
+
+                if serialized_kwargs.get("deployment_name"):
+                    deployment_name = serialized_kwargs.get("deployment_name")
+
+            if not isinstance(deployment_name, str):
+                return None
+
+            if not isinstance(deployment_version, str):
+                return deployment_name
+
+            return (
+                deployment_name + "-" + deployment_version
+                if deployment_version not in deployment_name
+                else deployment_name
+            )
+
+    # Third, for some models, we are unable to extract the model by a path in an object. Aeri provides us with a string representation of the model pbjects
+    # We use regex to extract the model from the repr string
+    models_by_pattern = [
+        ("Anthropic", "model", "anthropic"),
+        ("ChatAnthropic", "model", None),
+        ("ChatTongyi", "model_name", None),
+        ("ChatCohere", "model", None),
+        ("Cohere", "model", None),
+        ("HuggingFaceHub", "model", None),
+        ("ChatAnyscale", "model_name", None),
+        ("TextGen", "model", "text-gen"),
+        ("Ollama", "model", None),
+        ("OllamaLLM", "model", None),
+        ("ChatOllama", "model", None),
+        ("ChatFireworks", "model", None),
+        ("ChatPerplexity", "model", None),
+        ("VLLM", "model", None),
+        ("Xinference", "model_uid", None),
+        ("ChatOCIGenAI", "model_id", None),
+        ("DeepInfra", "model_id", None),
+    ]
+
+    for model_name, pattern, default in models_by_pattern:
+        model = _extract_model_from_repr_by_pattern(
+            model_name, serialized, pattern, default
+        )
+        if model:
+            return model
+
+    # Finally, we try to extract the most likely paths as a catch all
+    random_paths = [
+        ["kwargs", "model_name"],
+        ["kwargs", "model"],
+        ["invocation_params", "model_name"],
+        ["invocation_params", "model"],
+    ]
+    for select in ["kwargs", "serialized"]:
+        for path in random_paths:
+            model = _extract_model_by_path(
+                serialized, kwargs, path, cast(Literal["serialized", "kwargs"], select)
+            )
+            if model:
+                return str(model)
+
+    return None
+
+
+def _extract_model_from_repr_by_pattern(
+    id: str,
+    serialized: Optional[Dict[str, Any]],
+    pattern: str,
+    default: Optional[str] = None,
+) -> Optional[str]:
+    if serialized is None:
+        return None
+
+    serialized_id = serialized.get("id")
+    if (
+        serialized_id
+        and isinstance(serialized_id, list)
+        and len(serialized_id) > 0
+        and serialized_id[-1] == id
+    ):
+        repr_str = serialized.get("repr")
+        if repr_str and isinstance(repr_str, str):
+            extracted = _extract_model_with_regex(pattern, repr_str)
+            return extracted if extracted else default if default else None
+
+    return None
+
+
+def _extract_model_with_regex(pattern: str, text: str) -> Optional[str]:
+    match = re.search(rf"{pattern}='(.*?)'", text)
+    if match:
+        return match.group(1)
+    return None
+
+
+def _extract_model_by_path_for_id(
+    id: str,
+    serialized: Optional[Dict[str, Any]],
+    kwargs: Dict[str, Any],
+    keys: List[str],
+    select_from: Literal["serialized", "kwargs"],
+) -> Optional[str]:
+    if serialized is None and select_from == "serialized":
+        return None
+
+    if serialized:
+        serialized_id = serialized.get("id")
+        if (
+            serialized_id
+            and isinstance(serialized_id, list)
+            and len(serialized_id) > 0
+            and serialized_id[-1] == id
+        ):
+            result = _extract_model_by_path(serialized, kwargs, keys, select_from)
+            return str(result) if result is not None else None
+
+    return None
+
+
+def _extract_model_by_path(
+    serialized: Optional[Dict[str, Any]],
+    kwargs: dict,
+    keys: List[str],
+    select_from: Literal["serialized", "kwargs"],
+) -> Optional[str]:
+    if serialized is None and select_from == "serialized":
+        return None
+
+    current_obj = kwargs if select_from == "kwargs" else serialized
+
+    for key in keys:
+        if current_obj and isinstance(current_obj, dict):
+            current_obj = current_obj.get(key)
+        else:
+            return None
+        if not current_obj:
+            return None
+
+    return str(current_obj) if current_obj else None

aeri/logger.py

@@ -0,0 +1,28 @@
+"""Logger configuration for Aeri OpenTelemetry integration.
+
+This module initializes and configures loggers used by the Aeri OpenTelemetry integration.
+It sets up the main 'aeri' logger and configures the httpx logger to reduce noise.
+
+Log levels used throughout Aeri:
+- DEBUG: Detailed tracing information useful for development and diagnostics
+- INFO: Normal operational information confirming expected behavior
+- WARNING: Indication of potential issues that don't prevent operation
+- ERROR: Errors that prevent specific operations but allow continued execution
+- CRITICAL: Critical errors that may prevent further operation
+"""
+
+import logging
+
+# Create the main Aeri logger
+aeri_logger = logging.getLogger("aeri")
+aeri_logger.setLevel(logging.WARNING)
+
+# Configure httpx logger to reduce noise from HTTP requests
+httpx_logger = logging.getLogger("httpx")
+httpx_logger.setLevel(logging.WARNING)
+
+# Add console handler if no handlers exist
+console_handler = logging.StreamHandler()
+formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+console_handler.setFormatter(formatter)
+httpx_logger.addHandler(console_handler)

aeri/media.py

@@ -0,0 +1,352 @@
+"""This module contains the AeriMedia class, which is used to wrap media objects for upload to Aeri."""
+
+import base64
+import hashlib
+import os
+import re
+from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple, TypeVar, cast
+
+import httpx
+
+from aeri.api import MediaContentType
+from aeri.logger import aeri_logger as logger
+from aeri.types import ParsedMediaReference
+
+if TYPE_CHECKING:
+    from aeri._client.client import Aeri
+
+T = TypeVar("T")
+
+
+class AeriMedia:
+    """A class for wrapping media objects for upload to Aeri.
+
+    This class handles the preparation and formatting of media content for Aeri,
+    supporting both base64 data URIs and raw content bytes.
+
+    Args:
+        obj (Optional[object]): The source object to be wrapped. Can be accessed via the `obj` attribute.
+        base64_data_uri (Optional[str]): A base64-encoded data URI containing the media content
+            and content type (e.g., "data:image/jpeg;base64,/9j/4AAQ...").
+        content_type (Optional[str]): The MIME type of the media content when providing raw bytes.
+        content_bytes (Optional[bytes]): Raw bytes of the media content.
+        file_path (Optional[str]): The path to the file containing the media content. For relative paths,
+            the current working directory is used.
+
+    Raises:
+        ValueError: If neither base64_data_uri or the combination of content_bytes
+            and content_type is provided.
+    """
+
+    obj: object
+
+    _content_bytes: Optional[bytes]
+    _content_type: Optional[MediaContentType]
+    _source: Optional[str]
+    _media_id: Optional[str]
+
+    def __init__(
+        self,
+        *,
+        obj: Optional[object] = None,
+        base64_data_uri: Optional[str] = None,
+        content_type: Optional[MediaContentType] = None,
+        content_bytes: Optional[bytes] = None,
+        file_path: Optional[str] = None,
+    ):
+        """Initialize a AeriMedia object.
+
+        Args:
+            obj: The object to wrap.
+
+            base64_data_uri: A base64-encoded data URI containing the media content
+                and content type (e.g., "data:image/jpeg;base64,/9j/4AAQ...").
+            content_type: The MIME type of the media content when providing raw bytes or reading from a file.
+            content_bytes: Raw bytes of the media content.
+            file_path: The path to the file containing the media content. For relative paths,
+                the current working directory is used.
+        """
+        self.obj = obj
+
+        if base64_data_uri is not None:
+            parsed_data = self._parse_base64_data_uri(base64_data_uri)
+            self._content_bytes, self._content_type = parsed_data
+            self._source = "base64_data_uri"
+
+        elif content_bytes is not None and content_type is not None:
+            self._content_type = content_type
+            self._content_bytes = content_bytes
+            self._source = "bytes"
+        elif (
+            file_path is not None
+            and content_type is not None
+            and os.path.exists(file_path)
+        ):
+            self._content_bytes = self._read_file(file_path)
+            self._content_type = content_type if self._content_bytes else None
+            self._source = "file" if self._content_bytes else None
+        else:
+            logger.error(
+                "base64_data_uri, or content_bytes and content_type, or file_path must be provided to AeriMedia"
+            )
+
+            self._content_bytes = None
+            self._content_type = None
+            self._source = None
+
+        self._media_id = self._get_media_id()
+
+    def _read_file(self, file_path: str) -> Optional[bytes]:
+        try:
+            with open(file_path, "rb") as file:
+                return file.read()
+        except Exception as e:
+            logger.error(f"Error reading file at path {file_path}", exc_info=e)
+
+            return None
+
+    def _get_media_id(self) -> Optional[str]:
+        content_hash = self._content_sha256_hash
+
+        if content_hash is None:
+            return None
+
+        # Convert hash to base64Url
+        url_safe_content_hash = content_hash.replace("+", "-").replace("/", "_")
+
+        return url_safe_content_hash[:22]
+
+    @property
+    def _content_length(self) -> Optional[int]:
+        return len(self._content_bytes) if self._content_bytes else None
+
+    @property
+    def _content_sha256_hash(self) -> Optional[str]:
+        if self._content_bytes is None:
+            return None
+
+        sha256_hash_bytes = hashlib.sha256(self._content_bytes).digest()
+
+        return base64.b64encode(sha256_hash_bytes).decode("utf-8")
+
+    @property
+    def _reference_string(self) -> Optional[str]:
+        if self._content_type is None or self._source is None or self._media_id is None:
+            return None
+
+        return f"@@@aeriMedia:type={self._content_type}|id={self._media_id}|source={self._source}@@@"
+
+    @staticmethod
+    def parse_reference_string(reference_string: str) -> ParsedMediaReference:
+        """Parse a media reference string into a ParsedMediaReference.
+
+        Example reference string:
+            "@@@aeriMedia:type=image/jpeg|id=some-uuid|source=base64_data_uri@@@"
+
+        Args:
+            reference_string: The reference string to parse.
+
+        Returns:
+            A TypedDict with the media_id, source, and content_type.
+
+        Raises:
+            ValueError: If the reference string is empty or not a string.
+            ValueError: If the reference string does not start with "@@@aeriMedia:type=".
+            ValueError: If the reference string does not end with "@@@".
+            ValueError: If the reference string is missing required fields.
+        """
+        if not reference_string:
+            raise ValueError("Reference string is empty")
+
+        if not isinstance(reference_string, str):
+            raise ValueError("Reference string is not a string")
+
+        if not reference_string.startswith("@@@aeriMedia:type="):
+            raise ValueError(
+                "Reference string does not start with '@@@aeriMedia:type='"
+            )
+
+        if not reference_string.endswith("@@@"):
+            raise ValueError("Reference string does not end with '@@@'")
+
+        content = reference_string[len("@@@aeriMedia:") :].rstrip("@@@")
+
+        # Split into key-value pairs
+        pairs = content.split("|")
+        parsed_data = {}
+
+        for pair in pairs:
+            key, value = pair.split("=", 1)
+            parsed_data[key] = value
+
+        # Verify all required fields are present
+        if not all(key in parsed_data for key in ["type", "id", "source"]):
+            raise ValueError("Missing required fields in reference string")
+
+        return ParsedMediaReference(
+            media_id=parsed_data["id"],
+            source=parsed_data["source"],
+            content_type=cast(MediaContentType, parsed_data["type"]),
+        )
+
+    def _parse_base64_data_uri(
+        self, data: str
+    ) -> Tuple[Optional[bytes], Optional[MediaContentType]]:
+        # Example data URI: data:image/jpeg;base64,/9j/4AAQ...
+        try:
+            if not data or not isinstance(data, str):
+                raise ValueError("Data URI is not a string")
+
+            if not data.startswith("data:"):
+                raise ValueError("Data URI does not start with 'data:'")
+
+            header, actual_data = data[5:].split(",", 1)
+            if not header or not actual_data:
+                raise ValueError("Invalid URI")
+
+            # Split header into parts and check for base64
+            header_parts = header.split(";")
+            if "base64" not in header_parts:
+                raise ValueError("Data is not base64 encoded")
+
+            # Content type is the first part
+            content_type = header_parts[0]
+            if not content_type:
+                raise ValueError("Content type is empty")
+
+            return base64.b64decode(actual_data), cast(MediaContentType, content_type)
+
+        except Exception as e:
+            logger.error("Error parsing base64 data URI", exc_info=e)
+
+            return None, None
+
+    @staticmethod
+    def resolve_media_references(
+        *,
+        obj: T,
+        aeri_client: "Aeri",
+        resolve_with: Literal["base64_data_uri"],
+        max_depth: int = 10,
+        content_fetch_timeout_seconds: int = 10,
+    ) -> T:
+        """Replace media reference strings in an object with base64 data URIs.
+
+        This method recursively traverses an object (up to max_depth) looking for media reference strings
+        in the format "@@@aeriMedia:...@@@". When found, it (synchronously) fetches the actual media content using
+        the provided Aeri client and replaces the reference string with a base64 data URI.
+
+        If fetching media content fails for a reference string, a warning is logged and the reference
+        string is left unchanged.
+
+        Args:
+            obj: The object to process. Can be a primitive value, array, or nested object.
+                If the object has a __dict__ attribute, a dict will be returned instead of the original object type.
+            aeri_client: Aeri client instance used to fetch media content.
+            resolve_with: The representation of the media content to replace the media reference string with.
+                Currently only "base64_data_uri" is supported.
+            max_depth: Optional. Default is 10. The maximum depth to traverse the object.
+
+        Returns:
+            A deep copy of the input object with all media references replaced with base64 data URIs where possible.
+            If the input object has a __dict__ attribute, a dict will be returned instead of the original object type.
+
+        Example:
+            obj = {
+                "image": "@@@aeriMedia:type=image/jpeg|id=123|source=bytes@@@",
+                "nested": {
+                    "pdf": "@@@aeriMedia:type=application/pdf|id=456|source=bytes@@@"
+                }
+            }
+
+            result = await AeriMedia.resolve_media_references(obj, aeri_client)
+
+            # Result:
+            # {
+            #     "image": "data:image/jpeg;base64,/9j/4AAQSkZJRg...",
+            #     "nested": {
+            #         "pdf": "data:application/pdf;base64,JVBERi0xLjcK..."
+            #     }
+            # }
+        """
+
+        def traverse(obj: Any, depth: int) -> Any:
+            if depth > max_depth:
+                return obj
+
+            # Handle string
+            if isinstance(obj, str):
+                regex = r"@@@aeriMedia:.+?@@@"
+                reference_string_matches = re.findall(regex, obj)
+                if len(reference_string_matches) == 0:
+                    return obj
+
+                result = obj
+                reference_string_to_media_content = {}
+                httpx_client = (
+                    aeri_client._resources.httpx_client
+                    if aeri_client._resources is not None
+                    else None
+                )
+
+                for reference_string in reference_string_matches:
+                    try:
+                        parsed_media_reference = AeriMedia.parse_reference_string(
+                            reference_string
+                        )
+                        media_data = aeri_client.api.media.get(
+                            parsed_media_reference["media_id"]
+                        )
+                        media_content = (
+                            httpx_client.get(
+                                media_data.url,
+                                timeout=content_fetch_timeout_seconds,
+                            )
+                            if httpx_client is not None
+                            else httpx.get(
+                                media_data.url, timeout=content_fetch_timeout_seconds
+                            )
+                        )
+                        media_content.raise_for_status()
+
+                        base64_media_content = base64.b64encode(
+                            media_content.content
+                        ).decode()
+                        base64_data_uri = f"data:{media_data.content_type};base64,{base64_media_content}"
+
+                        reference_string_to_media_content[reference_string] = (
+                            base64_data_uri
+                        )
+                    except Exception as e:
+                        logger.warning(
+                            f"Error fetching media content for reference string {reference_string}: {e}"
+                        )
+                        # Do not replace the reference string if there's an error
+                        continue
+
+                for (
+                    ref_str,
+                    media_content_str,
+                ) in reference_string_to_media_content.items():
+                    result = result.replace(ref_str, media_content_str)
+
+                return result
+
+            # Handle arrays
+            if isinstance(obj, list):
+                return [traverse(item, depth + 1) for item in obj]
+
+            # Handle dictionaries
+            if isinstance(obj, dict):
+                return {key: traverse(value, depth + 1) for key, value in obj.items()}
+
+            # Handle objects:
+            if hasattr(obj, "__dict__"):
+                return {
+                    key: traverse(value, depth + 1)
+                    for key, value in obj.__dict__.items()
+                }
+
+            return obj
+
+        return cast(T, traverse(obj, 0))