Skip to content

Commit 3aa4af8

Browse files
nimarbnimarb
committed
from review
1 parent e7f34fe commit 3aa4af8

File tree

3 files changed

+70
-92
lines changed

3 files changed

+70
-92
lines changed

‎langfuse/_client/client.py‎

Lines changed: 49 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
ObservationTypeLiteral,
5252
ObservationTypeLiteralNoEvent,
5353
ObservationTypeGenerationLike,
54+
ObservationTypeSpanLike,
5455
get_observation_types_list,
5556
)
5657
from langfuse._client.resource_manager import LangfuseResourceManager
@@ -1100,8 +1101,7 @@ def start_as_current_observation(
11001101
status_message: Optional status message for the observation
11011102
end_on_exit (default: True): Whether to end the span automatically when leaving the context manager. If False, the span must be manually ended to avoid memory leaks.
11021103
1103-
The following parameters are available when as_type is: "generation", "agent",
1104-
"tool", "chain", "retriever", "evaluator" or "embedding".
1104+
The following parameters are available when as_type is: "generation" or "embedding".
11051105
completion_start_time: When the model started generating the response
11061106
model: Name/identifier of the AI model used (e.g., "gpt-4")
11071107
model_parameters: Parameters used for the model (e.g., temperature, max_tokens)
@@ -1142,7 +1142,7 @@ def start_as_current_observation(
11421142
generation.update(output=response)
11431143
```
11441144
"""
1145-
if as_type == "generation":
1145+
if as_type in get_observation_types_list(ObservationTypeGenerationLike):
11461146
if trace_context:
11471147
trace_id = trace_context.get("trace_id", None)
11481148
parent_span_id = trace_context.get("parent_span_id", None)
@@ -1153,9 +1153,12 @@ def start_as_current_observation(
11531153
)
11541154

11551155
return cast(
1156-
_AgnosticContextManager[LangfuseGeneration],
1156+
Union[
1157+
_AgnosticContextManager[LangfuseGeneration],
1158+
_AgnosticContextManager[LangfuseEmbedding],
1159+
],
11571160
self._create_span_with_parent_context(
1158-
as_type="generation",
1161+
as_type=as_type,
11591162
name=name,
11601163
remote_parent_span=remote_parent_span,
11611164
parent=None,
@@ -1176,9 +1179,12 @@ def start_as_current_observation(
11761179
)
11771180

11781181
return cast(
1179-
_AgnosticContextManager[LangfuseGeneration],
1182+
Union[
1183+
_AgnosticContextManager[LangfuseGeneration],
1184+
_AgnosticContextManager[LangfuseEmbedding],
1185+
],
11801186
self._start_as_current_otel_span_with_processed_media(
1181-
as_type="generation",
1187+
as_type=as_type,
11821188
name=name,
11831189
end_on_exit=end_on_exit,
11841190
input=input,
@@ -1196,7 +1202,7 @@ def start_as_current_observation(
11961202
),
11971203
)
11981204

1199-
if as_type == "span":
1205+
if as_type in get_observation_types_list(ObservationTypeSpanLike):
12001206
if trace_context:
12011207
trace_id = trace_context.get("trace_id", None)
12021208
parent_span_id = trace_context.get("parent_span_id", None)
@@ -1207,9 +1213,17 @@ def start_as_current_observation(
12071213
)
12081214

12091215
return cast(
1210-
_AgnosticContextManager[LangfuseSpan],
1216+
Union[
1217+
_AgnosticContextManager[LangfuseSpan],
1218+
_AgnosticContextManager[LangfuseAgent],
1219+
_AgnosticContextManager[LangfuseTool],
1220+
_AgnosticContextManager[LangfuseChain],
1221+
_AgnosticContextManager[LangfuseRetriever],
1222+
_AgnosticContextManager[LangfuseEvaluator],
1223+
_AgnosticContextManager[LangfuseGuardrail],
1224+
],
12111225
self._create_span_with_parent_context(
1212-
as_type="span",
1226+
as_type=as_type,
12131227
name=name,
12141228
remote_parent_span=remote_parent_span,
12151229
parent=None,
@@ -1224,9 +1238,17 @@ def start_as_current_observation(
12241238
)
12251239

12261240
return cast(
1227-
_AgnosticContextManager[LangfuseSpan],
1241+
Union[
1242+
_AgnosticContextManager[LangfuseSpan],
1243+
_AgnosticContextManager[LangfuseAgent],
1244+
_AgnosticContextManager[LangfuseTool],
1245+
_AgnosticContextManager[LangfuseChain],
1246+
_AgnosticContextManager[LangfuseRetriever],
1247+
_AgnosticContextManager[LangfuseEvaluator],
1248+
_AgnosticContextManager[LangfuseGuardrail],
1249+
],
12281250
self._start_as_current_otel_span_with_processed_media(
1229-
as_type="span",
1251+
as_type=as_type,
12301252
name=name,
12311253
end_on_exit=end_on_exit,
12321254
input=input,
@@ -1238,73 +1260,20 @@ def start_as_current_observation(
12381260
),
12391261
)
12401262

1241-
if trace_context:
1242-
trace_id = trace_context.get("trace_id", None)
1243-
parent_span_id = trace_context.get("parent_span_id", None)
1244-
1245-
if trace_id:
1246-
remote_parent_span = self._create_remote_parent_span(
1247-
trace_id=trace_id, parent_span_id=parent_span_id
1248-
)
1249-
1250-
return cast(
1251-
Union[
1252-
_AgnosticContextManager[LangfuseAgent],
1253-
_AgnosticContextManager[LangfuseTool],
1254-
_AgnosticContextManager[LangfuseChain],
1255-
_AgnosticContextManager[LangfuseRetriever],
1256-
_AgnosticContextManager[LangfuseEvaluator],
1257-
_AgnosticContextManager[LangfuseEmbedding],
1258-
_AgnosticContextManager[LangfuseGuardrail],
1259-
],
1260-
self._create_span_with_parent_context(
1261-
as_type=as_type,
1262-
name=name,
1263-
remote_parent_span=remote_parent_span,
1264-
parent=None,
1265-
end_on_exit=end_on_exit,
1266-
input=input,
1267-
output=output,
1268-
metadata=metadata,
1269-
version=version,
1270-
level=level,
1271-
status_message=status_message,
1272-
completion_start_time=completion_start_time,
1273-
model=model,
1274-
model_parameters=model_parameters,
1275-
usage_details=usage_details,
1276-
cost_details=cost_details,
1277-
prompt=prompt,
1278-
),
1279-
)
1280-
1281-
return cast(
1282-
Union[
1283-
_AgnosticContextManager[LangfuseAgent],
1284-
_AgnosticContextManager[LangfuseTool],
1285-
_AgnosticContextManager[LangfuseChain],
1286-
_AgnosticContextManager[LangfuseRetriever],
1287-
_AgnosticContextManager[LangfuseEvaluator],
1288-
_AgnosticContextManager[LangfuseEmbedding],
1289-
_AgnosticContextManager[LangfuseGuardrail],
1290-
],
1291-
self._start_as_current_otel_span_with_processed_media(
1292-
as_type=as_type,
1293-
name=name,
1294-
end_on_exit=end_on_exit,
1295-
input=input,
1296-
output=output,
1297-
metadata=metadata,
1298-
version=version,
1299-
level=level,
1300-
status_message=status_message,
1301-
completion_start_time=completion_start_time,
1302-
model=model,
1303-
model_parameters=model_parameters,
1304-
usage_details=usage_details,
1305-
cost_details=cost_details,
1306-
prompt=prompt,
1307-
),
1263+
# This should never be reached since all valid types are handled above
1264+
langfuse_logger.warning(
1265+
f"Unknown observation type: {as_type}, falling back to span"
1266+
)
1267+
return self._start_as_current_otel_span_with_processed_media(
1268+
as_type="span",
1269+
name=name,
1270+
end_on_exit=end_on_exit,
1271+
input=input,
1272+
output=output,
1273+
metadata=metadata,
1274+
version=version,
1275+
level=level,
1276+
status_message=status_message,
13081277
)
13091278

13101279
def _get_span_class(
@@ -1438,11 +1407,6 @@ def _start_as_current_otel_span_with_processed_media(
14381407

14391408
if span_class in [
14401409
LangfuseGeneration,
1441-
LangfuseAgent,
1442-
LangfuseTool,
1443-
LangfuseChain,
1444-
LangfuseRetriever,
1445-
LangfuseEvaluator,
14461410
LangfuseEmbedding,
14471411
]:
14481412
common_args.update(
@@ -1455,7 +1419,7 @@ def _start_as_current_otel_span_with_processed_media(
14551419
"prompt": prompt,
14561420
}
14571421
)
1458-
# For span and guardrail types, no generation properties needed
1422+
# For span-like types (span, agent, tool, chain, retriever, evaluator, guardrail), no generation properties needed
14591423

14601424
yield span_class(**common_args) # type: ignore[arg-type]
14611425

‎langfuse/_client/span.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -866,7 +866,10 @@ def start_observation(
866866

867867
observation_class = _OBSERVATION_CLASS_MAP.get(as_type)
868868
if not observation_class:
869-
raise ValueError(f"Unknown observation type: {as_type}")
869+
langfuse_logger.warning(
870+
f"Unknown observation type: {as_type}, falling back to LangfuseSpan"
871+
)
872+
observation_class = LangfuseSpan
870873

871874
with otel_trace_api.use_span(self._otel_span):
872875
new_otel_span = self._langfuse_client._otel_tracer.start_span(name=name)
@@ -1644,7 +1647,7 @@ def update(
16441647
cost_details: Optional[Dict[str, float]] = None,
16451648
prompt: Optional[PromptClient] = None,
16461649
**kwargs: Any,
1647-
) -> "LangfuseObservationWrapper":
1650+
) -> "LangfuseEvent":
16481651
"""Update is not allowed for LangfuseEvent because events cannot be updated.
16491652
16501653
This method logs a warning and returns self without making changes.

‎langfuse/langchain/CallbackHandler.py‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,14 @@
33
import pydantic
44

55
from langfuse._client.get_client import get_client
6+
from langfuse._client.attributes import LangfuseOtelSpanAttributes
67
from langfuse._client.span import (
78
LangfuseGeneration,
89
LangfuseSpan,
910
LangfuseAgent,
1011
LangfuseChain,
1112
LangfuseTool,
1213
LangfuseRetriever,
13-
LangfuseObservationWrapper,
1414
)
1515
from langfuse.logger import langfuse_logger
1616

@@ -299,8 +299,7 @@ def on_chain_start(
299299
self.runs[run_id] = span
300300
else:
301301
self.runs[run_id] = cast(
302-
# TODO: make this more precise (can be chain or agent here)
303-
LangfuseObservationWrapper,
302+
LangfuseChain,
304303
self.runs[parent_run_id],
305304
).start_observation(
306305
name=span_name,
@@ -365,7 +364,13 @@ def on_agent_action(
365364
if run_id not in self.runs:
366365
raise Exception("run not found")
367366

368-
self.runs[run_id].update(
367+
agent_run = self.runs[run_id]
368+
if hasattr(agent_run, '_otel_span'):
369+
agent_run._otel_span.set_attribute(
370+
LangfuseOtelSpanAttributes.OBSERVATION_TYPE, "agent"
371+
)
372+
373+
agent_run.update(
369374
output=action,
370375
input=kwargs.get("inputs"),
371376
).end()
@@ -388,7 +393,13 @@ def on_agent_finish(
388393
if run_id not in self.runs:
389394
raise Exception("run not found")
390395

391-
self.runs[run_id].update(
396+
agent_run = self.runs[run_id]
397+
if hasattr(agent_run, '_otel_span'):
398+
agent_run._otel_span.set_attribute(
399+
LangfuseOtelSpanAttributes.OBSERVATION_TYPE, "agent"
400+
)
401+
402+
agent_run.update(
392403
output=finish,
393404
input=kwargs.get("inputs"),
394405
).end()

0 commit comments

Comments
 (0)