Adding mypy check and test

.github/workflows/test_llm.yml

@@ -31,4 +31,4 @@ jobs:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
         run: |
-          uv run pytest tests/test_handlers_llm_provider.py -v --tb=short
+          uv run pytest tests/test_handlers_llm_provider*.py -v --tb=short

effectful/handlers/llm/encoding.py

@@ -37,7 +37,13 @@ def encode(cls, vl: T) -> U:
 
     @classmethod
     @abstractmethod
-    def decode(cls, vl: U) -> T:
+    def decode(cls, vl: U, template: typing.Any = None) -> T:
+        """Decode an encoded value back to the original type.
+
+        Args:
+            vl: The encoded value
+            template: Optional Template providing context (e.g., lexical scope)
+        """
         pass
 
     @classmethod
@@ -67,7 +73,7 @@ def encode(cls, vl: T) -> T:
             return vl
 
         @classmethod
-        def decode(cls, vl: T) -> T:
+        def decode(cls, vl: T, template: typing.Any = None) -> T:
             return vl
 
     return typing.cast(Encodable[T], BaseEncodable())
@@ -81,7 +87,7 @@ class EncodablePydanticBaseModel(EncodableAs[T, T]):
         t: type[T] = ty
 
         @classmethod
-        def decode(cls, vl: T) -> T:
+        def decode(cls, vl: T, template: typing.Any = None) -> T:
             return vl
 
         @classmethod
@@ -107,7 +113,9 @@ def encode(cls, image: Image.Image) -> ChatCompletionImageUrlObject:
         }
 
     @classmethod
-    def decode(cls, image: ChatCompletionImageUrlObject) -> Image.Image:
+    def decode(
+        cls, image: ChatCompletionImageUrlObject, template: typing.Any = None
+    ) -> Image.Image:
         image_url = image["url"]
         if not image_url.startswith("data:image/"):
             raise RuntimeError(
@@ -156,13 +164,14 @@ def encode(cls, t: T) -> typing.Any:
             return tuple([enc.encode(elem) for enc, elem in zip(element_encoders, t)])
 
         @classmethod
-        def decode(cls, t: typing.Any) -> T:
+        def decode(cls, t: typing.Any, template: typing.Any = None) -> T:
             if len(t) != len(element_encoders):
                 raise ValueError(
                     f"tuple length {len(t)} does not match expected length {len(element_encoders)}"
                 )
             decoded_elements: list[typing.Any] = [
-                enc.decode(elem) for enc, elem in zip(element_encoders, t)
+                enc.decode(elem, template=template)
+                for enc, elem in zip(element_encoders, t)
             ]
             return typing.cast(T, tuple(decoded_elements))
 
@@ -217,9 +226,9 @@ def encode(cls, t: T) -> typing.Any:
             return [element_encoder.encode(elem) for elem in t]
 
         @classmethod
-        def decode(cls, t: typing.Any) -> T:
+        def decode(cls, t: typing.Any, template: typing.Any = None) -> T:
             decoded_elements: list[typing.Any] = [
-                element_encoder.decode(elem) for elem in t
+                element_encoder.decode(elem, template=template) for elem in t
             ]
             return typing.cast(T, decoded_elements)
 

effectful/handlers/llm/providers.py

@@ -234,6 +234,23 @@ def tool_call[T](
     return tool(*args, **kwargs)
 
 
+@defop
+def type_check[T](value: T, template: Template) -> T:
+    """Validate that a value conforms to the template's expected return type.
+
+    Default behavior is pass-through (no validation).
+    Handlers can intercept this to perform type checking (e.g., mypy for Callables).
+
+    Args:
+        value: The value to validate
+        template: The template providing expected type and context
+
+    Returns:
+        The value unchanged if validation passes
+    """
+    return value
+
+
 class CacheLLMRequestHandler(ObjectInterpretation):
     """Caches LLM requests."""
 
@@ -437,7 +454,7 @@ def decode_response[**P, T](template: Callable[P, T], response: ModelResponse) -
         assert isinstance(result, Result)
         value = result.value  # type: ignore
 
-    return encodable_ty.decode(value)  # type: ignore
+    return encodable_ty.decode(value, template=template)  # type: ignore
 
 
 @defop
@@ -446,7 +463,6 @@ def format_model_input[**P, T](
 ) -> list[Any]:
     """Format a template applied to arguments into a sequence of input
     messages.
-
     """
     bound_args = template.__signature__.bind(*args, **kwargs)
     bound_args.apply_defaults()
@@ -489,4 +505,6 @@ def _call[**P, T](
     ) -> T:
         model_input = format_model_input(template, *args, **kwargs)
         resp = compute_response(template, model_input)
-        return decode_response(template, resp)
+        result = decode_response(template, resp)
+        # Validate the result against the expected return type
+        return type_check(result, template)