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

opik/integrations/langchain/provider_usage_extractors/__init__.py

@@ -0,0 +1,5 @@
+from .usage_extractor import try_extract_provider_usage_data
+
+__all__ = [
+    "try_extract_provider_usage_data",
+]

opik/integrations/langchain/provider_usage_extractors/anthropic_usage_extractor.py

@@ -0,0 +1,101 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, Optional
+import opik
+from opik import llm_usage, _logging as opik_logging, logging_messages
+from . import provider_usage_extractor_protocol
+from . import langchain_run_helpers
+from .langchain_run_helpers import langchain_usage
+
+if TYPE_CHECKING:
+    pass
+
+LOGGER = logging.getLogger(__name__)
+
+
+class AnthropicUsageExtractor(
+    provider_usage_extractor_protocol.ProviderUsageExtractorProtocol
+):
+    PROVIDER = opik.LLMProvider.ANTHROPIC
+
+    def is_provider_run(self, run_dict: Dict[str, Any]) -> bool:
+        try:
+            if run_dict.get("serialized") is None:
+                return False
+
+            serialized_kwargs = run_dict.get("serialized", {}).get("kwargs", {})
+            has_anthropic_key = "anthropic_api_key" in serialized_kwargs
+
+            return has_anthropic_key
+
+        except Exception:
+            LOGGER.debug(
+                "Failed to check if Run instance is from Anthropic LLM, returning False.",
+                exc_info=True,
+            )
+            return False
+
+    def get_llm_usage_info(self, run_dict: Dict[str, Any]) -> llm_usage.LLMUsageInfo:
+        usage_dict = _try_get_token_usage(run_dict)
+        model = _try_get_model_name(run_dict)
+
+        return llm_usage.LLMUsageInfo(
+            provider=self.PROVIDER, model=model, usage=usage_dict
+        )
+
+
+def _try_get_token_usage(run_dict: Dict[str, Any]) -> Optional[llm_usage.OpikUsage]:
+    try:
+        if token_usage := langchain_run_helpers.try_to_get_usage_by_search(
+            run_dict, candidate_keys=None
+        ):
+            if isinstance(token_usage, langchain_usage.LangChainUsage):
+                anthropic_usage_dict = token_usage.map_to_anthropic_usage()
+                return llm_usage.OpikUsage.from_anthropic_dict(anthropic_usage_dict)
+
+        opik_logging.log_once_at_level(
+            logging.WARNING,
+            "Failed to extract token usage from presumably Anthropic LLM langchain run. Run dict: %s",
+            LOGGER,
+            run_dict,
+        )
+
+        opik_logging.log_once_at_level(
+            logging_level=logging.WARNING,
+            message=logging_messages.WARNING_TOKEN_USAGE_DATA_IS_NOT_AVAILABLE,
+            logger=LOGGER,
+        )
+
+    except Exception:
+        LOGGER.warning(
+            "Failed to extract token usage from presumably Anthropic LLM langchain run.",
+            exc_info=True,
+        )
+
+    return None
+
+
+def _try_get_model_name(run_dict: Dict[str, Any]) -> Optional[str]:
+    POSSIBLE_MODEL_NAME_KEYS = [
+        "model",  # detected in langchain-anthropic 0.3.5
+        "model_name",  # detected in langchain-anthropic 0.3.17
+    ]
+    model = None
+    for model_name_key in POSSIBLE_MODEL_NAME_KEYS:
+        try:
+            if run_dict["outputs"]["llm_output"] is not None:
+                model = run_dict["outputs"]["llm_output"][model_name_key]
+            else:
+                # Handle the streaming mode
+                model = run_dict["outputs"]["generations"][-1][-1]["message"]["kwargs"][
+                    "response_metadata"
+                ][model_name_key]
+        except KeyError:
+            continue
+
+    if model is None:
+        LOGGER.error(
+            "Failed to extract model name from presumably Anthropic LLM langchain Run object: %s",
+            run_dict,
+        )
+
+    return model

opik/integrations/langchain/provider_usage_extractors/anthropic_vertexai_usage_extractor.py

@@ -0,0 +1,67 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+import opik
+from opik import llm_usage
+from . import provider_usage_extractor_protocol
+from . import langchain_run_helpers
+
+if TYPE_CHECKING:
+    pass
+
+LOGGER = logging.getLogger(__name__)
+
+
+class AnthropicVertexAIUsageExtractor(
+    provider_usage_extractor_protocol.ProviderUsageExtractorProtocol
+):
+    PROVIDER = opik.LLMProvider.ANTHROPIC_VERTEXAI
+
+    def is_provider_run(self, run_dict: Dict[str, Any]) -> bool:
+        try:
+            if run_dict.get("serialized") is None:
+                return False
+
+            invocation_params = run_dict["extra"].get("invocation_params", {})
+            provider = invocation_params.get("_type", "").lower()
+            is_anthropic_vertexai = (
+                "vertexai" in provider.lower() and "anthropic" in provider.lower()
+            )
+
+            return is_anthropic_vertexai
+
+        except Exception:
+            LOGGER.debug(
+                "Failed to check if Run instance is from Anthropic LLM vertexai, returning False.",
+                exc_info=True,
+            )
+            return False
+
+    def get_llm_usage_info(self, run_dict: Dict[str, Any]) -> llm_usage.LLMUsageInfo:
+        usage_dict = _try_get_token_usage(run_dict)
+        model = _try_get_model_name(run_dict)
+
+        return llm_usage.LLMUsageInfo(
+            provider=self.PROVIDER, model=model, usage=usage_dict
+        )
+
+
+def _try_get_token_usage(run_dict: Dict[str, Any]) -> Optional[llm_usage.OpikUsage]:
+    try:
+        langchain_usage = langchain_run_helpers.try_get_token_usage(run_dict)
+        anthropic_usage_dict = langchain_usage.map_to_anthropic_usage()
+
+        opik_usage = llm_usage.OpikUsage.from_anthropic_dict(anthropic_usage_dict)
+        return opik_usage
+    except Exception:
+        LOGGER.warning(
+            "Failed to extract token usage from presumably Anthropic LLM vertexai langchain run.",
+            exc_info=True,
+        )
+        return None
+
+
+def _try_get_model_name(run_dict: Dict[str, Any]) -> Optional[str]:
+    if invocation_params := run_dict["extra"].get("invocation_params"):
+        return invocation_params.get("model_name")
+    return None

opik/integrations/langchain/provider_usage_extractors/bedrock_usage_extractor.py

@@ -0,0 +1,94 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, Optional, List
+
+import opik
+from opik import llm_usage, _logging as opik_logging, logging_messages
+from . import provider_usage_extractor_protocol
+from . import langchain_run_helpers
+from .langchain_run_helpers import langchain_usage
+
+if TYPE_CHECKING:
+    pass
+
+LOGGER = logging.getLogger(__name__)
+
+
+class BedrockUsageExtractor(
+    provider_usage_extractor_protocol.ProviderUsageExtractorProtocol
+):
+    PROVIDER = opik.LLMProvider.BEDROCK
+
+    def is_provider_run(self, run_dict: Dict[str, Any]) -> bool:
+        try:
+            if run_dict.get("serialized") is None:
+                return False
+
+            class_id: List[str] = run_dict.get("serialized", {}).get("id", [])
+            if len(class_id) == 0:
+                return False
+
+            class_name = class_id[-1]
+            is_bedrock = "ChatBedrock" in class_name
+
+            return is_bedrock
+
+        except Exception:
+            LOGGER.debug(
+                "Failed to check if Run instance is from Bedrock LLM, returning False.",
+                exc_info=True,
+            )
+            return False
+
+    def get_llm_usage_info(self, run_dict: Dict[str, Any]) -> llm_usage.LLMUsageInfo:
+        usage_dict = _try_get_token_usage(run_dict)
+        model = _try_get_model_name(run_dict)
+
+        return llm_usage.LLMUsageInfo(
+            provider=self.PROVIDER, model=model, usage=usage_dict
+        )
+
+
+def _try_get_token_usage(run_dict: Dict[str, Any]) -> Optional[llm_usage.OpikUsage]:
+    try:
+        if token_usage := langchain_run_helpers.try_to_get_usage_by_search(
+            run_dict, candidate_keys=None
+        ):
+            if isinstance(token_usage, langchain_usage.LangChainUsage):
+                bedrock_usage_dict = token_usage.map_to_bedrock_usage()
+                return llm_usage.OpikUsage.from_bedrock_dict(bedrock_usage_dict)
+
+        opik_logging.log_once_at_level(
+            logging.WARNING,
+            "Failed to extract token usage from presumably Bedrock LLM langchain run. Run dict: %s",
+            LOGGER,
+            run_dict,
+        )
+
+        opik_logging.log_once_at_level(
+            logging_level=logging.WARNING,
+            message=logging_messages.WARNING_TOKEN_USAGE_DATA_IS_NOT_AVAILABLE,
+            logger=LOGGER,
+        )
+
+    except Exception:
+        LOGGER.warning(
+            "Failed to extract token usage from presumably Bedrock LLM langchain run. Run dict: %s",
+            run_dict,
+            exc_info=True,
+        )
+
+    return None
+
+
+def _try_get_model_name(run_dict: Dict[str, Any]) -> Optional[str]:
+    MODEL_NAME_KEY = "model_id"
+
+    model = run_dict.get("serialized", {}).get("kwargs", {}).get(MODEL_NAME_KEY, None)
+
+    if model is None:
+        LOGGER.error(
+            "Failed to extract model name from presumably Bedrock LLM langchain Run object: %s",
+            run_dict,
+        )
+
+    return model

opik/integrations/langchain/provider_usage_extractors/google_generative_ai_usage_extractor.py

@@ -0,0 +1,109 @@
+import logging
+from typing import Any, Dict, Optional
+
+import opik
+from opik import llm_usage, logging_messages, _logging as opik_logging
+from . import langchain_run_helpers
+from . import provider_usage_extractor_protocol
+from .langchain_run_helpers import langchain_usage
+
+LOGGER = logging.getLogger(__name__)
+
+
+class GoogleGenerativeAIUsageExtractor(
+    provider_usage_extractor_protocol.ProviderUsageExtractorProtocol
+):
+    PROVIDER = opik.LLMProvider.GOOGLE_AI
+
+    def is_provider_run(self, run_dict: Dict[str, Any]) -> bool:
+        try:
+            if run_dict.get("serialized") is None:
+                return False
+
+            if (
+                ls_metadata := langchain_run_helpers.try_get_ls_metadata(run_dict)
+            ) is not None:
+                if "google_genai" == ls_metadata.provider:
+                    return True
+
+            if (
+                invocation_params := run_dict["extra"].get("invocation_params")
+            ) is not None:
+                if _is_invocation_param_of_google_gen_ai_type(
+                    invocation_params.get("_type").lower()
+                ):
+                    return True
+
+            return False
+
+        except Exception as ex:
+            LOGGER.debug(
+                "Failed to check if Run instance is from Google Generative AI, returning False. Reason: %s",
+                ex,
+                exc_info=True,
+            )
+            return False
+
+    def get_llm_usage_info(self, run_dict: Dict[str, Any]) -> llm_usage.LLMUsageInfo:
+        usage_dict = _try_get_token_usage(run_dict)
+        model = _get_model_name(run_dict)
+
+        return llm_usage.LLMUsageInfo(
+            provider=self.PROVIDER, model=model, usage=usage_dict
+        )
+
+
+def _try_get_token_usage(run_dict: Dict[str, Any]) -> Optional[llm_usage.OpikUsage]:
+    try:
+        if token_usage := langchain_run_helpers.try_to_get_usage_by_search(
+            run_dict, candidate_keys=None
+        ):
+            if isinstance(token_usage, langchain_usage.LangChainUsage):
+                gemini_usage_dict = token_usage.map_to_google_gemini_usage()
+                return llm_usage.OpikUsage.from_google_dict(gemini_usage_dict)
+
+        opik_logging.log_once_at_level(
+            logging.WARNING,
+            logging_messages.FAILED_TO_EXTRACT_TOKEN_USAGE_FROM_PRESUMABLY_LANGCHAIN_GOOGLE_LLM_RUN,
+            LOGGER,
+            run_dict,
+        )
+
+        opik_logging.log_once_at_level(
+            logging_level=logging.WARNING,
+            message=logging_messages.WARNING_TOKEN_USAGE_DATA_IS_NOT_AVAILABLE,
+            logger=LOGGER,
+        )
+
+    except Exception:
+        LOGGER.warning(
+            logging_messages.FAILED_TO_EXTRACT_TOKEN_USAGE_FROM_PRESUMABLY_LANGCHAIN_GOOGLE_LLM_RUN,
+            run_dict,
+            exc_info=True,
+        )
+
+    return None
+
+
+def _get_model_name(run_dict: Dict[str, Any]) -> Optional[str]:
+    """
+    Extracts the model name from the run dictionary.
+    """
+    model = None
+
+    # try metadata first
+    if (ls_metadata := langchain_run_helpers.try_get_ls_metadata(run_dict)) is not None:
+        model = ls_metadata.model
+
+    elif (invocation_params := run_dict["extra"].get("invocation_params")) is not None:
+        model = invocation_params.get("model")
+
+    if model is not None:
+        # Gemini **may** add "models/" prefix to some model versions
+        model = model.split("/")[-1]
+
+    return model
+
+
+def _is_invocation_param_of_google_gen_ai_type(provider: str) -> bool:
+    return "google-generative-ai" == provider or "google_gemini" == provider

opik/integrations/langchain/provider_usage_extractors/groq_usage_extractor.py

@@ -0,0 +1,92 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+import opik
+from opik import _logging as opik_logging
+from opik import llm_usage, logging_messages
+from . import provider_usage_extractor_protocol, langchain_run_helpers
+from .langchain_run_helpers import langchain_usage
+
+if TYPE_CHECKING:
+    pass
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+GROQ_CANDIDATE_USAGE_KEYS = {"prompt_tokens", "completion_tokens", "total_tokens"}
+
+
+class GroqUsageExtractor(
+    provider_usage_extractor_protocol.ProviderUsageExtractorProtocol
+):
+    PROVIDER = opik.LLMProvider.GROQ
+
+    def is_provider_run(self, run_dict: Dict[str, Any]) -> bool:
+        try:
+            if run_dict.get("serialized") is None:
+                return False
+
+            serialized_kwargs = run_dict["serialized"].get("kwargs", {})
+            has_groq_key = "groq_api_key" in serialized_kwargs
+
+            return has_groq_key
+
+        except Exception:
+            LOGGER.debug(
+                "Failed to check if Run instance is from Groq LLM, returning False.",
+                exc_info=True,
+            )
+            return False
+
+    def get_llm_usage_info(self, run_dict: Dict[str, Any]) -> llm_usage.LLMUsageInfo:
+        opik_usage = _try_get_token_usage(run_dict)
+        model = _try_get_model_name(run_dict)
+
+        return llm_usage.LLMUsageInfo(
+            provider=self.PROVIDER, model=model, usage=opik_usage
+        )
+
+
+def _try_get_token_usage(run_dict: Dict[str, Any]) -> Optional[llm_usage.OpikUsage]:
+    """
+    Attempts to extract and return the token usage from the given run dictionary.
+
+    Depending on the execution type (invoke, streaming mode, async, etc.), or even the model name itself,
+    token usage info might be in different places, different formats, or completely missing.
+    """
+    try:
+        if token_usage := langchain_run_helpers.try_to_get_usage_by_search(
+            run_dict, GROQ_CANDIDATE_USAGE_KEYS
+        ):
+            if isinstance(token_usage, dict):
+                return llm_usage.OpikUsage.from_groq_completions_dict(token_usage)
+            elif isinstance(token_usage, langchain_usage.LangChainUsage):
+                # streaming mode handling
+                # token usage data MAY be available at the end of streaming
+                # in async mode may not provide token usage info
+                groq_usage_dict = token_usage.map_to_groq_completions_usage()
+                return llm_usage.OpikUsage.from_groq_completions_dict(groq_usage_dict)
+
+        opik_logging.log_once_at_level(
+            logging.WARNING,
+            "Failed to extract token usage from presumably Groq LLM langchain run. Run dict: %s",
+            LOGGER,
+            run_dict,
+        )
+
+    except Exception:
+        LOGGER.warning(
+            logging_messages.FAILED_TO_EXTRACT_TOKEN_USAGE_FROM_PRESUMABLY_LANGCHAIN_GROQ_LLM_RUN,
+            run_dict,
+            exc_info=True,
+        )
+    return None
+
+
+def _try_get_model_name(run_dict: Dict[str, Any]) -> Optional[str]:
+    model = run_dict["extra"].get("metadata", {}).get("ls_model_name")
+    if model is not None:
+        model = model.split("/")[-1]
+
+    return model

opik/integrations/langchain/provider_usage_extractors/langchain_run_helpers/__init__.py

@@ -0,0 +1,15 @@
+from .helpers import (
+    try_get_token_usage,
+    try_get_streaming_token_usage,
+    try_get_ls_metadata,
+    find_token_usage_dict,
+    try_to_get_usage_by_search,
+)
+
+__all__ = [
+    "try_get_token_usage",
+    "try_get_streaming_token_usage",
+    "try_get_ls_metadata",
+    "find_token_usage_dict",
+    "try_to_get_usage_by_search",
+]

opik/integrations/langchain/provider_usage_extractors/langchain_run_helpers/helpers.py

@@ -0,0 +1,134 @@
+from typing import Dict, Any, Optional, Set, NamedTuple, List, Union
+
+from . import langchain_usage
+
+
+class LSMetadata(NamedTuple):
+    model: str
+    provider: str
+    model_type: str
+
+
+def try_get_token_usage(run_dict: Dict[str, Any]) -> langchain_usage.LangChainUsage:
+    if (usage := try_get_streaming_token_usage(run_dict)) is not None:
+        return usage
+
+    # try generation_info
+    usage_metadata = run_dict["outputs"]["generations"][-1][-1]["generation_info"][
+        "usage_metadata"
+    ]
+    return langchain_usage.LangChainUsage.from_original_usage_dict(usage_metadata)
+
+
+def try_get_streaming_token_usage(
+    run_dict: Dict[str, Any],
+) -> Optional[langchain_usage.LangChainUsage]:
+    if "message" in run_dict["outputs"]["generations"][-1][-1]:
+        usage_metadata = run_dict["outputs"]["generations"][-1][-1]["message"][
+            "kwargs"
+        ].get("usage_metadata")
+
+        if usage_metadata is not None:
+            return langchain_usage.LangChainUsage.from_original_usage_dict(
+                usage_metadata
+            )
+
+    return None
+
+
+def try_get_ls_metadata(run_dict: Dict[str, Any]) -> Optional[LSMetadata]:
+    if metadata := run_dict["extra"].get("metadata"):
+        model = metadata.get("ls_model_name")
+        provider = metadata.get("ls_provider")
+        model_type = metadata.get("ls_model_type")
+        return LSMetadata(model=model, provider=provider, model_type=model_type)
+
+    return None
+
+
+def try_to_get_usage_by_search(
+    run_dict: Dict[str, Any], candidate_keys: Optional[Set[str]]
+) -> Optional[Union[Dict[str, Any], langchain_usage.LangChainUsage]]:
+    """
+    Attempts to extract usage data from the given dictionary.
+
+    This function updates the provided set of candidate keys with additional
+    keys specific to LangChain usage. It then searches for token usage data
+    within the provided dictionary using the updated set of keys. If the usage
+    data conforms to the LangChain usage format, it converts it accordingly.
+    Otherwise, it simply returns the found usage dictionary if available.
+
+    Args:
+        run_dict: Dictionary potentially containing usage data.
+        candidate_keys: Set of keys to consider when extracting usage
+            information. Additional LangChain-specific keys are added to this set
+            or will be used if not provided.
+
+    Returns:
+        Extracted usage data. Returns a LangChainUsage object if the data matches
+        LangChain usage format. Otherwise, returns a plain dictionary with the usage data,
+        or None if no usage data is found.
+    """
+    if candidate_keys is None:
+        all_keys_should_match = True
+        candidate_keys = set()
+    else:
+        all_keys_should_match = False
+    candidate_keys.update(langchain_usage.LANGCHAIN_USAGE_KEYS)
+
+    usage_dict = find_token_usage_dict(run_dict, candidate_keys, all_keys_should_match)
+    if usage_dict is None:
+        return None
+
+    if langchain_usage.is_langchain_usage(usage_dict):
+        return langchain_usage.LangChainUsage.from_original_usage_dict(usage_dict)
+
+    return usage_dict
+
+
+def find_token_usage_dict(
+    data: Union[Dict[str, Any], List[Any]],
+    candidate_keys: Set[str],
+    all_keys_should_match: bool,
+) -> Optional[Dict[str, Any]]:
+    """
+    Find the first dictionary containing any of the specified candidate keys within a nested data structure.
+
+    The function recursively traverses through the given nested data structure, which
+    can be a dictionary, list, or tuple. It searches for the first dictionary
+    that includes one or more keys from the specified candidate keys and returns it.
+    If no such dictionary is found, the function returns None.
+
+    Args:
+        all_keys_should_match: if True, all candidate keys must be present in the dictionary.
+        data: A nested data structure containing dictionaries, lists, or tuples to search through.
+        candidate_keys: A set of strings representing the keys to look for in the dictionaries.
+
+    Returns:
+        The first dictionary as a `dict` that contains at least one of the candidate keys.
+        Returns None if no such dictionary is found.
+    """
+    # Handle dictionary case
+    if isinstance(data, dict):
+        # Check if the current dictionary contains any or all of the candidate keys
+        matched_keys = candidate_keys.intersection(data.keys())
+        if all_keys_should_match and len(matched_keys) == len(candidate_keys):
+            return data
+        elif not all_keys_should_match and len(matched_keys) > 0:
+            return data
+
+        # Recursively search through dictionary values
+        for value in data.values():
+            result = find_token_usage_dict(value, candidate_keys, all_keys_should_match)
+            if result is not None:
+                return result
+
+    # Handle list and tuple cases
+    elif isinstance(data, (list, tuple)):
+        for item in data:
+            result = find_token_usage_dict(item, candidate_keys, all_keys_should_match)
+            if result is not None:
+                return result
+
+    # If data is neither dict, list, nor tuple, or no token usage found
+    return None