Skip to content

Commit 180cb96

Browse files
nimarbnimarb
committed
fix cases
1 parent 53a9bff commit 180cb96

3 files changed

Lines changed: 363 additions & 75 deletions

File tree

‎langfuse/_client/client.py‎

Lines changed: 312 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,10 @@
3636
LANGFUSE_TRACING_ENABLED,
3737
LANGFUSE_TRACING_ENVIRONMENT,
3838
)
39-
from langfuse._client.constants import ObservationTypeLiteral, ObservationTypeLiteralNoEvent
39+
from langfuse._client.constants import (
40+
ObservationTypeLiteral,
41+
ObservationTypeLiteralNoEvent,
42+
)
4043
from langfuse._client.resource_manager import LangfuseResourceManager
4144
from langfuse._client.span import (
4245
LangfuseEvent,
@@ -674,32 +677,326 @@ def start_as_current_generation(
674677
),
675678
)
676679

680+
@overload
681+
def start_as_current_observation(
682+
self,
683+
*,
684+
trace_context: Optional[TraceContext] = None,
685+
name: str,
686+
as_type: Literal["generation"],
687+
input: Optional[Any] = None,
688+
output: Optional[Any] = None,
689+
metadata: Optional[Any] = None,
690+
version: Optional[str] = None,
691+
level: Optional[SpanLevel] = None,
692+
status_message: Optional[str] = None,
693+
completion_start_time: Optional[datetime] = None,
694+
model: Optional[str] = None,
695+
model_parameters: Optional[Dict[str, MapValue]] = None,
696+
usage_details: Optional[Dict[str, int]] = None,
697+
cost_details: Optional[Dict[str, float]] = None,
698+
prompt: Optional[PromptClient] = None,
699+
end_on_exit: Optional[bool] = None,
700+
) -> _AgnosticContextManager[LangfuseGeneration]: ...
701+
702+
@overload
703+
def start_as_current_observation(
704+
self,
705+
*,
706+
trace_context: Optional[TraceContext] = None,
707+
name: str,
708+
as_type: Literal["span"] = "span",
709+
input: Optional[Any] = None,
710+
output: Optional[Any] = None,
711+
metadata: Optional[Any] = None,
712+
version: Optional[str] = None,
713+
level: Optional[SpanLevel] = None,
714+
status_message: Optional[str] = None,
715+
end_on_exit: Optional[bool] = None,
716+
) -> _AgnosticContextManager[LangfuseSpan]: ...
717+
718+
@overload
719+
def start_as_current_observation(
720+
self,
721+
*,
722+
trace_context: Optional[TraceContext] = None,
723+
name: str,
724+
as_type: Literal[
725+
"agent", "tool", "chain", "retriever", "evaluator", "embedding", "guardrail"
726+
],
727+
input: Optional[Any] = None,
728+
output: Optional[Any] = None,
729+
metadata: Optional[Any] = None,
730+
version: Optional[str] = None,
731+
level: Optional[SpanLevel] = None,
732+
status_message: Optional[str] = None,
733+
end_on_exit: Optional[bool] = None,
734+
) -> _AgnosticContextManager[
735+
Union[
736+
LangfuseAgent,
737+
LangfuseTool,
738+
LangfuseChain,
739+
LangfuseRetriever,
740+
LangfuseEvaluator,
741+
LangfuseEmbedding,
742+
LangfuseGuardrail,
743+
]
744+
]: ...
745+
746+
def start_as_current_observation(
747+
self,
748+
*,
749+
trace_context: Optional[TraceContext] = None,
750+
name: str,
751+
as_type: ObservationTypeLiteralNoEvent = "span",
752+
input: Optional[Any] = None,
753+
output: Optional[Any] = None,
754+
metadata: Optional[Any] = None,
755+
version: Optional[str] = None,
756+
level: Optional[SpanLevel] = None,
757+
status_message: Optional[str] = None,
758+
completion_start_time: Optional[datetime] = None,
759+
model: Optional[str] = None,
760+
model_parameters: Optional[Dict[str, MapValue]] = None,
761+
usage_details: Optional[Dict[str, int]] = None,
762+
cost_details: Optional[Dict[str, float]] = None,
763+
prompt: Optional[PromptClient] = None,
764+
end_on_exit: Optional[bool] = None,
765+
) -> _AgnosticContextManager[
766+
Union[
767+
LangfuseSpan,
768+
LangfuseGeneration,
769+
LangfuseAgent,
770+
LangfuseTool,
771+
LangfuseChain,
772+
LangfuseRetriever,
773+
LangfuseEvaluator,
774+
LangfuseEmbedding,
775+
LangfuseGuardrail,
776+
]
777+
]:
778+
"""Create a new observation and set it as the current span in a context manager.
779+
780+
This method creates a new observation of the specified type and sets it as the
781+
current span within a context manager. Use this method with a 'with' statement to
782+
automatically handle the observation lifecycle within a code block.
783+
784+
The created observation will be the child of the current span in the context.
785+
786+
Args:
787+
trace_context: Optional context for connecting to an existing trace
788+
name: Name of the observation (e.g., function or operation name)
789+
as_type: Type of observation to create (defaults to "span")
790+
input: Input data for the operation (can be any JSON-serializable object)
791+
output: Output data from the operation (can be any JSON-serializable object)
792+
metadata: Additional metadata to associate with the observation
793+
version: Version identifier for the code or component
794+
level: Importance level of the observation (info, warning, error)
795+
status_message: Optional status message for the observation
796+
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.
797+
798+
The following parameters are only available when as_type="generation":
799+
completion_start_time: When the model started generating the response
800+
model: Name/identifier of the AI model used (e.g., "gpt-4")
801+
model_parameters: Parameters used for the model (e.g., temperature, max_tokens)
802+
usage_details: Token usage information (e.g., prompt_tokens, completion_tokens)
803+
cost_details: Cost information for the model call
804+
prompt: Associated prompt template from Langfuse prompt management
805+
806+
Returns:
807+
A context manager that yields the appropriate observation type based on as_type
808+
809+
Example:
810+
```python
811+
# Create an agent observation
812+
with langfuse.start_as_current_observation(name="planning-agent", as_type="agent") as agent:
813+
# Do agent work
814+
plan = create_plan()
815+
agent.update(output=plan)
816+
817+
# Create a tool observation
818+
with langfuse.start_as_current_observation(name="web-search", as_type="tool") as tool:
819+
# Do tool work
820+
results = search_web(query)
821+
tool.update(output=results)
822+
823+
# Create a generation observation
824+
with langfuse.start_as_current_observation(
825+
name="answer-generation",
826+
as_type="generation",
827+
model="gpt-4"
828+
) as generation:
829+
# Generate answer
830+
response = llm.generate(...)
831+
generation.update(output=response)
832+
```
833+
"""
834+
# Delegate to existing methods for consistency
835+
if as_type == "span":
836+
return cast(
837+
_AgnosticContextManager[
838+
Union[
839+
LangfuseSpan,
840+
LangfuseGeneration,
841+
LangfuseAgent,
842+
LangfuseTool,
843+
LangfuseChain,
844+
LangfuseRetriever,
845+
LangfuseEvaluator,
846+
LangfuseEmbedding,
847+
LangfuseGuardrail,
848+
]
849+
],
850+
self.start_as_current_span(
851+
trace_context=trace_context,
852+
name=name,
853+
input=input,
854+
output=output,
855+
metadata=metadata,
856+
version=version,
857+
level=level,
858+
status_message=status_message,
859+
end_on_exit=end_on_exit,
860+
),
861+
)
862+
863+
if as_type == "generation":
864+
return cast(
865+
_AgnosticContextManager[
866+
Union[
867+
LangfuseSpan,
868+
LangfuseGeneration,
869+
LangfuseAgent,
870+
LangfuseTool,
871+
LangfuseChain,
872+
LangfuseRetriever,
873+
LangfuseEvaluator,
874+
LangfuseEmbedding,
875+
LangfuseGuardrail,
876+
]
877+
],
878+
self.start_as_current_generation(
879+
trace_context=trace_context,
880+
name=name,
881+
input=input,
882+
output=output,
883+
metadata=metadata,
884+
version=version,
885+
level=level,
886+
status_message=status_message,
887+
completion_start_time=completion_start_time,
888+
model=model,
889+
model_parameters=model_parameters,
890+
usage_details=usage_details,
891+
cost_details=cost_details,
892+
prompt=prompt,
893+
end_on_exit=end_on_exit,
894+
),
895+
)
896+
897+
if trace_context:
898+
trace_id = trace_context.get("trace_id", None)
899+
parent_span_id = trace_context.get("parent_span_id", None)
900+
901+
if trace_id:
902+
remote_parent_span = self._create_remote_parent_span(
903+
trace_id=trace_id, parent_span_id=parent_span_id
904+
)
905+
906+
return cast(
907+
_AgnosticContextManager[
908+
Union[
909+
LangfuseSpan,
910+
LangfuseGeneration,
911+
LangfuseAgent,
912+
LangfuseTool,
913+
LangfuseChain,
914+
LangfuseRetriever,
915+
LangfuseEvaluator,
916+
LangfuseEmbedding,
917+
LangfuseGuardrail,
918+
]
919+
],
920+
self._create_span_with_parent_context(
921+
as_type=as_type,
922+
name=name,
923+
remote_parent_span=remote_parent_span,
924+
parent=None,
925+
end_on_exit=end_on_exit,
926+
input=input,
927+
output=output,
928+
metadata=metadata,
929+
version=version,
930+
level=level,
931+
status_message=status_message,
932+
completion_start_time=completion_start_time,
933+
model=model,
934+
model_parameters=model_parameters,
935+
usage_details=usage_details,
936+
cost_details=cost_details,
937+
prompt=prompt,
938+
),
939+
)
940+
941+
return cast(
942+
_AgnosticContextManager[
943+
Union[
944+
LangfuseSpan,
945+
LangfuseGeneration,
946+
LangfuseAgent,
947+
LangfuseTool,
948+
LangfuseChain,
949+
LangfuseRetriever,
950+
LangfuseEvaluator,
951+
LangfuseEmbedding,
952+
LangfuseGuardrail,
953+
]
954+
],
955+
self._start_as_current_otel_span_with_processed_media(
956+
as_type=as_type,
957+
name=name,
958+
end_on_exit=end_on_exit,
959+
input=input,
960+
output=output,
961+
metadata=metadata,
962+
version=version,
963+
level=level,
964+
status_message=status_message,
965+
completion_start_time=completion_start_time,
966+
model=model,
967+
model_parameters=model_parameters,
968+
usage_details=usage_details,
969+
cost_details=cost_details,
970+
prompt=prompt,
971+
),
972+
)
973+
677974
def _get_span_class(
678975
self,
679976
as_type: ObservationTypeLiteral,
680977
) -> type:
681978
"""Get the appropriate span class based on as_type."""
682-
normalized_type = as_type.upper()
979+
normalized_type = as_type.lower()
683980

684-
if normalized_type == "AGENT":
981+
if normalized_type == "agent":
685982
return LangfuseAgent
686-
elif normalized_type == "TOOL":
983+
elif normalized_type == "tool":
687984
return LangfuseTool
688-
elif normalized_type == "CHAIN":
985+
elif normalized_type == "chain":
689986
return LangfuseChain
690-
elif normalized_type == "RETRIEVER":
987+
elif normalized_type == "retriever":
691988
return LangfuseRetriever
692-
elif normalized_type == "EVALUATOR":
989+
elif normalized_type == "evaluator":
693990
return LangfuseEvaluator
694-
elif normalized_type == "EMBEDDING":
991+
elif normalized_type == "embedding":
695992
return LangfuseEmbedding
696-
elif normalized_type == "GUARDRAIL":
993+
elif normalized_type == "guardrail":
697994
return LangfuseGuardrail
698-
elif normalized_type == "GENERATION":
995+
elif normalized_type == "generation":
699996
return LangfuseGeneration
700-
elif normalized_type == "EVENT":
997+
elif normalized_type == "event":
701998
return LangfuseEvent
702-
elif normalized_type == "SPAN":
999+
elif normalized_type == "span":
7031000
return LangfuseSpan
7041001
else:
7051002
# Default to LangfuseSpan for unrecognized types
@@ -778,7 +1075,9 @@ def _start_as_current_otel_span_with_processed_media(
7781075
name=name,
7791076
end_on_exit=end_on_exit if end_on_exit is not None else True,
7801077
) as otel_span:
781-
span_class = self._get_span_class(as_type or "generation") # default was "generation"
1078+
span_class = self._get_span_class(
1079+
as_type or "generation"
1080+
) # default was "generation"
7821081
common_args = {
7831082
"otel_span": otel_span,
7841083
"langfuse_client": self,

0 commit comments

Comments
 (0)