clean

Author: nimarb

langfuse/_client/client.py

@@ -669,6 +669,7 @@ def start_as_current_generation(
 
     def _get_span_class(self, as_type: str):
         """Get the appropriate span class based on as_type."""
+        # TODO: make it case insensitive
         if as_type == "AGENT":
             return LangfuseAgent
         elif as_type == "TOOL":
@@ -686,7 +687,6 @@ def _get_span_class(self, as_type: str):
         elif as_type in ("span", "SPAN"):
             return LangfuseSpan
         else:
-            # Default to LangfuseSpan for unknown types
             return LangfuseSpan
 
     @_agnosticcontextmanager
@@ -696,7 +696,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: Optional[str] = None,
+        as_type: Literal["generation", "span", "AGENT", "TOOL", "CHAIN", "RETRIEVER"],
         end_on_exit: Optional[bool] = None,
         input: Optional[Any] = None,
         output: Optional[Any] = None,
@@ -793,10 +793,9 @@ def _start_as_current_otel_span_with_processed_media(
                 LangfuseRetriever,
                 LangfuseEmbedding,
             ]:
-                # Graph observation classes set their specific observation_type internally
+                # set their type internally in the class
                 pass
             else:
-                # Regular spans (LangfuseSpan, LangfuseEvent) need as_type
                 if as_type is not None:
                     common_args["as_type"] = as_type
 

langfuse/_client/observe.py

@@ -121,7 +121,7 @@ def observe(
             name (Optional[str]): Custom name for the created trace or span. If not provided, the function name is used.
             as_type (Optional[Literal["generation"]]): Set to "generation" to create a specialized LLM generation span
                     with model metrics support, suitable for tracking language model outputs.
-            type (Optional[Literal]): Set the observation type for agentic workflows. Supported values: "SPAN", "EVENT",
+            type (Optional[Literal]): Set the observation type for agentic workflows. Supported values:
                     "GENERATION", "AGENT", "TOOL", "CHAIN", "RETRIEVER", "EMBEDDING". When specified, creates spans with
                     the specified type for graph visualization and filtering in the Langfuse UI.
 
@@ -182,7 +182,6 @@ def sub_process():
             - For async functions, the decorator returns an async function wrapper.
             - For sync functions, the decorator returns a synchronous wrapper.
         """
-        # Validate parameters
         if type is not None and type not in VALID_OBSERVATION_TYPES:
             raise ValueError(
                 f"Invalid observation type '{type}'. Valid types are: {', '.join(sorted(VALID_OBSERVATION_TYPES))}"
@@ -235,11 +234,11 @@ def decorator(func: F) -> F:
             )
 
         """Handle decorator with or without parentheses.
-        
+
         This logic enables the decorator to work both with and without parentheses:
         - @observe - Python passes the function directly to the decorator
         - @observe() - Python calls the decorator first, which must return a function decorator
-        
+
         When called without arguments (@observe), the func parameter contains the function to decorate,
         so we directly apply the decorator to it. When called with parentheses (@observe()),
         func is None, so we return the decorator function itself for Python to apply in the next step.