improve typing

Author: hassiebp

langfuse/_client/client.py

@@ -29,7 +29,10 @@
 from opentelemetry import trace
 from opentelemetry import trace as otel_trace_api
 from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
-from opentelemetry.util._decorator import _agnosticcontextmanager
+from opentelemetry.util._decorator import (
+    _AgnosticContextManager,
+    _agnosticcontextmanager,
+)
 
 from langfuse._client.attributes import (
     LangfuseOtelSpanAttributes,
@@ -310,7 +313,6 @@ def start_span(
             metadata=metadata,
         )
 
-    # TODO: add return typing
     def start_as_current_span(
         self,
         *,
@@ -323,7 +325,7 @@ def start_as_current_span(
         level: Optional[SpanLevel] = None,
         status_message: Optional[str] = None,
         end_on_exit: Optional[bool] = None,
-    ):
+    ) -> _AgnosticContextManager[LangfuseSpan]:
         """Create a new span and set it as the current span in a context manager.
 
         This method creates a new span and sets it as the current span within a context
@@ -529,7 +531,7 @@ def start_as_current_generation(
         cost_details: Optional[Dict[str, float]] = None,
         prompt: Optional[PromptClient] = None,
         end_on_exit: Optional[bool] = None,
-    ):
+    ) -> _AgnosticContextManager[LangfuseGeneration]:
         """Create a new generation span and set it as the current span in a context manager.
 
         This method creates a specialized span for AI model generations and sets it as the
@@ -954,7 +956,7 @@ def _is_valid_span_id(self, span_id):
 
         return bool(re.match(pattern, span_id))
 
-    def create_observation_id(self, *, seed: Optional[str] = None) -> str:
+    def _create_observation_id(self, *, seed: Optional[str] = None) -> str:
         """Create a unique observation ID for use with Langfuse.
 
         This method generates a unique observation ID (span ID in OpenTelemetry terms)
@@ -1167,7 +1169,7 @@ def create_score(
         if not self.tracing_enabled:
             return
 
-        score_id = score_id or self.create_observation_id()
+        score_id = score_id or self._create_observation_id()
 
         try:
             score_event = {

langfuse/_client/span.py

@@ -28,6 +28,7 @@
 )
 
 from opentelemetry import trace as otel_trace_api
+from opentelemetry.util._decorator import _AgnosticContextManager
 
 from langfuse.model import PromptClient
 
@@ -654,7 +655,7 @@ def start_as_current_span(
         version: Optional[str] = None,
         level: Optional[SpanLevel] = None,
         status_message: Optional[str] = None,
-    ):
+    ) -> _AgnosticContextManager["LangfuseSpan"]:
         """Create a new child span and set it as the current span in a context manager.
 
         This method creates a new child span and sets it as the current span within
@@ -698,15 +699,18 @@ def start_as_current_span(
             status_message=status_message,
         )
 
-        return self._langfuse_client._create_span_with_parent_context(
-            name=name,
-            attributes=attributes,
-            as_type="span",
-            remote_parent_span=None,
-            parent=self._otel_span,
-            input=input,
-            output=output,
-            metadata=metadata,
+        return cast(
+            _AgnosticContextManager["LangfuseSpan"],
+            self._langfuse_client._create_span_with_parent_context(
+                name=name,
+                attributes=attributes,
+                as_type="span",
+                remote_parent_span=None,
+                parent=self._otel_span,
+                input=input,
+                output=output,
+                metadata=metadata,
+            ),
         )
 
     def start_generation(
@@ -832,7 +836,7 @@ def start_as_current_generation(
         usage_details: Optional[Dict[str, int]] = None,
         cost_details: Optional[Dict[str, float]] = None,
         prompt: Optional[PromptClient] = None,
-    ):
+    ) -> _AgnosticContextManager["LangfuseGeneration"]:
         """Create a new child generation span and set it as the current span in a context manager.
 
         This method creates a new child generation span and sets it as the current span
@@ -900,15 +904,18 @@ def start_as_current_generation(
             prompt=prompt,
         )
 
-        return self._langfuse_client._create_span_with_parent_context(
-            name=name,
-            attributes=attributes,
-            as_type="generation",
-            remote_parent_span=None,
-            parent=self._otel_span,
-            input=input,
-            output=output,
-            metadata=metadata,
+        return cast(
+            _AgnosticContextManager["LangfuseGeneration"],
+            self._langfuse_client._create_span_with_parent_context(
+                name=name,
+                attributes=attributes,
+                as_type="generation",
+                remote_parent_span=None,
+                parent=self._otel_span,
+                input=input,
+                output=output,
+                metadata=metadata,
+            ),
         )