feat(langchain): [MLOB-1972] update langchain to handle vertex and gemini llm calls

[ncybul]

This PR updates the Langchain integration to handle LLM calls instrumented in the Vertex AI and Gemini integrations by checking for their respective provider names (vertexai and google_palm) and instrumenting the wrapper Langchain calls as workflow spans if these provider names are detected.

Importantly, because of the way Langchain invokes chat generations for Vertex AI and Gemini, our integrations will not capture these inner LLM calls (in both cases, the prediction client is called directly with the input instead of using the chat.send_message method which we have instrumented in both cases). Therefore, we will only capture the Langchain LLM call and hence leave it as an LLM span for chat generations.

Testing

As we work on a more stable way to test our langchain integration, I opted to manually verify these changes by submitting traces to staging. Below I have included the code that I ran and the resulting trace that appeared in the product. It is expected that LLM calls have a langchain workflow span with a Gemini/Vertex AI child LLM span whereas Chat calls have only the Langchain LLM span instrumented (see above for details on why this is the case).

Gemini LLM call

from langchain_google_genai import GoogleGenerativeAI
llm = GoogleGenerativeAI(model="gemini-pro")
print(
    llm.invoke(
        "What are some tips for improving sleep quality?"
    )
)

Gemini Chat call

from langchain_google_genai import ChatGoogleGenerativeAI
llm = ChatGoogleGenerativeAI(model="gemini-pro")
resp = llm.invoke("Tell me a joke.")
print(resp)

Vertex AI LLM call

from langchain_google_vertexai import VertexAI
model = VertexAI(model_name="gemini-pro")
message = "What is the optimal temperature for sleeping?"
model.invoke(message)

Vertex AI Chat call

from langchain_google_vertexai import ChatVertexAI
llm = ChatVertexAI(
    model="gemini-1.5-flash-001",
    temperature=0,
    max_tokens=None,
    max_retries=6,
    stop=None,
)
messages = [
    (
        "system",
        "You are a helpful assistant that translates English to French. Translate the user sentence.",
    ),
    ("human", "I love programming."),
]
ai_msg = llm.invoke(messages)
print(ai_msg)

Checklist

• PR author has checked that all the criteria below are met
• The PR description includes an overview of the change
• The PR description articulates the motivation for the change
• The change includes tests OR the PR description describes a testing strategy
• The PR description notes risks associated with the change, if any
• Newly-added code is easy to change
• The change follows the library release note guidelines
• The change includes or references documentation updates if necessary
• Backport labels are set (if applicable)

Reviewer Checklist

• Reviewer has checked that all the criteria below are met
• Title is accurate
• All changes are related to the pull request's stated goal
• Avoids breaking API changes
• Testing strategy adequately addresses listed risks
• Newly-added code is easy to change
• Release note makes sense to a user of the library
• If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment
• Backport labels are set in a manner that is consistent with the release branch maintenance policy

[github-actions]

CODEOWNERS have been resolved as:

ddtrace/llmobs/_integrations/langchain.py                               @DataDog/ml-observability

[datadog-dd-trace-py-rkomorn]

Datadog Report

Branch report: nicole-cybul/add-langchain-tests2
Commit report: e4d51b7
Test service: dd-trace-py

✅ 0 Failed, 20 Passed, 1448 Skipped, 1m 9.81s Total duration (35m 36.91s time saved)

[pr-commenter]

Benchmarks

Benchmark execution time: 2024-12-09 15:54:20

Comparing candidate commit e4d51b7 in PR branch nicole-cybul/add-langchain-tests2 with baseline commit 7320b6f in branch main.

Found 0 performance improvements and 2 performance regressions! Performance is the same for 392 metrics, 2 unstable metrics.

scenario:iast_aspects-ospathbasename_aspect

• 🟥 execution_time [+367.818ns; +418.590ns] or [+10.984%; +12.500%]

scenario:iast_aspects-ospathdirname_aspect

• 🟥 execution_time [+320.994ns; +389.810ns] or [+8.635%; +10.487%]