Use as_type

Author: nimarb

langfuse/_client/client.py

@@ -350,7 +350,7 @@ def start_as_current_span(
         level: Optional[SpanLevel] = None,
         status_message: Optional[str] = None,
         end_on_exit: Optional[bool] = None,
-        observation_type: Optional[str] = None,
+        as_type: Optional[str] = None,
     ) -> _AgnosticContextManager[LangfuseSpan]:
         """Create a new span and set it as the current span in a context manager.
 
@@ -410,14 +410,13 @@ def start_as_current_span(
                         version=version,
                         level=level,
                         status_message=status_message,
-                        observation_type=observation_type,
                     ),
                 )
 
         return cast(
             _AgnosticContextManager[LangfuseSpan],
             self._start_as_current_otel_span_with_processed_media(
-                as_type="span",
+                as_type=as_type,
                 name=name,
                 end_on_exit=end_on_exit,
                 input=input,
@@ -426,7 +425,6 @@ def start_as_current_span(
                 version=version,
                 level=level,
                 status_message=status_message,
-                observation_type=observation_type,
             ),
         )
 
@@ -669,27 +667,26 @@ def start_as_current_generation(
             ),
         )
 
-    def _get_span_class(self, as_type: str, observation_type: Optional[str] = None):
-        """Get the appropriate span class based on as_type and observation_type."""
-        if observation_type == "AGENT":
+    def _get_span_class(self, as_type: str):
+        """Get the appropriate span class based on as_type."""
+        if as_type == "AGENT":
             return LangfuseAgent
-        elif observation_type == "TOOL":
+        elif as_type == "TOOL":
             return LangfuseTool
-        elif observation_type == "CHAIN":
+        elif as_type == "CHAIN":
             return LangfuseChain
-        elif observation_type == "RETRIEVER":
+        elif as_type == "RETRIEVER":
             return LangfuseRetriever
-        elif observation_type == "EMBEDDING":
+        elif as_type == "EMBEDDING":
             return LangfuseEmbedding
-        if as_type == "generation" or observation_type == "GENERATION":
+        elif as_type in ("generation", "GENERATION"):
             return LangfuseGeneration
-        elif as_type == "event" or observation_type == "EVENT":
+        elif as_type in ("event", "EVENT"):
             return LangfuseEvent
-        elif as_type == "span" or observation_type == "SPAN":
+        elif as_type in ("span", "SPAN"):
             return LangfuseSpan
         else:
-            # TODO: this should never happen -> error out?
-            # for now default to LangfuseSpan for unknown types or "span" as_type
+            # Default to LangfuseSpan for unknown types
             return LangfuseSpan
 
     @_agnosticcontextmanager
@@ -699,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: Literal["generation", "span"],
+        as_type: Optional[str] = None,
         end_on_exit: Optional[bool] = None,
         input: Optional[Any] = None,
         output: Optional[Any] = None,
@@ -713,7 +710,6 @@ def _create_span_with_parent_context(
         usage_details: Optional[Dict[str, int]] = None,
         cost_details: Optional[Dict[str, float]] = None,
         prompt: Optional[PromptClient] = None,
-        observation_type: Optional[str] = None,
     ) -> Any:
         parent_span = parent or cast(otel_trace_api.Span, remote_parent_span)
 
@@ -734,7 +730,6 @@ def _create_span_with_parent_context(
                 usage_details=usage_details,
                 cost_details=cost_details,
                 prompt=prompt,
-                observation_type=observation_type,
             ) as langfuse_span:
                 if remote_parent_span is not None:
                     langfuse_span._otel_span.set_attribute(
@@ -748,7 +743,7 @@ def _start_as_current_otel_span_with_processed_media(
         self,
         *,
         name: str,
-        as_type: Optional[Literal["generation", "span"]] = None,
+        as_type: Optional[str] = None,
         end_on_exit: Optional[bool] = None,
         input: Optional[Any] = None,
         output: Optional[Any] = None,
@@ -762,13 +757,12 @@ def _start_as_current_otel_span_with_processed_media(
         usage_details: Optional[Dict[str, int]] = None,
         cost_details: Optional[Dict[str, float]] = None,
         prompt: Optional[PromptClient] = None,
-        observation_type: Optional[str] = None,
     ) -> Any:
         with self._otel_tracer.start_as_current_span(
             name=name,
             end_on_exit=end_on_exit if end_on_exit is not None else True,
         ) as otel_span:
-            span_class = self._get_span_class(as_type, observation_type)
+            span_class = self._get_span_class(as_type)
             common_args = {
                 "otel_span": otel_span,
                 "langfuse_client": self,
@@ -790,12 +784,21 @@ def _start_as_current_otel_span_with_processed_media(
                         "usage_details": usage_details,
                         "cost_details": cost_details,
                         "prompt": prompt,
-                        "observation_type": observation_type,
                     }
                 )
+            elif span_class in [
+                LangfuseAgent,
+                LangfuseTool,
+                LangfuseChain,
+                LangfuseRetriever,
+                LangfuseEmbedding,
+            ]:
+                # Graph observation classes set their specific observation_type internally
+                pass
             else:
-                if observation_type is not None:
-                    common_args["observation_type"] = observation_type
+                # Regular spans (LangfuseSpan, LangfuseEvent) need as_type
+                if as_type is not None:
+                    common_args["as_type"] = as_type
 
             yield span_class(**common_args)
 

langfuse/_client/observe.py

@@ -19,14 +19,15 @@
     overload,
 )
 
-from opentelemetry import trace
 from opentelemetry.util._decorator import _AgnosticContextManager
 from typing_extensions import ParamSpec
 
-from langfuse._client.attributes import LangfuseOtelSpanAttributes
 from langfuse._client.environment_variables import (
     LANGFUSE_OBSERVE_DECORATOR_IO_CAPTURE_ENABLED,
 )
+
+from langfuse._client.constants import VALID_OBSERVATION_TYPES
+from langfuse._client.get_client import _set_current_public_key, get_client
 from langfuse._client.span import LangfuseGeneration, LangfuseSpan
 from langfuse.types import TraceContext
 
@@ -67,7 +68,18 @@ def observe(
         *,
         name: Optional[str] = None,
         as_type: Optional[Literal["generation"]] = None,
-        type: Optional[Literal["SPAN", "EVENT", "GENERATION", "AGENT", "TOOL", "CHAIN", "RETRIEVER", "EMBEDDING"]] = None,
+        type: Optional[
+            Literal[
+                "SPAN",
+                "EVENT",
+                "GENERATION",
+                "AGENT",
+                "TOOL",
+                "CHAIN",
+                "RETRIEVER",
+                "EMBEDDING",
+            ]
+        ] = None,
         capture_input: Optional[bool] = None,
         capture_output: Optional[bool] = None,
         transform_to_string: Optional[Callable[[Iterable], str]] = None,
@@ -79,7 +91,18 @@ def observe(
         *,
         name: Optional[str] = None,
         as_type: Optional[Literal["generation"]] = None,
-        type: Optional[Literal["SPAN", "EVENT", "GENERATION", "AGENT", "TOOL", "CHAIN", "RETRIEVER", "EMBEDDING"]] = None,
+        type: Optional[
+            Literal[
+                "SPAN",
+                "EVENT",
+                "GENERATION",
+                "AGENT",
+                "TOOL",
+                "CHAIN",
+                "RETRIEVER",
+                "EMBEDDING",
+            ]
+        ] = None,
         capture_input: Optional[bool] = None,
         capture_output: Optional[bool] = None,
         transform_to_string: Optional[Callable[[Iterable], str]] = None,
@@ -98,8 +121,8 @@ 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 directly. Supported values: "SPAN", "EVENT", 
-                    "GENERATION", "AGENT", "TOOL", "CHAIN", "RETRIEVER", "EMBEDDING". When specified, creates spans with 
+            type (Optional[Literal]): Set the observation type for agentic workflows. Supported values: "SPAN", "EVENT",
+                    "GENERATION", "AGENT", "TOOL", "CHAIN", "RETRIEVER", "EMBEDDING". When specified, creates spans with
                     the specified type for graph visualization and filtering in the Langfuse UI.
 
         Returns:
@@ -130,14 +153,7 @@ async def generate_answer(query):
             ```python
             @observe(type="AGENT")
             def planning_agent():
-                # Creates a span with observation type "AGENT" for graph visualization
                 return create_plan()
-
-            @observe(type="AGENT") 
-            def execution_agent():
-                # Creates a span with observation type "AGENT" 
-                # Parent relationships inferred from OpenTelemetry span hierarchy
-                return execute_plan()
             ```
 
             For trace context propagation between functions:
@@ -166,13 +182,19 @@ def sub_process():
             - For async functions, the decorator returns an async function wrapper.
             - For sync functions, the decorator returns a synchronous wrapper.
         """
-        # Validate type parameter if provided
-        if type is not None:
-            from langfuse._client.constants import VALID_OBSERVATION_TYPES
-            if type not in VALID_OBSERVATION_TYPES:
-                raise ValueError(
-                    f"Invalid observation type '{type}'. Valid types are: {', '.join(sorted(VALID_OBSERVATION_TYPES))}"
-                )
+        # 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))}"
+            )
+        if as_type is not None and as_type.upper() not in VALID_OBSERVATION_TYPES:
+            valid_values = sorted(
+                list(VALID_OBSERVATION_TYPES)
+                + [t.lower() for t in VALID_OBSERVATION_TYPES]
+            )
+            raise ValueError(
+                f"Invalid as_type '{as_type}'. Valid values are: {', '.join(valid_values)}"
+            )
 
         function_io_capture_enabled = os.environ.get(
             LANGFUSE_OBSERVE_DECORATOR_IO_CAPTURE_ENABLED, "True"
@@ -189,12 +211,14 @@ def sub_process():
         )
 
         def decorator(func: F) -> F:
+            # Merge as_type and type parameters - type takes precedence for graph observations
+            final_as_type = type or as_type
+
             return (
                 self._async_observe(
                     func,
                     name=name,
-                    as_type=as_type,
-                    observation_type=type,
+                    as_type=final_as_type,
                     capture_input=should_capture_input,
                     capture_output=should_capture_output,
                     transform_to_string=transform_to_string,
@@ -203,8 +227,7 @@ def decorator(func: F) -> F:
                 else self._sync_observe(
                     func,
                     name=name,
-                    as_type=as_type,
-                    observation_type=type,
+                    as_type=final_as_type,
                     capture_input=should_capture_input,
                     capture_output=should_capture_output,
                     transform_to_string=transform_to_string,
@@ -231,8 +254,7 @@ def _async_observe(
         func: F,
         *,
         name: Optional[str],
-        as_type: Optional[Literal["generation"]],
-        observation_type: Optional[str],
+        as_type: Optional[str],
         capture_input: bool,
         capture_output: bool,
         transform_to_string: Optional[Callable[[Iterable], str]] = None,
@@ -252,7 +274,6 @@ async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any:
                 else None
             )
             final_name = name or func.__name__
-            
             input = (
                 self._get_input_from_func_args(
                     is_method=self._is_method(func),
@@ -264,9 +285,9 @@ async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any:
             )
             public_key = cast(str, kwargs.pop("langfuse_public_key", None))
             langfuse_client = get_client(public_key=public_key)
-            
-            # Determine final observation type and create appropriate span
-            final_obs_type = observation_type or as_type
+
+            # Use consolidated as_type parameter
+            final_obs_type = as_type
 
             context_manager: Optional[
                 Union[
@@ -281,13 +302,13 @@ async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any:
                         input=input,
                         end_on_exit=False,  # when returning a generator, closing on exit would be to early
                     )
-                    if final_obs_type == "generation" or observation_type == "GENERATION"
+                    if final_obs_type in ("generation", "GENERATION")
                     else langfuse_client.start_as_current_span(
                         name=final_name,
                         trace_context=trace_context,
                         input=input,
                         end_on_exit=False,  # when returning a generator, closing on exit would be to early
-                        observation_type=observation_type,
+                        as_type=final_obs_type,
                     )
                 )
                 if langfuse_client
@@ -298,7 +319,6 @@ async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any:
                 return await func(*args, **kwargs)
 
             with context_manager as langfuse_span_or_generation:
-                
                 is_return_type_generator = False
 
                 try:
@@ -343,8 +363,7 @@ def _sync_observe(
         func: F,
         *,
         name: Optional[str],
-        as_type: Optional[Literal["generation"]],
-        observation_type: Optional[str],
+        as_type: Optional[str],
         capture_input: bool,
         capture_output: bool,
         transform_to_string: Optional[Callable[[Iterable], str]] = None,
@@ -362,7 +381,6 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                 else None
             )
             final_name = name or func.__name__
-            
             input = (
                 self._get_input_from_func_args(
                     is_method=self._is_method(func),
@@ -374,9 +392,9 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
             )
             public_key = kwargs.pop("langfuse_public_key", None)
             langfuse_client = get_client(public_key=public_key)
-            
-            # Determine final observation type and create appropriate span
-            final_obs_type = observation_type or as_type
+
+            # Use consolidated as_type parameter
+            final_obs_type = as_type
 
             context_manager: Optional[
                 Union[
@@ -391,13 +409,13 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                         input=input,
                         end_on_exit=False,  # when returning a generator, closing on exit would be to early
                     )
-                    if final_obs_type == "generation" or observation_type == "GENERATION"
+                    if final_obs_type in ("generation", "GENERATION")
                     else langfuse_client.start_as_current_span(
                         name=final_name,
                         trace_context=trace_context,
                         input=input,
                         end_on_exit=False,  # when returning a generator, closing on exit would be to early
-                        observation_type=observation_type,
+                        as_type=final_obs_type,
                     )
                 )
                 if langfuse_client
@@ -454,7 +472,6 @@ def _is_method(func: Callable) -> bool:
             or "cls" in inspect.signature(func).parameters
         )
 
-
     def _get_input_from_func_args(
         self,
         *,