How to Extend the API Server
In this guide, we will explore how to extend the API server provided by the Agents SDK to add custom endpoints and functionality. This allows you to tailor the server to your specific needs and integrate additional features.
Prerequisites
Before you begin, make sure you have completed the Getting Started with the Agents SDK tutorial and have a basic understanding of how create and run agents.
Overview
The API server provided by the Agents SDK is built using FastAPI, a modern, fast (high-performance) web framework for building APIs with Python. FastAPI makes it easy to extend the server by adding new routes, middleware, and other customizations. For more information on FastAPI, check out the official documentation.
The Agents SDK implements a hexagonal architecture design pattern. Refer to the Agents SDK Architecture article for more information. The book Architecture Patterns with Python is also a nice general introduction to the topic. However, you can extend the API server following any design pattern you prefer. In this guide, we provide examples that follow the FastAPI documentation since it requires less background knowledge.
Step 1: Setting Up the Project
First, create a new Python file for your custom API server.
from fastapi import APIRouter
from fastapi import APIRouter
from zav.logging import logger
from zav.agents_sdk import AgentSetupRetrieverFromFile, setup_app
# This is the path to the agents project directory
from .agents import ChatAgentFactory
agent_setup_retriever = AgentSetupRetrieverFromFile(
file_path="./agents/agent_setups.json"
)
app = setup_app(
agent_setup_retriever=agent_setup_retriever,
chat_agent_factory=ChatAgentFactory,
debug_backend=logger.info,
)
# Create a new router for custom endpoints
custom_router = APIRouter()
# Declare a new endpoint
@custom_router.get("/custom-endpoint")
async def custom_endpoint():
return {"message": "This is a custom endpoint"}
# Include the custom router in the app
app.include_router(custom_router)
# Run the app
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
In this example, we create a new FastAPI router called custom_router and define a custom endpoint /custom-endpoint. We then include this router in the main app using app.include_router(custom_router).
Step 2: Adding Custom Middleware
You can also add custom middleware to the API server. Middleware is a function that runs before and after each request. It can be used for tasks such as logging, authentication, and modifying the request or response.
Here is an example of adding custom middleware to log the request and response:
from fastapi.middleware import Middleware
from fastapi.middleware.cors import CORSMiddleware
# Define custom middleware
class CustomMiddleware:
async def __call__(self, request, call_next):
logger.info(f"Request: {request.method} {request.url}")
response = await call_next(request)
logger.info(f"Response: {response.status_code}")
return response
# Add the custom middleware to the app
app.add_middleware(CustomMiddleware)
# Optionally, add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
In this example, we define a custom middleware class CustomMiddleware that logs the request and response. We then add this middleware to the app using app.add_middleware(CustomMiddleware). Additionally, we add CORS middleware to allow cross-origin requests.
Step 3: Adding Custom Dependencies
You can also add custom dependencies to the API server. Dependencies are functions that are called before the endpoint handler and can be used to inject data or perform actions.
Here is an example of adding a custom dependency to inject a database connection:
from fastapi import Depends
# Define a custom dependency
def get_db():
db = DatabaseConnection()
try:
yield db
finally:
db.close()
# Use the custom dependency in an endpoint
@custom_router.get("/items/")
async def read_items(db: DatabaseConnection = Depends(get_db)):
items = db.query("SELECT * FROM items")
return items
In this example, we define a custom dependency get_db that provides a database connection. We then use this dependency in an endpoint by adding db: DatabaseConnection = Depends(get_db) to the endpoint's parameters.
Step 4: Running the Custom API Server
To run the custom API server, use the following command:
python custom_app.py
This will start the API server on 0.0.0.0 at port 8000 by default. You can access the API documentation at http://<host>:<port>/docs.
Adding New Endpoints to the Chat Agents
For this section, it is recommended that you are familiar with the Agents SDK Architecture. The lifecycle of a request in the Agents SDK is as follows:
- The user makes an HTTP request to the API server.
- The API server routes the request to the corresponding FastAPI route.
- The FastAPI route Transform the request parameters into a command object.
- The command is sent to the message bus (which is injected to the route as a dependency).
- The message bus routes the command to the corresponding handler.
- The handler processes the command and returns some domain object.
- The FastAPI route parses the domain object into a response object.
- The API server returns the response payload to the user.
Since the command and message bus are decoupled from the transport layer (the FastAPI routers), it means that it's very simple to add new endpoints for the same command, or even create new endpoints with new commands. In this section we focus on the former.
Adding a new endpoint that reuses an existing command is useful if you want to make the API compatible with a different client, or if you want to add a new feature to the API.
The following example demonstrates how to add a new endpoint with different body payload and response payload for the chat command.
from typing import Optional, List, Dict, Any
from fastapi import APIRouter
from fastapi import APIRouter, Depends
from pydantic import BaseModel
from zav.logging import logger
from zav.api.dependencies import get_message_bus
from zav.message_bus import MessageBus
from zav.agents_sdk.domain import ChatRequest, ChatMessage
from zav.agents_sdk.handlers import commands
from zav.agents_sdk import AgentSetupRetrieverFromFile, setup_app
# This is the path to the agents project directory
from .agents import ChatAgentFactory
agent_setup_retriever = AgentSetupRetrieverFromFile(
file_path="./agents/agent_setups.json"
)
app = setup_app(
agent_setup_retriever=agent_setup_retriever,
chat_agent_factory=ChatAgentFactory,
debug_backend=logger.info,
)
# Create a new router for custom endpoints
custom_router = APIRouter()
# Define new response model
class Message(BaseModel):
sender: str
content: str
class ChatResponseForm(BaseModel):
agent_identifier: str
agent_params: Optional[Dict[str, Any]] = None
messages: List[Message]
class ChatResponseItem(ChatResponseForm):
class Config:
orm_mode = True
# Declare a new endpoint
@custom_router.post("/custom-chat", response_model=ChatResponseItem)
async def custom_endpoint(
body: ChatResponseForm,
message_bus: MessageBus = Depends(get_message_bus),
):
results = await message_bus.handle(
# Translate the request to the appropriate command
commands.CreateChatResponse(
chat_request=ChatRequest(
agent_identifier=body.agent_identifier,
bot_params=body.agent_params,
conversation=[
ChatMessage(
sender=message.sender,
content=message.content,
)
for message in body.messages
],
),
)
)
result: ChatRequest = results.pop(0)
# Translate the handler output to the response model
return ChatResponseItem(
agent_identifier=result.agent_identifier,
agent_params=result.bot_params,
messages=[
Message(
sender=message.sender,
content=message.content,
)
for message in result.conversation
],
)
# Include the custom router in the app
app.include_router(custom_router)
# Run the app
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
For more details, refer to the Agents SDK API Reference and the controllers source code.