chore: added context rendering in the templates

Author: hendrixmar

py/packages/genkit/src/genkit/blocks/prompt.py

@@ -192,7 +192,7 @@ async def __call__(
         """
         return await generate_action(
             self._registry,
-            await self.render(input=input, config=config),
+            await self.render(input=input, config=config, context=context),
             on_chunk=on_chunk,
             middleware=self._use,
             context=context if context else ActionRunContext._current_context(),
@@ -236,6 +236,7 @@ async def render(
         self,
         input: dict[str, Any] | None = None,
         config: GenerationCommonConfig | dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
     ) -> GenerateActionOptions:
         """Renders the prompt with the given input and configuration.
 
@@ -277,18 +278,26 @@ async def render(
                 input,
                 options,
                 self._cache_prompt,
-                ActionRunContext._current_context() or {}
+                context
             )
             resolved_msgs.append(result)
         if options.messages:
-            resolved_msgs += options.messages
+            resolved_msgs.extend(
+                await render_message_prompt(
+                    self._registry,
+                    input,
+                    options,
+                    self._cache_prompt,
+                    context
+                )
+            )
         if options.prompt:
             result = await render_user_prompt(
                 self._registry,
                 input,
                 options,
                 self._cache_prompt,
-                ActionRunContext._current_context() or {}
+                context
             )
             resolved_msgs.append(result)
 
@@ -437,14 +446,7 @@ async def to_generate_action_options(
         raise Exception('No model configured.')
     resolved_msgs: list[Message] = []
     if options.system:
-        result = await render_system_prompt(
-            registry,
-            None,
-            options,
-            PromptCache(),
-            ActionRunContext._current_context() or {}
-        )
-        resolved_msgs.append(result)
+        resolved_msgs.append(Message(role=Role.SYSTEM, content=_normalize_prompt_arg(options.system)))
     if options.messages:
         resolved_msgs += options.messages
     if options.prompt:
@@ -551,10 +553,9 @@ async def render_system_prompt(
                 context,
                 prompt_cache.system,
                 input,
-
                 PromptMetadata(
                     input=PromptInputConfig(
-
+                        schema=options.input_schema,
                     )
                 ),
             )
@@ -565,6 +566,7 @@ async def render_system_prompt(
         content=_normalize_prompt_arg(options.system)
     )
 
+
 async def render_dotprompt_to_parts(
     context: dict[str, Any],
     prompt_function: PromptFunction,
@@ -585,9 +587,10 @@ async def render_dotprompt_to_parts(
     Raises:
         Exception: If the template produces more than one message.
     """
+    merged_input = input_
     rendered = await prompt_function(
         data=DataArgument[dict[str, Any]](
-            input=input_,
+            input=merged_input,
             context=context,
         ),
         options=options,
@@ -606,6 +609,50 @@ async def render_dotprompt_to_parts(
 
 
 
+async def render_message_prompt(
+    registry: Registry,
+    input: dict[str, Any],
+    options: PromptConfig,
+    prompt_cache: PromptCache,
+    context: dict[str, Any] | None = None,
+) -> list[Message]:
+    """Renders the user prompt for a prompt action."""
+    if isinstance(options.messages, str):
+        if prompt_cache.messages is None:
+            prompt_cache.messages = await registry.dotprompt.compile(options.messages)
+
+        if options.metadata:
+            context = {**context, "state": options.metadata.get("state")}
+
+        messages_ = None
+        if isinstance(options.messages, list):
+            messages_ = [e.model_dump() for e in options.messages]
+
+        rendered = await prompt_cache.messages(
+            data=DataArgument[dict[str, Any]](
+                input=input,
+                context=context,
+                messages=messages_,
+            ),
+            options=PromptMetadata(
+                input=PromptInputConfig(
+                )
+            ),
+        )
+        return [Message.model_validate(e.model_dump()) for e in rendered.messages]
+
+    elif isinstance(options.messages, list):
+        return options.messages
+
+    return [
+        Message(
+            role=Role.USER,
+            content=_normalize_prompt_arg(options.prompt)
+        )
+    ]
+
+
+
 async def render_user_prompt(
     registry: Registry,
     input: dict[str, Any],
@@ -615,17 +662,17 @@ async def render_user_prompt(
 ) -> Message:
     """Renders the user prompt for a prompt action."""
     if isinstance(options.prompt, str):
-        if prompt_cache.prompt is None:
-            prompt_cache.prompt = await registry.dotprompt.compile(options.prompt)
+        if prompt_cache.user_prompt is None:
+            prompt_cache.user_prompt = await registry.dotprompt.compile(options.prompt)
 
         if options.metadata:
             context = {**context, "state": options.metadata.get("state")}
 
         return Message(
-            role=Role.SYSTEM,
+            role=Role.USER,
             content=await render_dotprompt_to_parts(
                 context,
-                prompt_cache.system,
+                prompt_cache.user_prompt,
                 input,
                 PromptMetadata(
                     input=PromptInputConfig(

py/packages/genkit/tests/genkit/blocks/prompt_test.py

@@ -124,7 +124,7 @@ def test_tool(input: ToolInput):
         tools=['testTool'],
         tool_choice=ToolChoice.REQUIRED,
         max_turns=5,
-        input_schema=PromptInput,
+        input_schema=PromptInput.model_json_schema(),
         output_constrained=True,
         output_format='json',
         description='a prompt descr',
@@ -143,7 +143,7 @@ def test_tool(input: ToolInput):
 
 test_cases_parse_partial_json = [
     (
-        "renders user prompt",
+        "renders system prompt",
         {
             "model": "echoModel",
             "config": {"banana": "ripe"},
@@ -158,28 +158,68 @@ def test_tool(input: ToolInput):
         },
         {"name": "foo"},
         GenerationCommonConfig.model_validate({"temperature": 11}),
+        {},
         """[ECHO] system: "hello foo (bar)" {"temperature":11.0}"""
+    ),
+    (
+        "renders user prompt",
+        {
+            "model": "echoModel",
+            "config": {"banana": "ripe"},
+            "input_schema": {
+                'type': 'object',
+                'properties': {
+                    'name': {'type': 'string'},
+                },
+            },  # Note: Schema representation might need adjustment
+            "prompt": "hello {{name}} ({{@state.name}})",
+            "metadata": {"state": {"name": "bar_system"}}
+        },
+        {"name": "foo"},
+        GenerationCommonConfig.model_validate({"temperature": 11}),
+        {},
+        """[ECHO] user: "hello foo (bar_system)" {"temperature":11.0}"""
+    ),
+    (
+        "renders user prompt with context",
+        {
+            "model": "echoModel",
+            "config": {"banana": "ripe"},
+            "input_schema": {
+                'type': 'object',
+                'properties': {
+                    'name': {'type': 'string'},
+                },
+            },  # Note: Schema representation might need adjustment
+            "prompt": "hello {{name}} ({{@state.name}}, {{@auth.email}})",
+            "metadata": {"state": {"name": "bar"}}
+        },
+        {"name": "foo"},
+        GenerationCommonConfig.model_validate({"temperature": 11}),
+        { "auth": { "email": '[email protected]' } },
+        """[ECHO] user: "hello foo (bar, [email protected])" {"temperature":11.0}"""
     )
 ]
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    'test_case, prompt, input, input_option, want_rendered',
+    'test_case, prompt, input, input_option, context, want_rendered',
     test_cases_parse_partial_json,
     ids=[tc[0] for tc in test_cases_parse_partial_json],
 )
-async def test_prompt_with_system(
+async def test_prompt_rendering_dotprompt(
     test_case: str,
     prompt: dict[str, Any],
     input: dict[str, Any],
     input_option: GenerationCommonConfig,
+    context: dict[str, Any],
     want_rendered: str
 ) -> None:
     """Test system prompt rendering."""
     ai, *_ = setup_test()
 
     my_prompt = ai.define_prompt(**prompt)
 
-    response = await my_prompt(input, input_option)
+    response = await my_prompt(input, input_option, context=context)
 
     assert response.text == want_rendered