Skip to main content

A lightweight, extensible framework for building LLM agents with Model Context Protocol (MCP) support

Project description

Pocket-Agent

Pocket Agent

A lightweight, extensible framework for building LLM agents with Model Context Protocol (MCP) support

License


Why Pocket Agent?

Most agent frameworks are severely over-bloated. The reason for this is that they are trying to support too many things at once and make every possible agent implementation "simple". This only works until it doesn't and you are stuck having to understand the enormous code base to implement what should be a very simple feature.

Pocket Agent takes the opposite approach by handling only the basic functions of an LLM agent and working with the MCP protocol. That way you don't give up any flexibility when building your agent but a lot of the lower level implementation details are taken care of.

Design Principles

🚀 Lightweight & Simple

  • Minimal dependencies - just fastmcp and litellm
  • Clean abstractions that separate agent logic from MCP client details
  • < 500 lines of code

🎯 Developer-Friendly

  • Abstract base class design for easy extension
  • Clear separation of concerns between agents and clients
  • Built-in logging and event system

🌐 Multi-Model Support

  • Works with any endpoint supported by LiteLLM without requiring code changes
  • Easy model switching and configuration

💡 Extensible

  • Use any custom logging implementation
  • Easily integrate custom frontends using the built-in event system
  • Easily create fully custom agent implementations

🧑‍🍳 Cookbook

Refer to the Cookbook to find example implementations and try out PocketAgent without any implementation overhead

Installation

Install with uv (Recommended):

uv add pocket-agent

Install with pip:

pip install pocket-agent

Creating Your First Pocket-Agent (Quick Start)

To build a Pocket-Agent, all you need to implement is the agent's run method:

class SimpleAgent(PocketAgent):
    async def run(self):
        """Simple conversation loop"""

        while True:
            # Accept user message
            user_input = input("Your input: ")
            if user_input.lower() == 'quit':
                break
                
            # Add user message
            await self.add_user_message(user_input)
            
             # Generates response and executes any tool calls
            step_result = await self.step()
            while step_result["llm_message"].tool_calls is not None:
                step_result = await self.step()
    
        return {"status": "completed"}

To run the agent, you only need to pass your JSON MCP config and your agent configuration:

mcp_config = {
    "mcpServers": {
        "weather": {
            "transport": "stdio",
            "command": "python",
            "args": ["server.py"],
            "cwd": os.path.dirname(os.path.abspath(__file__))
        }
    }
}
# Configure agent  
config = AgentConfig(
    llm_model="gpt-5-nano",
    system_prompt="You are a helpful assistant who answers user questions and uses provided tools when applicable"
)
# Create and run agent
agent = SimpleAgent(
    agent_config=config,
    mcp_config=mcp_config
)

await agent.run()

Core Concepts

🏗️ PocketAgent Base Class

The PocketAgent is an abstract base class that provides the foundation for building custom agents. You inherit from this class and implement the run() method to define your agent's behavior.

from pocket_agent import PocketAgent, AgentConfig

class MyAgent(PocketAgent):
    async def run(self):
        # Your agent logic here
        return {"status": "completed"}

PocketAgent Parameters:

agent = PocketAgent(
    agent_config,   # Required: Instance of the AgentConfig class
    mcp_config,     # Required: JSON MCP server configuration to pass tools to the agent
    router,         # Optional: A litellm router to manage llm rate limits
    logger,         # Optional: A logger instance to capture logs
    hooks,          # Optional: Instance of AgentHooks to optionally define custom behavior at common junction points
    **client_kwargs # Optional: additional kwargs passed to the PocketAgentClient

)

AgentConfig Parameters:

config = AgentConfig(
    llm_model="gpt-4",                    # Required: LLM model to use
    system_prompt="You are helpful...",   # Optional: System prompt for the agent
    agent_id="my-agent-123",              # Optional: Custom context ID
    allow_images=False,                   # Optional: Enable image input support (default: False)
    messages=[],                          # Optional: Initial conversation history (default: [])
    completion_kwargs={                   # Optional: Additional LLM parameters (default: {"tool_choice": "auto"})
        "tool_choice": "auto",
        "temperature": 0.7
    }
)

🔄 The Step Method

The step() method is the core execution unit that:

  1. Gets an LLM response with available tools
  2. Executes any tool calls in parallel
  3. Updates conversation history

The output of calling the step() method is the StepResult

@dataclass
class StepResult:
    llm_message: LitellmMessage                                 # The message generated by the llm including str content, tool calls, images, etc.
    tool_execution_results: Optional[list[ToolResult]] = None   # Results of any executed tools 
# Single step execution
step_result = await agent.step()

# continue until no more tool calls
while step_result.llm_message.tool_calls is not None:
    step_result = await agent.step()

Step Result Structure:

{
    "llm_message": LitellmMessage,           # The LLM response
    "tool_execution_results": [ToolResult]   # Results from tool calls (if any)
}

💬 Message Management

Pocket Agent automatically adds llm generated messages and tool result messages in the step() function. Input provided by a user can easily be managed using add_user_message() and should be done before calling the step() method:

class Agent(PocketAgent)
    async def run(self):
        # Add user messages (with optional images)
        await agent.add_user_message("Hello!", image_base64s=["base64_image_data"])
        await self.step()

# Clear all messages except the system promp `reset_messages` function
agent.reset_messages()

🪝 Hook System

Customize agent behavior at key execution points:

@dataclass
class HookContext:
    """Context object passed to all hooks"""
    agent: 'PocketAgent'                                    # provides hooks access to the Agent instance
    metadata: Dict[str, Any] = field(default_factory=dict)  # additional metadata (default is empty)

class CustomHooks(AgentHooks):
    async def pre_step(self, context: HookContext):
        # executed before the llm response is generated in the step() method
        print("About to execute step")
    
    async def post_step(self, context: HookContext):
        # executed after all tool results (if any) are retrieved; This runs even if tool calling results in an error
        print("Step completed")
    
    async def pre_tool_call(self, context: HookContext, tool_call):
        # executed right before a tool is run
        print(f"Calling tool: {tool_call.name}")
        # Return modified tool_call or None
    
    async def post_tool_call(self, context: HookContext, tool_call, result):
        # executed right after a tool call result is retrieved from the PocketAgentClient
        print(f"Tool {tool_call.name} completed")
        return result  # Return modified result
    
    async def on_llm_response(self, context: HookContext, response):
        # executed right after a response message has been generated by the llm
        print("Got LLM response")
    
    async def on_event(self, event: AgentEvent):
        # Custom publishing of events useful for frontend integration

    async def on_tool_error(self, context: HookContext, tool_call: MCPCallToolRequestParams, error: Exception) -> Union[str, False]:
        # custom error handling described in more detail in PocketAgentClient docs 

    async def on_tool_result(self, context: HookContext, tool_call: ChatCompletionMessageToolCall, tool_result: FastMCPCallToolResult) -> ToolResult:
        # custom parser for tool results described in more detail in PocketAgentClient docs

# Use custom hooks
agent = MyAgent(
    agent_config=config,
    mcp_config=mcp_config,
    hooks=CustomHooks()
)

By Default, the HookContext is created with the Agent instance and empty metadata but this behavior can be customized by implementing the _create_hook_context method in your custom agent:

class Agent(PocketAgent):
    async def _create_hook_context(self) -> HookContext:
        return HookContext(
            agent=self,
            metadata={
                # custom metadata
            }
        )

📡 Event System

PocketAgent includes an AgentEvent type:

@dataclass
class AgentEvent:
    event_type: str  # e.g., "new_message"
    data: dict       # Event-specific data

By default, events are automatically emitted when any new message is added to the message history:

  • llm message
  • tool result message
  • user message

You can easily add on_event calls with custom AgentEvents in other hooks if necessary:

class CustomHooks(AgentHooks):
    async def pre_tool_call(self, context, tool_call):
        event = AgentEvent(
            event_type="tool_call",
            data=tool_call
        )

🔧 Multi-Model Support

Works seamlessly with any LiteLLM-supported model:

# OpenAI
config = AgentConfig(llm_model="gpt-4")

# Anthropic
config = AgentConfig(llm_model="anthropic/claude-3-sonnet-20240229")

# Local models
config = AgentConfig(llm_model="ollama/llama2")

# Azure OpenAI
config = AgentConfig(llm_model="azure/gpt-4")

🚏 LiteLLM Router Integration

To easily set rate limits or implement load balancing with multiple LLM API providers you can pass a LiteLLM Router instance to PocketAgent:

from litellm import Router
router_info = {
    "models": [
        {
            "model_name": "gpt-5-nano",
            "litellm_params": {
                "model": "gpt-5-nano",
                "tpm": 3000000,
                "rpm": 5000
            }
        }
    ]
}

litellm_router = Router(model_list=router_info["models"])

agent = PocketAgent(
    router=litellm_router,
    # other args
)

PocketAgentClient

Each PocketAgent instance creates initializes a PocketAgentClient which acts as a wrapper for the FastMCP Client to implement the standard mcp protocol features and some additional features.

Custom Query Params

  • Sending metadata such as a custom id to MCP servers is not handled well by the protocol (until this is merged). For now a workaround is to send metadata via query params to mcp servers using an http transport.

    agent = PocketAgent(
        mcp_server_query_params = {
            "context_id": "1111"    # context_id will automatically be added to the server endpoint when sending a request
        }
    )
    

    Note: Query params must use custom MCP middleware to be interpreted by servers

on_tool_error (hook)

  • If a tool call fails and is not handled within the tool itself, it will result in a ToolError result. You can add custom handling of such errors using the on_tool_error hook method in AgentHooks. Any custom functionality should either return a string or False. If the method returns a string, the contents will be sent to the agent as the tool result, if the method returns False the ToolError will be raised an execution of the agent will stop. The following handler is implemented by default to handle a common scenario where LLMs pass invalid parameters to tools resulting in an error:

    class AgentHooks:
        async def on_tool_error(self, context: HookContext, tool_call: MCPCallToolRequestParams, error: Exception) -> Union[str, False]:
            if "unexpected_keyword_argument" in str(error):
                tool_call_name = tool_call.name
                tool_format = await context.agent.mcp_client.get_tool_input_format(tool_call_name)
                return "You supplied an unexpected keyword argument to the tool. \
                    Try again with the correct arguments as specified in expected format: \n" + tool_format
            return False
    

tool_result_handler (hook)

  • When a tool is called successfully it results in a CallToolResult. Most of the time, you will likely just want to parse the content which is a list of MCP content objects (i.e. TextContent, ImageContent, etc). For this reason, the PocketAgentClient uses its default parser to parse these objects into content that can directly be fed to the agent as a message. Specifically, the default parser will return a ToolResult object:

    return ToolResult(
        tool_call_id=tool_call.id,                              # ID of the original tool call (needed by most apis when passing tool results)
        tool_call_name=tool_call.name,                          # Name of the tool which the result is for
        tool_result_content=tool_result_content,                # Tool result content compatible with LiteLLM message format
        _extra={
            "tool_result_raw_content": tool_result_raw_content  # Unprocessed MCP tool result (unused by default)
        }
    )
    

    However, in some cases you may want to specifically parse structured content from a known tool in which case you can override the default parser by implementing the on_tool_result hook method in AgentHooks:

    class CustomHooks(AgentHooks):
        async def on_tool_result(self, context: HookContext, tool_call: ChatCompletionMessageToolCall, tool_result: FastMCPCallToolResult) -> ToolResult:
            # your custom tool result parsing
    

Server-initiated Events

By default, PocketAgent only implements the logging handler.

To define custom behavior for any other server initiated events they can be passed as additional arguments to the agent:

agent = PocketAgent(
    elicitation_handler=your_elicitation_handler,
    log_handler=your_log_handler,
    progress_handler=your_progress_handler,
    sampling_handler=your_sampling_handler,
    message_handler=your_message_handler
)

You can also provide a custom LiteLLM Router for advanced model routing and fallback logic.

Feature Roadmap

Core Features

Feature Status Priority Description
Agent Abstraction ✅ Implemented - Basic agent abstraction with PocketAgent base class
MCP Protocol Support ✅ Implemented - Full integration with Model Context Protocol via fastmcp
Multi-Model Support ✅ Implemented - Support for any LiteLLM compatible model/endpoint
Tool Execution ✅ Implemented - Automatic parallel tool calling and results handling
Hook System ✅ Implemented - Allow configurable hooks to inject functionality during agent execution
Logging Integration ✅ Implemented - Built-in logging with custom logger support
Streaming Responses 📋 Planned Medium Real-time response streaming support
**Define Defaults for standard MCP Client handlers 📋 Planned Medium Standard MCP client methods (i.e. sampling, progress, etc) may benefit from default implementations if custom behavior is not often needed
Multi-Agent Integration 📋 Planned High Allow a PocketAgent to accept other PocketAgents as Sub Agents and automatically set up Sub Agents as tools for the Agent to use
Resources Integration 📋 Planned Medium Automatically set up mcp read_resource functionality as a tool

Modality support

Modality Status Priority Description
Text ✅ Implemented - Multi-modal input support for vision models
Images ✅ Implemented - Multi-modal input support for VLMs with option to enable/disable
Audio 📋 Planned Low Multi-modal input support for LLMs which allow audio inputs

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

pocket_agent-0.1.1.tar.gz (1.9 MB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

pocket_agent-0.1.1-py3-none-any.whl (16.5 kB view details)

Uploaded Python 3

File details

Details for the file pocket_agent-0.1.1.tar.gz.

File metadata

  • Download URL: pocket_agent-0.1.1.tar.gz
  • Upload date:
  • Size: 1.9 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.8

File hashes

Hashes for pocket_agent-0.1.1.tar.gz
Algorithm Hash digest
SHA256 a7b43b9054bfd0162f45f04618dff5dcde89821c2789299b50e9dd12b98efbf9
MD5 bd3109ce65060eb23ec8ba6e600e32cd
BLAKE2b-256 c1919da7de981dd6604e503c4c4c9166604b9ad1f23e44aa1f10cf20fc862c7a

See more details on using hashes here.

File details

Details for the file pocket_agent-0.1.1-py3-none-any.whl.

File metadata

  • Download URL: pocket_agent-0.1.1-py3-none-any.whl
  • Upload date:
  • Size: 16.5 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.8

File hashes

Hashes for pocket_agent-0.1.1-py3-none-any.whl
Algorithm Hash digest
SHA256 c16d11de7af08926751fef209955a62a6caf57c413ee366d82bf21186605ed16
MD5 d0ee3c4e74f392b0564a42b812bc1391
BLAKE2b-256 d9f40fff5f88ff344ed643b1e603f70a62571c33127cd27b979ee45e70f21bb2

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page