Build session-based agent tool loop

tests/test_injected_context.py

@@ -0,0 +1,26 @@
+import unittest
+from typing import Annotated
+
+from micro_agents import Injected, OpenAIBasedLLM, SessionContext, Tool, ToolCallMessage, ToolManager
+
+
+class InjectedContextTests(unittest.IsolatedAsyncioTestCase):
+    async def test_context_is_injected_but_not_exposed_to_the_llm(self):
+        def who_am_i(text: str, context: Annotated[SessionContext, Injected]) -> str:
+            return f"{context['user_id'].__str__()}: {text}"
+
+        tool = Tool(who_am_i)
+        schema = OpenAIBasedLLM._infer_parameters_schema(who_am_i)
+        response = await ToolManager([tool]).execute(
+            ToolCallMessage("call-1", "who_am_i", {"text": "hello"}),
+            {"user_id": 42},
+        )
+
+        properties = schema["properties"]
+        assert isinstance(properties, dict)
+        self.assertEqual({"text"}, set(properties))
+        self.assertEqual(["text"], schema["required"])
+        self.assertEqual("42: hello", response.response.text())
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_runtime.py

@@ -0,0 +1,187 @@
+import unittest
+
+from micro_agents import (
+    Aggregator,
+    Completed,
+    DynamicToolManager,
+    LLM,
+    LlmMessage,
+    Media,
+    Tool,
+    ToolCallFailed,
+    ToolCallMessage,
+    ToolCallResponseMessage,
+    ToolManager,
+    UserMessage,
+)
+from micro_agents.conversation import ReasoningMessage
+from micro_agents.llm import OpenAIBasedLLM
+
+
+class ScriptedLLM(LLM):
+    def __init__(self, responses, after_first_call=None):
+        self.responses = list(responses)
+        self.after_first_call = after_first_call
+        self.instructions = []
+        self.tools = []
+        self.thinking = []
+
+    async def ask(self, conversation, tools, *, instructions=None, thinking=None):
+        self.instructions.append(instructions)
+        self.tools.append([tool.name for tool in tools])
+        self.thinking.append(thinking)
+        if len(self.instructions) == 1 and self.after_first_call is not None:
+            self.after_first_call()
+        return self.responses.pop(0)
+
+
+class ContextualToolManager(ToolManager):
+    async def is_allowed(self, tool, context):
+        return bool(context.get("can_echo"))
+
+
+class SessionRuntimeTests(unittest.IsolatedAsyncioTestCase):
+    async def test_master_prompt_is_rebuilt_on_each_llm_round(self):
+        async def echo(text: str):
+            return {"echo": text}
+
+        tool = Tool(echo, name="echo")
+        llm = ScriptedLLM(
+            [
+                [ToolCallMessage("call-1", "echo", {"text": "hello"})],
+                [LlmMessage("done")],
+            ]
+        )
+        contexts = []
+
+        def prompt(prompt_context):
+            contexts.append(prompt_context)
+            return f"user={prompt_context['user_id']}"
+
+        context = {"user_id": 42, "can_echo": True}
+        session = Aggregator(llm, ToolManager([tool]), master_prompt=prompt).create_session(context=context)
+
+        result = await session.run("start")
+
+        self.assertTrue(result.succeeded)
+        final_message = result.final_message
+        assert final_message is not None
+        self.assertEqual("done", final_message.message())
+        self.assertEqual(["user=42", "user=42"], llm.instructions)
+        self.assertEqual([context, context], contexts)
+        responses = [m for m in session.conversation.messages if isinstance(m, ToolCallResponseMessage)]
+        self.assertEqual('{"echo":"hello"}', responses[0].text())
+
+    async def test_permission_is_checked_again_before_execution(self):
+        def restricted():
+            raise AssertionError("A denied tool must not execute")
+
+        context = {"can_echo": True}
+        llm = ScriptedLLM(
+            [
+                [ToolCallMessage("call-1", "restricted", {})],
+                [LlmMessage("I cannot access that tool")],
+            ],
+            after_first_call=lambda: context.update(can_echo=False),
+        )
+        session = Aggregator(llm, ContextualToolManager([Tool(restricted, name="restricted")])).create_session(
+            context=context
+        )
+
+        events = [event async for event in session.stream("do it")]
+        result = next(event.result for event in events if isinstance(event, Completed))
+
+        self.assertTrue(result.succeeded)
+        self.assertEqual([['restricted'], []], llm.tools)
+        response = next(m for m in session.conversation.messages if isinstance(m, ToolCallResponseMessage))
+        self.assertTrue(response.is_error())
+        self.assertEqual("tool_not_allowed", response.error_code())
+        self.assertTrue(any(isinstance(event, ToolCallFailed) for event in events))
+
+    async def test_unknown_tool_becomes_a_recoverable_tool_response(self):
+        llm = ScriptedLLM(
+            [
+                [ToolCallMessage("call-1", "missing", {})],
+                [LlmMessage("I do not have that capability")],
+            ]
+        )
+        session = Aggregator(llm).create_session()
+
+        result = await session.run("use missing")
+
+        self.assertTrue(result.succeeded)
+        response = next(m for m in session.conversation.messages if isinstance(m, ToolCallResponseMessage))
+        self.assertEqual("tool_not_found", response.error_code())
+
+    async def test_dynamic_tool_manager_updates_running_sessions(self):
+        manager = DynamicToolManager()
+        session = Aggregator(
+            ScriptedLLM(
+                [
+                    [ToolCallMessage("call-1", "echo", {"text": "hello"})],
+                    [LlmMessage("done")],
+                ]
+            ),
+            manager,
+        ).create_session()
+        manager.add_tool(Tool(lambda text: text, name="echo"))
+
+        result = await session.run("start")
+
+        self.assertTrue(result.succeeded)
+        manager.remove_tool("echo")
+        self.assertEqual((), manager.registered_tools())
+
+    def test_dynamic_tool_manager_requires_explicit_replace(self):
+        manager = DynamicToolManager([Tool(lambda: "first", name="status")])
+
+        with self.assertRaises(ValueError):
+            manager.add_tool(Tool(lambda: "second", name="status"))
+
+        replacement = Tool(lambda: "second", name="status")
+        manager.add_tool(replacement, replace=True)
+        self.assertEqual((replacement,), manager.registered_tools())
+
+    def test_openai_reasoning_state_is_retained_for_the_next_request(self):
+        raw_reasoning = {"type": "reasoning", "id": "rs_123", "encrypted_content": "opaque"}
+        message = ReasoningMessage("", provider_item=raw_reasoning)
+
+        self.assertEqual(raw_reasoning, OpenAIBasedLLM._serialize_reasoning_message(message))
+
+    async def test_session_can_override_llm_and_thinking(self):
+        default_llm = ScriptedLLM([[LlmMessage("default")]])
+        session_llm = ScriptedLLM([[LlmMessage("override")]])
+        session = Aggregator(default_llm, thinking=False).create_session(llm=session_llm, thinking=True)
+
+        result = await session.run("hello")
+
+        self.assertEqual("override", result.final_message.message() if result.final_message else None)
+        self.assertEqual([True], session_llm.thinking)
+        self.assertEqual([], default_llm.thinking)
+
+    async def test_session_accepts_media(self):
+        llm = ScriptedLLM([[LlmMessage("seen")]])
+        session = Aggregator(llm).create_session()
+        media = Media("image", data=b"image", mime_type="image/png", filename="image.png")
+
+        await session.run("describe", media=[media])
+
+        user_message = session.conversation.messages[0]
+        self.assertIsInstance(user_message, UserMessage)
+        assert isinstance(user_message, UserMessage)
+        self.assertEqual([media], user_message.media())
+
+    def test_openai_serializes_inline_image(self):
+        media = Media("image", data=b"image", mime_type="image/png", filename="image.png")
+
+        serialized = OpenAIBasedLLM._serialize_user_message(UserMessage("describe", [media]))
+        content = serialized["content"]
+        assert isinstance(content, list)
+
+        self.assertEqual("user", serialized["role"])
+        self.assertEqual("input_text", content[0]["type"])
+        self.assertEqual("data:image/png;base64,aW1hZ2U=", content[1].get("image_url"))
+
+
+if __name__ == "__main__":
+    unittest.main()