websocket_server.py
"""
Background Execution WebSocket Server
=====================================
Demonstrates running background workflows and streaming workflow events over WebSocket.
"""
import json
import os
from typing import Dict
import uvicorn
from agno.agent import Agent
from agno.db.sqlite import SqliteDb
from agno.models.openai import OpenAIChat
from agno.tools.hackernews import HackerNewsTools
from agno.tools.websearch import WebSearchTools
from agno.workflow.step import Step
from agno.workflow.workflow import Workflow
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
# ---------------------------------------------------------------------------
# Create Configuration
# ---------------------------------------------------------------------------
SECURITY_KEY = os.getenv("SECURITY_KEY", "your-secret-key") # Set your key here
# ---------------------------------------------------------------------------
# Create Agents
# ---------------------------------------------------------------------------
hackernews_agent = Agent(
name="HackerNews Researcher",
model=OpenAIChat(id="gpt-4o-mini"),
tools=[HackerNewsTools()],
instructions="Research tech news and trends from HackerNews",
)
search_agent = Agent(
name="Search Agent",
model=OpenAIChat(id="gpt-4o-mini"),
tools=[WebSearchTools()],
instructions="Search for additional information on the web",
)
# ---------------------------------------------------------------------------
# Create WebSocket App
# ---------------------------------------------------------------------------
app = FastAPI(title="Background Workflow WebSocket Server")
# Store active WebSocket connections and their auth status
active_connections: Dict[str, WebSocket] = {}
authenticated_connections: Dict[str, bool] = {} # {connection_id: is_authenticated}
# ---------------------------------------------------------------------------
# Define Helpers
# ---------------------------------------------------------------------------
def validate_token(token: str) -> bool:
"""Validate authentication token"""
# If no security key set, allow all connections
if not SECURITY_KEY or SECURITY_KEY == "your-secret-key":
return True
return token == SECURITY_KEY
@app.get("/")
async def get():
"""API status endpoint"""
return {
"status": "running",
"message": "Background Workflow WebSocket Server",
"endpoints": {
"websocket": "/ws",
"start-workflow": "/workflow/start",
},
"connections": len(active_connections),
"authenticated": len([c for c in authenticated_connections.values() if c]),
}
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
"""WebSocket endpoint for background workflow events"""
await websocket.accept()
connection_id = f"conn_{len(active_connections)}"
active_connections[connection_id] = websocket
authenticated_connections[connection_id] = False # Start unauthenticated
print(f"[CONN] Client connected: {connection_id}")
try:
# Send connection confirmation
await websocket.send_text(
json.dumps(
{
"event": "connected",
"connection_id": connection_id,
"message": "Connected to workflow events. Please authenticate to continue.",
"requires_auth": True,
}
)
)
# Keep connection alive
while True:
try:
data = await websocket.receive_text()
message_data = json.loads(data)
action = message_data.get("action") or message_data.get("type")
# Handle authentication
if action == "authenticate":
token = message_data.get("token")
if not token:
await websocket.send_text(
json.dumps(
{"event": "auth_error", "error": "Token is required"}
)
)
continue
if validate_token(token):
authenticated_connections[connection_id] = True
await websocket.send_text(
json.dumps(
{
"event": "authenticated",
"message": "Authentication successful. You can now send commands.",
}
)
)
print(f"[AUTH] Client authenticated: {connection_id}")
else:
await websocket.send_text(
json.dumps(
{"event": "auth_error", "error": "Invalid token"}
)
)
continue
# Check authentication for other actions
if not authenticated_connections.get(connection_id, False):
await websocket.send_text(
json.dumps(
{
"event": "auth_required",
"error": "Authentication required. Send authenticate action with valid token.",
}
)
)
continue
# Handle authenticated actions
if action == "start-workflow":
await handle_start_workflow(websocket, message_data)
elif action == "ping":
await websocket.send_text(json.dumps({"event": "pong"}))
else:
# Echo back for testing
await websocket.send_text(
json.dumps({"event": "echo", "original_message": message_data})
)
except WebSocketDisconnect:
break
except Exception as e:
await websocket.send_text(
json.dumps(
{
"event": "error",
"message": f"Error processing message: {str(e)}",
}
)
)
except WebSocketDisconnect:
pass
finally:
if connection_id in active_connections:
del active_connections[connection_id]
if connection_id in authenticated_connections:
del authenticated_connections[connection_id]
print(f"[CONN] Client disconnected: {connection_id}")
async def handle_start_workflow(websocket: WebSocket, message_data: dict):
"""Handle workflow start request via WebSocket"""
message = message_data.get("message", "AI trends 2024")
session_id = message_data.get("session_id", f"ws-session-{len(active_connections)}")
workflow = Workflow(
name="Tech Research Pipeline",
steps=[
Step(name="hackernews_research", agent=hackernews_agent),
Step(name="web_search", agent=search_agent),
],
db=SqliteDb(
db_file="tmp/workflow_bg.db",
session_table="workflow_bg",
),
)
try:
# Send acknowledgment
await websocket.send_text(
json.dumps(
{
"event": "workflow_starting",
"message": f"Starting workflow with message: {message}",
"session_id": session_id,
}
)
)
# Execute workflow in background with streaming and WebSocket
result = await workflow.arun(
input=message,
session_id=session_id,
stream=True,
stream_events=True,
background=True,
websocket=websocket,
)
# Send completion notification
await websocket.send_text(
json.dumps(
{
"event": "workflow_initiated",
"run_id": result.run_id,
"session_id": result.session_id,
"message": "Background streaming workflow initiated successfully",
}
)
)
except Exception as e:
await websocket.send_text(
json.dumps(
{
"event": "workflow_error",
"error": str(e),
"message": "Failed to start workflow",
}
)
)
# ---------------------------------------------------------------------------
# Run Server
# ---------------------------------------------------------------------------
if __name__ == "__main__":
print("[START] Starting Background Workflow WebSocket Server...")
print("[CONN] WebSocket: ws://localhost:8000/ws")
print("[API] HTTP API: http://localhost:8000")
print("[DOCS] API Docs: http://localhost:8000/docs")
print(f"[AUTH] Security Key: {SECURITY_KEY}")
uvicorn.run(
app,
host="0.0.0.0",
port=8000,
log_level="info",
)
Run the Example
1
Set up your virtual environment
uv venv --python 3.12
source .venv/bin/activate
uv venv --python 3.12
.venv\Scripts\activate
2
Install dependencies
uv pip install -U agno ddgs fastapi openai sqlalchemy uvicorn
3
Export environment variables
export OPENAI_API_KEY="your_openai_api_key_here"
export SECURITY_KEY="your_security_key_here"
$Env:OPENAI_API_KEY="your_openai_api_key_here"
$Env:SECURITY_KEY="your_security_key_here"
4
Run the example
Save the code above as
websocket_server.py, then run:python websocket_server.py