Message return packaging created

tests/unit/tools/summary/summary_test.py

@@ -0,0 +1,33 @@
+from pydantic_ai import Agent
+
+from app.core.pii import PiiScrubber, PiiAuditStore
+from tools.summary.summary import ConversationSummaryTool
+
+
+async def main() -> None:
+    agent = Agent("google-vertex:gemini-2.5-flash")
+
+    # Generate a real conversation so we get genuine ModelMessage objects.
+    conversation = await agent.run(
+        "My name is John Smith. "
+        "My phone number is 555-123-4567. "
+        "Orphan from Morocco with family at 123 Main St, Columbus, Ohio"
+    )
+
+    tool = ConversationSummaryTool(
+        agent=agent,
+        scrubber=PiiScrubber(),
+        audit_store=PiiAuditStore(),
+    )
+
+    summary = await tool.run(conversation.all_messages())
+
+    print("\nSummary:")
+    print(summary)
+
+
+import asyncio
+
+if __name__ == "__main__":
+    print("Calling main")
+    asyncio.run(main())

tools/summary/summary.py

@@ -0,0 +1,68 @@
+from collections.abc import Sequence
+
+from pydantic_ai import Agent
+from pydantic_ai.messages import ModelMessage
+
+from app.core.pii import PiiScrubber, PiiAuditStore
+
+
+class ConversationSummaryTool:
+    def __init__(
+        self,
+        agent: Agent,
+        scrubber: PiiScrubber,
+        audit_store: PiiAuditStore,
+    ):
+        self._agent = agent
+        self._scrubber = scrubber
+        self._audit_store = audit_store
+
+    async def run(
+        self,
+        messages: Sequence[ModelMessage],
+    ) -> str:
+        transcript = self._build_transcript(messages)
+
+        scrubbed = self._scrubber.scrub(
+            transcript,
+            audit_store=self._audit_store,
+        )
+
+        return await self._summarize(scrubbed)
+
+    def _build_transcript(
+        self,
+        messages: Sequence[ModelMessage],
+    ) -> str:
+        return "\n".join(
+            message.model_dump_json(exclude_none=True)
+            for message in messages
+        )
+
+    async def _summarize(
+        self,
+        scrubbed_conversation: str,
+    ) -> str:
+        prompt = f"""
+            Summarize the following conversation history.
+
+            Requirements:
+            - The conversation has already been scrubbed of PII.
+            - Include:
+                • User goals
+                • Important decisions
+                • Completed work
+                • Remaining work
+                • Important tool calls, if relevant
+            - Maximum 200 words.
+            - Do not speculate.
+            - Return plain text.
+
+            Conversation:
+
+            {scrubbed_conversation}
+        """
+
+        result = await self._agent.run(prompt)
+
+        return result.output

tools/summary_tool.py

@@ -0,0 +1,22 @@
+from collections.abc import Sequence
+
+from summary.summary import ConversationSummaryTool
+from pydantic_ai import Agent, ModelMessage
+from app.core.pii import PiiScrubber, PiiAuditStore
+
+
+async def handle(value: Sequence[ModelMessage]) -> str:
+    """Return a summary of the user's conversation"""
+    """Parameters : ModelMessages for complete conversation"""
+    agent = Agent("google-vertex:gemini-2.5-flash")
+
+    tool = ConversationSummaryTool(
+        agent=agent,
+        scrubber=PiiScrubber(),
+        audit_store=PiiAuditStore(),
+    )
+
+    summary = await tool.run(value)
+
+    return summary
+

tools/summary_tool.yaml

@@ -0,0 +1,13 @@
+# tools/summary_tool.yaml
+name: summary_tool
+description: "Tool for a complete summary of the user's conversation"
+input_schema:
+  type: object
+  properties:
+    value:
+      type: string
+      description: A tool used to summarize the conversation and pass back a brief summary to the user."
+  required: []
+implementation:
+  type: function
+  handler: tools.summary_tool.handle