from mirix.agent import AgentWrapper

# Initialize agent with configuration
agent = AgentWrapper("./mirix/configs/mirix.yaml")

# Send basic text information
agent.send_message(
    message="The moon now has a president.",
    memorizing=True,
    force_absorb_content=True
)

from mirix import Mirix

# Initialize memory agent (defaults to Google Gemini 2.0 Flash)
memory_agent = Mirix(api_key="your-google-api-key")

# Add memories
memory_agent.add("The moon now has a president")
memory_agent.add("John loves Italian food and is allergic to peanuts")

# Chat with memory context
response = memory_agent.chat("Does the moon have a president?")
print(response)  # "Yes, according to my memory, the moon has a president."

response = memory_agent.chat("What does John like to eat?") 
print(response)  # "John loves Italian food. However, he's allergic to peanuts."

#!/usr/bin/env python3
import os
import asyncio
from claude_agent_sdk import query, ClaudeAgentOptions, AssistantMessage
from mirix import Mirix
from collections import deque
from dotenv import load_dotenv
load_dotenv()

# Configuration
MEMORY_UPDATE_INTERVAL = 3  # Update memory every N turns
REINIT_INTERVAL = 1  # Rebuild system prompt (retrieve from Mirix) every N turns
KEEP_LAST_N_TURNS = 50  # Keep last N turns in memory buffer

def build_system_prompt(mirix_agent=None, user_id=None, conversation_buffer=""):
    """Build system prompt with optional Mirix memory context"""
    system_prompt = """You are a helpful assistant."""
    
    # Add Mirix memory context if available
    if mirix_agent and user_id and conversation_buffer:
        memory_context = mirix_agent.extract_memory_for_system_prompt(
            conversation_buffer, user_id
        )
        if memory_context:
            system_prompt += "\n\nRelevant Memory Context:\n" + memory_context
    
    return system_prompt

async def run_agent():
    """Run Claude Agent SDK with Mirix memory integration"""
    
    # Initialize Mirix memory agent
    mirix_agent = Mirix(
        model_name="gemini-2.0-flash",
        api_key=os.getenv("GEMINI_API_KEY"),
    )
    user = mirix_agent.create_user(user_name="Alice")
    
    # Track conversation for memory updates
    conversation_history = deque(maxlen=KEEP_LAST_N_TURNS)
    turn_count = 0
    turns_since_reinit = 0
    session_id = None
    
    while True:
        # Build system prompt with memory context
        options = ClaudeAgentOptions(
            resume=session_id,
            allowed_tools=["Task", "Bash", "Read", "Edit", "Write", "WebSearch"],
            system_prompt=build_system_prompt(
                mirix_agent, user.id, 
                "\n".join([f"User: {u}\nAssistant: {a}" for u, a in conversation_history])
            ),
            model="claude-sonnet-4-5",
            max_turns=50
        )
        
        user_input = input("User: ").strip()
        if user_input.lower() in ['exit', 'quit', 'bye']:
            break
        
        # Get Claude's response
        assistant_response = ""
        async for message in query(prompt=user_input, options=options):
            if hasattr(message, 'subtype') and message.subtype == 'init':
                session_id = message.data.get('session_id')
            
            if isinstance(message, AssistantMessage):
                for block in message.content:
                    if hasattr(block, 'text'):
                        assistant_response += block.text
                        print(block.text, flush=True)
        
        # Update conversation history
        conversation_history.append((user_input, assistant_response))
        turn_count += 1
        turns_since_reinit += 1
        
        # Periodically update Mirix memory
        if turn_count % MEMORY_UPDATE_INTERVAL == 0:
            combined = "\n".join([
                f"[User] {u}\n[Assistant] {a}" 
                for u, a in conversation_history
            ])
            await asyncio.to_thread(mirix_agent.add, combined, user_id=user.id)
            print("✅ Memory updated!")

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