format

Author: nimarb

langfuse/_client/client.py

@@ -355,7 +355,20 @@ def start_as_current_span(
         level: Optional[SpanLevel] = None,
         status_message: Optional[str] = None,
         end_on_exit: Optional[bool] = None,
-        as_type: Optional[Literal["generation", "span", "GENERATION", "EVENT", "SPAN", "AGENT", "TOOL", "CHAIN", "RETRIEVER", "EMBEDDING"]],
+        as_type: Optional[
+            Literal[
+                "generation",
+                "span",
+                "GENERATION",
+                "EVENT",
+                "SPAN",
+                "AGENT",
+                "TOOL",
+                "CHAIN",
+                "RETRIEVER",
+                "EMBEDDING",
+            ]
+        ],
     ) -> _AgnosticContextManager[LangfuseSpan]:
         """Create a new span and set it as the current span in a context manager.
 
@@ -672,7 +685,19 @@ def start_as_current_generation(
             ),
         )
 
-    def _get_span_class(self, as_type: Literal["span", "generation", "event", "AGENT", "TOOL", "CHAIN", "RETRIEVER", "EMBEDDING"]):
+    def _get_span_class(
+        self,
+        as_type: Literal[
+            "span",
+            "generation",
+            "event",
+            "AGENT",
+            "TOOL",
+            "CHAIN",
+            "RETRIEVER",
+            "EMBEDDING",
+        ],
+    ):
         """Get the appropriate span class based on as_type."""
         # TODO: make it case insensitive
         if as_type == "AGENT":
@@ -701,7 +726,7 @@ def _create_span_with_parent_context(
         name: str,
         parent: Optional[otel_trace_api.Span] = None,
         remote_parent_span: Optional[otel_trace_api.Span] = None,
-        as_type: Literal["generation", "span", "AGENT", "TOOL", "CHAIN", "RETRIEVER"],
+        as_type: Literal["generation", "span", "agent", "tool", "chain", "retriever"],
         end_on_exit: Optional[bool] = None,
         input: Optional[Any] = None,
         output: Optional[Any] = None,
@@ -748,7 +773,11 @@ def _start_as_current_otel_span_with_processed_media(
         self,
         *,
         name: str,
-        as_type: Optional[Literal["generation", "span", "event", "AGENT", "TOOL", "CHAIN", "RETRIEVER"]] = None,
+        as_type: Optional[
+            Literal[
+                "generation", "span", "event", "AGENT", "TOOL", "CHAIN", "RETRIEVER"
+            ]
+        ] = None,
         end_on_exit: Optional[bool] = None,
         input: Optional[Any] = None,
         output: Optional[Any] = None,

langfuse/_client/observe.py

@@ -257,7 +257,19 @@ def _async_observe(
         func: F,
         *,
         name: Optional[str],
-        as_type: Optional[Literal["generation", "GENERATION", "EVENT", "SPAN", "AGENT", "TOOL", "CHAIN", "RETRIEVER", "EMBEDDING"]],
+        as_type: Optional[
+            Literal[
+                "generation",
+                "GENERATION",
+                "EVENT",
+                "SPAN",
+                "AGENT",
+                "TOOL",
+                "CHAIN",
+                "RETRIEVER",
+                "EMBEDDING",
+            ]
+        ],
         capture_input: bool,
         capture_output: bool,
         transform_to_string: Optional[Callable[[Iterable], str]] = None,

langfuse/_client/span.py

@@ -64,7 +64,16 @@ def __init__(
         *,
         otel_span: otel_trace_api.Span,
         langfuse_client: "Langfuse",
-        as_type: Literal["span", "generation", "event", "agent", "tool", "chain", "retriever", "embedding"],
+        as_type: Literal[
+            "span",
+            "generation",
+            "event",
+            "agent",
+            "tool",
+            "chain",
+            "retriever",
+            "embedding",
+        ],
         input: Optional[Any] = None,
         output: Optional[Any] = None,
         metadata: Optional[Any] = None,
@@ -531,7 +540,18 @@ def __init__(
         version: Optional[str] = None,
         level: Optional[SpanLevel] = None,
         status_message: Optional[str] = None,
-        as_type: Optional[Literal["span", "generation", "event", "agent", "tool", "chain", "retriever", "embedding"]] = None,
+        as_type: Optional[
+            Literal[
+                "span",
+                "generation",
+                "event",
+                "agent",
+                "tool",
+                "chain",
+                "retriever",
+                "embedding",
+            ]
+        ] = None,
     ):
         """Initialize a new LangfuseSpan.