feat(traces): add more observation types

langfuse/_client/attributes.py

@@ -93,9 +93,10 @@ def create_span_attributes(
     level: Optional[SpanLevel] = None,
     status_message: Optional[str] = None,
     version: Optional[str] = None,
+    observation_type: Optional[str] = None,
 ) -> dict:
     attributes = {
-        LangfuseOtelSpanAttributes.OBSERVATION_TYPE: "span",
+        LangfuseOtelSpanAttributes.OBSERVATION_TYPE: observation_type or "span",
         LangfuseOtelSpanAttributes.OBSERVATION_LEVEL: level,
         LangfuseOtelSpanAttributes.OBSERVATION_STATUS_MESSAGE: status_message,
         LangfuseOtelSpanAttributes.VERSION: version,

langfuse/_client/client.py

@@ -41,6 +41,11 @@
     LangfuseEvent,
     LangfuseGeneration,
     LangfuseSpan,
+    LangfuseAgent,
+    LangfuseTool,
+    LangfuseChain,
+    LangfuseRetriever,
+    LangfuseEmbedding,
 )
 from langfuse._utils import _get_timestamp
 from langfuse._utils.parse_error import handle_fern_exception
@@ -350,6 +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",
+            ]
+        ],
     ) -> _AgnosticContextManager[LangfuseSpan]:
         """Create a new span and set it as the current span in a context manager.
 
@@ -415,7 +434,7 @@ def start_as_current_span(
         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,
@@ -666,14 +685,48 @@ def start_as_current_generation(
             ),
         )
 
+    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":
+            return LangfuseAgent
+        elif as_type == "TOOL":
+            return LangfuseTool
+        elif as_type == "CHAIN":
+            return LangfuseChain
+        elif as_type == "RETRIEVER":
+            return LangfuseRetriever
+        elif as_type == "EMBEDDING":
+            return LangfuseEmbedding
+        elif as_type in ("generation", "GENERATION"):
+            return LangfuseGeneration
+        elif as_type in ("event", "EVENT"):
+            return LangfuseEvent
+        elif as_type in ("span", "SPAN"):
+            return LangfuseSpan
+        else:
+            return LangfuseSpan
+
     @_agnosticcontextmanager
     def _create_span_with_parent_context(
         self,
         *,
         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: Literal["generation", "span", "agent", "tool", "chain", "retriever"],
         end_on_exit: Optional[bool] = None,
         input: Optional[Any] = None,
         output: Optional[Any] = None,
@@ -720,7 +773,11 @@ def _start_as_current_otel_span_with_processed_media(
         self,
         *,
         name: str,
-        as_type: Optional[Literal["generation", "span"]] = 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,
@@ -739,37 +796,44 @@ def _start_as_current_otel_span_with_processed_media(
             name=name,
             end_on_exit=end_on_exit if end_on_exit is not None else True,
         ) as otel_span:
-            yield (
-                LangfuseSpan(
-                    otel_span=otel_span,
-                    langfuse_client=self,
-                    environment=self._environment,
-                    input=input,
-                    output=output,
-                    metadata=metadata,
-                    version=version,
-                    level=level,
-                    status_message=status_message,
-                )
-                if as_type == "span"
-                else LangfuseGeneration(
-                    otel_span=otel_span,
-                    langfuse_client=self,
-                    environment=self._environment,
-                    input=input,
-                    output=output,
-                    metadata=metadata,
-                    version=version,
-                    level=level,
-                    status_message=status_message,
-                    completion_start_time=completion_start_time,
-                    model=model,
-                    model_parameters=model_parameters,
-                    usage_details=usage_details,
-                    cost_details=cost_details,
-                    prompt=prompt,
+            span_class = self._get_span_class(as_type or "generation")
+            common_args = {
+                "otel_span": otel_span,
+                "langfuse_client": self,
+                "environment": self._environment,
+                "input": input,
+                "output": output,
+                "metadata": metadata,
+                "version": version,
+                "level": level,
+                "status_message": status_message,
+            }
+
+            if span_class == LangfuseGeneration:
+                common_args.update(
+                    {
+                        "completion_start_time": completion_start_time,
+                        "model": model,
+                        "model_parameters": model_parameters,
+                        "usage_details": usage_details,
+                        "cost_details": cost_details,
+                        "prompt": prompt,
+                    }
                 )
-            )
+            elif span_class in [
+                LangfuseAgent,
+                LangfuseTool,
+                LangfuseChain,
+                LangfuseRetriever,
+                LangfuseEmbedding,
+            ]:
+                # set their type internally in the class
+                pass
+            else:
+                if as_type is not None:
+                    common_args["as_type"] = as_type
+
+            yield span_class(**common_args)
 
     def _get_current_otel_span(self) -> Optional[otel_trace_api.Span]:
         current_span = otel_trace_api.get_current_span()

langfuse/_client/constants.py

@@ -4,3 +4,15 @@
 """
 
 LANGFUSE_TRACER_NAME = "langfuse-sdk"
+
+# Valid observation types for the @observe decorator
+VALID_OBSERVATION_TYPES = {
+    "SPAN",
+    "EVENT",
+    "GENERATION",
+    "AGENT",
+    "TOOL",
+    "CHAIN",
+    "RETRIEVER",
+    "EMBEDDING",
+}