agent.py

"""RSS Feed Reader Agent - PydanticAI Agent with MCP Server Integration."""

import argparse
import asyncio
import os
import traceback
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

import logfire
from dotenv import load_dotenv
from pydantic_ai import Agent
from pydantic_ai.agent import AgentRunResult
from pydantic_ai.mcp import MCPServerStdio
from pydantic_ai.messages import (
    ModelMessage,
    SystemPromptPart,
    ToolCallPart,
    ToolReturnPart,
    UserPromptPart,
)
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider

load_dotenv()

logfire.configure(token=os.getenv("LOGFIRE_API_KEY"))
logfire.instrument_openai()


# ============================================================================
# Command Line Arguments
# ============================================================================


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="RSS Feed Reader Agent - Discover and read RSS/Atom feeds"
    )
    parser.add_argument(
        "--model",
        type=str,
        default="anthropic/claude-3.7-sonnet",
        help="Model identifier to use with OpenRouter (default: anthropic/claude-3.7-sonnet)",
    )
    return parser.parse_args()


# ============================================================================
# Model Setup
# ============================================================================

args = parse_args()

API_KEY = os.getenv("OPENROUTER_API_KEY")
if not API_KEY:
    raise ValueError("OPENROUTER_API_KEY environment variable is required")

model = OpenAIModel(
    args.model,
    provider=OpenAIProvider(base_url="https://openrouter.ai/api/v1", api_key=API_KEY),
)

# ============================================================================
# MCP Server Setup
# ============================================================================

# No special environment variables needed for RSS client
env = {}

mcp_servers = [
    MCPServerStdio("python", ["./mcp_server.py"], env=env),
]


# ============================================================================
# Message History Filtering
# ============================================================================


def filtered_message_history(
    result: Optional[AgentRunResult],
    limit: Optional[int] = None,
    include_tool_messages: bool = True,
) -> Optional[List[Dict[str, Any]]]:
    """Filter and limit the message history from an AgentRunResult.

    Args:
        result: The AgentRunResult object with message history
        limit: Optional int, if provided returns only system message + last N messages
        include_tool_messages: Whether to include tool messages in the history

    Returns:
        Filtered list of messages in the format expected by the agent
    """
    if result is None:
        return None

    # Get all messages
    messages: list[ModelMessage] = result.all_messages()

    # Extract system message (always the first one with role="system")
    system_message = next(
        (msg for msg in messages if type(msg.parts[0]) == SystemPromptPart), None
    )

    # Filter non-system messages
    non_system_messages = [
        msg for msg in messages if type(msg.parts[0]) != SystemPromptPart
    ]

    # Apply tool message filtering if requested
    if not include_tool_messages:
        non_system_messages = [
            msg
            for msg in non_system_messages
            if not any(
                isinstance(part, ToolCallPart) or isinstance(part, ToolReturnPart)
                for part in msg.parts
            )
        ]

    # Find the most recent UserPromptPart before applying limit
    latest_user_prompt_part = None
    latest_user_prompt_index = -1
    for i, msg in enumerate(non_system_messages):
        for part in msg.parts:
            if isinstance(part, UserPromptPart):
                latest_user_prompt_part = part
                latest_user_prompt_index = i

    # Apply limit if specified, but ensure paired tool calls and returns stay together
    if limit is not None and limit > 0:
        # Identify tool call IDs and their corresponding return parts
        tool_call_ids = {}
        tool_return_ids = set()

        for i, msg in enumerate(non_system_messages):
            for part in msg.parts:
                if isinstance(part, ToolCallPart):
                    tool_call_ids[part.tool_call_id] = i
                elif isinstance(part, ToolReturnPart):
                    tool_return_ids.add(part.tool_call_id)

        # Take the last 'limit' messages but ensure we include paired messages
        if len(non_system_messages) > limit:
            included_indices = set(
                range(len(non_system_messages) - limit, len(non_system_messages))
            )

            # Include any missing tool call messages for tool returns that are included
            for i, msg in enumerate(non_system_messages):
                if i in included_indices:
                    for part in msg.parts:
                        if (
                            isinstance(part, ToolReturnPart)
                            and part.tool_call_id in tool_call_ids
                        ):
                            included_indices.add(tool_call_ids[part.tool_call_id])

            # Check if the latest UserPromptPart would be excluded by the limit
            if (
                latest_user_prompt_index >= 0
                and latest_user_prompt_index not in included_indices
                and latest_user_prompt_part is not None
                and system_message is not None
            ):
                # Find if system_message already has a UserPromptPart
                user_prompt_index = next(
                    (
                        i
                        for i, part in enumerate(system_message.parts)
                        if isinstance(part, UserPromptPart)
                    ),
                    None,
                )

                if user_prompt_index is not None:
                    # Replace existing UserPromptPart
                    system_message.parts[user_prompt_index] = latest_user_prompt_part
                else:
                    # Add new UserPromptPart to system message
                    system_message.parts.append(latest_user_prompt_part)

            # Create a new list with only the included messages
            non_system_messages = [
                msg for i, msg in enumerate(non_system_messages) if i in included_indices
            ]

    # Combine system message with other messages
    result_messages = []
    if system_message:
        result_messages.append(system_message)
    result_messages.extend(non_system_messages)

    return result_messages


# ============================================================================
# Agent Setup
# ============================================================================

agent_name = "rss"


def load_agent_prompt(agent: str) -> str:
    """Load the agent system prompt and replace time_now variable.

    Args:
        agent: Name of the agent (matches the .md filename)

    Returns:
        The agent system prompt with time_now replaced
    """
    print(f"Loading {agent} agent...")
    time_now = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
    prompt_path = f"./agents/{agent}.md"

    if not os.path.exists(prompt_path):
        raise FileNotFoundError(f"Agent prompt file not found: {prompt_path}")

    with open(prompt_path, "r") as f:
        agent_prompt = f.read()

    agent_prompt = agent_prompt.replace("{time_now}", time_now)
    return agent_prompt


# Load the agent system prompt
agent_prompt = load_agent_prompt(agent_name)

# Display the selected model
print(f"Using model: {args.model}")

# Create the agent
agent = Agent(model, mcp_servers=mcp_servers, system_prompt=agent_prompt)


# ============================================================================
# Main CLI Loop
# ============================================================================


async def main():
    """CLI testing in a conversation with the agent."""
    async with agent.run_mcp_servers():
        result: AgentRunResult = None

        print("\n" + "=" * 60)
        print("RSS Feed Reader Agent")
        print("=" * 60)
        print("Type your message or 'exit' to quit\n")

        # Chat Loop
        while True:
            if result:
                print(f"\n{result.output}")

            user_input = input("\n> ")

            if user_input.lower() in ["exit", "quit", "q"]:
                print("\nGoodbye!")
                break

            err = None
            for i in range(0, 2):
                try:
                    # Use the filtered message history
                    result: AgentRunResult = await agent.run(
                        user_input,
                        message_history=filtered_message_history(
                            result,
                            limit=24,  # Last 24 non-system messages
                            include_tool_messages=True,  # Include tool messages
                        ),
                    )
                    break
                except Exception as e:
                    err = e
                    traceback.print_exc()
                    await asyncio.sleep(2)

            if result is None:
                print(f"\nError: {err}. Try again...\n")
                continue
            elif len(result.output) == 0:
                continue


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