todo clarifty case sensitivity

Author: nimarb

langfuse/_client/client.py

@@ -355,7 +355,7 @@ 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", "event", "agent", "tool", "chain", "retriever"]] = None,
+        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.
 

langfuse/_client/observe.py

@@ -182,6 +182,10 @@ def sub_process():
             - For async functions, the decorator returns an async function wrapper.
             - For sync functions, the decorator returns a synchronous wrapper.
         """
+        if as_type is not None and type is not None:
+            raise ValueError(
+                "Cannot specify both 'as_type' and 'type' parameters. Please use only one of them."
+            )
         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))}"
@@ -210,14 +214,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
+            # TODO: merge smartly with manually set as_type
+            consolidated_type = type or as_type
 
             return (
                 self._async_observe(
                     func,
                     name=name,
-                    as_type=final_as_type,
+                    as_type=consolidated_type,
                     capture_input=should_capture_input,
                     capture_output=should_capture_output,
                     transform_to_string=transform_to_string,
@@ -226,7 +230,7 @@ def decorator(func: F) -> F:
                 else self._sync_observe(
                     func,
                     name=name,
-                    as_type=final_as_type,
+                    as_type=consolidated_type,
                     capture_input=should_capture_input,
                     capture_output=should_capture_output,
                     transform_to_string=transform_to_string,
@@ -253,7 +257,7 @@ def _async_observe(
         func: F,
         *,
         name: Optional[str],
-        as_type: Optional[str],
+        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,
@@ -283,77 +287,80 @@ async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any:
                 else None
             )
             public_key = cast(str, kwargs.pop("langfuse_public_key", None))
-            langfuse_client = get_client(public_key=public_key)
 
-            # Use consolidated as_type parameter
-            final_obs_type = as_type
-
-            context_manager: Optional[
-                Union[
-                    _AgnosticContextManager[LangfuseGeneration],
-                    _AgnosticContextManager[LangfuseSpan],
-                ]
-            ] = (
-                (
-                    langfuse_client.start_as_current_generation(
-                        name=final_name,
-                        trace_context=trace_context,
-                        input=input,
-                        end_on_exit=False,  # when returning a generator, closing on exit would be to early
-                    )
-                    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
-                        as_type=final_obs_type,
+            # Set public key in execution context for nested decorated functions
+            with _set_current_public_key(public_key):
+                langfuse_client = get_client(public_key=public_key)
+
+                # Use consolidated as_type parameter
+                final_obs_type = as_type
+
+                context_manager: Optional[
+                    Union[
+                        _AgnosticContextManager[LangfuseGeneration],
+                        _AgnosticContextManager[LangfuseSpan],
+                    ]
+                ] = (
+                    (
+                        langfuse_client.start_as_current_generation(
+                            name=final_name,
+                            trace_context=trace_context,
+                            input=input,
+                            end_on_exit=False,  # when returning a generator, closing on exit would be to early
+                        )
+                        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
+                            as_type=final_obs_type,
+                        )
                     )
+                    if langfuse_client
+                    else None
                 )
-                if langfuse_client
-                else None
-            )
 
-            if context_manager is None:
-                return await func(*args, **kwargs)
+                if context_manager is None:
+                    return await func(*args, **kwargs)
 
-            with context_manager as langfuse_span_or_generation:
-                is_return_type_generator = False
+                with context_manager as langfuse_span_or_generation:
+                    is_return_type_generator = False
 
-                try:
-                    result = await func(*args, **kwargs)
+                    try:
+                        result = await func(*args, **kwargs)
 
-                    if capture_output is True:
-                        if inspect.isgenerator(result):
-                            is_return_type_generator = True
+                        if capture_output is True:
+                            if inspect.isgenerator(result):
+                                is_return_type_generator = True
 
-                            return self._wrap_sync_generator_result(
-                                langfuse_span_or_generation,
-                                result,
-                                transform_to_string,
-                            )
+                                return self._wrap_sync_generator_result(
+                                    langfuse_span_or_generation,
+                                    result,
+                                    transform_to_string,
+                                )
 
-                        if inspect.isasyncgen(result):
-                            is_return_type_generator = True
+                            if inspect.isasyncgen(result):
+                                is_return_type_generator = True
 
-                            return self._wrap_async_generator_result(
-                                langfuse_span_or_generation,
-                                result,
-                                transform_to_string,
-                            )
+                                return self._wrap_async_generator_result(
+                                    langfuse_span_or_generation,
+                                    result,
+                                    transform_to_string,
+                                )
 
-                        langfuse_span_or_generation.update(output=result)
+                            langfuse_span_or_generation.update(output=result)
 
-                    return result
-                except Exception as e:
-                    langfuse_span_or_generation.update(
-                        level="ERROR", status_message=str(e)
-                    )
+                        return result
+                    except Exception as e:
+                        langfuse_span_or_generation.update(
+                            level="ERROR", status_message=str(e)
+                        )
 
-                    raise e
-                finally:
-                    if not is_return_type_generator:
-                        langfuse_span_or_generation.end()
+                        raise e
+                    finally:
+                        if not is_return_type_generator:
+                            langfuse_span_or_generation.end()
 
         return cast(F, async_wrapper)