cleanup

Author: nimarb

langfuse/_client/attributes.py

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

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,
-        as_type: Optional[str] = None,
+        as_type: Optional[Literal["generation", "span", "event", "agent", "tool", "chain", "retriever"]] = None,
     ) -> _AgnosticContextManager[LangfuseSpan]:
         """Create a new span and set it as the current span in a context manager.
 
@@ -743,7 +743,7 @@ def _start_as_current_otel_span_with_processed_media(
         self,
         *,
         name: str,
-        as_type: Optional[str] = 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,
@@ -762,7 +762,7 @@ 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:
-            span_class = self._get_span_class(as_type)
+            span_class = self._get_span_class(as_type or "generation")
             common_args = {
                 "otel_span": otel_span,
                 "langfuse_client": self,