Zentry as an Agentic Tool

pip install Zentryai pydantic openai-agents

export Zentry_API_KEY="your_Zentry_api_key"

import os
os.environ["Zentry_API_KEY"] = "your_Zentry_api_key"

from __future__ import annotations
import os
import asyncio
from pydantic import BaseModel
try:
    from Zentry import AsyncMemoryClient
except ImportError:
    raise ImportError("Zentry is not installed. Please install it using 'pip install Zentryai'.")
from agents import (
    Agent,
    ItemHelpers,
    MessageOutputItem,
    RunContextWrapper,
    Runner,
    ToolCallItem,
    ToolCallOutputItem,
    TResponseInputItem,
    function_tool,
)

class ZentryContext(BaseModel):
    user_id: str | None = None

client = AsyncMemoryClient(api_key=os.getenv("Zentry_API_KEY"))

@function_tool
async def add_to_memory(
    context: RunContextWrapper[ZentryContext],
    content: str,
) -> str:
    """
    Add a message to Zentry
    Args:
        content: The content to store in memory.
    """
    messages = [{"role": "user", "content": content}]
    user_id = context.context.user_id or "default_user"
    await client.add(messages, user_id=user_id)
    return f"Stored message: {content}"

@function_tool
async def search_memory(
    context: RunContextWrapper[ZentryContext],
    query: str,
) -> str:
    """
    Search for memories in Zentry
    Args:
        query: The search query.
    """
    user_id = context.context.user_id or "default_user"
    memories = await client.search(query, user_id=user_id, output_format="v1.1")
    results = '\n'.join([result["memory"] for result in memories["results"]])
    return str(results)

@function_tool
async def get_all_memory(
    context: RunContextWrapper[ZentryContext],
) -> str:
    """Retrieve all memories from Zentry"""
    user_id = context.context.user_id or "default_user"
    memories = await client.get_all(user_id=user_id, output_format="v1.1")
    results = '\n'.join([result["memory"] for result in memories["results"]])
    return str(results)

memory_agent = Agent[ZentryContext](
    name="Memory Assistant",
    instructions="""You are a helpful assistant with memory capabilities. You can:
    1. Store new information using add_to_memory
    2. Search existing information using search_memory
    3. Retrieve all stored information using get_all_memory
    When users ask questions:
    - If they want to store information, use add_to_memory
    - If they're searching for specific information, use search_memory
    - If they want to see everything stored, use get_all_memory""",
    tools=[add_to_memory, search_memory, get_all_memory],
)

async def main():
    current_agent: Agent[ZentryContext] = memory_agent
    input_items: list[TResponseInputItem] = []
    context = ZentryContext()
    while True:
        user_input = input("Enter your message (or 'quit' to exit): ")
        if user_input.lower() == 'quit':
            break
        input_items.append({"content": user_input, "role": "user"})
        result = await Runner.run(current_agent, input_items, context=context)
        for new_item in result.new_items:
            agent_name = new_item.agent.name
            if isinstance(new_item, MessageOutputItem):
                print(f"{agent_name}: {ItemHelpers.text_message_output(new_item)}")
            elif isinstance(new_item, ToolCallItem):
                print(f"{agent_name}: Calling a tool")
            elif isinstance(new_item, ToolCallOutputItem):
                print(f"{agent_name}: Tool call output: {new_item.output}")
            else:
                print(f"{agent_name}: Skipping item: {new_item.__class__.__name__}")
        input_items = result.to_input_list()

if __name__ == "__main__":
    asyncio.run(main())

User: Remember that my favorite color is blue
Agent: Calling a tool
Agent: Tool call output: Stored message: my favorite color is blue
Agent: I've stored that your favorite color is blue in my memory. I'll remember that for future conversations.

User: What's my favorite color?
Agent: Calling a tool
Agent: Tool call output: my favorite color is blue
Agent: Your favorite color is blue, based on what you've told me earlier.

User: What do you know about me?
Agent: Calling a tool
Agent: Tool call output: favorite color is blue
my birthday is on March 15
Agent: Based on our previous conversations, I know that:
1. Your favorite color is blue
2. Your birthday is on March 15

context = ZentryContext(user_id="user123")