Adds context parameter to encode and decode calls (closes #499)

effectful/handlers/llm/completions.py

@@ -5,7 +5,7 @@
 import string
 import traceback
 import typing
-from collections.abc import Callable, Hashable
+from collections.abc import Callable, Hashable, Mapping
 from typing import Any
 
 import litellm
@@ -176,7 +176,9 @@ def function_definition(tool: Tool) -> OpenAIChatCompletionToolParam:
     }
 
 
-def call_with_json_args(tool: Tool, json_str: str) -> OpenAIMessageContent:
+def call_with_json_args(
+    tool: Tool, context: Mapping[str, Any], json_str: str
+) -> OpenAIMessageContent:
     """Implements a roundtrip call to a python function. Input is a json
     string representing an LLM tool call request parameters. The output is
     the serialised response to the model.
@@ -192,7 +194,7 @@ def call_with_json_args(tool: Tool, json_str: str) -> OpenAIMessageContent:
         params: dict[str, Any] = {
             param_name: type_to_encodable_type(
                 sig.parameters[param_name].annotation
-            ).decode(getattr(raw_args, param_name))
+            ).decode(getattr(raw_args, param_name), context)
             for param_name in raw_args.model_fields_set
         }
 
@@ -201,7 +203,7 @@ def call_with_json_args(tool: Tool, json_str: str) -> OpenAIMessageContent:
 
         # serialize back to U using encoder for return type
         encoded_ty = type_to_encodable_type(sig.return_annotation)
-        encoded_value = encoded_ty.encode(result)
+        encoded_value = encoded_ty.encode(result, context)
 
         # serialise back to Json
         return encoded_ty.serialize(encoded_value)
@@ -246,7 +248,9 @@ def compute_response(template: Template, model_input: list[Any]) -> ModelRespons
             function_name = function.name
             assert function_name is not None
             tool = tools[function_name]
-            tool_result = call_with_json_args(tool, function.arguments)
+            tool_result = call_with_json_args(
+                tool, template.__context__, function.arguments
+            )
             model_input.append(
                 {
                     "role": "tool",
@@ -280,7 +284,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.__context__)  # type: ignore
 
 
 @defop
@@ -299,7 +303,7 @@ def format_model_input[**P, T](
         encoder = type_to_encodable_type(
             template.__signature__.parameters[param].annotation
         )
-        encoded = encoder.encode(bound_args.arguments[param])
+        encoded = encoder.encode(bound_args.arguments[param], template.__context__)
         arguments[param] = encoder.serialize(encoded)
 
     prompt = _OpenAIPromptFormatter().format_as_messages(

effectful/handlers/llm/encoding.py

@@ -2,7 +2,8 @@
 import io
 import typing
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
+from typing import Any
 
 import pydantic
 from litellm import (
@@ -33,12 +34,12 @@ def __init__(self, *args, **kwargs):
 
     @classmethod
     @abstractmethod
-    def encode(cls, vl: T) -> U:
+    def encode(cls, vl: T, env: Mapping[str, Any] | None = None) -> U:
         pass
 
     @classmethod
     @abstractmethod
-    def decode(cls, vl: U) -> T:
+    def decode(cls, vl: U, env: Mapping[str, Any] | None = None) -> T:
         pass
 
     @classmethod
@@ -64,11 +65,11 @@ class BaseEncodable(EncodableAs[T, T]):
         t: type[T] = ty
 
         @classmethod
-        def encode(cls, vl: T) -> T:
+        def encode(cls, vl: T, env: Mapping[str, Any] | None = None) -> T:
             return vl
 
         @classmethod
-        def decode(cls, vl: T) -> T:
+        def decode(cls, vl: T, env: Mapping[str, Any] | None = None) -> T:
             return vl
 
     return typing.cast(Encodable[T], BaseEncodable())
@@ -92,16 +93,16 @@ class EncodablePydanticBaseModel(EncodableAs[T, T]):
         t: type[T] = ty
 
         @classmethod
-        def decode(cls, vl: T) -> T:
+        def decode(cls, vl: T, env: Mapping[str, Any] | None = None) -> T:
             return vl
 
         @classmethod
-        def encode(cls, vl: T) -> T:
+        def encode(cls, vl: T, env: Mapping[str, Any] | None = None) -> T:
             return vl
 
         @classmethod
-        def serialize(cls, vl: T) -> list[OpenAIMessageContentListBlock]:
-            return [{"type": "text", "text": vl.model_dump_json()}]
+        def serialize(cls, value: T) -> list[OpenAIMessageContentListBlock]:
+            return [{"type": "text", "text": value.model_dump_json()}]
 
     return typing.cast(Encodable[T], EncodablePydanticBaseModel())
 
@@ -111,14 +112,18 @@ class EncodableImage(EncodableAs[Image.Image, ChatCompletionImageUrlObject]):
     t = ChatCompletionImageUrlObject
 
     @classmethod
-    def encode(cls, image: Image.Image) -> ChatCompletionImageUrlObject:
+    def encode(
+        cls, image: Image.Image, env: Mapping[str, Any] | None = None
+    ) -> ChatCompletionImageUrlObject:
         return {
             "detail": "auto",
             "url": _pil_image_to_base64_data_uri(image),
         }
 
     @classmethod
-    def decode(cls, image: ChatCompletionImageUrlObject) -> Image.Image:
+    def decode(
+        cls, image: ChatCompletionImageUrlObject, env: Mapping[str, Any] | None = None
+    ) -> Image.Image:
         image_url = image["url"]
         if not image_url.startswith("data:image/"):
             raise RuntimeError(
@@ -157,23 +162,27 @@ class TupleEncodable(EncodableAs[T, typing.Any]):
         t: type[typing.Any] = encoded_ty
 
         @classmethod
-        def encode(cls, t: T) -> typing.Any:
+        def encode(
+            cls, t: T, env: typing.Mapping[str, Any] | None = None
+        ) -> typing.Any:
             if not isinstance(t, tuple):
                 raise TypeError(f"Expected tuple, got {type(t)}")
             if len(t) != len(element_encoders):
                 raise ValueError(
                     f"Tuple length {len(t)} does not match expected length {len(element_encoders)}"
                 )
-            return tuple([enc.encode(elem) for enc, elem in zip(element_encoders, t)])
+            return tuple(
+                [enc.encode(elem, env) for enc, elem in zip(element_encoders, t)]
+            )
 
         @classmethod
-        def decode(cls, t: typing.Any) -> T:
+        def decode(cls, t: typing.Any, env: Mapping[str, Any] | None = 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, env) for enc, elem in zip(element_encoders, t)
             ]
             return typing.cast(T, tuple(decoded_elements))
 
@@ -222,15 +231,15 @@ class ListEncodable(EncodableAs[T, typing.Any]):
         t: type[typing.Any] = encoded_ty
 
         @classmethod
-        def encode(cls, t: T) -> typing.Any:
+        def encode(cls, t: T, env: Mapping[str, Any] | None = None) -> typing.Any:
             if not isinstance(t, list):
                 raise TypeError(f"Expected list, got {type(t)}")
-            return [element_encoder.encode(elem) for elem in t]
+            return [element_encoder.encode(elem, env) for elem in t]
 
         @classmethod
-        def decode(cls, t: typing.Any) -> T:
+        def decode(cls, t: typing.Any, env: Mapping[str, Any] | None = None) -> T:
             decoded_elements: list[typing.Any] = [
-                element_encoder.decode(elem) for elem in t
+                element_encoder.decode(elem, env) for elem in t
             ]
             return typing.cast(T, decoded_elements)