update langchain to handle vertex and gemini llm calls

DataDog · Dec 9, 2024 · e4d51b7 · e4d51b7
1 parent 7320b6f
commit e4d51b7
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/ddtrace/llmobs/_integrations/langchain.py b/ddtrace/llmobs/_integrations/langchain.py
@@ -44,6 +44,8 @@
 ANTHROPIC_PROVIDER_NAME = "anthropic"
 BEDROCK_PROVIDER_NAME = "amazon_bedrock"
 OPENAI_PROVIDER_NAME = "openai"
+VERTEXAI_PROVIDER_NAME = "vertexai"
+GEMINI_PROVIDER_NAME = "google_palm"
 
 ROLE_MAPPING = {
     "human": "user",
@@ -81,6 +83,12 @@ def _llmobs_set_tags(
         if model_provider:
             if model_provider.startswith(BEDROCK_PROVIDER_NAME):
                 llmobs_integration = "bedrock"
+            # only the llm interface for Vertex AI will get instrumented
+            elif model_provider.startswith(VERTEXAI_PROVIDER_NAME) and operation == "llm":
+                llmobs_integration = "vertexai"
+            # only the llm interface for Gemini will get instrumented
+            elif model_provider.startswith(GEMINI_PROVIDER_NAME) and operation == "llm":
+                llmobs_integration = "google_generativeai"
             elif model_provider.startswith(OPENAI_PROVIDER_NAME):
                 llmobs_integration = "openai"
             elif operation == "chat" and model_provider.startswith(ANTHROPIC_PROVIDER_NAME):