mirror of
https://github.com/hanzoai/insights-python.git
synced 2026-10-30 18:33:17 +00:00
zeekay
3ca982343e
- Remove Posthog = Insights alias from __init__.py - Remove PosthogContextMiddleware alias from django.py - Remove PostHogTracingProcessor alias from processor.py - Change $lib from "posthog-python" to "insights-python" - Change ingestion URLs from posthog.com to insights.hanzo.ai - Rename all posthog_* kwargs to insights_* across AI wrappers - Rename __posthog_exception_captured to __insights_exception_captured - Rename posthog_context_stack contextvar to insights_context_stack - Rename posthog🎏 Redis prefix to insights🎏 - Rename $$_posthog_redacted_* sentinels to $$_insights_redacted_* - Remove POSTHOG_MW_* Django settings fallback, X-POSTHOG-* headers - Rename Prompts(posthog=) param to Prompts(client=) - Update APP_ENDPOINT to us.insights.hanzo.ai - Update all tests, examples, docs, mypy config
2647 lines
93 KiB
Python
2647 lines
93 KiB
Python
import asyncio
|
|
import logging
|
|
import math
|
|
import os
|
|
import time
|
|
import uuid
|
|
from typing import List, Literal, Optional, TypedDict, Union
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
import pytest
|
|
|
|
try:
|
|
from langchain_anthropic.chat_models import ChatAnthropic
|
|
from langchain_community.chat_models.fake import FakeMessagesListChatModel
|
|
from langchain_community.llms.fake import FakeListLLM, FakeStreamingListLLM
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
from langchain_core.runnables import RunnableLambda
|
|
from langchain_core.tools import tool
|
|
from langchain_openai.chat_models import ChatOpenAI
|
|
from langgraph.graph.state import END, START, StateGraph
|
|
from langgraph.prebuilt import create_react_agent
|
|
|
|
from hanzo_insights.ai.langchain import CallbackHandler
|
|
from hanzo_insights.ai.langchain.callbacks import GenerationMetadata, SpanMetadata
|
|
|
|
LANGCHAIN_AVAILABLE = True
|
|
except ImportError:
|
|
|
|
class FakeListLLM:
|
|
pass
|
|
|
|
class FakeStreamingListLLM:
|
|
pass
|
|
|
|
class HumanMessage:
|
|
pass
|
|
|
|
class AIMessage:
|
|
pass
|
|
|
|
LANGCHAIN_AVAILABLE = False
|
|
|
|
|
|
# Skip all tests if LangChain is not available
|
|
pytestmark = pytest.mark.skipif(
|
|
not LANGCHAIN_AVAILABLE, reason="LangChain package is not available"
|
|
)
|
|
|
|
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
|
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def mock_client():
|
|
with patch("hanzo_insights.client.Client") as mock_client:
|
|
mock_client.privacy_mode = False
|
|
logging.getLogger("hanzo_insights").setLevel(logging.DEBUG)
|
|
yield mock_client
|
|
|
|
|
|
def test_parent_capture(mock_client):
|
|
callbacks = CallbackHandler(mock_client)
|
|
parent_run_id = uuid.uuid4()
|
|
run_id = uuid.uuid4()
|
|
callbacks._set_parent_of_run(run_id, parent_run_id)
|
|
assert callbacks._parent_tree == {run_id: parent_run_id}
|
|
callbacks._pop_parent_of_run(run_id)
|
|
assert callbacks._parent_tree == {}
|
|
callbacks._pop_parent_of_run(parent_run_id) # should not raise
|
|
|
|
|
|
def test_find_root_run(mock_client):
|
|
callbacks = CallbackHandler(mock_client)
|
|
root_run_id = uuid.uuid4()
|
|
parent_run_id = uuid.uuid4()
|
|
run_id = uuid.uuid4()
|
|
callbacks._set_parent_of_run(run_id, parent_run_id)
|
|
callbacks._set_parent_of_run(parent_run_id, root_run_id)
|
|
assert callbacks._find_root_run(run_id) == root_run_id
|
|
new_run_id = uuid.uuid4()
|
|
assert callbacks._find_root_run(new_run_id) == new_run_id
|
|
|
|
|
|
def test_trace_id_generation(mock_client):
|
|
callbacks = CallbackHandler(mock_client)
|
|
run_id = uuid.uuid4()
|
|
with patch("uuid.uuid4", return_value=run_id):
|
|
assert callbacks._get_trace_id(run_id) == run_id
|
|
run_id = uuid.uuid4()
|
|
callbacks = CallbackHandler(mock_client, trace_id=run_id)
|
|
assert callbacks._get_trace_id(uuid.uuid4()) == run_id
|
|
|
|
|
|
def test_metadata_capture(mock_client):
|
|
callbacks = CallbackHandler(mock_client)
|
|
run_id = uuid.uuid4()
|
|
with patch("time.time", return_value=1234567890):
|
|
callbacks._set_llm_metadata(
|
|
{"kwargs": {"openai_api_base": "https://us.posthog.com"}},
|
|
run_id,
|
|
messages=[{"role": "user", "content": "Who won the world series in 2020?"}],
|
|
invocation_params={"temperature": 0.5},
|
|
metadata={"ls_model_name": "hog-mini", "ls_provider": "hanzo_insights"},
|
|
name="test",
|
|
)
|
|
expected = GenerationMetadata(
|
|
model="hog-mini",
|
|
input=[{"role": "user", "content": "Who won the world series in 2020?"}],
|
|
start_time=1234567890,
|
|
model_params={"temperature": 0.5},
|
|
provider="hanzo_insights",
|
|
base_url="https://us.posthog.com",
|
|
name="test",
|
|
end_time=None,
|
|
insights_properties=None,
|
|
)
|
|
assert callbacks._runs[run_id] == expected
|
|
with patch("time.time", return_value=1234567891):
|
|
run = callbacks._pop_run_metadata(run_id)
|
|
expected.end_time = 1234567891
|
|
assert run == expected
|
|
assert callbacks._runs == {}
|
|
callbacks._pop_run_metadata(uuid.uuid4()) # should not raise
|
|
|
|
|
|
def test_run_metadata_capture(mock_client):
|
|
callbacks = CallbackHandler(mock_client)
|
|
run_id = uuid.uuid4()
|
|
with patch("time.time", return_value=1234567890):
|
|
callbacks._set_trace_or_span_metadata(None, 1, run_id)
|
|
expected = SpanMetadata(
|
|
name="trace",
|
|
input=1,
|
|
start_time=1234567890,
|
|
end_time=None,
|
|
)
|
|
assert callbacks._runs[run_id] == expected
|
|
with patch("time.time", return_value=1234567890):
|
|
callbacks._set_trace_or_span_metadata(None, 1, run_id, uuid.uuid4())
|
|
expected = SpanMetadata(
|
|
name="span",
|
|
input=1,
|
|
start_time=1234567890,
|
|
end_time=None,
|
|
)
|
|
assert callbacks._runs[run_id] == expected
|
|
|
|
with patch("time.time", return_value=1234567890):
|
|
callbacks._set_trace_or_span_metadata({"name": "test"}, 1, run_id)
|
|
expected = SpanMetadata(
|
|
name="test",
|
|
input=1,
|
|
start_time=1234567890,
|
|
end_time=None,
|
|
)
|
|
assert callbacks._runs[run_id] == expected
|
|
|
|
|
|
@pytest.mark.parametrize("stream", [True, False])
|
|
def test_basic_chat_chain(mock_client, stream):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", "You are a helpful assistant."),
|
|
("user", "Who won the world series in 2020?"),
|
|
]
|
|
)
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="The Los Angeles Dodgers won the World Series in 2020.",
|
|
usage_metadata={
|
|
"input_tokens": 10,
|
|
"output_tokens": 10,
|
|
"total_tokens": 20,
|
|
},
|
|
)
|
|
]
|
|
)
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
if stream:
|
|
result = [m for m in chain.stream({}, config={"callbacks": callbacks})][0]
|
|
else:
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "The Los Angeles Dodgers won the World Series in 2020."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
span_args = mock_client.capture.call_args_list[0][1]
|
|
span_props = span_args["properties"]
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
trace_args = mock_client.capture.call_args_list[2][1]
|
|
trace_props = trace_args["properties"]
|
|
|
|
# Span is first
|
|
assert span_args["event"] == "$ai_span"
|
|
assert span_props["$ai_trace_id"] == generation_props["$ai_trace_id"]
|
|
assert span_props["$ai_parent_id"] == trace_props["$ai_trace_id"]
|
|
assert "$ai_span_id" in span_props
|
|
|
|
# Generation is second
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert "distinct_id" in generation_args
|
|
assert generation_props["$ai_framework"] == "langchain"
|
|
assert "$ai_model" in generation_props
|
|
assert "$ai_provider" in generation_props
|
|
assert generation_props["$ai_input"] == [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "Who won the world series in 2020?"},
|
|
]
|
|
assert generation_props["$ai_output_choices"] == [
|
|
{
|
|
"role": "assistant",
|
|
"content": "The Los Angeles Dodgers won the World Series in 2020.",
|
|
}
|
|
]
|
|
assert generation_props["$ai_input_tokens"] == 10
|
|
assert generation_props["$ai_output_tokens"] == 10
|
|
assert generation_props["$ai_http_status"] == 200
|
|
assert isinstance(generation_props["$ai_latency"], float)
|
|
assert "$ai_span_id" in generation_props
|
|
assert generation_props["$ai_parent_id"] == trace_props["$ai_trace_id"]
|
|
assert generation_props["$ai_trace_id"] == trace_props["$ai_trace_id"]
|
|
assert generation_props["$ai_span_name"] == "FakeMessagesListChatModel"
|
|
|
|
# Trace is last
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert "$ai_trace_id" in trace_props
|
|
assert "$ai_parent_id" not in trace_props
|
|
|
|
|
|
@pytest.mark.parametrize("stream", [True, False])
|
|
async def test_async_basic_chat_chain(mock_client, stream):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", "You are a helpful assistant."),
|
|
("user", "Who won the world series in 2020?"),
|
|
]
|
|
)
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="The Los Angeles Dodgers won the World Series in 2020.",
|
|
usage_metadata={
|
|
"input_tokens": 10,
|
|
"output_tokens": 10,
|
|
"total_tokens": 20,
|
|
},
|
|
)
|
|
]
|
|
)
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
if stream:
|
|
result = [m async for m in chain.astream({}, config={"callbacks": callbacks})][
|
|
0
|
|
]
|
|
else:
|
|
result = await chain.ainvoke({}, config={"callbacks": callbacks})
|
|
assert result.content == "The Los Angeles Dodgers won the World Series in 2020."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
span_args = mock_client.capture.call_args_list[0][1]
|
|
span_props = span_args["properties"]
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
trace_args = mock_client.capture.call_args_list[2][1]
|
|
trace_props = trace_args["properties"]
|
|
|
|
# Span is first
|
|
assert span_args["event"] == "$ai_span"
|
|
assert span_props["$ai_trace_id"] == generation_props["$ai_trace_id"]
|
|
assert span_props["$ai_parent_id"] == trace_props["$ai_trace_id"]
|
|
assert "$ai_span_id" in span_props
|
|
|
|
# Generation is second
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert "distinct_id" in generation_args
|
|
assert "$ai_model" in generation_props
|
|
assert "$ai_provider" in generation_props
|
|
assert generation_props["$ai_input"] == [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "Who won the world series in 2020?"},
|
|
]
|
|
assert generation_props["$ai_output_choices"] == [
|
|
{
|
|
"role": "assistant",
|
|
"content": "The Los Angeles Dodgers won the World Series in 2020.",
|
|
}
|
|
]
|
|
assert generation_props["$ai_input_tokens"] == 10
|
|
assert generation_props["$ai_output_tokens"] == 10
|
|
assert generation_props["$ai_http_status"] == 200
|
|
assert isinstance(generation_props["$ai_latency"], float)
|
|
assert "$ai_span_id" in generation_props
|
|
assert generation_props["$ai_parent_id"] == trace_props["$ai_trace_id"]
|
|
assert generation_props["$ai_trace_id"] == trace_props["$ai_trace_id"]
|
|
|
|
# Trace is last
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert "distinct_id" in generation_args
|
|
assert trace_props["$ai_trace_id"] == generation_props["$ai_trace_id"]
|
|
assert "$ai_parent_id" not in trace_props
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"Model,stream",
|
|
[
|
|
(FakeListLLM, True),
|
|
(FakeListLLM, False),
|
|
(FakeStreamingListLLM, True),
|
|
(FakeStreamingListLLM, False),
|
|
],
|
|
)
|
|
def test_basic_llm_chain(mock_client, Model, stream):
|
|
model = Model(responses=["The Los Angeles Dodgers won the World Series in 2020."])
|
|
callbacks: List[CallbackHandler] = [CallbackHandler(mock_client)]
|
|
|
|
if stream:
|
|
result = "".join(
|
|
[
|
|
m
|
|
for m in model.stream(
|
|
"Who won the world series in 2020?", config={"callbacks": callbacks}
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
result = model.invoke(
|
|
"Who won the world series in 2020?", config={"callbacks": callbacks}
|
|
)
|
|
assert result == "The Los Angeles Dodgers won the World Series in 2020."
|
|
|
|
assert mock_client.capture.call_count == 1
|
|
args = mock_client.capture.call_args_list[0][1]
|
|
props = args["properties"]
|
|
|
|
assert args["event"] == "$ai_generation"
|
|
assert "distinct_id" in args
|
|
assert "$ai_model" in props
|
|
assert "$ai_provider" in props
|
|
assert props["$ai_input"] == ["Who won the world series in 2020?"]
|
|
assert props["$ai_output_choices"] == [
|
|
"The Los Angeles Dodgers won the World Series in 2020."
|
|
]
|
|
assert props["$ai_http_status"] == 200
|
|
assert props["$ai_trace_id"] is not None
|
|
assert isinstance(props["$ai_latency"], float)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"Model,stream",
|
|
[
|
|
(FakeListLLM, True),
|
|
(FakeListLLM, False),
|
|
(FakeStreamingListLLM, True),
|
|
(FakeStreamingListLLM, False),
|
|
],
|
|
)
|
|
async def test_async_basic_llm_chain(mock_client, Model, stream):
|
|
model = Model(responses=["The Los Angeles Dodgers won the World Series in 2020."])
|
|
callbacks: List[CallbackHandler] = [CallbackHandler(mock_client)]
|
|
|
|
if stream:
|
|
result = "".join(
|
|
[
|
|
m
|
|
async for m in model.astream(
|
|
"Who won the world series in 2020?", config={"callbacks": callbacks}
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
result = await model.ainvoke(
|
|
"Who won the world series in 2020?", config={"callbacks": callbacks}
|
|
)
|
|
assert result == "The Los Angeles Dodgers won the World Series in 2020."
|
|
|
|
assert mock_client.capture.call_count == 1
|
|
args = mock_client.capture.call_args_list[0][1]
|
|
props = args["properties"]
|
|
|
|
assert args["event"] == "$ai_generation"
|
|
assert "distinct_id" in args
|
|
assert "$ai_model" in props
|
|
assert "$ai_provider" in props
|
|
assert props["$ai_input"] == ["Who won the world series in 2020?"]
|
|
assert props["$ai_output_choices"] == [
|
|
"The Los Angeles Dodgers won the World Series in 2020."
|
|
]
|
|
assert props["$ai_http_status"] == 200
|
|
assert props["$ai_trace_id"] is not None
|
|
assert isinstance(props["$ai_latency"], float)
|
|
|
|
|
|
def test_trace_id_and_inputs_for_multiple_chains(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("user", "Foo {var}"),
|
|
]
|
|
)
|
|
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model | RunnableLambda(lambda x: [x]) | model
|
|
result = chain.invoke({"var": "bar"}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Bar"
|
|
# span, generation, span, generation, trace
|
|
assert mock_client.capture.call_count == 5
|
|
|
|
first_span_args = mock_client.capture.call_args_list[0][1]
|
|
first_span_props = first_span_args["properties"]
|
|
|
|
first_generation_args = mock_client.capture.call_args_list[1][1]
|
|
first_generation_props = first_generation_args["properties"]
|
|
|
|
second_span_args = mock_client.capture.call_args_list[2][1]
|
|
second_span_props = second_span_args["properties"]
|
|
|
|
second_generation_args = mock_client.capture.call_args_list[3][1]
|
|
second_generation_props = second_generation_args["properties"]
|
|
|
|
trace_args = mock_client.capture.call_args_list[4][1]
|
|
trace_props = trace_args["properties"]
|
|
|
|
# Prompt span
|
|
assert first_span_args["event"] == "$ai_span"
|
|
assert first_span_props["$ai_input_state"] == {"var": "bar"}
|
|
assert first_span_props["$ai_trace_id"] == trace_props["$ai_trace_id"]
|
|
assert first_span_props["$ai_parent_id"] == trace_props["$ai_trace_id"]
|
|
assert "$ai_span_id" in first_span_props
|
|
assert first_span_props["$ai_output_state"] == ChatPromptTemplate(
|
|
messages=[HumanMessage(content="Foo bar")]
|
|
).invoke({})
|
|
|
|
# first model
|
|
assert first_generation_args["event"] == "$ai_generation"
|
|
assert "distinct_id" in first_generation_args
|
|
assert "$ai_model" in first_generation_props
|
|
assert "$ai_provider" in first_generation_props
|
|
assert first_generation_props["$ai_input"] == [
|
|
{"role": "user", "content": "Foo bar"}
|
|
]
|
|
assert first_generation_props["$ai_output_choices"] == [
|
|
{"role": "assistant", "content": "Bar"}
|
|
]
|
|
assert first_generation_props["$ai_http_status"] == 200
|
|
assert isinstance(first_generation_props["$ai_latency"], float)
|
|
assert "$ai_span_id" in first_generation_props
|
|
assert first_generation_props["$ai_parent_id"] == trace_props["$ai_trace_id"]
|
|
assert first_generation_props["$ai_trace_id"] == trace_props["$ai_trace_id"]
|
|
|
|
# lambda span
|
|
assert second_span_args["event"] == "$ai_span"
|
|
assert second_span_props["$ai_input_state"].content == "Bar"
|
|
assert second_span_props["$ai_trace_id"] == trace_props["$ai_trace_id"]
|
|
assert second_span_props["$ai_parent_id"] == trace_props["$ai_trace_id"]
|
|
assert "$ai_span_id" in second_span_props
|
|
assert second_span_props["$ai_output_state"][0].content == "Bar"
|
|
|
|
# second model
|
|
assert second_generation_args["event"] == "$ai_generation"
|
|
assert "distinct_id" in second_generation_args
|
|
assert "$ai_model" in second_generation_props
|
|
assert "$ai_provider" in second_generation_props
|
|
assert second_generation_props["$ai_input"] == [
|
|
{"role": "assistant", "content": "Bar"}
|
|
]
|
|
assert second_generation_props["$ai_output_choices"] == [
|
|
{"role": "assistant", "content": "Bar"}
|
|
]
|
|
assert second_generation_props["$ai_http_status"] == 200
|
|
assert second_generation_props["$ai_trace_id"] is not None
|
|
assert isinstance(second_generation_props["$ai_latency"], float)
|
|
|
|
# trace
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert "distinct_id" in trace_args
|
|
assert trace_props["$ai_input_state"] == {"var": "bar"}
|
|
assert isinstance(trace_props["$ai_output_state"], AIMessage)
|
|
assert trace_props["$ai_output_state"].content == "Bar"
|
|
assert trace_props["$ai_trace_id"] is not None
|
|
assert trace_props["$ai_span_name"] == "RunnableSequence"
|
|
|
|
|
|
def test_personless_mode(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
|
chain = prompt | FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
|
chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client)]})
|
|
assert mock_client.capture.call_count == 3
|
|
span_args = mock_client.capture.call_args_list[0][1]
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
trace_args = mock_client.capture.call_args_list[2][1]
|
|
|
|
# span
|
|
assert span_args["event"] == "$ai_span"
|
|
assert span_args["properties"]["$process_person_profile"] is False
|
|
# generation
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert generation_args["properties"]["$process_person_profile"] is False
|
|
# trace
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert trace_args["properties"]["$process_person_profile"] is False
|
|
|
|
id = uuid.uuid4()
|
|
chain.invoke(
|
|
{}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]}
|
|
)
|
|
assert mock_client.capture.call_count == 6
|
|
span_args = mock_client.capture.call_args_list[3][1]
|
|
generation_args = mock_client.capture.call_args_list[4][1]
|
|
trace_args = mock_client.capture.call_args_list[5][1]
|
|
|
|
# span
|
|
assert "$process_person_profile" not in span_args["properties"]
|
|
assert span_args["distinct_id"] == id
|
|
# generation
|
|
assert "$process_person_profile" not in generation_args["properties"]
|
|
assert generation_args["distinct_id"] == id
|
|
# trace
|
|
assert "$process_person_profile" not in trace_args["properties"]
|
|
assert trace_args["distinct_id"] == id
|
|
|
|
|
|
def test_personless_mode_exception(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
|
chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
|
|
callbacks = CallbackHandler(mock_client)
|
|
with pytest.raises(Exception):
|
|
chain.invoke({}, config={"callbacks": [callbacks]})
|
|
assert mock_client.capture.call_count == 3
|
|
span_args = mock_client.capture.call_args_list[0][1]
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
trace_args = mock_client.capture.call_args_list[2][1]
|
|
|
|
# span
|
|
assert span_args["event"] == "$ai_span"
|
|
assert span_args["properties"]["$process_person_profile"] is False
|
|
# generation
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert generation_args["properties"]["$process_person_profile"] is False
|
|
# trace
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert trace_args["properties"]["$process_person_profile"] is False
|
|
|
|
id = uuid.uuid4()
|
|
with pytest.raises(Exception):
|
|
chain.invoke(
|
|
{}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]}
|
|
)
|
|
assert mock_client.capture.call_count == 6
|
|
span_args = mock_client.capture.call_args_list[3][1]
|
|
generation_args = mock_client.capture.call_args_list[4][1]
|
|
trace_args = mock_client.capture.call_args_list[5][1]
|
|
|
|
# span
|
|
assert span_args["event"] == "$ai_span"
|
|
assert "$process_person_profile" not in span_args["properties"]
|
|
assert span_args["distinct_id"] == id
|
|
|
|
# generation
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert "$process_person_profile" not in generation_args["properties"]
|
|
assert generation_args["distinct_id"] == id
|
|
|
|
# trace
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert "$process_person_profile" not in trace_args["properties"]
|
|
assert trace_args["distinct_id"] == id
|
|
|
|
|
|
def test_metadata(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("user", "Foo"),
|
|
]
|
|
)
|
|
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
|
callbacks = [
|
|
CallbackHandler(
|
|
mock_client,
|
|
trace_id="test-trace-id",
|
|
distinct_id="test_id",
|
|
properties={"foo": "bar"},
|
|
)
|
|
]
|
|
chain = prompt | model
|
|
result = chain.invoke({"plan": None}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Bar"
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
span_call_args = mock_client.capture.call_args_list[0][1]
|
|
span_call_props = span_call_args["properties"]
|
|
assert span_call_args["distinct_id"] == "test_id"
|
|
assert span_call_args["event"] == "$ai_span"
|
|
assert span_call_props["$ai_trace_id"] == "test-trace-id"
|
|
assert span_call_props["foo"] == "bar"
|
|
assert "$ai_parent_id" in span_call_props
|
|
assert "$ai_span_id" in span_call_props
|
|
|
|
generation_call_args = mock_client.capture.call_args_list[1][1]
|
|
generation_call_props = generation_call_args["properties"]
|
|
assert generation_call_args["distinct_id"] == "test_id"
|
|
assert generation_call_args["event"] == "$ai_generation"
|
|
assert generation_call_props["$ai_trace_id"] == "test-trace-id"
|
|
assert generation_call_props["foo"] == "bar"
|
|
assert generation_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}]
|
|
assert generation_call_props["$ai_output_choices"] == [
|
|
{"role": "assistant", "content": "Bar"}
|
|
]
|
|
assert generation_call_props["$ai_http_status"] == 200
|
|
assert isinstance(generation_call_props["$ai_latency"], float)
|
|
|
|
trace_call_args = mock_client.capture.call_args_list[2][1]
|
|
trace_call_props = trace_call_args["properties"]
|
|
assert trace_call_args["distinct_id"] == "test_id"
|
|
assert trace_call_args["event"] == "$ai_trace"
|
|
assert trace_call_props["$ai_trace_id"] == "test-trace-id"
|
|
assert trace_call_props["$ai_span_name"] == "RunnableSequence"
|
|
assert trace_call_props["foo"] == "bar"
|
|
assert trace_call_props["$ai_input_state"] == {"plan": None}
|
|
assert isinstance(trace_call_props["$ai_output_state"], AIMessage)
|
|
assert trace_call_props["$ai_output_state"].content == "Bar"
|
|
|
|
|
|
class FakeGraphState(TypedDict):
|
|
messages: List[Union[HumanMessage, AIMessage]]
|
|
xyz: Optional[str]
|
|
|
|
|
|
def test_graph_state(mock_client):
|
|
config = {"callbacks": [CallbackHandler(mock_client)]}
|
|
|
|
graph = StateGraph(FakeGraphState)
|
|
graph.add_node(
|
|
"fake_plain",
|
|
lambda state: {
|
|
"messages": [
|
|
*state["messages"],
|
|
AIMessage(content="Let's explore bar."),
|
|
],
|
|
"xyz": "abc",
|
|
},
|
|
)
|
|
intermediate_chain = ChatPromptTemplate.from_messages(
|
|
[("user", "Question: What's a bar?")]
|
|
) | FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(content="It's a type of greeble."),
|
|
]
|
|
)
|
|
graph.add_node(
|
|
"fake_llm",
|
|
lambda state: {
|
|
"messages": [
|
|
*state["messages"],
|
|
intermediate_chain.invoke(state),
|
|
],
|
|
"xyz": state["xyz"],
|
|
},
|
|
)
|
|
graph.add_edge(START, "fake_plain")
|
|
graph.add_edge("fake_plain", "fake_llm")
|
|
graph.add_edge("fake_llm", END)
|
|
|
|
initial_state = {"messages": [HumanMessage(content="What's a bar?")], "xyz": None}
|
|
result = graph.compile().invoke(initial_state, config=config)
|
|
|
|
assert len(result["messages"]) == 3
|
|
assert isinstance(result["messages"][0], HumanMessage)
|
|
assert result["messages"][0].content == "What's a bar?"
|
|
assert isinstance(result["messages"][1], AIMessage)
|
|
assert result["messages"][1].content == "Let's explore bar."
|
|
assert isinstance(result["messages"][2], AIMessage)
|
|
assert result["messages"][2].content == "It's a type of greeble."
|
|
|
|
assert mock_client.capture.call_count == 6
|
|
calls = [call[1] for call in mock_client.capture.call_args_list]
|
|
|
|
# The trace event is captured at the end
|
|
trace_args = calls[-1]
|
|
trace_props = calls[-1]["properties"]
|
|
|
|
# Events are captured in the reverse order.
|
|
# Check all trace_ids
|
|
for call in calls:
|
|
assert call["properties"]["$ai_trace_id"] == trace_props["$ai_trace_id"]
|
|
|
|
# 1. Span, finish initialization
|
|
second_state = {
|
|
"messages": [
|
|
HumanMessage(content="What's a bar?"),
|
|
AIMessage(content="Let's explore bar."),
|
|
],
|
|
"xyz": "abc",
|
|
}
|
|
|
|
# 1. Span - the fake_plain node, which doesn't do anything
|
|
assert calls[0]["event"] == "$ai_span"
|
|
assert calls[0]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
|
|
assert "$ai_span_id" in calls[0]["properties"]
|
|
assert calls[0]["properties"]["$ai_span_name"] == "fake_plain"
|
|
assert calls[0]["properties"]["$ai_input_state"] == initial_state
|
|
assert calls[0]["properties"]["$ai_output_state"] == second_state
|
|
|
|
# 2. Span - the ChatPromptTemplate within fake_llm's FakeMessagesListChatModel
|
|
assert calls[1]["event"] == "$ai_span"
|
|
assert (
|
|
calls[1]["properties"]["$ai_parent_id"] == calls[3]["properties"]["$ai_span_id"]
|
|
)
|
|
assert "$ai_span_id" in calls[1]["properties"]
|
|
assert calls[1]["properties"]["$ai_span_name"] == "ChatPromptTemplate"
|
|
|
|
# 3. Generation - the FakeMessagesListChatModel within fake_llm's RunnableSequence
|
|
assert calls[2]["event"] == "$ai_generation"
|
|
assert (
|
|
calls[2]["properties"]["$ai_parent_id"] == calls[3]["properties"]["$ai_span_id"]
|
|
)
|
|
assert "$ai_span_id" in calls[2]["properties"]
|
|
assert calls[2]["properties"]["$ai_span_name"] == "FakeMessagesListChatModel"
|
|
|
|
# 4. Span - RunnableSequence within fake_llm
|
|
assert calls[3]["event"] == "$ai_span"
|
|
assert (
|
|
calls[3]["properties"]["$ai_parent_id"] == calls[4]["properties"]["$ai_span_id"]
|
|
)
|
|
assert "$ai_span_id" in calls[3]["properties"]
|
|
assert calls[3]["properties"]["$ai_span_name"] == "RunnableSequence"
|
|
|
|
# 5. Span - the fake_llm node
|
|
assert calls[4]["event"] == "$ai_span"
|
|
assert calls[4]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
|
|
assert "$ai_span_id" in calls[4]["properties"]
|
|
assert calls[4]["properties"]["$ai_span_name"] == "fake_llm"
|
|
|
|
# 6. Trace
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert trace_props["$ai_span_name"] == "LangGraph"
|
|
|
|
assert len(trace_props["$ai_input_state"]["messages"]) == 1
|
|
assert isinstance(trace_props["$ai_input_state"]["messages"][0], HumanMessage)
|
|
assert trace_props["$ai_input_state"]["messages"][0].content == "What's a bar?"
|
|
assert trace_props["$ai_input_state"]["messages"][0].type == "human"
|
|
assert trace_props["$ai_input_state"]["xyz"] is None
|
|
assert len(trace_props["$ai_output_state"]["messages"]) == 3
|
|
|
|
assert isinstance(trace_props["$ai_output_state"]["messages"][0], HumanMessage)
|
|
assert trace_props["$ai_output_state"]["messages"][0].content == "What's a bar?"
|
|
assert isinstance(trace_props["$ai_output_state"]["messages"][1], AIMessage)
|
|
assert (
|
|
trace_props["$ai_output_state"]["messages"][1].content == "Let's explore bar."
|
|
)
|
|
assert isinstance(trace_props["$ai_output_state"]["messages"][2], AIMessage)
|
|
assert (
|
|
trace_props["$ai_output_state"]["messages"][2].content
|
|
== "It's a type of greeble."
|
|
)
|
|
assert trace_args["properties"]["$ai_output_state"]["xyz"] == "abc"
|
|
|
|
|
|
def test_callbacks_logic(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
|
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
|
callbacks = CallbackHandler(
|
|
mock_client,
|
|
trace_id="test-trace-id",
|
|
distinct_id="test_id",
|
|
properties={"foo": "bar"},
|
|
)
|
|
chain = prompt | model
|
|
|
|
chain.invoke({}, config={"callbacks": [callbacks]})
|
|
assert callbacks._runs == {}
|
|
assert callbacks._parent_tree == {}
|
|
|
|
def assert_intermediary_run(m):
|
|
assert len(callbacks._runs) != 0
|
|
run = next(iter(callbacks._runs.values()))
|
|
assert run.name == "RunnableSequence"
|
|
assert run.input == {}
|
|
assert run.start_time is not None
|
|
assert len(callbacks._parent_tree.items()) == 1
|
|
return [m]
|
|
|
|
(chain | RunnableLambda(assert_intermediary_run) | model).invoke(
|
|
{}, config={"callbacks": [callbacks]}
|
|
)
|
|
assert callbacks._runs == {}
|
|
assert callbacks._parent_tree == {}
|
|
|
|
|
|
def test_exception_in_chain(mock_client):
|
|
def runnable(_):
|
|
raise ValueError("test")
|
|
|
|
callbacks = CallbackHandler(mock_client)
|
|
with pytest.raises(ValueError):
|
|
RunnableLambda(runnable).invoke({}, config={"callbacks": [callbacks]})
|
|
|
|
assert callbacks._runs == {}
|
|
assert callbacks._parent_tree == {}
|
|
assert mock_client.capture.call_count == 1
|
|
trace_call_args = mock_client.capture.call_args_list[0][1]
|
|
assert trace_call_args["event"] == "$ai_trace"
|
|
assert trace_call_args["properties"]["$ai_span_name"] == "runnable"
|
|
|
|
|
|
def test_openai_error(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
|
chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
|
|
callbacks = CallbackHandler(mock_client)
|
|
|
|
# 401
|
|
with pytest.raises(Exception):
|
|
chain.invoke({}, config={"callbacks": [callbacks]})
|
|
|
|
assert callbacks._runs == {}
|
|
assert callbacks._parent_tree == {}
|
|
assert mock_client.capture.call_count == 3
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
props = generation_args["properties"]
|
|
assert props["$ai_http_status"] == 401
|
|
assert props["$ai_input"] == [{"role": "user", "content": "Foo"}]
|
|
assert "$ai_output_choices" not in props
|
|
|
|
|
|
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
|
|
def test_openai_chain(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", 'You must always answer with "Bar".'),
|
|
("user", "Foo"),
|
|
]
|
|
)
|
|
chain = prompt | ChatOpenAI(
|
|
api_key=OPENAI_API_KEY,
|
|
model="gpt-4o-mini",
|
|
temperature=0,
|
|
max_tokens=1,
|
|
)
|
|
callbacks = CallbackHandler(
|
|
mock_client,
|
|
trace_id="test-trace-id",
|
|
distinct_id="test_id",
|
|
properties={"foo": "bar"},
|
|
)
|
|
start_time = time.time()
|
|
result = chain.invoke({}, config={"callbacks": [callbacks]})
|
|
approximate_latency = math.floor(time.time() - start_time)
|
|
|
|
assert result.content == "Bar"
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
gen_args = mock_client.capture.call_args_list[1][1]
|
|
gen_props = gen_args["properties"]
|
|
assert gen_args["event"] == "$ai_generation"
|
|
assert gen_props["$ai_trace_id"] == "test-trace-id"
|
|
assert gen_props["$ai_provider"] == "openai"
|
|
assert gen_props["$ai_model"] == "gpt-4o-mini"
|
|
assert gen_props["foo"] == "bar"
|
|
|
|
# langchain-openai for langchain v3
|
|
if "max_completion_tokens" in gen_props["$ai_model_parameters"]:
|
|
assert gen_props["$ai_model_parameters"] == {
|
|
"temperature": 0.0,
|
|
"max_completion_tokens": 1,
|
|
"stream": False,
|
|
}
|
|
else:
|
|
assert gen_props["$ai_model_parameters"] == {
|
|
"temperature": 0.0,
|
|
"max_tokens": 1,
|
|
"n": 1,
|
|
"stream": False,
|
|
}
|
|
assert gen_props["$ai_input"] == [
|
|
{"role": "system", "content": 'You must always answer with "Bar".'},
|
|
{"role": "user", "content": "Foo"},
|
|
]
|
|
assert gen_props["$ai_output_choices"] == [
|
|
{"role": "assistant", "content": "Bar", "refusal": None}
|
|
]
|
|
assert gen_props["$ai_http_status"] == 200
|
|
assert isinstance(gen_props["$ai_latency"], float)
|
|
assert (
|
|
min(approximate_latency - 1, 0)
|
|
<= math.floor(gen_props["$ai_latency"])
|
|
<= approximate_latency
|
|
)
|
|
assert gen_props["$ai_input_tokens"] == 20
|
|
assert gen_props["$ai_output_tokens"] == 1
|
|
|
|
|
|
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
|
|
def test_openai_captures_multiple_generations(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", 'You must always answer with "Bar".'),
|
|
("user", "Foo"),
|
|
]
|
|
)
|
|
chain = prompt | ChatOpenAI(
|
|
api_key=OPENAI_API_KEY,
|
|
model="gpt-4o-mini",
|
|
temperature=0,
|
|
max_tokens=1,
|
|
n=2,
|
|
)
|
|
callbacks = CallbackHandler(mock_client)
|
|
result = chain.invoke({}, config={"callbacks": [callbacks]})
|
|
|
|
assert result.content == "Bar"
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
gen_args = mock_client.capture.call_args_list[1][1]
|
|
gen_props = gen_args["properties"]
|
|
trace_args = mock_client.capture.call_args_list[2][1]
|
|
trace_props = trace_args["properties"]
|
|
|
|
assert gen_args["event"] == "$ai_generation"
|
|
assert gen_props["$ai_input"] == [
|
|
{"role": "system", "content": 'You must always answer with "Bar".'},
|
|
{"role": "user", "content": "Foo"},
|
|
]
|
|
assert gen_props["$ai_output_choices"] == [
|
|
{"role": "assistant", "content": "Bar", "refusal": None},
|
|
{
|
|
"role": "assistant",
|
|
"content": "Bar",
|
|
},
|
|
]
|
|
|
|
# langchain-openai for langchain v3
|
|
if "max_completion_tokens" in gen_props["$ai_model_parameters"]:
|
|
assert gen_props["$ai_model_parameters"] == {
|
|
"temperature": 0.0,
|
|
"max_completion_tokens": 1,
|
|
"stream": False,
|
|
"n": 2,
|
|
}
|
|
else:
|
|
assert gen_props["$ai_model_parameters"] == {
|
|
"temperature": 0.0,
|
|
"max_tokens": 1,
|
|
"stream": False,
|
|
"n": 2,
|
|
}
|
|
assert gen_props["$ai_http_status"] == 200
|
|
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert trace_props["$ai_input_state"] == {}
|
|
assert isinstance(trace_props["$ai_output_state"], AIMessage)
|
|
|
|
|
|
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
|
|
def test_openai_streaming(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", 'You must always answer with "Bar".'),
|
|
("user", "Foo"),
|
|
]
|
|
)
|
|
chain = prompt | ChatOpenAI(
|
|
api_key=OPENAI_API_KEY,
|
|
model="gpt-4o-mini",
|
|
temperature=0,
|
|
max_tokens=1,
|
|
stream=True,
|
|
stream_usage=True,
|
|
)
|
|
callbacks = CallbackHandler(mock_client)
|
|
result = [m for m in chain.stream({}, config={"callbacks": [callbacks]})]
|
|
result = sum(result[1:], result[0])
|
|
|
|
assert result.content == "Bar"
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
gen_args = mock_client.capture.call_args_list[1][1]
|
|
gen_props = gen_args["properties"]
|
|
trace_args = mock_client.capture.call_args_list[2][1]
|
|
trace_props = trace_args["properties"]
|
|
|
|
assert gen_args["event"] == "$ai_generation"
|
|
assert gen_props["$ai_model_parameters"]["stream"]
|
|
assert gen_props["$ai_input"] == [
|
|
{"role": "system", "content": 'You must always answer with "Bar".'},
|
|
{"role": "user", "content": "Foo"},
|
|
]
|
|
assert gen_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
|
assert gen_props["$ai_http_status"] == 200
|
|
assert gen_props["$ai_input_tokens"] == 20
|
|
assert gen_props["$ai_output_tokens"] == 1
|
|
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert trace_props["$ai_input_state"] == {"input": ""}
|
|
assert isinstance(trace_props["$ai_output_state"], AIMessage)
|
|
|
|
|
|
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
|
|
async def test_async_openai_streaming(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", 'You must always answer with "Bar".'),
|
|
("user", "Foo"),
|
|
]
|
|
)
|
|
chain = prompt | ChatOpenAI(
|
|
api_key=OPENAI_API_KEY,
|
|
model="gpt-4o-mini",
|
|
temperature=0,
|
|
max_tokens=1,
|
|
stream=True,
|
|
stream_usage=True,
|
|
)
|
|
callbacks = CallbackHandler(mock_client)
|
|
result = [m async for m in chain.astream({}, config={"callbacks": [callbacks]})]
|
|
result = sum(result[1:], result[0])
|
|
|
|
assert result.content == "Bar"
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
gen_args = mock_client.capture.call_args_list[1][1]
|
|
gen_props = gen_args["properties"]
|
|
trace_args = mock_client.capture.call_args_list[2][1]
|
|
trace_props = trace_args["properties"]
|
|
|
|
assert gen_args["event"] == "$ai_generation"
|
|
assert gen_props["$ai_model_parameters"]["stream"]
|
|
assert gen_props["$ai_input"] == [
|
|
{"role": "system", "content": 'You must always answer with "Bar".'},
|
|
{"role": "user", "content": "Foo"},
|
|
]
|
|
assert gen_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
|
assert gen_props["$ai_http_status"] == 200
|
|
assert gen_props["$ai_input_tokens"] == 20
|
|
assert gen_props["$ai_output_tokens"] == 1
|
|
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert trace_props["$ai_input_state"] == {"input": ""}
|
|
assert isinstance(trace_props["$ai_output_state"], AIMessage)
|
|
|
|
|
|
def test_base_url_retrieval(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
|
chain = prompt | ChatOpenAI(
|
|
api_key="test",
|
|
model="insights-mini",
|
|
base_url="https://test.posthog.com",
|
|
)
|
|
callbacks = CallbackHandler(mock_client)
|
|
with pytest.raises(Exception):
|
|
chain.invoke({}, config={"callbacks": [callbacks]})
|
|
|
|
assert mock_client.capture.call_count == 3
|
|
generation_call = mock_client.capture.call_args_list[1][1]
|
|
assert generation_call["properties"]["$ai_base_url"] == "https://test.posthog.com"
|
|
|
|
|
|
def test_groups(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", 'You must always answer with "Bar".'),
|
|
("user", "Foo"),
|
|
]
|
|
)
|
|
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
|
chain = prompt | model
|
|
callbacks = CallbackHandler(mock_client, groups={"company": "test_company"})
|
|
chain.invoke({}, config={"callbacks": [callbacks]})
|
|
|
|
assert mock_client.capture.call_count == 3
|
|
generation_call = mock_client.capture.call_args_list[1][1]
|
|
assert generation_call["groups"] == {"company": "test_company"}
|
|
|
|
|
|
def test_privacy_mode_local(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", 'You must always answer with "Bar".'),
|
|
("user", "Foo"),
|
|
]
|
|
)
|
|
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
|
chain = prompt | model
|
|
callbacks = CallbackHandler(mock_client, privacy_mode=True)
|
|
chain.invoke({}, config={"callbacks": [callbacks]})
|
|
|
|
assert mock_client.capture.call_count == 3
|
|
generation_call = mock_client.capture.call_args_list[1][1]
|
|
assert generation_call["properties"]["$ai_input"] is None
|
|
assert generation_call["properties"]["$ai_output_choices"] is None
|
|
|
|
|
|
def test_privacy_mode_global(mock_client):
|
|
mock_client.privacy_mode = True
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", 'You must always answer with "Bar".'),
|
|
("user", "Foo"),
|
|
]
|
|
)
|
|
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
|
chain = prompt | model
|
|
callbacks = CallbackHandler(mock_client)
|
|
chain.invoke({}, config={"callbacks": [callbacks]})
|
|
|
|
assert mock_client.capture.call_count == 3
|
|
generation_call = mock_client.capture.call_args_list[1][1]
|
|
assert generation_call["properties"]["$ai_input"] is None
|
|
assert generation_call["properties"]["$ai_output_choices"] is None
|
|
|
|
|
|
@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
|
|
def test_anthropic_chain(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", 'You must always answer with "Bar".'),
|
|
("user", "Foo"),
|
|
]
|
|
)
|
|
chain = prompt | ChatAnthropic(
|
|
api_key=ANTHROPIC_API_KEY,
|
|
model="claude-sonnet-4-5-20250929",
|
|
temperature=0,
|
|
max_tokens=1024,
|
|
)
|
|
callbacks = CallbackHandler(
|
|
mock_client,
|
|
trace_id="test-trace-id",
|
|
distinct_id="test_id",
|
|
properties={"foo": "bar"},
|
|
)
|
|
start_time = time.time()
|
|
result = chain.invoke({}, config={"callbacks": [callbacks]})
|
|
approximate_latency = math.floor(time.time() - start_time)
|
|
|
|
assert result.content == "Bar"
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
gen_args = mock_client.capture.call_args_list[1][1]
|
|
gen_props = gen_args["properties"]
|
|
trace_args = mock_client.capture.call_args_list[2][1]
|
|
trace_props = trace_args["properties"]
|
|
|
|
assert gen_args["event"] == "$ai_generation"
|
|
assert gen_props["$ai_trace_id"] == "test-trace-id"
|
|
assert gen_props["$ai_provider"] == "anthropic"
|
|
assert gen_props["$ai_model"] == "claude-sonnet-4-5-20250929"
|
|
assert gen_props["foo"] == "bar"
|
|
|
|
assert gen_props["$ai_model_parameters"] == {
|
|
"temperature": 0.0,
|
|
"max_tokens": 1024,
|
|
"streaming": False,
|
|
}
|
|
assert gen_props["$ai_input"] == [
|
|
{"role": "system", "content": 'You must always answer with "Bar".'},
|
|
{"role": "user", "content": "Foo"},
|
|
]
|
|
assert gen_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
|
assert gen_props["$ai_http_status"] == 200
|
|
assert isinstance(gen_props["$ai_latency"], float)
|
|
assert (
|
|
min(approximate_latency - 1, 0)
|
|
<= math.floor(gen_props["$ai_latency"])
|
|
<= approximate_latency
|
|
)
|
|
assert gen_props["$ai_input_tokens"] == 17
|
|
assert gen_props["$ai_output_tokens"] == 4
|
|
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert trace_props["$ai_input_state"] == {}
|
|
assert isinstance(trace_props["$ai_output_state"], AIMessage)
|
|
|
|
|
|
@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
|
|
async def test_async_anthropic_streaming(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", 'You must always answer with "Bar".'),
|
|
("user", "Foo"),
|
|
]
|
|
)
|
|
chain = prompt | ChatAnthropic(
|
|
api_key=ANTHROPIC_API_KEY,
|
|
model="claude-sonnet-4-5-20250929",
|
|
temperature=0,
|
|
max_tokens=1024,
|
|
streaming=True,
|
|
stream_usage=True,
|
|
)
|
|
callbacks = CallbackHandler(mock_client)
|
|
result = [m async for m in chain.astream({}, config={"callbacks": [callbacks]})]
|
|
result = sum(result[1:], result[0])
|
|
|
|
assert result.content == "Bar"
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
gen_args = mock_client.capture.call_args_list[1][1]
|
|
gen_props = gen_args["properties"]
|
|
trace_args = mock_client.capture.call_args_list[2][1]
|
|
trace_props = trace_args["properties"]
|
|
|
|
assert gen_args["event"] == "$ai_generation"
|
|
assert gen_props["$ai_model_parameters"]["streaming"]
|
|
assert gen_props["$ai_input"] == [
|
|
{"role": "system", "content": 'You must always answer with "Bar".'},
|
|
{"role": "user", "content": "Foo"},
|
|
]
|
|
assert gen_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
|
assert gen_props["$ai_http_status"] == 200
|
|
assert gen_props["$ai_input_tokens"] == 17
|
|
assert gen_props["$ai_output_tokens"] is not None
|
|
|
|
assert trace_args["event"] == "$ai_trace"
|
|
assert trace_props["$ai_input_state"] == {
|
|
"input": "",
|
|
}
|
|
assert isinstance(trace_props["$ai_output_state"], AIMessage)
|
|
|
|
|
|
def test_metadata_tools(mock_client):
|
|
callbacks = CallbackHandler(mock_client)
|
|
run_id = uuid.uuid4()
|
|
tools = [
|
|
[
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "foo",
|
|
"description": "The foo.",
|
|
"parameters": {
|
|
"properties": {
|
|
"bar": {
|
|
"description": "The bar of foo.",
|
|
"type": "string",
|
|
},
|
|
},
|
|
"required": ["query_description", "query_kind"],
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
},
|
|
"strict": True,
|
|
},
|
|
}
|
|
]
|
|
]
|
|
|
|
with patch("time.time", return_value=1234567890):
|
|
callbacks._set_llm_metadata(
|
|
{"kwargs": {"openai_api_base": "https://us.posthog.com"}},
|
|
run_id,
|
|
messages=[{"role": "user", "content": "What's the weather like in SF?"}],
|
|
invocation_params={"temperature": 0.5, "tools": tools},
|
|
metadata={"ls_model_name": "hog-mini", "ls_provider": "hanzo_insights"},
|
|
name="test",
|
|
)
|
|
expected = GenerationMetadata(
|
|
model="hog-mini",
|
|
input=[{"role": "user", "content": "What's the weather like in SF?"}],
|
|
start_time=1234567890,
|
|
model_params={"temperature": 0.5},
|
|
provider="hanzo_insights",
|
|
base_url="https://us.posthog.com",
|
|
name="test",
|
|
tools=tools,
|
|
end_time=None,
|
|
insights_properties=None,
|
|
)
|
|
assert callbacks._runs[run_id] == expected
|
|
with patch("time.time", return_value=1234567891):
|
|
run = callbacks._pop_run_metadata(run_id)
|
|
expected.end_time = 1234567891
|
|
assert run == expected
|
|
assert callbacks._runs == {}
|
|
|
|
|
|
def test_tool_calls(mock_client):
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Bar",
|
|
additional_kwargs={
|
|
"tool_calls": [
|
|
{
|
|
"type": "function",
|
|
"id": "123",
|
|
"function": {
|
|
"name": "test",
|
|
"args": '{"a": 1}',
|
|
},
|
|
}
|
|
]
|
|
},
|
|
)
|
|
]
|
|
)
|
|
chain = prompt | model
|
|
callbacks = CallbackHandler(mock_client)
|
|
chain.invoke({}, config={"callbacks": [callbacks]})
|
|
|
|
assert mock_client.capture.call_count == 3
|
|
generation_call = mock_client.capture.call_args_list[1][1]
|
|
assert generation_call["properties"]["$ai_output_choices"][0]["tool_calls"] == [
|
|
{
|
|
"type": "function",
|
|
"id": "123",
|
|
"function": {
|
|
"name": "test",
|
|
"args": '{"a": 1}',
|
|
},
|
|
}
|
|
]
|
|
assert (
|
|
"additional_kwargs"
|
|
not in generation_call["properties"]["$ai_output_choices"][0]
|
|
)
|
|
|
|
|
|
async def test_async_traces(mock_client):
|
|
async def sleep(x): # -> Any:
|
|
await asyncio.sleep(0.1)
|
|
return x
|
|
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
|
chain1 = RunnableLambda(sleep)
|
|
chain2 = prompt | FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
|
|
|
cb = CallbackHandler(mock_client)
|
|
|
|
start_time = time.time()
|
|
await asyncio.gather(
|
|
chain1.ainvoke({}, config={"callbacks": [cb]}),
|
|
chain2.ainvoke({}, config={"callbacks": [cb]}),
|
|
)
|
|
approximate_latency = math.floor(time.time() - start_time)
|
|
assert mock_client.capture.call_count == 4
|
|
|
|
first_call, second_call, third_call, fourth_call = (
|
|
mock_client.capture.call_args_list
|
|
)
|
|
assert first_call[1]["event"] == "$ai_span"
|
|
assert second_call[1]["event"] == "$ai_generation"
|
|
assert third_call[1]["event"] == "$ai_trace"
|
|
assert third_call[1]["properties"]["$ai_span_name"] == "RunnableSequence"
|
|
assert fourth_call[1]["event"] == "$ai_trace"
|
|
assert fourth_call[1]["properties"]["$ai_span_name"] == "sleep"
|
|
assert (
|
|
min(approximate_latency - 1, 0)
|
|
<= math.floor(third_call[1]["properties"]["$ai_latency"])
|
|
<= approximate_latency
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OPENAI_API_KEY is not set")
|
|
def test_langgraph_agent(mock_client):
|
|
@tool
|
|
def get_weather(city: Literal["nyc", "sf"]):
|
|
"""
|
|
Use this to get weather information.
|
|
|
|
Args:
|
|
city: The city to get weather information for.
|
|
"""
|
|
if city == "sf":
|
|
return "It's always sunny in sf"
|
|
return "No info"
|
|
|
|
tools = [get_weather]
|
|
model = ChatOpenAI(api_key=OPENAI_API_KEY, model="gpt-4o-mini", temperature=0)
|
|
graph = create_react_agent(model, tools=tools)
|
|
inputs = {"messages": [("user", "what is the weather in sf")]}
|
|
cb = CallbackHandler(
|
|
mock_client, trace_id="test-trace-id", distinct_id="test-distinct-id"
|
|
)
|
|
graph.invoke(inputs, config={"callbacks": [cb]})
|
|
calls = [call[1] for call in mock_client.capture.call_args_list]
|
|
assert len(calls) == 15
|
|
for call in calls:
|
|
assert call["properties"]["$ai_trace_id"] == "test-trace-id"
|
|
assert len([call for call in calls if call["event"] == "$ai_generation"]) == 2
|
|
assert len([call for call in calls if call["event"] == "$ai_span"]) == 12
|
|
assert len([call for call in calls if call["event"] == "$ai_trace"]) == 1
|
|
|
|
|
|
@pytest.mark.parametrize("trace_id", ["test-trace-id", None])
|
|
def test_span_set_parent_ids(mock_client, trace_id):
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", "You are a helpful assistant."),
|
|
("user", "Who won the world series in 2020?"),
|
|
]
|
|
)
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(content="The Los Angeles Dodgers won the World Series in 2020.")
|
|
]
|
|
)
|
|
callbacks = [CallbackHandler(mock_client, trace_id=trace_id)]
|
|
chain = prompt | model
|
|
chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
span_props = mock_client.capture.call_args_list[0][1]
|
|
assert (
|
|
span_props["properties"]["$ai_trace_id"]
|
|
== span_props["properties"]["$ai_parent_id"]
|
|
)
|
|
|
|
generation_props = mock_client.capture.call_args_list[1][1]
|
|
assert (
|
|
generation_props["properties"]["$ai_trace_id"]
|
|
== generation_props["properties"]["$ai_parent_id"]
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("trace_id", ["test-trace-id", None])
|
|
def test_span_set_parent_ids_for_third_level_run(mock_client, trace_id):
|
|
def span_1(_):
|
|
def span_2(_):
|
|
def span_3(_):
|
|
return "span 3"
|
|
|
|
return RunnableLambda(span_3)
|
|
|
|
return RunnableLambda(span_2)
|
|
|
|
callbacks = [CallbackHandler(mock_client, trace_id=trace_id)]
|
|
chain = RunnableLambda(span_1)
|
|
chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
calls = mock_client.capture.call_args_list
|
|
span_props_2 = calls[0][1]["properties"]
|
|
span_props_1 = calls[1][1]["properties"]
|
|
trace_props = calls[2][1]["properties"]
|
|
|
|
assert span_props_2["$ai_parent_id"] == span_props_1["$ai_span_id"]
|
|
assert span_props_1["$ai_parent_id"] == trace_props["$ai_trace_id"]
|
|
|
|
|
|
def test_captures_error_with_details_in_span(mock_client):
|
|
def span(_):
|
|
raise ValueError("test")
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = RunnableLambda(span) | RunnableLambda(lambda _: "foo")
|
|
try:
|
|
chain.invoke({}, config={"callbacks": callbacks})
|
|
except ValueError:
|
|
pass
|
|
|
|
assert mock_client.capture.call_count == 2
|
|
assert (
|
|
mock_client.capture.call_args_list[1][1]["properties"]["$ai_error"]
|
|
== "ValueError: test"
|
|
)
|
|
assert mock_client.capture.call_args_list[1][1]["properties"]["$ai_is_error"]
|
|
|
|
|
|
def test_captures_error_without_details_in_span(mock_client):
|
|
def span(_):
|
|
raise ValueError
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = RunnableLambda(span) | RunnableLambda(lambda _: "foo")
|
|
try:
|
|
chain.invoke({}, config={"callbacks": callbacks})
|
|
except ValueError:
|
|
pass
|
|
|
|
assert mock_client.capture.call_count == 2
|
|
assert (
|
|
mock_client.capture.call_args_list[1][1]["properties"]["$ai_error"]
|
|
== "ValueError"
|
|
)
|
|
assert mock_client.capture.call_args_list[1][1]["properties"]["$ai_is_error"]
|
|
|
|
|
|
def test_openai_reasoning_tokens(mock_client):
|
|
"""Test that OpenAI reasoning tokens are captured correctly."""
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[("user", "Think step by step about this problem")]
|
|
)
|
|
|
|
# Mock response with reasoning tokens in output_token_details
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Let me think through this step by step...",
|
|
usage_metadata={
|
|
"input_tokens": 10,
|
|
"output_tokens": 25,
|
|
"total_tokens": 35,
|
|
"output_token_details": {"reasoning": 15}, # 15 reasoning tokens
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Let me think through this step by step..."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert generation_props["$ai_input_tokens"] == 10
|
|
assert generation_props["$ai_output_tokens"] == 25
|
|
assert generation_props["$ai_reasoning_tokens"] == 15
|
|
|
|
|
|
def test_anthropic_cache_write_and_read_tokens(mock_client):
|
|
"""Test that Anthropic cache creation and read tokens are captured correctly."""
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Analyze this large document")])
|
|
|
|
# First call with cache creation
|
|
model_write = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="I've analyzed the document and cached the context.",
|
|
usage_metadata={
|
|
"total_tokens": 1050,
|
|
"input_tokens": 1000,
|
|
"output_tokens": 50,
|
|
"cache_creation_input_tokens": 800, # Anthropic cache write
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model_write
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "I've analyzed the document and cached the context."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert generation_props["$ai_input_tokens"] == 1000
|
|
assert generation_props["$ai_output_tokens"] == 50
|
|
assert generation_props["$ai_cache_creation_input_tokens"] == 800
|
|
assert generation_props["$ai_cache_read_input_tokens"] == 0
|
|
assert generation_props["$ai_reasoning_tokens"] == 0
|
|
|
|
# Reset mock for second call
|
|
mock_client.reset_mock()
|
|
|
|
# Second call with cache read
|
|
model_read = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Using cached analysis to provide quick response.",
|
|
usage_metadata={
|
|
"input_tokens": 1200,
|
|
"output_tokens": 30,
|
|
"total_tokens": 1230,
|
|
"cache_read_input_tokens": 800, # Anthropic cache read
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
chain = prompt | model_read
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Using cached analysis to provide quick response."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert (
|
|
generation_props["$ai_input_tokens"] == 1200
|
|
) # No provider metadata, no subtraction
|
|
assert generation_props["$ai_output_tokens"] == 30
|
|
assert generation_props["$ai_cache_creation_input_tokens"] == 0
|
|
assert generation_props["$ai_cache_read_input_tokens"] == 800
|
|
assert generation_props["$ai_reasoning_tokens"] == 0
|
|
|
|
|
|
def test_anthropic_provider_subtracts_cache_tokens(mock_client):
|
|
"""Test that Anthropic provider correctly subtracts cache tokens from input tokens."""
|
|
from langchain_core.outputs import LLMResult, ChatGeneration
|
|
from langchain_core.messages import AIMessage
|
|
from uuid import uuid4
|
|
|
|
cb = CallbackHandler(mock_client)
|
|
run_id = uuid4()
|
|
|
|
# Set up with Anthropic provider
|
|
cb._set_llm_metadata(
|
|
serialized={},
|
|
run_id=run_id,
|
|
messages=[{"role": "user", "content": "test"}],
|
|
metadata={"ls_provider": "anthropic", "ls_model_name": "claude-3-sonnet"},
|
|
)
|
|
|
|
# Response with cache tokens: 1200 input (includes 800 cached)
|
|
response = LLMResult(
|
|
generations=[
|
|
[
|
|
ChatGeneration(
|
|
message=AIMessage(content="Response"),
|
|
generation_info={
|
|
"usage_metadata": {
|
|
"input_tokens": 1200,
|
|
"output_tokens": 50,
|
|
"cache_read_input_tokens": 800,
|
|
}
|
|
},
|
|
)
|
|
]
|
|
],
|
|
llm_output={},
|
|
)
|
|
|
|
cb._pop_run_and_capture_generation(run_id, None, response)
|
|
|
|
generation_args = mock_client.capture.call_args_list[0][1]
|
|
assert generation_args["properties"]["$ai_input_tokens"] == 400 # 1200 - 800
|
|
assert generation_args["properties"]["$ai_cache_read_input_tokens"] == 800
|
|
|
|
|
|
def test_anthropic_provider_subtracts_cache_write_tokens(mock_client):
|
|
"""Test that Anthropic provider correctly subtracts cache write tokens from input tokens."""
|
|
from langchain_core.outputs import LLMResult, ChatGeneration
|
|
from langchain_core.messages import AIMessage
|
|
from uuid import uuid4
|
|
|
|
cb = CallbackHandler(mock_client)
|
|
run_id = uuid4()
|
|
|
|
# Set up with Anthropic provider
|
|
cb._set_llm_metadata(
|
|
serialized={},
|
|
run_id=run_id,
|
|
messages=[{"role": "user", "content": "test"}],
|
|
metadata={"ls_provider": "anthropic", "ls_model_name": "claude-3-sonnet"},
|
|
)
|
|
|
|
# Response with cache creation: 1000 input (includes 800 being written to cache)
|
|
response = LLMResult(
|
|
generations=[
|
|
[
|
|
ChatGeneration(
|
|
message=AIMessage(content="Response"),
|
|
generation_info={
|
|
"usage_metadata": {
|
|
"input_tokens": 1000,
|
|
"output_tokens": 50,
|
|
"cache_creation_input_tokens": 800,
|
|
}
|
|
},
|
|
)
|
|
]
|
|
],
|
|
llm_output={},
|
|
)
|
|
|
|
cb._pop_run_and_capture_generation(run_id, None, response)
|
|
|
|
generation_args = mock_client.capture.call_args_list[0][1]
|
|
assert generation_args["properties"]["$ai_input_tokens"] == 200 # 1000 - 800
|
|
assert generation_args["properties"]["$ai_cache_creation_input_tokens"] == 800
|
|
|
|
|
|
def test_anthropic_provider_subtracts_both_cache_read_and_write_tokens(mock_client):
|
|
"""Test that Anthropic provider correctly subtracts both cache read and write tokens."""
|
|
from langchain_core.outputs import LLMResult, ChatGeneration
|
|
from langchain_core.messages import AIMessage
|
|
from uuid import uuid4
|
|
|
|
cb = CallbackHandler(mock_client)
|
|
run_id = uuid4()
|
|
|
|
# Set up with Anthropic provider
|
|
cb._set_llm_metadata(
|
|
serialized={},
|
|
run_id=run_id,
|
|
messages=[{"role": "user", "content": "test"}],
|
|
metadata={"ls_provider": "anthropic", "ls_model_name": "claude-3-sonnet"},
|
|
)
|
|
|
|
# Response with both cache read and creation
|
|
response = LLMResult(
|
|
generations=[
|
|
[
|
|
ChatGeneration(
|
|
message=AIMessage(content="Response"),
|
|
generation_info={
|
|
"usage_metadata": {
|
|
"input_tokens": 2000,
|
|
"output_tokens": 50,
|
|
"cache_read_input_tokens": 800,
|
|
"cache_creation_input_tokens": 500,
|
|
}
|
|
},
|
|
)
|
|
]
|
|
],
|
|
llm_output={},
|
|
)
|
|
|
|
cb._pop_run_and_capture_generation(run_id, None, response)
|
|
|
|
generation_args = mock_client.capture.call_args_list[0][1]
|
|
# 2000 - 800 (read) - 500 (write) = 700
|
|
assert generation_args["properties"]["$ai_input_tokens"] == 700
|
|
assert generation_args["properties"]["$ai_cache_read_input_tokens"] == 800
|
|
assert generation_args["properties"]["$ai_cache_creation_input_tokens"] == 500
|
|
|
|
|
|
def test_openai_cache_read_tokens(mock_client):
|
|
"""Test that OpenAI cache read tokens are captured correctly."""
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[("user", "Use the cached prompt for this request")]
|
|
)
|
|
|
|
# Mock response with cache read tokens in input_token_details
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Response using cached prompt context.",
|
|
usage_metadata={
|
|
"input_tokens": 150,
|
|
"output_tokens": 40,
|
|
"total_tokens": 190,
|
|
"input_token_details": {
|
|
"cache_read": 100, # 100 tokens read from cache
|
|
"cache_creation": 0,
|
|
},
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Response using cached prompt context."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert generation_props["$ai_input_tokens"] == 150 # No subtraction for OpenAI
|
|
assert generation_props["$ai_output_tokens"] == 40
|
|
assert generation_props["$ai_cache_read_input_tokens"] == 100
|
|
assert generation_props["$ai_cache_creation_input_tokens"] == 0
|
|
assert generation_props["$ai_reasoning_tokens"] == 0
|
|
|
|
|
|
def test_openai_cache_creation_tokens(mock_client):
|
|
"""Test that OpenAI cache creation tokens are captured correctly."""
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[("user", "Create a cache for this large prompt context")]
|
|
)
|
|
|
|
# Mock response with cache creation tokens in input_token_details
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Created cache for the prompt context.",
|
|
usage_metadata={
|
|
"input_tokens": 2000,
|
|
"output_tokens": 25,
|
|
"total_tokens": 2025,
|
|
"input_token_details": {
|
|
"cache_creation": 1500, # 1500 tokens written to cache
|
|
"cache_read": 0,
|
|
},
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Created cache for the prompt context."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert generation_props["$ai_input_tokens"] == 2000
|
|
assert generation_props["$ai_output_tokens"] == 25
|
|
assert generation_props["$ai_cache_creation_input_tokens"] == 1500
|
|
assert generation_props["$ai_cache_read_input_tokens"] == 0
|
|
assert generation_props["$ai_reasoning_tokens"] == 0
|
|
|
|
|
|
def test_combined_reasoning_and_cache_tokens(mock_client):
|
|
"""Test that both reasoning tokens and cache tokens can be captured together."""
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[("user", "Think through this cached problem")]
|
|
)
|
|
|
|
# Mock response with both reasoning and cache tokens
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Let me reason through this using cached context...",
|
|
usage_metadata={
|
|
"input_tokens": 500,
|
|
"output_tokens": 100,
|
|
"total_tokens": 600,
|
|
"input_token_details": {"cache_read": 300, "cache_creation": 0},
|
|
"output_token_details": {"reasoning": 60}, # 60 reasoning tokens
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Let me reason through this using cached context..."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
assert generation_props["$ai_input_tokens"] == 500 # No subtraction for OpenAI
|
|
assert generation_props["$ai_output_tokens"] == 100
|
|
assert generation_props["$ai_cache_read_input_tokens"] == 300
|
|
assert generation_props["$ai_cache_creation_input_tokens"] == 0
|
|
assert generation_props["$ai_reasoning_tokens"] == 60
|
|
|
|
|
|
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OPENAI_API_KEY is not set")
|
|
def test_openai_reasoning_tokens_o4_mini(mock_client):
|
|
model = ChatOpenAI(
|
|
api_key=OPENAI_API_KEY, model="o4-mini", max_completion_tokens=10
|
|
)
|
|
cb = CallbackHandler(
|
|
mock_client, trace_id="test-trace-id", distinct_id="test-distinct-id"
|
|
)
|
|
model.invoke("what is the weather in sf", config={"callbacks": [cb]})
|
|
call = mock_client.capture.call_args_list[0][1]
|
|
assert call["properties"]["$ai_reasoning_tokens"] is not None
|
|
assert call["properties"]["$ai_input_tokens"] is not None
|
|
assert call["properties"]["$ai_output_tokens"] is not None
|
|
|
|
|
|
def test_callback_handler_without_client():
|
|
"""Test that CallbackHandler works properly when no Insights client is passed."""
|
|
with patch("hanzo_insights.ai.langchain.callbacks.setup") as mock_setup:
|
|
mock_client = mock_setup.return_value
|
|
|
|
callbacks = CallbackHandler()
|
|
|
|
# Verify that setup() was called
|
|
mock_setup.assert_called_once()
|
|
|
|
# Verify that the client was set to the result of setup()
|
|
assert callbacks._ph_client == mock_client
|
|
|
|
# Test that the callback handler works with a simple chain
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
|
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
|
chain = prompt | model
|
|
|
|
# This should work and call the mock client
|
|
result = chain.invoke({}, config={"callbacks": [callbacks]})
|
|
assert result.content == "Bar"
|
|
|
|
# Verify that the mock client was used for capturing events
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
|
|
def test_convert_message_to_dict_tool_calls():
|
|
"""Test that _convert_message_to_dict properly converts tool calls in AIMessage."""
|
|
from hanzo_insights.ai.langchain.callbacks import _convert_message_to_dict
|
|
from langchain_core.messages import AIMessage
|
|
from langchain_core.messages.tool import ToolCall
|
|
|
|
# Create an AIMessage with tool calls
|
|
tool_calls = [
|
|
ToolCall(
|
|
id="call_123",
|
|
name="get_weather",
|
|
args={"city": "San Francisco", "units": "celsius"},
|
|
)
|
|
]
|
|
|
|
ai_message = AIMessage(
|
|
content="I'll check the weather for you.", tool_calls=tool_calls
|
|
)
|
|
|
|
# Convert to dict
|
|
result = _convert_message_to_dict(ai_message)
|
|
|
|
# Verify the conversion
|
|
assert result["role"] == "assistant"
|
|
assert result["content"] == "I'll check the weather for you."
|
|
assert result["tool_calls"] == [
|
|
{
|
|
"type": "function",
|
|
"id": "call_123",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"city": "San Francisco", "units": "celsius"}',
|
|
},
|
|
}
|
|
]
|
|
|
|
|
|
def test_tool_definition(mock_client):
|
|
"""Test that tools defined in invocation parameters are captured in $ai_tools property"""
|
|
callbacks = CallbackHandler(mock_client)
|
|
run_id = uuid.uuid4()
|
|
|
|
# Define tools to be passed to the invocation parameters
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"description": "Get the current weather for a specific location",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {
|
|
"type": "string",
|
|
"description": "The city or location name to get weather for",
|
|
}
|
|
},
|
|
"required": ["location"],
|
|
},
|
|
},
|
|
}
|
|
]
|
|
|
|
with patch("time.time", return_value=1234567890):
|
|
callbacks._set_llm_metadata(
|
|
{"kwargs": {"openai_api_base": "https://api.openai.com/v1"}},
|
|
run_id,
|
|
messages=[{"role": "user", "content": "hey"}],
|
|
invocation_params={"temperature": 0.7, "tools": tools},
|
|
metadata={"ls_model_name": "gpt-4o-mini", "ls_provider": "openai"},
|
|
name="test",
|
|
)
|
|
|
|
expected = GenerationMetadata(
|
|
model="gpt-4o-mini",
|
|
input=[{"role": "user", "content": "hey"}],
|
|
start_time=1234567890,
|
|
model_params={"temperature": 0.7},
|
|
provider="openai",
|
|
base_url="https://api.openai.com/v1",
|
|
name="test",
|
|
tools=tools,
|
|
end_time=None,
|
|
)
|
|
assert callbacks._runs[run_id] == expected
|
|
|
|
with patch("time.time", return_value=1234567891):
|
|
run = callbacks._pop_run_metadata(run_id)
|
|
expected.end_time = 1234567891
|
|
assert run == expected
|
|
assert callbacks._runs == {}
|
|
|
|
# Now test that the tools are properly captured in the Insights event
|
|
mock_response = MagicMock()
|
|
mock_response.generations = [[MagicMock()]]
|
|
|
|
callbacks._capture_generation(
|
|
trace_id=run_id,
|
|
run_id=run_id,
|
|
run=run,
|
|
output=mock_response,
|
|
parent_run_id=None,
|
|
)
|
|
|
|
assert mock_client.capture.call_count == 1
|
|
call_args = mock_client.capture.call_args[1]
|
|
props = call_args["properties"]
|
|
|
|
assert call_args["distinct_id"] == run_id
|
|
assert call_args["event"] == "$ai_generation"
|
|
assert props["$ai_provider"] == "openai"
|
|
assert props["$ai_model"] == "gpt-4o-mini"
|
|
assert props["$ai_input"] == [{"role": "user", "content": "hey"}]
|
|
assert props["$ai_model_parameters"] == {"temperature": 0.7}
|
|
assert props["$ai_base_url"] == "https://api.openai.com/v1"
|
|
assert props["$ai_span_name"] == "test"
|
|
assert props["$ai_span_id"] == run_id
|
|
assert props["$ai_trace_id"] == run_id
|
|
assert props["$ai_latency"] == 1.0
|
|
# Verify that tools are captured in the $ai_tools property
|
|
assert props["$ai_tools"] == tools
|
|
|
|
|
|
def test_cache_read_tokens_subtraction_from_input_tokens(mock_client):
|
|
"""Test that cache_read_tokens are properly subtracted from input_tokens.
|
|
|
|
This tests the logic in callbacks.py lines 757-758:
|
|
if normalized_usage.input_tokens and normalized_usage.cache_read_tokens:
|
|
normalized_usage.input_tokens = max(normalized_usage.input_tokens - normalized_usage.cache_read_tokens, 0)
|
|
"""
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[("user", "Use the cached prompt for this request")]
|
|
)
|
|
|
|
# Scenario 1: input_tokens includes cache_read_tokens (typical case)
|
|
# input_tokens=150 includes 100 cache_read tokens, so actual input is 50
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Response using cached prompt context.",
|
|
usage_metadata={
|
|
"input_tokens": 150, # Total includes cache reads
|
|
"output_tokens": 40,
|
|
"total_tokens": 190,
|
|
"cache_read_input_tokens": 100, # 100 tokens read from cache
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Response using cached prompt context."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
# Input tokens not reduced without provider metadata
|
|
assert generation_props["$ai_input_tokens"] == 150
|
|
assert generation_props["$ai_output_tokens"] == 40
|
|
assert generation_props["$ai_cache_read_input_tokens"] == 100
|
|
|
|
|
|
def test_cache_read_tokens_subtraction_prevents_negative(mock_client):
|
|
"""Test that cache_read_tokens subtraction doesn't result in negative input_tokens.
|
|
|
|
This tests the max(..., 0) part of the logic in callbacks.py lines 757-758.
|
|
"""
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[("user", "Edge case with large cache read")]
|
|
)
|
|
|
|
# Edge case: cache_read_tokens >= input_tokens
|
|
# This could happen in some API responses where accounting differs
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Response with edge case token counts.",
|
|
usage_metadata={
|
|
"input_tokens": 80,
|
|
"output_tokens": 20,
|
|
"total_tokens": 100,
|
|
"cache_read_input_tokens": 100, # More than input_tokens
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Response with edge case token counts."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
# Input tokens not reduced without provider metadata
|
|
assert generation_props["$ai_input_tokens"] == 80
|
|
assert generation_props["$ai_output_tokens"] == 20
|
|
assert generation_props["$ai_cache_read_input_tokens"] == 100
|
|
|
|
|
|
def test_no_cache_read_tokens_no_subtraction(mock_client):
|
|
"""Test that when there are no cache_read_tokens, input_tokens remain unchanged.
|
|
|
|
This tests the conditional check before the subtraction in callbacks.py line 757.
|
|
"""
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[("user", "Normal request without cache")]
|
|
)
|
|
|
|
# No cache usage - input_tokens should remain as-is
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Response without cache.",
|
|
usage_metadata={
|
|
"input_tokens": 100,
|
|
"output_tokens": 30,
|
|
"total_tokens": 130,
|
|
# No cache_read_input_tokens
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Response without cache."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
# Input tokens should remain unchanged at 100
|
|
assert generation_props["$ai_input_tokens"] == 100
|
|
assert generation_props["$ai_output_tokens"] == 30
|
|
assert generation_props["$ai_cache_read_input_tokens"] == 0
|
|
|
|
|
|
def test_zero_input_tokens_with_cache_read(mock_client):
|
|
"""Test edge case where input_tokens is 0 but cache_read_tokens exist.
|
|
|
|
This tests the falsy check in the conditional (line 757).
|
|
"""
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Edge case query")])
|
|
|
|
# Edge case: input_tokens is 0 (falsy), should skip subtraction
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Response.",
|
|
usage_metadata={
|
|
"input_tokens": 0,
|
|
"output_tokens": 10,
|
|
"total_tokens": 10,
|
|
"cache_read_input_tokens": 50,
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Response."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
# Input tokens should remain 0 (no subtraction because input_tokens is falsy)
|
|
assert generation_props["$ai_input_tokens"] == 0
|
|
assert generation_props["$ai_output_tokens"] == 10
|
|
assert generation_props["$ai_cache_read_input_tokens"] == 50
|
|
|
|
|
|
def test_non_anthropic_cache_write_tokens_not_subtracted_from_input(mock_client):
|
|
"""Test that cache_creation_input_tokens do NOT affect input_tokens for non-Anthropic providers.
|
|
|
|
When no provider metadata is set (or for non-Anthropic providers), cache tokens should
|
|
NOT be subtracted from input_tokens. This is because different providers report tokens
|
|
differently - only Anthropic's LangChain integration requires subtraction.
|
|
"""
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Create cache")])
|
|
|
|
# Cache creation without cache read
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="Creating cache.",
|
|
usage_metadata={
|
|
"input_tokens": 1000,
|
|
"output_tokens": 20,
|
|
"total_tokens": 1020,
|
|
"cache_creation_input_tokens": 800, # Cache write, not read
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
result = chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
assert result.content == "Creating cache."
|
|
assert mock_client.capture.call_count == 3
|
|
|
|
generation_args = mock_client.capture.call_args_list[1][1]
|
|
generation_props = generation_args["properties"]
|
|
|
|
assert generation_args["event"] == "$ai_generation"
|
|
# Input tokens should NOT be reduced by cache_creation_input_tokens
|
|
assert generation_props["$ai_input_tokens"] == 1000
|
|
assert generation_props["$ai_output_tokens"] == 20
|
|
assert generation_props["$ai_cache_creation_input_tokens"] == 800
|
|
assert generation_props["$ai_cache_read_input_tokens"] == 0
|
|
|
|
|
|
def test_agent_action_and_finish_imports():
|
|
"""
|
|
Regression test for LangChain 1.0+ compatibility (Issue #362).
|
|
Verifies that AgentAction and AgentFinish can be imported and used.
|
|
This test ensures the imports work with both LangChain 0.x and 1.0+.
|
|
"""
|
|
# Import the types that caused the compatibility issue
|
|
try:
|
|
from langchain_core.agents import AgentAction, AgentFinish
|
|
except (ImportError, ModuleNotFoundError):
|
|
from langchain.schema.agent import AgentAction, AgentFinish # type: ignore
|
|
|
|
# Verify they're available in the callbacks module
|
|
from hanzo_insights.ai.langchain.callbacks import CallbackHandler
|
|
|
|
# Test on_agent_action with mock data
|
|
mock_client = MagicMock()
|
|
callbacks = CallbackHandler(mock_client)
|
|
run_id = uuid.uuid4()
|
|
parent_run_id = uuid.uuid4()
|
|
|
|
# Create mock AgentAction
|
|
action = AgentAction(tool="test_tool", tool_input="test_input", log="test_log")
|
|
|
|
# Should not raise an exception
|
|
callbacks.on_agent_action(action, run_id=run_id, parent_run_id=parent_run_id)
|
|
|
|
# Verify parent was set
|
|
assert run_id in callbacks._parent_tree
|
|
assert callbacks._parent_tree[run_id] == parent_run_id
|
|
|
|
# Test on_agent_finish with mock data
|
|
finish = AgentFinish(return_values={"output": "test_output"}, log="finish_log")
|
|
|
|
# Should not raise an exception
|
|
callbacks.on_agent_finish(finish, run_id=run_id, parent_run_id=parent_run_id)
|
|
|
|
# Verify capture was called
|
|
assert mock_client.capture.call_count == 1
|
|
call_args = mock_client.capture.call_args[1]
|
|
assert call_args["event"] == "$ai_span"
|
|
|
|
|
|
def test_insights_properties_field_in_generation_metadata(mock_client):
|
|
"""Test that insights_properties is properly stored in GenerationMetadata."""
|
|
callbacks = CallbackHandler(mock_client)
|
|
run_id = uuid.uuid4()
|
|
|
|
# Test with billable=True
|
|
with patch("time.time", return_value=1234567890):
|
|
callbacks._set_llm_metadata(
|
|
{"kwargs": {"openai_api_base": "https://api.openai.com"}},
|
|
run_id,
|
|
messages=[{"role": "user", "content": "Test message"}],
|
|
invocation_params={"temperature": 0.5},
|
|
metadata={
|
|
"ls_model_name": "gpt-4o",
|
|
"ls_provider": "openai",
|
|
"insights_properties": {"$ai_billable": True},
|
|
},
|
|
name="test",
|
|
)
|
|
|
|
expected = GenerationMetadata(
|
|
model="gpt-4o",
|
|
input=[{"role": "user", "content": "Test message"}],
|
|
start_time=1234567890,
|
|
model_params={"temperature": 0.5},
|
|
provider="openai",
|
|
base_url="https://api.openai.com",
|
|
name="test",
|
|
insights_properties={"$ai_billable": True},
|
|
end_time=None,
|
|
)
|
|
assert callbacks._runs[run_id] == expected
|
|
assert callbacks._runs[run_id].insights_properties == {"$ai_billable": True}
|
|
|
|
callbacks._pop_run_metadata(run_id)
|
|
|
|
# Test with billable=False (explicit)
|
|
run_id2 = uuid.uuid4()
|
|
with patch("time.time", return_value=1234567890):
|
|
callbacks._set_llm_metadata(
|
|
{"kwargs": {"openai_api_base": "https://api.openai.com"}},
|
|
run_id2,
|
|
messages=[{"role": "user", "content": "Test message"}],
|
|
invocation_params={"temperature": 0.5},
|
|
metadata={
|
|
"ls_model_name": "gpt-4o",
|
|
"ls_provider": "openai",
|
|
"insights_properties": {"$ai_billable": False},
|
|
},
|
|
name="test",
|
|
)
|
|
|
|
assert callbacks._runs[run_id2].insights_properties == {"$ai_billable": False}
|
|
callbacks._pop_run_metadata(run_id2)
|
|
|
|
# Test when insights_properties not provided
|
|
run_id3 = uuid.uuid4()
|
|
with patch("time.time", return_value=1234567890):
|
|
callbacks._set_llm_metadata(
|
|
{"kwargs": {"openai_api_base": "https://api.openai.com"}},
|
|
run_id3,
|
|
messages=[{"role": "user", "content": "Test message"}],
|
|
invocation_params={"temperature": 0.5},
|
|
metadata={"ls_model_name": "gpt-4o", "ls_provider": "openai"},
|
|
name="test",
|
|
)
|
|
|
|
assert callbacks._runs[run_id3].insights_properties is None
|
|
|
|
|
|
def test_billable_property_in_generation_event(mock_client):
|
|
"""Test that the billable property is captured in the $ai_generation event."""
|
|
callbacks = CallbackHandler(mock_client)
|
|
|
|
# We need to test the _set_llm_metadata directly since FakeMessagesListChatModel
|
|
# doesn't support metadata in the same way as real models
|
|
run_id = uuid.uuid4()
|
|
with patch("time.time", return_value=1234567890):
|
|
callbacks._set_llm_metadata(
|
|
{},
|
|
run_id,
|
|
messages=[{"role": "user", "content": "Test"}],
|
|
metadata={
|
|
"insights_properties": {"$ai_billable": True},
|
|
"ls_model_name": "test-model",
|
|
},
|
|
invocation_params={},
|
|
)
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.generations = [[MagicMock()]]
|
|
|
|
with patch("time.time", return_value=1234567891):
|
|
run = callbacks._pop_run_metadata(run_id)
|
|
|
|
callbacks._capture_generation(
|
|
trace_id=run_id,
|
|
run_id=run_id,
|
|
run=run,
|
|
output=mock_response,
|
|
parent_run_id=None,
|
|
)
|
|
|
|
assert mock_client.capture.call_count == 1
|
|
call_args = mock_client.capture.call_args[1]
|
|
props = call_args["properties"]
|
|
|
|
assert call_args["event"] == "$ai_generation"
|
|
assert props["$ai_billable"] is True
|
|
|
|
|
|
def test_billable_defaults_to_false_in_event(mock_client):
|
|
"""Test that $ai_billable is not present when not specified."""
|
|
prompt = ChatPromptTemplate.from_messages([("user", "Test query")])
|
|
model = FakeMessagesListChatModel(
|
|
responses=[AIMessage(content="Test response")],
|
|
)
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = prompt | model
|
|
chain.invoke({}, config={"callbacks": callbacks})
|
|
|
|
generation_call = None
|
|
for call in mock_client.capture.call_args_list:
|
|
if call[1]["event"] == "$ai_generation":
|
|
generation_call = call
|
|
break
|
|
|
|
assert generation_call is not None
|
|
props = generation_call[1]["properties"]
|
|
assert "$ai_billable" not in props
|
|
|
|
|
|
def test_billable_with_real_chain(mock_client):
|
|
"""Test billable tracking through a complete chain execution with mocked metadata."""
|
|
callbacks = CallbackHandler(mock_client)
|
|
run_id = uuid.uuid4()
|
|
|
|
with patch("time.time", return_value=1000.0):
|
|
callbacks._set_llm_metadata(
|
|
{},
|
|
run_id,
|
|
messages=[{"role": "user", "content": "What's the weather?"}],
|
|
metadata={
|
|
"ls_model_name": "fake-model",
|
|
"ls_provider": "fake",
|
|
"insights_properties": {"$ai_billable": True},
|
|
},
|
|
invocation_params={"temperature": 0.7},
|
|
)
|
|
|
|
assert callbacks._runs[run_id].insights_properties == {"$ai_billable": True}
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.generations = [[MagicMock()]]
|
|
|
|
with patch("time.time", return_value=1001.0):
|
|
run = callbacks._pop_run_metadata(run_id)
|
|
|
|
callbacks._capture_generation(
|
|
trace_id=run_id,
|
|
run_id=run_id,
|
|
run=run,
|
|
output=mock_response,
|
|
parent_run_id=None,
|
|
)
|
|
|
|
assert mock_client.capture.call_count == 1
|
|
call_args = mock_client.capture.call_args[1]
|
|
props = call_args["properties"]
|
|
|
|
assert call_args["event"] == "$ai_generation"
|
|
assert props["$ai_billable"] is True
|
|
assert props["$ai_model"] == "fake-model"
|
|
assert props["$ai_provider"] == "fake"
|
|
|
|
|
|
# Exception Capture Integration Tests
|
|
|
|
|
|
def test_exception_autocapture_on_span_error():
|
|
"""Test that capture_exception is called when a span errors and autocapture is enabled."""
|
|
mock_client = MagicMock()
|
|
mock_client.privacy_mode = False
|
|
mock_client.enable_exception_autocapture = True
|
|
mock_client.capture_exception.return_value = "exception-uuid-123"
|
|
|
|
def failing_span(_):
|
|
raise ValueError("test error")
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = RunnableLambda(failing_span)
|
|
|
|
try:
|
|
chain.invoke({}, config={"callbacks": callbacks})
|
|
except ValueError:
|
|
pass
|
|
|
|
# Verify capture_exception was called
|
|
assert mock_client.capture_exception.call_count == 1
|
|
exception_call = mock_client.capture_exception.call_args
|
|
assert isinstance(exception_call[0][0], ValueError)
|
|
assert str(exception_call[0][0]) == "test error"
|
|
|
|
|
|
def test_exception_autocapture_adds_exception_id_to_span_event():
|
|
"""Test that $exception_event_id is added to the span event properties."""
|
|
mock_client = MagicMock()
|
|
mock_client.privacy_mode = False
|
|
mock_client.enable_exception_autocapture = True
|
|
mock_client.capture_exception.return_value = "exception-uuid-456"
|
|
|
|
def failing_span(_):
|
|
raise ValueError("test error")
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = RunnableLambda(failing_span)
|
|
|
|
try:
|
|
chain.invoke({}, config={"callbacks": callbacks})
|
|
except ValueError:
|
|
pass
|
|
|
|
# Find the span event (should have $ai_is_error=True)
|
|
span_calls = [
|
|
call
|
|
for call in mock_client.capture.call_args_list
|
|
if call[1].get("properties", {}).get("$ai_is_error") is True
|
|
]
|
|
assert len(span_calls) >= 1
|
|
|
|
span_props = span_calls[0][1]["properties"]
|
|
assert span_props["$exception_event_id"] == "exception-uuid-456"
|
|
assert span_props["$ai_error"] == "ValueError: test error"
|
|
|
|
|
|
def test_exception_autocapture_disabled_does_not_capture():
|
|
"""Test that capture_exception is NOT called when autocapture is disabled."""
|
|
mock_client = MagicMock()
|
|
mock_client.privacy_mode = False
|
|
mock_client.enable_exception_autocapture = False
|
|
|
|
def failing_span(_):
|
|
raise ValueError("test error")
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = RunnableLambda(failing_span)
|
|
|
|
try:
|
|
chain.invoke({}, config={"callbacks": callbacks})
|
|
except ValueError:
|
|
pass
|
|
|
|
# Verify capture_exception was NOT called
|
|
assert mock_client.capture_exception.call_count == 0
|
|
|
|
# But the span event should still have error info
|
|
span_calls = [
|
|
call
|
|
for call in mock_client.capture.call_args_list
|
|
if call[1].get("properties", {}).get("$ai_is_error") is True
|
|
]
|
|
assert len(span_calls) >= 1
|
|
|
|
span_props = span_calls[0][1]["properties"]
|
|
assert "$exception_event_id" not in span_props
|
|
assert span_props["$ai_error"] == "ValueError: test error"
|
|
|
|
|
|
def test_exception_autocapture_on_llm_generation_error(mock_client):
|
|
"""Test that capture_exception is called when an LLM generation fails."""
|
|
mock_client.privacy_mode = False
|
|
mock_client.enable_exception_autocapture = True
|
|
mock_client.capture_exception.return_value = "exception-uuid-789"
|
|
|
|
callbacks = CallbackHandler(mock_client)
|
|
run_id = uuid.uuid4()
|
|
|
|
# Simulate LLM start
|
|
callbacks.on_llm_start(
|
|
serialized={"kwargs": {"openai_api_base": "https://api.openai.com"}},
|
|
prompts=["Hello"],
|
|
run_id=run_id,
|
|
)
|
|
|
|
# Simulate LLM error
|
|
error = Exception("API rate limit exceeded")
|
|
callbacks.on_llm_error(error, run_id=run_id)
|
|
|
|
# Verify capture_exception was called
|
|
assert mock_client.capture_exception.call_count == 1
|
|
exception_call = mock_client.capture_exception.call_args
|
|
assert exception_call[0][0] is error
|
|
|
|
# Verify the generation event has $exception_event_id
|
|
generation_calls = [
|
|
call
|
|
for call in mock_client.capture.call_args_list
|
|
if call[1].get("event") == "$ai_generation"
|
|
]
|
|
assert len(generation_calls) == 1
|
|
|
|
gen_props = generation_calls[0][1]["properties"]
|
|
assert gen_props["$exception_event_id"] == "exception-uuid-789"
|
|
assert gen_props["$ai_is_error"] is True
|
|
|
|
|
|
def test_exception_autocapture_passes_ai_properties_to_exception():
|
|
"""Test that AI properties are passed to the exception event."""
|
|
mock_client = MagicMock()
|
|
mock_client.privacy_mode = False
|
|
mock_client.enable_exception_autocapture = True
|
|
mock_client.capture_exception.return_value = "exception-uuid-abc"
|
|
|
|
callbacks = CallbackHandler(
|
|
mock_client,
|
|
distinct_id="user-123",
|
|
properties={"custom_prop": "custom_value"},
|
|
)
|
|
run_id = uuid.uuid4()
|
|
|
|
# Simulate LLM start
|
|
callbacks.on_llm_start(
|
|
serialized={"kwargs": {"openai_api_base": "https://api.openai.com"}},
|
|
prompts=["Hello"],
|
|
run_id=run_id,
|
|
)
|
|
|
|
# Simulate LLM error
|
|
error = Exception("API error")
|
|
callbacks.on_llm_error(error, run_id=run_id)
|
|
|
|
# Verify capture_exception received the properties
|
|
exception_call = mock_client.capture_exception.call_args
|
|
props = exception_call[1]["properties"]
|
|
|
|
# Should have AI-related properties
|
|
assert "$ai_trace_id" in props
|
|
assert "$ai_is_error" in props
|
|
assert props["$ai_is_error"] is True
|
|
|
|
# Should have distinct_id passed through
|
|
assert exception_call[1]["distinct_id"] == "user-123"
|
|
|
|
|
|
def test_exception_autocapture_none_return_no_exception_id():
|
|
"""Test that when capture_exception returns None, no $exception_event_id is added."""
|
|
mock_client = MagicMock()
|
|
mock_client.privacy_mode = False
|
|
mock_client.enable_exception_autocapture = True
|
|
mock_client.capture_exception.return_value = (
|
|
None # e.g., exception already captured
|
|
)
|
|
|
|
def failing_span(_):
|
|
raise ValueError("test error")
|
|
|
|
callbacks = [CallbackHandler(mock_client)]
|
|
chain = RunnableLambda(failing_span)
|
|
|
|
try:
|
|
chain.invoke({}, config={"callbacks": callbacks})
|
|
except ValueError:
|
|
pass
|
|
|
|
# capture_exception was called but returned None
|
|
assert mock_client.capture_exception.call_count == 1
|
|
|
|
# Span event should NOT have $exception_event_id
|
|
span_calls = [
|
|
call
|
|
for call in mock_client.capture.call_args_list
|
|
if call[1].get("properties", {}).get("$ai_is_error") is True
|
|
]
|
|
assert len(span_calls) >= 1
|
|
|
|
span_props = span_calls[0][1]["properties"]
|
|
assert "$exception_event_id" not in span_props
|