cleanup

Author: nimarb

langfuse/_client/client.py

@@ -667,7 +667,7 @@ def start_as_current_generation(
             ),
         )
 
-    def _get_span_class(self, as_type: str):
+    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":

langfuse/_client/span.py

@@ -152,16 +152,10 @@ def __init__(
                     observation_type=as_type,
                 )
 
-            # Keep OBSERVATION_TYPE in attributes - don't remove it
-            # attributes.pop(LangfuseOtelSpanAttributes.OBSERVATION_TYPE, None)
+            attributes.pop(LangfuseOtelSpanAttributes.OBSERVATION_TYPE, None)
 
-            final_attributes = {k: v for k, v in attributes.items() if v is not None}
-            self._otel_span.set_attributes(final_attributes)
-
-            # Check what actually got set
-            actual_attributes = dict(self._otel_span.attributes)
-            obs_type_value = actual_attributes.get(
-                "langfuse.observation.type", "NOT_FOUND"
+            self._otel_span.set_attributes(
+                {k: v for k, v in attributes.items() if v is not None}
             )
 
     def end(self, *, end_time: Optional[int] = None) -> "LangfuseSpanWrapper":
@@ -552,11 +546,9 @@ def __init__(
             level: Importance level of the span (info, warning, error)
             status_message: Optional status message for the span
         """
-        # Default to span if not provided mimicing previous behavior
-        final_as_type = as_type or "span"
         super().__init__(
             otel_span=otel_span,
-            as_type=final_as_type,
+            as_type="span",
             langfuse_client=langfuse_client,
             input=input,
             output=output,

langfuse/openai.py

@@ -570,10 +570,7 @@ def _extract_streamed_openai_response(resource: Any, chunks: Any) -> Any:
                         )
                         curr["arguments"] += getattr(tool_call_chunk, "arguments", "")
 
-                elif (
-                    delta.get("tool_calls", None) is not None
-                    and len(delta.get("tool_calls")) > 0
-                ):
+                elif delta.get("tool_calls", None) is not None and len(delta.get("tool_calls")) > 0:
                     curr = completion["tool_calls"]
                     tool_call_chunk = getattr(
                         delta.get("tool_calls", None)[0], "function", None