Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
52df246a3e | ||
|
|
f1f9ecf7a4 | ||
|
|
9db1b7e9f3 | ||
|
|
01751d1205 |
@@ -0,0 +1,17 @@
|
||||
# This workflow is used to call the flags-project-board workflow when a pull request is opened, ready for review, review requested, synchronized, converted to draft, or reopened.
|
||||
# It is used to update the feature flags project board with the pull request information.
|
||||
|
||||
name: Call Feature Flags Project Workflow
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [opened, ready_for_review, review_requested, synchronize, converted_to_draft, reopened]
|
||||
|
||||
jobs:
|
||||
call-flags-project:
|
||||
uses: PostHog/.github/.github/workflows/flags-project-board.yml@main
|
||||
with:
|
||||
pr_number: ${{ github.event.pull_request.number }}
|
||||
pr_node_id: ${{ github.event.pull_request.node_id }}
|
||||
is_draft: ${{ github.event.pull_request.draft }}
|
||||
secrets: inherit
|
||||
@@ -1,3 +1,16 @@
|
||||
## 4.9.0 - 2025-06-13
|
||||
|
||||
- feat(ai): track reasoning and cache tokens in the LangChain callback
|
||||
|
||||
## 4.8.0 - 2025-06-10
|
||||
|
||||
- fix: export scoped, rather than tracked, decorator
|
||||
- feat: allow use of contexts without error tracking
|
||||
|
||||
## 4.7.0 - 2025-06-10
|
||||
|
||||
- feat: add support for parse endpoint in responses API (no longer beta)
|
||||
|
||||
## 4.6.2 - 2025-06-09
|
||||
|
||||
- fix: replace `import posthog` with direct method imports
|
||||
|
||||
+1
-1
@@ -15,7 +15,7 @@ new_context = new_context
|
||||
tag = tag
|
||||
get_tags = get_tags
|
||||
clear_tags = clear_tags
|
||||
tracked = scoped
|
||||
scoped = scoped
|
||||
|
||||
"""Settings."""
|
||||
api_key = None # type: Optional[str]
|
||||
|
||||
@@ -14,7 +14,6 @@ from typing import (
|
||||
List,
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
@@ -569,9 +568,14 @@ class CallbackHandler(BaseCallbackHandler):
|
||||
event_properties["$ai_is_error"] = True
|
||||
else:
|
||||
# Add usage
|
||||
input_tokens, output_tokens = _parse_usage(output)
|
||||
event_properties["$ai_input_tokens"] = input_tokens
|
||||
event_properties["$ai_output_tokens"] = output_tokens
|
||||
usage = _parse_usage(output)
|
||||
event_properties["$ai_input_tokens"] = usage.input_tokens
|
||||
event_properties["$ai_output_tokens"] = usage.output_tokens
|
||||
event_properties["$ai_cache_creation_input_tokens"] = (
|
||||
usage.cache_write_tokens
|
||||
)
|
||||
event_properties["$ai_cache_read_input_tokens"] = usage.cache_read_tokens
|
||||
event_properties["$ai_reasoning_tokens"] = usage.reasoning_tokens
|
||||
|
||||
# Generation results
|
||||
generation_result = output.generations[-1]
|
||||
@@ -647,9 +651,18 @@ def _convert_message_to_dict(message: BaseMessage) -> Dict[str, Any]:
|
||||
return message_dict
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelUsage:
|
||||
input_tokens: Optional[int]
|
||||
output_tokens: Optional[int]
|
||||
cache_write_tokens: Optional[int]
|
||||
cache_read_tokens: Optional[int]
|
||||
reasoning_tokens: Optional[int]
|
||||
|
||||
|
||||
def _parse_usage_model(
|
||||
usage: Union[BaseModel, Dict],
|
||||
) -> Tuple[Union[int, None], Union[int, None]]:
|
||||
usage: Union[BaseModel, dict],
|
||||
) -> ModelUsage:
|
||||
if isinstance(usage, BaseModel):
|
||||
usage = usage.__dict__
|
||||
|
||||
@@ -657,15 +670,23 @@ def _parse_usage_model(
|
||||
# https://pypi.org/project/langchain-anthropic/ (works also for Bedrock-Anthropic)
|
||||
("input_tokens", "input"),
|
||||
("output_tokens", "output"),
|
||||
("cache_creation_input_tokens", "cache_write"),
|
||||
("cache_read_input_tokens", "cache_read"),
|
||||
# https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/get-token-count
|
||||
("prompt_token_count", "input"),
|
||||
("candidates_token_count", "output"),
|
||||
("cached_content_token_count", "cache_read"),
|
||||
("thoughts_token_count", "reasoning"),
|
||||
# Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/monitoring-cw.html#runtime-cloudwatch-metrics
|
||||
("inputTokenCount", "input"),
|
||||
("outputTokenCount", "output"),
|
||||
("cacheCreationInputTokenCount", "cache_write"),
|
||||
("cacheReadInputTokenCount", "cache_read"),
|
||||
# Bedrock Anthropic
|
||||
("prompt_tokens", "input"),
|
||||
("completion_tokens", "output"),
|
||||
("cache_creation_input_tokens", "cache_write"),
|
||||
("cache_read_input_tokens", "cache_read"),
|
||||
# langchain-ibm https://pypi.org/project/langchain-ibm/
|
||||
("input_token_count", "input"),
|
||||
("generated_token_count", "output"),
|
||||
@@ -683,13 +704,45 @@ def _parse_usage_model(
|
||||
|
||||
parsed_usage[type_key] = final_count
|
||||
|
||||
return parsed_usage.get("input"), parsed_usage.get("output")
|
||||
# Caching (OpenAI & langchain 0.3.9+)
|
||||
if "input_token_details" in usage and isinstance(
|
||||
usage["input_token_details"], dict
|
||||
):
|
||||
parsed_usage["cache_write"] = usage["input_token_details"].get("cache_creation")
|
||||
parsed_usage["cache_read"] = usage["input_token_details"].get("cache_read")
|
||||
|
||||
# Reasoning (OpenAI & langchain 0.3.9+)
|
||||
if "output_token_details" in usage and isinstance(
|
||||
usage["output_token_details"], dict
|
||||
):
|
||||
parsed_usage["reasoning"] = usage["output_token_details"].get("reasoning")
|
||||
|
||||
field_mapping = {
|
||||
"input": "input_tokens",
|
||||
"output": "output_tokens",
|
||||
"cache_write": "cache_write_tokens",
|
||||
"cache_read": "cache_read_tokens",
|
||||
"reasoning": "reasoning_tokens",
|
||||
}
|
||||
return ModelUsage(
|
||||
**{
|
||||
dataclass_key: parsed_usage.get(mapped_key) or 0
|
||||
for mapped_key, dataclass_key in field_mapping.items()
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _parse_usage(response: LLMResult):
|
||||
def _parse_usage(response: LLMResult) -> ModelUsage:
|
||||
# langchain-anthropic uses the usage field
|
||||
llm_usage_keys = ["token_usage", "usage"]
|
||||
llm_usage: Tuple[Union[int, None], Union[int, None]] = (None, None)
|
||||
llm_usage: ModelUsage = ModelUsage(
|
||||
input_tokens=None,
|
||||
output_tokens=None,
|
||||
cache_write_tokens=None,
|
||||
cache_read_tokens=None,
|
||||
reasoning_tokens=None,
|
||||
)
|
||||
|
||||
if response.llm_output is not None:
|
||||
for key in llm_usage_keys:
|
||||
if response.llm_output.get(key):
|
||||
|
||||
@@ -230,6 +230,42 @@ class WrappedResponses:
|
||||
groups=posthog_groups,
|
||||
)
|
||||
|
||||
def parse(
|
||||
self,
|
||||
posthog_distinct_id: Optional[str] = None,
|
||||
posthog_trace_id: Optional[str] = None,
|
||||
posthog_properties: Optional[Dict[str, Any]] = None,
|
||||
posthog_privacy_mode: bool = False,
|
||||
posthog_groups: Optional[Dict[str, Any]] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
Parse structured output using OpenAI's 'responses.parse' method, but also track usage in PostHog.
|
||||
|
||||
Args:
|
||||
posthog_distinct_id: Optional ID to associate with the usage event.
|
||||
posthog_trace_id: Optional trace UUID for linking events.
|
||||
posthog_properties: Optional dictionary of extra properties to include in the event.
|
||||
posthog_privacy_mode: Whether to anonymize the input and output.
|
||||
posthog_groups: Optional dictionary of groups to associate with the event.
|
||||
**kwargs: Any additional parameters for the OpenAI Responses Parse API.
|
||||
|
||||
Returns:
|
||||
The response from OpenAI's responses.parse call.
|
||||
"""
|
||||
return call_llm_and_track_usage(
|
||||
posthog_distinct_id,
|
||||
self._client._ph_client,
|
||||
"openai",
|
||||
posthog_trace_id,
|
||||
posthog_properties,
|
||||
posthog_privacy_mode,
|
||||
posthog_groups,
|
||||
self._client.base_url,
|
||||
self._original.parse,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
class WrappedChat:
|
||||
"""Wrapper for OpenAI chat that tracks usage in PostHog."""
|
||||
|
||||
@@ -230,6 +230,42 @@ class WrappedResponses:
|
||||
groups=posthog_groups,
|
||||
)
|
||||
|
||||
async def parse(
|
||||
self,
|
||||
posthog_distinct_id: Optional[str] = None,
|
||||
posthog_trace_id: Optional[str] = None,
|
||||
posthog_properties: Optional[Dict[str, Any]] = None,
|
||||
posthog_privacy_mode: bool = False,
|
||||
posthog_groups: Optional[Dict[str, Any]] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
Parse structured output using OpenAI's 'responses.parse' method, but also track usage in PostHog.
|
||||
|
||||
Args:
|
||||
posthog_distinct_id: Optional ID to associate with the usage event.
|
||||
posthog_trace_id: Optional trace UUID for linking events.
|
||||
posthog_properties: Optional dictionary of extra properties to include in the event.
|
||||
posthog_privacy_mode: Whether to anonymize the input and output.
|
||||
posthog_groups: Optional dictionary of groups to associate with the event.
|
||||
**kwargs: Any additional parameters for the OpenAI Responses Parse API.
|
||||
|
||||
Returns:
|
||||
The response from OpenAI's responses.parse call.
|
||||
"""
|
||||
return await call_llm_and_track_usage_async(
|
||||
posthog_distinct_id,
|
||||
self._client._ph_client,
|
||||
"openai",
|
||||
posthog_trace_id,
|
||||
posthog_properties,
|
||||
posthog_privacy_mode,
|
||||
posthog_groups,
|
||||
self._client.base_url,
|
||||
self._original.parse,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
class WrappedChat:
|
||||
"""Async wrapper for OpenAI chat that tracks usage in PostHog."""
|
||||
|
||||
+28
-23
@@ -12,31 +12,34 @@ def _get_current_context() -> Dict[str, Any]:
|
||||
|
||||
|
||||
@contextmanager
|
||||
def new_context(fresh=False):
|
||||
def new_context(fresh=False, capture_exceptions=True):
|
||||
"""
|
||||
Create a new context scope that will be active for the duration of the with block.
|
||||
Any tags set within this scope will be isolated to this context. Any exceptions raised
|
||||
Create a new context scope that will be active for the duration of the with block.
|
||||
Any tags set within this scope will be isolated to this context. Any exceptions raised
|
||||
or events captured within the context will be tagged with the context tags.
|
||||
|
||||
Args:
|
||||
fresh: Whether to start with a fresh context (default: False).
|
||||
If False, inherits tags from parent context.
|
||||
If True, starts with no tags.
|
||||
Args:
|
||||
fresh: Whether to start with a fresh context (default: False).
|
||||
If False, inherits tags from parent context.
|
||||
If True, starts with no tags.
|
||||
capture_exceptions: Whether to capture exceptions raised within the context (default: True).
|
||||
If True, captures exceptions and tags them with the context tags before propagating them.
|
||||
If False, exceptions will propagate without being tagged or captured.
|
||||
|
||||
Examples:
|
||||
# Inherit parent context tags
|
||||
with posthog.new_context():
|
||||
posthog.tag("request_id", "123")
|
||||
# Both this event and the exception will be tagged with the context tags
|
||||
posthog.capture("event_name", {"property": "value"})
|
||||
raise ValueError("Something went wrong")
|
||||
Examples:
|
||||
# Inherit parent context tags
|
||||
with posthog.new_context():
|
||||
posthog.tag("request_id", "123")
|
||||
# Both this event and the exception will be tagged with the context tags
|
||||
posthog.capture("event_name", {"property": "value"})
|
||||
raise ValueError("Something went wrong")
|
||||
|
||||
# Start with fresh context (no inherited tags)
|
||||
with posthog.new_context(fresh=True):
|
||||
posthog.tag("request_id", "123")
|
||||
# Both this event and the exception will be tagged with the context tags
|
||||
posthog.capture("event_name", {"property": "value"})
|
||||
raise ValueError("Something went wrong")
|
||||
# Start with fresh context (no inherited tags)
|
||||
with posthog.new_context(fresh=True):
|
||||
posthog.tag("request_id", "123")
|
||||
# Both this event and the exception will be tagged with the context tags
|
||||
posthog.capture("event_name", {"property": "value"})
|
||||
raise ValueError("Something went wrong")
|
||||
|
||||
"""
|
||||
from posthog import capture_exception
|
||||
@@ -49,7 +52,8 @@ def new_context(fresh=False):
|
||||
try:
|
||||
yield
|
||||
except Exception as e:
|
||||
capture_exception(e)
|
||||
if capture_exceptions:
|
||||
capture_exception(e)
|
||||
raise
|
||||
finally:
|
||||
_context_stack.reset(token)
|
||||
@@ -88,13 +92,14 @@ def clear_tags() -> None:
|
||||
F = TypeVar("F", bound=Callable[..., Any])
|
||||
|
||||
|
||||
def scoped(fresh=False):
|
||||
def scoped(fresh=False, capture_exceptions=True):
|
||||
"""
|
||||
Decorator that creates a new context for the function. Simply wraps
|
||||
the function in a with posthog.new_context(): block.
|
||||
|
||||
Args:
|
||||
fresh: Whether to start with a fresh context (default: False)
|
||||
capture_exceptions: Whether to capture and track exceptions with posthog error tracking (default: True)
|
||||
|
||||
Example:
|
||||
@posthog.scoped()
|
||||
@@ -114,7 +119,7 @@ def scoped(fresh=False):
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
with new_context(fresh=fresh):
|
||||
with new_context(fresh=fresh, capture_exceptions=capture_exceptions):
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return cast(F, wrapper)
|
||||
|
||||
@@ -1378,11 +1378,11 @@ def test_langgraph_agent(mock_client):
|
||||
)
|
||||
graph.invoke(inputs, config={"callbacks": [cb]})
|
||||
calls = [call[1] for call in mock_client.capture.call_args_list]
|
||||
assert len(calls) == 21
|
||||
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"]) == 18
|
||||
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
|
||||
|
||||
|
||||
@@ -1435,11 +1435,13 @@ def test_span_set_parent_ids_for_third_level_run(mock_client, trace_id):
|
||||
|
||||
assert mock_client.capture.call_count == 3
|
||||
|
||||
span2, span1, trace = [
|
||||
call[1]["properties"] for call in mock_client.capture.call_args_list
|
||||
]
|
||||
assert span2["$ai_parent_id"] == span1["$ai_span_id"]
|
||||
assert span1["$ai_parent_id"] == trace["$ai_trace_id"]
|
||||
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):
|
||||
@@ -1478,3 +1480,250 @@ def test_captures_error_without_details_in_span(mock_client):
|
||||
== "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": 200,
|
||||
"output_tokens": 30,
|
||||
"total_tokens": 1030,
|
||||
"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"] == 200
|
||||
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_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
|
||||
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
|
||||
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(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
|
||||
|
||||
@@ -26,6 +26,11 @@ try:
|
||||
ResponseOutputMessage,
|
||||
ResponseOutputText,
|
||||
ResponseUsage,
|
||||
ParsedResponse,
|
||||
)
|
||||
from openai.types.responses.parsed_response import (
|
||||
ParsedResponseOutputMessage,
|
||||
ParsedResponseOutputText,
|
||||
)
|
||||
|
||||
from posthog.ai.openai import OpenAI
|
||||
@@ -115,6 +120,59 @@ def mock_openai_response_with_responses_api():
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_parsed_response():
|
||||
return ParsedResponse(
|
||||
id="test",
|
||||
model="gpt-4o-2024-08-06",
|
||||
object="response",
|
||||
created_at=1741476542,
|
||||
status="completed",
|
||||
error=None,
|
||||
incomplete_details=None,
|
||||
instructions=None,
|
||||
max_output_tokens=None,
|
||||
tools=[],
|
||||
tool_choice="auto",
|
||||
output=[
|
||||
ParsedResponseOutputMessage(
|
||||
id="msg_123",
|
||||
type="message",
|
||||
role="assistant",
|
||||
status="completed",
|
||||
content=[
|
||||
ParsedResponseOutputText(
|
||||
type="output_text",
|
||||
text='{"name": "Science Fair", "date": "Friday", "participants": ["Alice", "Bob"]}',
|
||||
annotations=[],
|
||||
parsed={
|
||||
"name": "Science Fair",
|
||||
"date": "Friday",
|
||||
"participants": ["Alice", "Bob"],
|
||||
},
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
output_parsed={
|
||||
"name": "Science Fair",
|
||||
"date": "Friday",
|
||||
"participants": ["Alice", "Bob"],
|
||||
},
|
||||
parallel_tool_calls=True,
|
||||
previous_response_id=None,
|
||||
usage=ResponseUsage(
|
||||
input_tokens=15,
|
||||
output_tokens=20,
|
||||
input_tokens_details={"prompt_tokens": 15, "cached_tokens": 0},
|
||||
output_tokens_details={"reasoning_tokens": 5},
|
||||
total_tokens=35,
|
||||
),
|
||||
user=None,
|
||||
metadata={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_embedding_response():
|
||||
return CreateEmbeddingResponse(
|
||||
@@ -646,3 +704,73 @@ def test_responses_api(mock_client, mock_openai_response_with_responses_api):
|
||||
assert props["$ai_http_status"] == 200
|
||||
assert props["foo"] == "bar"
|
||||
assert isinstance(props["$ai_latency"], float)
|
||||
|
||||
|
||||
def test_responses_parse(mock_client, mock_parsed_response):
|
||||
with patch(
|
||||
"openai.resources.responses.Responses.parse",
|
||||
return_value=mock_parsed_response,
|
||||
):
|
||||
client = OpenAI(api_key="test-key", posthog_client=mock_client)
|
||||
response = client.responses.parse(
|
||||
model="gpt-4o-2024-08-06",
|
||||
input=[
|
||||
{"role": "system", "content": "Extract the event information."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Alice and Bob are going to a science fair on Friday.",
|
||||
},
|
||||
],
|
||||
text={
|
||||
"format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "event",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"date": {"type": "string"},
|
||||
"participants": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
},
|
||||
},
|
||||
"required": ["name", "date", "participants"],
|
||||
},
|
||||
},
|
||||
}
|
||||
},
|
||||
posthog_distinct_id="test-id",
|
||||
posthog_properties={"foo": "bar"},
|
||||
)
|
||||
|
||||
assert response == mock_parsed_response
|
||||
assert mock_client.capture.call_count == 1
|
||||
|
||||
call_args = mock_client.capture.call_args[1]
|
||||
props = call_args["properties"]
|
||||
|
||||
assert call_args["distinct_id"] == "test-id"
|
||||
assert call_args["event"] == "$ai_generation"
|
||||
assert props["$ai_provider"] == "openai"
|
||||
assert props["$ai_model"] == "gpt-4o-2024-08-06"
|
||||
assert props["$ai_input"] == [
|
||||
{"role": "system", "content": "Extract the event information."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Alice and Bob are going to a science fair on Friday.",
|
||||
},
|
||||
]
|
||||
assert props["$ai_output_choices"] == [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": '{"name": "Science Fair", "date": "Friday", "participants": ["Alice", "Bob"]}',
|
||||
}
|
||||
]
|
||||
assert props["$ai_input_tokens"] == 15
|
||||
assert props["$ai_output_tokens"] == 20
|
||||
assert props["$ai_reasoning_tokens"] == 5
|
||||
assert props["$ai_http_status"] == 200
|
||||
assert props["foo"] == "bar"
|
||||
assert isinstance(props["$ai_latency"], float)
|
||||
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
VERSION = "4.6.2"
|
||||
VERSION = "4.9.0"
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(VERSION, end="") # noqa: T201
|
||||
|
||||
+5
-4
@@ -68,10 +68,11 @@ test = [
|
||||
"django",
|
||||
"openai",
|
||||
"anthropic",
|
||||
"langgraph",
|
||||
"langchain-community>=0.2.0",
|
||||
"langchain-openai>=0.2.0",
|
||||
"langchain-anthropic>=0.2.0",
|
||||
"langgraph>=0.4.8",
|
||||
"langchain-core>=0.3.65",
|
||||
"langchain-community>=0.3.25",
|
||||
"langchain-openai>=0.3.22",
|
||||
"langchain-anthropic>=0.3.15",
|
||||
"google-genai",
|
||||
"pydantic",
|
||||
"parameterized>=0.8.1",
|
||||
|
||||
Reference in New Issue
Block a user