cleanup

Author: nimarb

langfuse/_client/client.py

@@ -707,13 +707,11 @@ def _get_span_class(
             return LangfuseRetriever
         elif as_type == "embedding":
             return LangfuseEmbedding
-        elif as_type in ("generation"):
+        elif as_type == "generation":
             return LangfuseGeneration
-        elif as_type in ("event"):
+        elif as_type == "event":
             return LangfuseEvent
-        elif as_type in ("span"):
-            return LangfuseSpan
-        else:
+        elif as_type == "span":
             return LangfuseSpan
 
     @_agnosticcontextmanager
@@ -825,6 +823,7 @@ def _start_as_current_otel_span_with_processed_media(
                     }
                 )
             elif span_class in [
+                LangfuseSpan,
                 LangfuseAgent,
                 LangfuseTool,
                 LangfuseChain,
@@ -833,9 +832,6 @@ def _start_as_current_otel_span_with_processed_media(
             ]:
                 # set their type internally in the class
                 pass
-            else:
-                if as_type is not None:
-                    common_args["as_type"] = as_type
 
             yield span_class(**common_args)
 

langfuse/_client/observe.py

@@ -190,13 +190,9 @@ def sub_process():
             raise ValueError(
                 f"Invalid observation type '{type}'. Valid types are: {', '.join(sorted(VALID_OBSERVATION_TYPES))}"
             )
-        if as_type is not None and as_type.upper() not in VALID_OBSERVATION_TYPES:
-            valid_values = sorted(
-                list(VALID_OBSERVATION_TYPES)
-                + [t.lower() for t in VALID_OBSERVATION_TYPES]
-            )
+        if as_type is not None and as_type not in VALID_OBSERVATION_TYPES:
             raise ValueError(
-                f"Invalid as_type '{as_type}'. Valid values are: {', '.join(valid_values)}"
+                f"Invalid as_type '{as_type}'. Valid types are: {', '.join(sorted(VALID_OBSERVATION_TYPES))}"
             )
 
         function_io_capture_enabled = os.environ.get(

tests/test_manual_graph_instrumentation.py

@@ -66,7 +66,7 @@ def end_agent():
     agent_observations = [
         obs
         for obs in all_observations
-        if obs.type in ["agent", "tool", "retriever", "chain", "embedding"]
+        if obs.type in ["agent", "tool", "retriever", "chain", "embedding", "AGENT", "TOOL", "RETRIEVER", "CHAIN", "EMBEDDING", "GENERATION", "generation"]
     ]
 
     assert (