Skip to content

Commit dc5e199

Browse files
nimarbnimarb
committed
fix types
1 parent 5c3a510 commit dc5e199

2 files changed

Lines changed: 78 additions & 94 deletions

File tree

‎langfuse/_client/client.py‎

Lines changed: 58 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -731,16 +731,14 @@ def start_as_current_observation(
731731
level: Optional[SpanLevel] = None,
732732
status_message: Optional[str] = None,
733733
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-
]
734+
) -> Union[
735+
_AgnosticContextManager[LangfuseAgent],
736+
_AgnosticContextManager[LangfuseTool],
737+
_AgnosticContextManager[LangfuseChain],
738+
_AgnosticContextManager[LangfuseRetriever],
739+
_AgnosticContextManager[LangfuseEvaluator],
740+
_AgnosticContextManager[LangfuseEmbedding],
741+
_AgnosticContextManager[LangfuseGuardrail],
744742
]: ...
745743

746744
def start_as_current_observation(
@@ -762,18 +760,16 @@ def start_as_current_observation(
762760
cost_details: Optional[Dict[str, float]] = None,
763761
prompt: Optional[PromptClient] = None,
764762
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-
]
763+
) -> Union[
764+
_AgnosticContextManager[LangfuseGeneration],
765+
_AgnosticContextManager[LangfuseSpan],
766+
_AgnosticContextManager[LangfuseAgent],
767+
_AgnosticContextManager[LangfuseTool],
768+
_AgnosticContextManager[LangfuseChain],
769+
_AgnosticContextManager[LangfuseRetriever],
770+
_AgnosticContextManager[LangfuseEvaluator],
771+
_AgnosticContextManager[LangfuseEmbedding],
772+
_AgnosticContextManager[LangfuseGuardrail],
777773
]:
778774
"""Create a new observation and set it as the current span in a context manager.
779775
@@ -834,18 +830,16 @@ def start_as_current_observation(
834830
# Delegate to existing methods for consistency
835831
if as_type == "span":
836832
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-
]
833+
Union[
834+
_AgnosticContextManager[LangfuseGeneration],
835+
_AgnosticContextManager[LangfuseSpan],
836+
_AgnosticContextManager[LangfuseAgent],
837+
_AgnosticContextManager[LangfuseTool],
838+
_AgnosticContextManager[LangfuseChain],
839+
_AgnosticContextManager[LangfuseRetriever],
840+
_AgnosticContextManager[LangfuseEvaluator],
841+
_AgnosticContextManager[LangfuseEmbedding],
842+
_AgnosticContextManager[LangfuseGuardrail],
849843
],
850844
self.start_as_current_span(
851845
trace_context=trace_context,
@@ -862,18 +856,16 @@ def start_as_current_observation(
862856

863857
if as_type == "generation":
864858
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-
]
859+
Union[
860+
_AgnosticContextManager[LangfuseGeneration],
861+
_AgnosticContextManager[LangfuseSpan],
862+
_AgnosticContextManager[LangfuseAgent],
863+
_AgnosticContextManager[LangfuseTool],
864+
_AgnosticContextManager[LangfuseChain],
865+
_AgnosticContextManager[LangfuseRetriever],
866+
_AgnosticContextManager[LangfuseEvaluator],
867+
_AgnosticContextManager[LangfuseEmbedding],
868+
_AgnosticContextManager[LangfuseGuardrail],
877869
],
878870
self.start_as_current_generation(
879871
trace_context=trace_context,
@@ -904,18 +896,16 @@ def start_as_current_observation(
904896
)
905897

906898
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-
]
899+
Union[
900+
_AgnosticContextManager[LangfuseGeneration],
901+
_AgnosticContextManager[LangfuseSpan],
902+
_AgnosticContextManager[LangfuseAgent],
903+
_AgnosticContextManager[LangfuseTool],
904+
_AgnosticContextManager[LangfuseChain],
905+
_AgnosticContextManager[LangfuseRetriever],
906+
_AgnosticContextManager[LangfuseEvaluator],
907+
_AgnosticContextManager[LangfuseEmbedding],
908+
_AgnosticContextManager[LangfuseGuardrail],
919909
],
920910
self._create_span_with_parent_context(
921911
as_type=as_type,
@@ -939,18 +929,16 @@ def start_as_current_observation(
939929
)
940930

941931
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-
]
932+
Union[
933+
_AgnosticContextManager[LangfuseGeneration],
934+
_AgnosticContextManager[LangfuseSpan],
935+
_AgnosticContextManager[LangfuseAgent],
936+
_AgnosticContextManager[LangfuseTool],
937+
_AgnosticContextManager[LangfuseChain],
938+
_AgnosticContextManager[LangfuseRetriever],
939+
_AgnosticContextManager[LangfuseEvaluator],
940+
_AgnosticContextManager[LangfuseEmbedding],
941+
_AgnosticContextManager[LangfuseGuardrail],
954942
],
955943
self._start_as_current_otel_span_with_processed_media(
956944
as_type=as_type,

‎langfuse/_client/observe.py‎

Lines changed: 20 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -266,18 +266,16 @@ async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any:
266266
langfuse_client = get_client(public_key=public_key)
267267

268268
context_manager: Optional[
269-
_AgnosticContextManager[
270-
Union[
271-
LangfuseSpan,
272-
LangfuseGeneration,
273-
LangfuseAgent,
274-
LangfuseTool,
275-
LangfuseChain,
276-
LangfuseRetriever,
277-
LangfuseEvaluator,
278-
LangfuseEmbedding,
279-
LangfuseGuardrail,
280-
]
269+
Union[
270+
_AgnosticContextManager[LangfuseGeneration],
271+
_AgnosticContextManager[LangfuseSpan],
272+
_AgnosticContextManager[LangfuseAgent],
273+
_AgnosticContextManager[LangfuseTool],
274+
_AgnosticContextManager[LangfuseChain],
275+
_AgnosticContextManager[LangfuseRetriever],
276+
_AgnosticContextManager[LangfuseEvaluator],
277+
_AgnosticContextManager[LangfuseEmbedding],
278+
_AgnosticContextManager[LangfuseGuardrail],
281279
]
282280
] = (
283281
langfuse_client.start_as_current_observation(
@@ -373,18 +371,16 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
373371
langfuse_client = get_client(public_key=public_key)
374372

375373
context_manager: Optional[
376-
_AgnosticContextManager[
377-
Union[
378-
LangfuseSpan,
379-
LangfuseGeneration,
380-
LangfuseAgent,
381-
LangfuseTool,
382-
LangfuseChain,
383-
LangfuseRetriever,
384-
LangfuseEvaluator,
385-
LangfuseEmbedding,
386-
LangfuseGuardrail,
387-
]
374+
Union[
375+
_AgnosticContextManager[LangfuseGeneration],
376+
_AgnosticContextManager[LangfuseSpan],
377+
_AgnosticContextManager[LangfuseAgent],
378+
_AgnosticContextManager[LangfuseTool],
379+
_AgnosticContextManager[LangfuseChain],
380+
_AgnosticContextManager[LangfuseRetriever],
381+
_AgnosticContextManager[LangfuseEvaluator],
382+
_AgnosticContextManager[LangfuseEmbedding],
383+
_AgnosticContextManager[LangfuseGuardrail],
388384
]
389385
] = (
390386
langfuse_client.start_as_current_observation(

0 commit comments

Comments
 (0)