openlit 1.32.2__py3-none-any.whl → 1.32.4__py3-none-any.whl

openlit/__init__.py

@@ -48,6 +48,7 @@ from openlit.instrumentation.chroma import ChromaInstrumentor
 from openlit.instrumentation.pinecone import PineconeInstrumentor
 from openlit.instrumentation.qdrant import QdrantInstrumentor
 from openlit.instrumentation.milvus import MilvusInstrumentor
+from openlit.instrumentation.astra import AstraInstrumentor
 from openlit.instrumentation.transformers import TransformersInstrumentor
 from openlit.instrumentation.litellm import LiteLLMInstrumentor
 from openlit.instrumentation.crewai import CrewAIInstrumentor
@@ -56,6 +57,7 @@ from openlit.instrumentation.multion import MultiOnInstrumentor
 from openlit.instrumentation.dynamiq import DynamiqInstrumentor
 from openlit.instrumentation.phidata import PhidataInstrumentor
 from openlit.instrumentation.julep import JulepInstrumentor
+from openlit.instrumentation.ai21 import AI21Instrumentor
 from openlit.instrumentation.gpu import GPUInstrumentor
 import openlit.guard
 import openlit.evals
@@ -250,6 +252,8 @@ def init(environment="default", application_name="default", tracer=None, otlp_en
         "reka-api": "reka",
         "premai": "premai",
         "julep": "julep",
+        "astra": "astrapy",
+        "ai21": "ai21",
     }
 
     invalid_instrumentors = [
@@ -339,6 +343,8 @@ def init(environment="default", application_name="default", tracer=None, otlp_en
             "reka-api": RekaInstrumentor(),
             "premai": PremAIInstrumentor(),
             "julep": JulepInstrumentor(),
+            "astra": AstraInstrumentor(),
+            "ai21": AI21Instrumentor(),
         }
 
         # Initialize and instrument only the enabled instrumentors

openlit/instrumentation/ai21/__init__.py

@@ -0,0 +1,66 @@
+# pylint: disable=useless-return, bad-staticmethod-argument, disable=duplicate-code
+"""Initializer of Auto Instrumentation of AI21 Functions"""
+
+from typing import Collection
+import importlib.metadata
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from wrapt import wrap_function_wrapper
+
+from openlit.instrumentation.ai21.ai21 import (
+    chat, chat_rag
+)
+from openlit.instrumentation.ai21.async_ai21 import (
+    async_chat, async_chat_rag
+)
+
+_instruments = ("ai21 >= 3.0.0",)
+
+class AI21Instrumentor(BaseInstrumentor):
+    """
+    An instrumentor for AI21's client library.
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        application_name = kwargs.get("application_name", "default_application")
+        environment = kwargs.get("environment", "default_environment")
+        tracer = kwargs.get("tracer")
+        metrics = kwargs.get("metrics_dict")
+        pricing_info = kwargs.get("pricing_info", {})
+        trace_content = kwargs.get("trace_content", False)
+        disable_metrics = kwargs.get("disable_metrics")
+        version = importlib.metadata.version("ai21")
+
+        #sync
+        wrap_function_wrapper(
+            "ai21.clients.studio.resources.chat.chat_completions",
+            "ChatCompletions.create",
+            chat("ai21.chat.completions", version, environment, application_name,
+                  tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+        wrap_function_wrapper(
+            "ai21.clients.studio.resources.studio_conversational_rag",
+            "StudioConversationalRag.create",
+            chat_rag("ai21.conversational_rag", version, environment, application_name,
+                  tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        #Async
+        wrap_function_wrapper(
+            "ai21.clients.studio.resources.chat.async_chat_completions",
+            "AsyncChatCompletions.create",
+            async_chat("ai21.chat.completions", version, environment, application_name,
+                  tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+        wrap_function_wrapper(
+            "ai21.clients.studio.resources.studio_conversational_rag",
+            "AsyncStudioConversationalRag.create",
+            async_chat_rag("ai21.conversational_rag", version, environment, application_name,
+                  tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+    def _uninstrument(self, **kwargs):
+        # Proper uninstrumentation logic to revert patched methods
+        pass