Merge branch 'webb/langchain/tool-test' into webb/langchain/text-completion

Author: alexander-alderman-webb

tests/integrations/langchain/test_langchain.py

@@ -211,7 +211,7 @@ def test_langchain_create_agent(
     )
     events = capture_events()
 
-    model_repsonse = get_model_response(
+    model_response = get_model_response(
         nonstreaming_responses_model_response,
         serialize_pydantic=True,
         request_headers={
@@ -235,7 +235,7 @@ def test_langchain_create_agent(
     with patch.object(
         llm.client._client._client,
         "send",
-        return_value=model_repsonse,
+        return_value=model_response,
     ) as _:
         with start_transaction():
             agent.invoke(
@@ -254,12 +254,9 @@ def test_langchain_create_agent(
     assert len(chat_spans) == 1
     assert chat_spans[0]["origin"] == "auto.ai.langchain"
 
-    # Token usage is only available in newer versions of langchain (v0.2+)
-    # where usage_metadata is supported on AIMessageChunk
-    if "gen_ai.usage.input_tokens" in chat_spans[0]["data"]:
-        assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 10
-        assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 20
-        assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 30
+    assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 10
+    assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 20
+    assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 30
 
     if send_default_pii and include_prompts:
         assert (
@@ -411,17 +408,13 @@ def test_tool_execution_span(
     assert chat_spans[1]["origin"] == "auto.ai.langchain"
     assert tool_exec_span["origin"] == "auto.ai.langchain"
 
-    # Token usage is only available in newer versions of langchain (v0.2+)
-    # where usage_metadata is supported on AIMessageChunk
-    if "gen_ai.usage.input_tokens" in chat_spans[0]["data"]:
-        assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 142
-        assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 50
-        assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 192
+    assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 142
+    assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 50
+    assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 192
 
-    if "gen_ai.usage.input_tokens" in chat_spans[1]["data"]:
-        assert chat_spans[1]["data"]["gen_ai.usage.input_tokens"] == 89
-        assert chat_spans[1]["data"]["gen_ai.usage.output_tokens"] == 28
-        assert chat_spans[1]["data"]["gen_ai.usage.total_tokens"] == 117
+    assert chat_spans[1]["data"]["gen_ai.usage.input_tokens"] == 89
+    assert chat_spans[1]["data"]["gen_ai.usage.output_tokens"] == 28
+    assert chat_spans[1]["data"]["gen_ai.usage.total_tokens"] == 117
 
     if send_default_pii and include_prompts:
         assert "word" in tool_exec_span["data"][SPANDATA.GEN_AI_TOOL_INPUT]