Implements Polymorphic Templates (#504)

* implemented polymorphic templates

switch to using types of dynamic values instead of parameters

* fixed tests

Author: kiranandcode

docs/source/agent.py

@@ -41,12 +41,12 @@ def wrapper(self, *args, **kwargs):
 
     def _format_model_input(self, template, other, *args, **kwargs):
         # update prompt with previous list of messages
-        prompt = fwd()
+        prompt, enc_ty = fwd()
         if Agent.current_agent() is self:
             assert self is other
             self.state.extend(prompt)
             prompt = self.state
-        return prompt
+        return (prompt, enc_ty)
 
     def _compute_response(self, *args, **kwargs):
         # save response into persisted state

effectful/handlers/llm/completions.py

@@ -21,7 +21,11 @@
 from litellm.types.utils import ModelResponse
 
 from effectful.handlers.llm import Template, Tool
-from effectful.handlers.llm.encoding import type_to_encodable_type
+from effectful.handlers.llm.encoding import (
+    Encodable,
+    type_to_encodable_type,
+)
+from effectful.internals.unification import TypeExpression, nested_type
 from effectful.ops.semantics import fwd, handler
 from effectful.ops.syntax import ObjectInterpretation, defop, implements
 from effectful.ops.types import Operation
@@ -213,16 +217,18 @@ def call_with_json_args(
 
 
 @defop
-def compute_response(template: Template, model_input: list[Any]) -> ModelResponse:
+def compute_response[**P, T](
+    template: Template[P, T], model_input_pair: tuple[list[Any], Encodable[T]]
+) -> ModelResponse:
     """Produce a complete model response for an input message sequence. This may
     involve multiple API requests if tools are invoked by the model.
 
     """
-    ret_type = template.__signature__.return_annotation
+    model_input, ret_type_encoder = model_input_pair
     tools = template.tools
 
     tool_schemas = [function_definition(t) for t in tools.values()]
-    response_encoding_type: type | None = type_to_encodable_type(ret_type).t
+    response_encoding_type: type | None = ret_type_encoder.t
     if response_encoding_type == str:
         response_encoding_type = None
 
@@ -262,7 +268,12 @@ def compute_response(template: Template, model_input: list[Any]) -> ModelRespons
             )
 
 
-def decode_response[**P, T](template: Callable[P, T], response: ModelResponse) -> T:
+type ModelInput[T] = tuple[list[Any], Encodable[T]]
+
+
+def decode_response[**P, T](
+    template: Callable[P, T], model_input_pair: ModelInput[T], response: ModelResponse
+) -> T:
     """Decode an LLM response into an instance of the template return type. This
     operation should raise if the output cannot be decoded.
     """
@@ -273,8 +284,7 @@ def decode_response[**P, T](template: Callable[P, T], response: ModelResponse) -
     result_str = last_resp.content or last_resp.reasoning_content
     assert result_str
 
-    ret_type = template.__signature__.return_annotation
-    encodable_ty = type_to_encodable_type(ret_type)
+    encodable_ty = model_input_pair[1]
 
     if encodable_ty.t == str:
         # if encoding as a type, value is just directly what the llm returned
@@ -291,28 +301,29 @@ def decode_response[**P, T](template: Callable[P, T], response: ModelResponse) -
 @defop
 def format_model_input[**P, T](
     template: Template[P, T], *args: P.args, **kwargs: P.kwargs
-) -> list[Any]:
+) -> ModelInput[T]:
     """Format a template applied to arguments into a sequence of input
     messages.
 
     """
-    bound_args = template.__signature__.bind(*args, **kwargs)
+    signature = template.__signature__
+
+    bound_args = signature.bind(*args, **kwargs)
     bound_args.apply_defaults()
     # encode arguments
     arguments = {}
-    for param in bound_args.arguments:
-        encoder = type_to_encodable_type(
-            template.__signature__.parameters[param].annotation
-        )
-        encoded = encoder.encode(bound_args.arguments[param], template.__context__)
+    for param, value in bound_args.arguments.items():
+        ty: TypeExpression = nested_type(value).value
+        encoder = type_to_encodable_type(ty)  # type: ignore
+        encoded = encoder.encode(value, template.__context__)
         arguments[param] = encoder.serialize(encoded)
 
     prompt = _OpenAIPromptFormatter().format_as_messages(
         template.__prompt_template__, **arguments
     )
 
     # install prefix if the return type has a return annotation
-    ret_type = template.__signature__.return_annotation
+    ret_type = template.__type_rule__(*args, **kwargs)
     origin = typing.get_origin(ret_type)
     ret_type = ret_type if origin is None else origin
     ret_type_encoder = type_to_encodable_type(ret_type)
@@ -327,7 +338,7 @@ def format_model_input[**P, T](
     # Note: The OpenAI api only seems to accept images in the 'user' role. The
     # effect of different roles on the model's response is currently unclear.
     messages = [{"type": "message", "content": prompt, "role": "user"}]
-    return messages
+    return (messages, ret_type_encoder)
 
 
 class InstructionHandler(ObjectInterpretation):
@@ -347,12 +358,16 @@ def __init__(self, instruction: str):
         self.instruction = instruction
 
     @implements(format_model_input)
-    def _inject_instruction(self, template: Template, *args, **kwargs) -> list[Any]:
+    def _inject_instruction(
+        self, template: Template, *args, **kwargs
+    ) -> tuple[list[Any], Any]:
         """Append instruction message to the formatted model input."""
-        messages = fwd()
-        return messages + [
-            {"type": "message", "content": self.instruction, "role": "user"}
-        ]
+        (messages, decode_ty) = fwd()
+        return (
+            messages
+            + [{"type": "message", "content": self.instruction, "role": "user"}],
+            decode_ty,
+        )
 
 
 class RetryLLMHandler(ObjectInterpretation):
@@ -421,4 +436,4 @@ def _call[**P, T](
     ) -> T:
         model_input = format_model_input(template, *args, **kwargs)
         resp = compute_response(template, model_input)
-        return decode_response(template, resp)
+        return decode_response(template, model_input, resp)

tests/fixtures/tests_test_handlers_llm_provider.py__test_polymorphic_templates[model_input0].json

@@ -0,0 +1,44 @@
+{
+  "choices": [
+    {
+      "finish_reason": "stop",
+      "index": 0,
+      "message": {
+        "annotations": [],
+        "content": "{\"value\":3}",
+        "function_call": null,
+        "provider_specific_fields": {
+          "refusal": null
+        },
+        "role": "assistant",
+        "tool_calls": null
+      },
+      "provider_specific_fields": {}
+    }
+  ],
+  "created": 1769197760,
+  "id": "chatcmpl-D1HRgEJGEJucGvOOBa4PSeUiXRhDs",
+  "model": "gpt-4o-2024-08-06",
+  "object": "chat.completion",
+  "service_tier": "default",
+  "system_fingerprint": "fp_cbf1785567",
+  "usage": {
+    "completion_tokens": 9,
+    "completion_tokens_details": {
+      "accepted_prediction_tokens": 0,
+      "audio_tokens": 0,
+      "image_tokens": null,
+      "reasoning_tokens": 0,
+      "rejected_prediction_tokens": 0,
+      "text_tokens": null
+    },
+    "prompt_tokens": 378,
+    "prompt_tokens_details": {
+      "audio_tokens": 0,
+      "cached_tokens": 0,
+      "image_tokens": null,
+      "text_tokens": null
+    },
+    "total_tokens": 387
+  }
+}

tests/fixtures/tests_test_handlers_llm_provider.py__test_polymorphic_templates[model_input1].json

@@ -0,0 +1,44 @@
+{
+  "choices": [
+    {
+      "finish_reason": "stop",
+      "index": 0,
+      "message": {
+        "annotations": [],
+        "content": "Here is another instance of one of the items in the list: \"banana\".",
+        "function_call": null,
+        "provider_specific_fields": {
+          "refusal": null
+        },
+        "role": "assistant",
+        "tool_calls": null
+      },
+      "provider_specific_fields": {}
+    }
+  ],
+  "created": 1769197761,
+  "id": "chatcmpl-D1HRhDz8ooFuBQH2N51NmO4zn0g8r",
+  "model": "gpt-4o-2024-08-06",
+  "object": "chat.completion",
+  "service_tier": "default",
+  "system_fingerprint": "fp_deacdd5f6f",
+  "usage": {
+    "completion_tokens": 17,
+    "completion_tokens_details": {
+      "accepted_prediction_tokens": 0,
+      "audio_tokens": 0,
+      "image_tokens": null,
+      "reasoning_tokens": 0,
+      "rejected_prediction_tokens": 0,
+      "text_tokens": null
+    },
+    "prompt_tokens": 271,
+    "prompt_tokens_details": {
+      "audio_tokens": 0,
+      "cached_tokens": 0,
+      "image_tokens": null,
+      "text_tokens": null
+    },
+    "total_tokens": 288
+  }
+}

tests/fixtures/tests_test_handlers_llm_provider.py__test_polymorphic_templates[model_input2].json

@@ -0,0 +1,44 @@
+{
+  "choices": [
+    {
+      "finish_reason": "stop",
+      "index": 0,
+      "message": {
+        "annotations": [],
+        "content": "{\"value\":{\"title\":\"Brave New World\",\"author\":\"Aldous Huxley\"}}",
+        "function_call": null,
+        "provider_specific_fields": {
+          "refusal": null
+        },
+        "role": "assistant",
+        "tool_calls": null
+      },
+      "provider_specific_fields": {}
+    }
+  ],
+  "created": 1769197762,
+  "id": "chatcmpl-D1HRi9XFRJF8gmevEAZ7xh0hxjgki",
+  "model": "gpt-4o-2024-08-06",
+  "object": "chat.completion",
+  "service_tier": "default",
+  "system_fingerprint": "fp_cbf1785567",
+  "usage": {
+    "completion_tokens": 23,
+    "completion_tokens_details": {
+      "accepted_prediction_tokens": 0,
+      "audio_tokens": 0,
+      "image_tokens": null,
+      "reasoning_tokens": 0,
+      "rejected_prediction_tokens": 0,
+      "text_tokens": null
+    },
+    "prompt_tokens": 373,
+    "prompt_tokens_details": {
+      "audio_tokens": 0,
+      "cached_tokens": 0,
+      "image_tokens": null,
+      "text_tokens": null
+    },
+    "total_tokens": 396
+  }
+}

tests/test_handlers_llm_provider.py

@@ -8,7 +8,7 @@
 import json
 import logging
 import os
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from enum import Enum
 from pathlib import Path
 
@@ -346,6 +346,39 @@ def review_book(plot: str) -> BookReview:
     raise NotImplementedError
 
 
+@Template.define
+def generate_similar[T](examples: Sequence[T]) -> T:
+    "Generate another instance of one of the items in {examples}"
+    raise NotHandled
+
+
+@dataclass
+class Book:
+    title: str
+    author: str
+
+
+@requires_openai
+@pytest.mark.parametrize(
+    "model_input",
+    [
+        [1, 2, 3, 4],
+        ["apple", "orange", "pear"],
+        [
+            Book("1984", "George Orwell"),
+            Book("To Kill a Mockingbird", "Harper Lee"),
+            Book("The Great Gatsby", "F. Scott Fitzgerald"),
+        ],
+    ],
+)
+def test_polymorphic_templates(request, model_input):
+    with (
+        handler(ReplayLiteLLMProvider(request, model_name="gpt-4o")),
+        handler(LimitLLMCallsHandler(max_calls=5)),
+    ):
+        assert isinstance(generate_similar(model_input), type(model_input[0]))
+
+
 class TestPydanticBaseModelReturn:
     @requires_openai
     def test_pydantic_basemodel_return(self, request):

tests/test_handlers_llm_template.py

@@ -155,7 +155,7 @@ class TemplateStringIntp(ObjectInterpretation):
     def _[**P, T](
         self, template: Template[P, T], *args: P.args, **kwargs: P.kwargs
     ) -> T:
-        model_input = format_model_input(template, *args, **kwargs)
+        model_input, _ = format_model_input(template, *args, **kwargs)
         template_result = model_input[0]["content"]
         assert len(template_result) == 1
         return template_result[0]["text"]