happy mypy

Author: nimarb

langfuse/_client/constants.py

@@ -3,7 +3,7 @@
 This module defines constants used throughout the Langfuse OpenTelemetry integration.
 """
 
-from typing import Literal, List, get_args, Union
+from typing import Literal, List, get_args, Union, Any
 from typing_extensions import TypeAlias
 
 LANGFUSE_TRACER_NAME = "langfuse-sdk"
@@ -39,16 +39,12 @@
 
 
 def get_observation_types_list(
-    literal_type: Union[
-        ObservationTypeGenerationLike,
-        ObservationTypeLiteralNoEvent,
-        ObservationTypeLiteral,
-    ],
+    literal_type: Any,
 ) -> List[str]:
     """Flattens the Literal type to provide a list of strings.
 
     Args:
-        literal_type: A Literal type or union of Literals to flatten
+        literal_type: A Literal type, TypeAlias, or union of Literals to flatten
 
     Returns:
         Flat list of all string values contained in the Literal type

langfuse/_client/observe.py

@@ -16,7 +16,6 @@
     Union,
     cast,
     overload,
-    get_args,
 )
 
 from opentelemetry.util._decorator import _AgnosticContextManager
@@ -28,6 +27,7 @@
 
 from langfuse._client.constants import (
     ObservationTypeLiteralNoEvent,
+    get_observation_types_list,
 )
 from langfuse._client.get_client import _set_current_public_key, get_client
 from langfuse._client.span import (
@@ -169,7 +169,7 @@ def sub_process():
             - For async functions, the decorator returns an async function wrapper.
             - For sync functions, the decorator returns a synchronous wrapper.
         """
-        valid_types = set(get_args(ObservationTypeLiteralNoEvent))
+        valid_types = set(get_observation_types_list(ObservationTypeLiteralNoEvent))
         if as_type is not None and as_type not in valid_types:
             self._log.warning(
                 f"Invalid as_type '{as_type}'. Valid types are: {', '.join(sorted(valid_types))}. Defaulting to 'span'."