fix lint

Author: nimarb

langfuse/_client/client.py

@@ -366,7 +366,7 @@ def start_as_current_span(
                 "retriever",
                 "embedding",
             ]
-        ],
+        ] = None,
     ) -> _AgnosticContextManager[LangfuseSpan]:
         """Create a new span and set it as the current span in a context manager.
 
@@ -695,9 +695,8 @@ def _get_span_class(
             "retriever",
             "embedding",
         ],
-    ):
+    ) -> type:
         """Get the appropriate span class based on as_type."""
-        # TODO: make it case insensitive
         if as_type == "agent":
             return LangfuseAgent
         elif as_type == "tool":

langfuse/_client/observe.py

@@ -303,9 +303,6 @@ async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any:
             with _set_current_public_key(public_key):
                 langfuse_client = get_client(public_key=public_key)
 
-                # Use consolidated as_type parameter
-                final_obs_type = as_type
-
                 context_manager: Optional[
                     Union[
                         _AgnosticContextManager[LangfuseGeneration],
@@ -319,13 +316,13 @@ async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any:
                             input=input,
                             end_on_exit=False,  # when returning a generator, closing on exit would be to early
                         )
-                        if final_obs_type in ("generation", "GENERATION")
+                        if as_type == "generation"
                         else langfuse_client.start_as_current_span(
                             name=final_name,
                             trace_context=trace_context,
                             input=input,
                             end_on_exit=False,  # when returning a generator, closing on exit would be to early
-                            as_type=final_obs_type,
+                            as_type=as_type,
                         )
                     )
                     if langfuse_client
@@ -413,9 +410,6 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
             with _set_current_public_key(public_key):
                 langfuse_client = get_client(public_key=public_key)
 
-                # Use consolidated as_type parameter
-                final_obs_type = as_type
-
                 context_manager: Optional[
                     Union[
                         _AgnosticContextManager[LangfuseGeneration],
@@ -429,13 +423,13 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                             input=input,
                             end_on_exit=False,  # when returning a generator, closing on exit would be to early
                         )
-                        if final_obs_type in ("generation", "GENERATION")
+                        if as_type == "generation"
                         else langfuse_client.start_as_current_span(
                             name=final_name,
                             trace_context=trace_context,
                             input=input,
                             end_on_exit=False,  # when returning a generator, closing on exit would be to early
-                            as_type=final_obs_type,
+                            as_type=as_type,
                         )
                     )
                     if langfuse_client

langfuse/_client/span.py

@@ -1333,7 +1333,7 @@ def update(
 class LangfuseAgent(LangfuseGraphObservation):
     """Specialized span for agent observations in agentic workflows."""
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs) -> None:
         """Initialize a new LangfuseAgent span."""
         kwargs["observation_type"] = "agent"
         super().__init__(**kwargs)
@@ -1342,7 +1342,7 @@ def __init__(self, **kwargs):
 class LangfuseTool(LangfuseGraphObservation):
     """Specialized span for tool observations in agentic workflows."""
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs) -> None:
         """Initialize a new LangfuseTool span."""
         kwargs["observation_type"] = "tool"
         super().__init__(**kwargs)
@@ -1351,7 +1351,7 @@ def __init__(self, **kwargs):
 class LangfuseChain(LangfuseGraphObservation):
     """Specialized span for chain observations in agentic workflows."""
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs) -> None:
         """Initialize a new LangfuseChain span."""
         kwargs["observation_type"] = "chain"
         super().__init__(**kwargs)
@@ -1360,7 +1360,7 @@ def __init__(self, **kwargs):
 class LangfuseRetriever(LangfuseGraphObservation):
     """Specialized span for retriever observations in agentic workflows."""
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs) -> None:
         """Initialize a new LangfuseRetriever span."""
         kwargs["observation_type"] = "retriever"
         super().__init__(**kwargs)
@@ -1369,7 +1369,7 @@ def __init__(self, **kwargs):
 class LangfuseEmbedding(LangfuseGraphObservation):
     """Specialized span for embedding observations in agentic workflows."""
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs) -> None:
         """Initialize a new LangfuseEmbedding span."""
         kwargs["observation_type"] = "embedding"
         super().__init__(**kwargs)

tests/test_manual_graph_instrumentation.py

@@ -1,7 +1,6 @@
 import time
 
 from langfuse import Langfuse, observe
-from tests.api_wrapper import LangfuseAPI
 from tests.utils import create_uuid, get_api
 
 
@@ -35,15 +34,15 @@ def end_agent():
     # Run the workflow within a trace context
     with langfuse.start_as_current_span(
         name="agent_workflow", as_type="agent"
-    ) as root_span:
+    ):
         langfuse.update_current_trace(name=trace_name)
 
         start_result = start_agent()
         process_result = process_agent()
         tool_result = tool_call()
         end_result = end_agent()
 
-        workflow_result = {
+        _workflow_result = {
             "start": start_result,
             "process": process_result,
             "tool": tool_result,
@@ -122,7 +121,7 @@ def end_agent():
     # Execute the parallel workflow
     with langfuse.start_as_current_span(
         name="parallel_workflow", as_type="span"
-    ) as root_span:
+    ):
         langfuse.update_current_trace(name=trace_name)
         start_result = start_agent()
 
@@ -143,7 +142,7 @@ def end_agent():
 
         end_result = end_agent()
 
-        workflow_result = {
+        _workflow_result = {
             "start": start_result,
             "tools": tool_results,
             "end": end_result,
@@ -171,8 +170,6 @@ def end_agent():
         for obs in all_observations
         if obs.type in ["agent", "tool", "AGENT", "TOOL"]
     ]
-    print(all_observations)
-    print(graph_observations)
 
     # Should have: start_agent (1) + 3 tools (3) + end_agent (1) = 5 total
     expected_count = 5
@@ -194,6 +191,3 @@ def end_agent():
     ), f"Expected 3 tool observations, got {len(tool_observations)}"
 
 
-if __name__ == "__main__":
-    test_observe_type_agent_instrumentation()
-    test_observe_type_parallel_tool_execution()