diff --git a/tests/unit/tools/summary/__init__.py b/tests/unit/tools/summary/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/tools/summary/summary_test.py b/tests/unit/tools/summary/summary_test.py new file mode 100644 index 0000000..a7e87c2 --- /dev/null +++ b/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()) \ No newline at end of file diff --git a/tools/summary/__init__.py b/tools/summary/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tools/summary/summary.py b/tools/summary/summary.py new file mode 100644 index 0000000..21cdd5e --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tools/summary_tool.py b/tools/summary_tool.py new file mode 100644 index 0000000..036f9f4 --- /dev/null +++ b/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 + diff --git a/tools/summary_tool.yaml b/tools/summary_tool.yaml new file mode 100644 index 0000000..52bca64 --- /dev/null +++ b/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 \ No newline at end of file