openlit 1.1.2__py3-none-any.whl → 1.2.0__py3-none-any.whl

openlit/__init__.py

@@ -18,6 +18,7 @@ from openlit.instrumentation.anthropic import AnthropicInstrumentor
 from openlit.instrumentation.cohere import CohereInstrumentor
 from openlit.instrumentation.mistral import MistralInstrumentor
 from openlit.instrumentation.bedrock import BedrockInstrumentor
+from openlit.instrumentation.vertexai import VertexAIInstrumentor
 from openlit.instrumentation.langchain import LangChainInstrumentor
 from openlit.instrumentation.chroma import ChromaInstrumentor
 from openlit.instrumentation.pinecone import PineconeInstrumentor
@@ -139,7 +140,6 @@ def init(environment="default", application_name="default", tracer=None, otlp_en
         disable_metrics (bool): Flag to disable metrics (Optional)
     """
     disabled_instrumentors = disabled_instrumentors if disabled_instrumentors else []
-
     # Check for invalid instrumentor names
 
     module_name_map = {
@@ -148,9 +148,10 @@ def init(environment="default", application_name="default", tracer=None, otlp_en
         "cohere": "cohere",  
         "mistral": "mistralai",
         "bedrock": "boto3",
+        "vertexai": "vertexai",
         "langchain": "langchain",
         "chroma": "chromadb",
-        "pinecone": "pinecone-client",
+        "pinecone": "pinecone",
         "transformers": "transformers"
     }
 
@@ -195,6 +196,7 @@ def init(environment="default", application_name="default", tracer=None, otlp_en
             "cohere": CohereInstrumentor(),
             "mistral": MistralInstrumentor(),
             "bedrock": BedrockInstrumentor(),
+            "vertexai": VertexAIInstrumentor(),
             "langchain": LangChainInstrumentor(),
             "chroma": ChromaInstrumentor(),
             "pinecone": PineconeInstrumentor(),

openlit/instrumentation/vertexai/__init__.py

@@ -0,0 +1,147 @@
+# pylint: disable=useless-return, bad-staticmethod-argument, disable=duplicate-code
+"""Initializer of Auto Instrumentation of VertexAI Functions"""
+
+from typing import Collection
+import importlib.metadata
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from wrapt import wrap_function_wrapper
+
+from openlit.instrumentation.vertexai.vertexai import (
+    generate_content, predict, predict_streaming,
+    send_message, start_chat, start_chat_streaming,
+    embeddings
+)
+from openlit.instrumentation.vertexai.async_vertexai import (
+    generate_content_async, predict_async,
+    predict_streaming_async,
+    send_message_async,
+    start_chat_async, start_chat_streaming_async,
+    embeddings_async
+)
+
+
+_instruments = ("google-cloud-aiplatform >= 1.38.1",)
+
+class VertexAIInstrumentor(BaseInstrumentor):
+    """
+    An instrumentor for VertexAI's client library.
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        application_name = kwargs.get("application_name", "default")
+        environment = kwargs.get("environment", "default")
+        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("google-cloud-aiplatform")
+
+        #sync
+        wrap_function_wrapper(
+            "vertexai.generative_models",
+            "GenerativeModel.generate_content",
+            generate_content("vertexai.generate_content", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.generative_models",
+            "ChatSession.send_message",
+            send_message("vertexai.send_message", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.language_models",  
+            "TextGenerationModel.predict",  
+            predict("vertexai.predict", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.language_models",
+            "TextGenerationModel.predict_streaming",
+            predict_streaming("vertexai.predict", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.language_models",
+            "ChatSession.send_message",
+            start_chat("vertexai.send_message", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.language_models",
+            "ChatSession.send_message_streaming",
+            start_chat_streaming("vertexai.send_message", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.language_models",
+            "TextEmbeddingModel.get_embeddings",
+            embeddings("vertexai.get_embeddings", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        #async
+        wrap_function_wrapper(
+            "vertexai.generative_models",
+            "GenerativeModel.generate_content_async",
+            generate_content_async("vertexai.generate_content", version, environment,
+                                   application_name, tracer, pricing_info, trace_content,
+                                   metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.generative_models",
+            "ChatSession.send_message_async",
+            send_message_async("vertexai.send_message", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.language_models",
+            "TextGenerationModel.predict_async",
+            predict_async("vertexai.predict", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.language_models",
+            "TextGenerationModel.predict_streaming_async",
+            predict_streaming_async("vertexai.predict", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.language_models",
+            "ChatSession.send_message_async",
+            start_chat_async("vertexai.send_message", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.language_models",
+            "ChatSession.send_message_streaming_async",
+            start_chat_streaming_async("vertexai.send_message", version, environment,
+                                       application_name, tracer, pricing_info, trace_content,
+                                       metrics, disable_metrics),
+        )
+
+        wrap_function_wrapper(
+            "vertexai.language_models",
+            "TextEmbeddingModel.get_embeddings_async",
+            embeddings_async("vertexai.get_embeddings", version, environment, application_name,
+                     tracer, pricing_info, trace_content, metrics, disable_metrics),
+        )
+
+    def _uninstrument(self, **kwargs):
+        # Proper uninstrumentation logic to revert patched methods
+        pass