Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 79 additions & 16 deletions langfuse/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -598,6 +598,12 @@ def _get_langfuse_data_from_kwargs(resource: OpenAiDefinition, kwargs: Any) -> A

parsed_n = kwargs.get("n", 1) if not isinstance(kwargs.get("n", 1), NotGiven) else 1

parsed_service_tier = (
kwargs.get("service_tier", None)
if not isinstance(kwargs.get("service_tier", None), NotGiven)
else None
)

if resource.type == "embedding":
parsed_dimensions = (
kwargs.get("dimensions", None)
Expand Down Expand Up @@ -634,6 +640,9 @@ def _get_langfuse_data_from_kwargs(resource: OpenAiDefinition, kwargs: Any) -> A
if parsed_seed is not None:
modelParameters["seed"] = parsed_seed

if parsed_service_tier is not None:
modelParameters["service_tier"] = parsed_service_tier

langfuse_prompt = kwargs.get("langfuse_prompt", None)

return {
Expand All @@ -657,6 +666,7 @@ def _create_langfuse_update(
model: Optional[str] = None,
usage: Optional[Any] = None,
metadata: Optional[Any] = None,
model_parameters: Optional[Any] = None,
) -> Any:
update = {
"output": completion,
Expand All @@ -668,6 +678,9 @@ def _create_langfuse_update(
if metadata is not None:
update["metadata"] = metadata

if model_parameters is not None:
update["model_parameters"] = model_parameters

if usage is not None:
update["usage_details"] = _parse_usage(usage)
update["cost_details"] = _parse_cost(usage)
Expand Down Expand Up @@ -721,7 +734,7 @@ def _parse_cost(usage: Optional[Any] = None) -> Any:


def _extract_streamed_response_api_response(chunks: Any) -> Any:
completion, model, usage = None, None, None
completion, model, usage, service_tier = None, None, None, None
metadata = {}

for raw_chunk in chunks:
Expand All @@ -731,6 +744,7 @@ def _extract_streamed_response_api_response(chunks: Any) -> Any:

response = raw_response.__dict__
model = response.get("model")
service_tier = response.get("service_tier", None) or service_tier

for key, val in response.items():
if key not in ["created_at", "model", "output", "usage", "text"]:
Expand All @@ -739,18 +753,19 @@ def _extract_streamed_response_api_response(chunks: Any) -> Any:
if key == "output":
completion = _extract_response_api_completion(val)

return (model, completion, usage, metadata)
return (model, completion, usage, metadata, service_tier)


def _extract_streamed_openai_response(resource: Any, chunks: Any) -> Any:
completion: Any = defaultdict(lambda: None) if resource.type == "chat" else ""
model, usage, finish_reason = None, None, None
model, usage, finish_reason, service_tier = None, None, None, None

for chunk in chunks:
if _is_openai_v1():
chunk = chunk.__dict__

model = model or chunk.get("model", None) or None
service_tier = service_tier or chunk.get("service_tier", None) or None
chunk_usage = chunk.get("usage", None)
if chunk_usage is not None:
usage = chunk_usage
Expand Down Expand Up @@ -884,16 +899,18 @@ def get_response_for_chat() -> Any:
get_response_for_chat() if resource.type == "chat" else completion,
usage,
{"finish_reason": finish_reason} if finish_reason is not None else None,
service_tier,
)


def _get_langfuse_data_from_default_response(
resource: OpenAiDefinition, response: Any
) -> Any:
if response is None:
return None, "<NoneType response returned from OpenAI>", None
return None, "<NoneType response returned from OpenAI>", None, None

model = response.get("model", None) or None
service_tier = response.get("service_tier", None) or None

completion = None

Expand Down Expand Up @@ -942,7 +959,23 @@ def _get_langfuse_data_from_default_response(

usage = _parse_usage(response.get("usage", None))

return (model, completion, usage)
return (model, completion, usage, service_tier)


def _merge_service_tier_into_model_parameters(
model_parameters: Optional[Any], service_tier: Optional[Any]
) -> Optional[Any]:
"""Merge the response-side service tier into the request-side model parameters.

The response value is authoritative because OpenAI returns the tier that
actually processed the request (e.g. when the request specified "auto").
Returns None when there is nothing to update so callers can skip the
update and keep the request-side model parameters untouched.
"""
if service_tier is None:
return None

return {**(model_parameters or {}), "service_tier": service_tier}


def _is_openai_v1() -> bool:
Expand Down Expand Up @@ -999,9 +1032,10 @@ def _finalize_stream_response(
items: list[Any],
generation: LangfuseGeneration,
completion_start_time: Optional[datetime],
model_parameters: Optional[Any] = None,
) -> None:
try:
model, completion, usage, metadata = (
model, completion, usage, metadata, service_tier = (
_extract_streamed_response_api_response(items)
if resource.object == "Responses" or resource.object == "AsyncResponses"
else _extract_streamed_openai_response(resource, items)
Expand All @@ -1014,6 +1048,9 @@ def _finalize_stream_response(
model=model,
usage=usage,
metadata=metadata,
model_parameters=_merge_service_tier_into_model_parameters(
model_parameters, service_tier
),
)
except Exception:
pass
Expand All @@ -1026,12 +1063,14 @@ def _instrument_openai_stream(
resource: OpenAiDefinition,
response: Any,
generation: LangfuseGeneration,
model_parameters: Optional[Any] = None,
) -> Any:
if not hasattr(response, "_iterator"):
return LangfuseResponseGeneratorSync(
resource=resource,
response=response,
generation=generation,
model_parameters=model_parameters,
)

items: list[Any] = []
Expand All @@ -1051,6 +1090,7 @@ def finalize_once() -> None:
items=items,
generation=generation,
completion_start_time=completion_start_time,
model_parameters=model_parameters,
)

response._langfuse_finalize_once = finalize_once # type: ignore[attr-defined]
Expand Down Expand Up @@ -1085,12 +1125,14 @@ def _instrument_openai_async_stream(
resource: OpenAiDefinition,
response: Any,
generation: LangfuseGeneration,
model_parameters: Optional[Any] = None,
) -> Any:
if not hasattr(response, "_iterator"):
return LangfuseResponseGeneratorAsync(
resource=resource,
response=response,
generation=generation,
model_parameters=model_parameters,
)

items: list[Any] = []
Expand All @@ -1110,6 +1152,7 @@ async def finalize_once() -> None:
items=items,
generation=generation,
completion_start_time=completion_start_time,
model_parameters=model_parameters,
)

response._langfuse_finalize_once = finalize_once # type: ignore[attr-defined]
Expand Down Expand Up @@ -1243,21 +1286,25 @@ def _wrap(
resource=open_ai_resource,
response=openai_response,
generation=generation,
model_parameters=langfuse_data.get("model_parameters", None),
)
elif _is_streaming_response(openai_response):
return LangfuseResponseGeneratorSync(
resource=open_ai_resource,
response=openai_response,
generation=generation,
model_parameters=langfuse_data.get("model_parameters", None),
)

else:
parsed_response = _unwrap_raw_response(openai_response)
model, completion, usage = _get_langfuse_data_from_default_response(
open_ai_resource,
(parsed_response and parsed_response.__dict__)
if _is_openai_v1()
else parsed_response,
model, completion, usage, service_tier = (
_get_langfuse_data_from_default_response(
open_ai_resource,
(parsed_response and parsed_response.__dict__)
if _is_openai_v1()
else parsed_response,
)
)

generation.update(
Expand All @@ -1267,6 +1314,9 @@ def _wrap(
cost_details=_parse_cost(parsed_response.usage)
if hasattr(parsed_response, "usage")
else None,
model_parameters=_merge_service_tier_into_model_parameters(
langfuse_data.get("model_parameters", None), service_tier
),
).end()

return openai_response
Expand Down Expand Up @@ -1325,21 +1375,25 @@ async def _wrap_async(
resource=open_ai_resource,
response=openai_response,
generation=generation,
model_parameters=langfuse_data.get("model_parameters", None),
)
elif _is_streaming_response(openai_response):
return LangfuseResponseGeneratorAsync(
resource=open_ai_resource,
response=openai_response,
generation=generation,
model_parameters=langfuse_data.get("model_parameters", None),
)

else:
parsed_response = _unwrap_raw_response(openai_response)
model, completion, usage = _get_langfuse_data_from_default_response(
open_ai_resource,
(parsed_response and parsed_response.__dict__)
if _is_openai_v1()
else parsed_response,
model, completion, usage, service_tier = (
_get_langfuse_data_from_default_response(
open_ai_resource,
(parsed_response and parsed_response.__dict__)
if _is_openai_v1()
else parsed_response,
)
)
generation.update(
model=model,
Expand All @@ -1349,6 +1403,9 @@ async def _wrap_async(
cost_details=_parse_cost(parsed_response.usage)
if hasattr(parsed_response, "usage")
else None,
model_parameters=_merge_service_tier_into_model_parameters(
langfuse_data.get("model_parameters", None), service_tier
),
).end()

return openai_response
Expand Down Expand Up @@ -1397,12 +1454,14 @@ def __init__(
resource: Any,
response: Any,
generation: Any,
model_parameters: Optional[Any] = None,
) -> None:
self.items: list[Any] = []

self.resource = resource
self.response = response
self.generation = generation
self.model_parameters = model_parameters
self.completion_start_time: Optional[datetime] = None
self._is_finalized = False

Expand Down Expand Up @@ -1458,6 +1517,7 @@ def _finalize(self) -> None:
items=self.items,
generation=self.generation,
completion_start_time=self.completion_start_time,
model_parameters=self.model_parameters,
)


Expand All @@ -1468,12 +1528,14 @@ def __init__(
resource: Any,
response: Any,
generation: Any,
model_parameters: Optional[Any] = None,
) -> None:
self.items: list[Any] = []

self.resource = resource
self.response = response
self.generation = generation
self.model_parameters = model_parameters
self.completion_start_time: Optional[datetime] = None
self._is_finalized = False

Expand Down Expand Up @@ -1520,6 +1582,7 @@ async def _finalize(self) -> None:
items=self.items,
generation=self.generation,
completion_start_time=self.completion_start_time,
model_parameters=self.model_parameters,
)

async def close(self) -> None:
Expand Down
1 change: 1 addition & 0 deletions tests/e2e/test_decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -940,6 +940,7 @@ async def level_1_function(*args, **kwargs):
assert generation.end_time is not None
assert generation.start_time < generation.end_time
assert generation.model_parameters == {
"service_tier": "default",
"temperature": 0,
"top_p": 1,
"frequency_penalty": 0,
Expand Down
Loading