diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py index b1345d89..1fbd33c2 100644 --- a/src/langtrace_python_sdk/langtrace.py +++ b/src/langtrace_python_sdk/langtrace.py @@ -16,8 +16,7 @@ import os import sys -from typing import Optional -import importlib.util +from typing import Any, Optional from colorama import Fore from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME, SENTRY_DSN from opentelemetry import trace @@ -62,7 +61,12 @@ InstrumentationMethods, InstrumentationType, ) -from langtrace_python_sdk.utils import check_if_sdk_is_outdated, get_sdk_version +from langtrace_python_sdk.utils import ( + check_if_sdk_is_outdated, + get_sdk_version, + is_package_installed, + validate_instrumentations, +) from langtrace_python_sdk.utils.langtrace_sampler import LangtraceSampler import sentry_sdk @@ -193,10 +197,10 @@ def init( def init_instrumentations( disable_instrumentations: Optional[DisableInstrumentations], - all_instrumentations: dict + all_instrumentations: dict, ): if disable_instrumentations is None: - for idx, (name, v) in enumerate(all_instrumentations.items()): + for name, v in all_instrumentations.items(): if is_package_installed(name): v.instrument() @@ -205,61 +209,19 @@ def init_instrumentations( validate_instrumentations(disable_instrumentations) for key in disable_instrumentations: - for vendor in disable_instrumentations[key]: - if key == "only": - filtered_dict = { - k: v - for k, v in all_instrumentations.items() - if k != vendor.value - } - for _, v in filtered_dict.items(): - v.instrument() - else: - filtered_dict = { - k: v - for k, v in all_instrumentations.items() - if k == vendor.value - } - - for _, v in filtered_dict.items(): - v.instrument() - - -def validate_instrumentations(disable_instrumentations): - if disable_instrumentations is not None: - for key, value in disable_instrumentations.items(): - if isinstance(value, str): - # Convert single string to list of enum values - disable_instrumentations[key] = [InstrumentationType.from_string(value)] - elif isinstance(value, list): - # Convert list of strings to list of enum values - disable_instrumentations[key] = [ - ( - InstrumentationType.from_string(item) - if isinstance(item, str) - else item - ) - for item in value - ] - # Validate all items are of enum type - if not all( - isinstance(item, InstrumentationType) - for item in disable_instrumentations[key] - ): - raise TypeError( - f"All items in {key} must be of type InstrumentationType" - ) - if ( - disable_instrumentations.get("all_except") is not None - and disable_instrumentations.get("only") is not None - ): - raise ValueError( - "Cannot specify both only and all_except in disable_instrumentations" - ) - - -def is_package_installed(package_name): - import pkg_resources - - installed_packages = {p.key for p in pkg_resources.working_set} - return package_name in installed_packages + vendors = [k.value for k in disable_instrumentations[key]] + + key = next(iter(disable_instrumentations)) + filtered_dict = {} + if key == "all_except": + filtered_dict = { + k: v for k, v in all_instrumentations.items() if k in vendors + } + elif key == "only": + filtered_dict = { + k: v for k, v in all_instrumentations.items() if k not in vendors + } + + for name, v in filtered_dict.items(): + if is_package_installed(name): + v.instrument() diff --git a/src/langtrace_python_sdk/types/__init__.py b/src/langtrace_python_sdk/types/__init__.py index 414aebff..3173efae 100644 --- a/src/langtrace_python_sdk/types/__init__.py +++ b/src/langtrace_python_sdk/types/__init__.py @@ -8,16 +8,22 @@ class InstrumentationType(Enum): ANTHROPIC = "anthropic" GROQ = "groq" MISTRAL = "mistral" - PINECONE = "pinecone" - LLAMAINDEX = "llamaindex" + PINECONE = "pinecone-client" + LLAMAINDEX = "llama-index" CHROMADB = "chromadb" QDRANT = "qdrant" LANGCHAIN = "langchain" - LANGCHAIN_CORE = "langchain_core" - LANGCHAIN_COMMUNITY = "langchain_community" + LANGCHAIN_CORE = "langchain-core" + LANGCHAIN_COMMUNITY = "langchain-community" LANGGRAPH = "langgraph" WEAVIATE = "weaviate" OLLAMA = "ollama" + AUTOGEN = "autogen" + DSPY = "dspy-ai" + CREWAI = "crewai" + GEMINI = "google-generativeai" + VERTEXAI = "google-cloud-aiplatform" + MISTRALAI = "mistralai" @staticmethod def from_string(value: str): @@ -112,7 +118,10 @@ class InstrumentationMethods(TypedDict): cohere: List[VendorMethods.CohereMethods] weaviate: List[str] + _T = TypeVar("_T") + + class NotGiven: """ A sentinel singleton class used to distinguish omitted keyword arguments @@ -139,4 +148,4 @@ def __repr__(self) -> str: NotGivenOr = Union[_T, NotGiven] -NOT_GIVEN = NotGiven() \ No newline at end of file +NOT_GIVEN = NotGiven() diff --git a/src/langtrace_python_sdk/utils/__init__.py b/src/langtrace_python_sdk/utils/__init__.py index d970ba3c..5aab83ed 100644 --- a/src/langtrace_python_sdk/utils/__init__.py +++ b/src/langtrace_python_sdk/utils/__init__.py @@ -1,4 +1,4 @@ -from langtrace_python_sdk.types import NOT_GIVEN +from langtrace_python_sdk.types import NOT_GIVEN, InstrumentationType from .sdk_version_checker import SDKVersionChecker from opentelemetry.trace import Span from langtrace.trace_attributes import SpanAttributes @@ -49,3 +49,44 @@ def check_if_sdk_is_outdated(): def get_sdk_version(): return SDKVersionChecker().get_sdk_version() + + +def validate_instrumentations(disable_instrumentations): + if disable_instrumentations is not None: + if ( + disable_instrumentations.get("all_except") is not None + and disable_instrumentations.get("only") is not None + ): + raise ValueError( + "Cannot specify both only and all_except in disable_instrumentations" + ) + + for key, value in disable_instrumentations.items(): + if isinstance(value, str): + # Convert single string to list of enum values + disable_instrumentations[key] = [InstrumentationType.from_string(value)] + elif isinstance(value, list): + # Convert list of strings to list of enum values + disable_instrumentations[key] = [ + ( + InstrumentationType.from_string(item) + if isinstance(item, str) + else item + ) + for item in value + ] + # Validate all items are of enum type + if not all( + isinstance(item, InstrumentationType) + for item in disable_instrumentations[key] + ): + raise TypeError( + f"All items in {key} must be of type InstrumentationType" + ) + + +def is_package_installed(package_name): + import pkg_resources + + installed_packages = {p.key for p in pkg_resources.working_set} + return package_name in installed_packages diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py index e4f37b40..d77091b9 100644 --- a/src/langtrace_python_sdk/version.py +++ b/src/langtrace_python_sdk/version.py @@ -1 +1 @@ -__version__ = "2.3.16" +__version__ = "2.3.17"