init commit

Author: nimarb

langfuse/_client/attributes.py

@@ -93,9 +93,11 @@ def create_span_attributes(
     level: Optional[SpanLevel] = None,
     status_message: Optional[str] = None,
     version: Optional[str] = None,
+    observation_type: Optional[str] = None,
 ) -> dict:
+    final_observation_type = observation_type or "span"
     attributes = {
-        LangfuseOtelSpanAttributes.OBSERVATION_TYPE: "span",
+        LangfuseOtelSpanAttributes.OBSERVATION_TYPE: final_observation_type,
         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
@@ -345,6 +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,
     ) -> _AgnosticContextManager[LangfuseSpan]:
         """Create a new span and set it as the current span in a context manager.
 
@@ -404,6 +410,7 @@ def start_as_current_span(
                         version=version,
                         level=level,
                         status_message=status_message,
+                        observation_type=observation_type,
                     ),
                 )
 
@@ -419,6 +426,7 @@ def start_as_current_span(
                 version=version,
                 level=level,
                 status_message=status_message,
+                observation_type=observation_type,
             ),
         )
 
@@ -661,6 +669,29 @@ 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":
+            return LangfuseAgent
+        elif observation_type == "TOOL":
+            return LangfuseTool
+        elif observation_type == "CHAIN":
+            return LangfuseChain
+        elif observation_type == "RETRIEVER":
+            return LangfuseRetriever
+        elif observation_type == "EMBEDDING":
+            return LangfuseEmbedding
+        if as_type == "generation" or observation_type == "GENERATION":
+            return LangfuseGeneration
+        elif as_type == "event" or observation_type == "EVENT":
+            return LangfuseEvent
+        elif as_type == "span" or observation_type == "SPAN":
+            return LangfuseSpan
+        else:
+            # TODO: this should never happen -> error out?
+            # for now default to LangfuseSpan for unknown types or "span" as_type
+            return LangfuseSpan
+
     @_agnosticcontextmanager
     def _create_span_with_parent_context(
         self,
@@ -682,6 +713,7 @@ 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)
 
@@ -702,6 +734,7 @@ 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(
@@ -729,42 +762,42 @@ 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:
-            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, observation_type)
+            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,
+                        "observation_type": observation_type,
+                    }
                 )
-            )
+            else:
+                if observation_type is not None:
+                    common_args["observation_type"] = observation_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",
+}