server.py

import asyncio
import json
import logging
import os
from contextlib import asynccontextmanager

from dotenv import load_dotenv
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from openai import AsyncAzureOpenAI
from pydantic import BaseModel

from agents import Agent, Runner, SQLiteSession, function_tool, set_default_openai_api, set_default_openai_client, set_tracing_disabled
from context_manager import trim_and_summarize

load_dotenv()

AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY", "")
AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT", "")
AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION", "2025-03-01-preview")
AZURE_OPENAI_DEPLOYMENT = os.getenv("AZURE_OPENAI_DEPLOYMENT", "gpt-4.1-mini")
AZURE_OPENAI_SUMMARY_DEPLOYMENT = os.getenv("AZURE_OPENAI_SUMMARY_DEPLOYMENT", "gpt-4.1-mini")

logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(name)-20s  %(levelname)-7s  %(message)s")
log = logging.getLogger("server")

azure_client = AsyncAzureOpenAI(
    api_key=AZURE_OPENAI_API_KEY,
    azure_endpoint=AZURE_OPENAI_ENDPOINT,
    api_version=AZURE_OPENAI_API_VERSION,
)

set_default_openai_client(client=azure_client, use_for_tracing=False)
set_default_openai_api("chat_completions")
set_tracing_disabled(disabled=True)

SYSTEM_PROMPT = (
    "You are a helpful, friendly assistant. "
    "When you see a '[Previous conversation summary]' message, treat it as "
    "a reliable recap of the earlier parts of the conversation that were "
    "trimmed to save context space. Use it to maintain continuity.\n"
    "You can check the weather using the get_weather tool."
)

@function_tool
def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is 22 °C and sunny."

agent = Agent(
    name="ContextManagedAssistant",
    instructions=SYSTEM_PROMPT,
    model=AZURE_OPENAI_DEPLOYMENT,
    tools=[get_weather],
)

session = SQLiteSession("demo_session", "conversation_history.db")

class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)

    async def broadcast(self, message: dict):
        for connection in self.active_connections:
            try:
                await connection.send_json(message)
            except Exception as e:
                log.error(f"Error broadcasting message: {e}")

manager = ConnectionManager()

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Setup
    yield
    # Teardown

app = FastAPI(lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

async def trigger_event(event: dict):
    await manager.broadcast(event)

class ChatMessage(BaseModel):
    message: str

@app.post("/api/chat")
async def chat_endpoint(msg: ChatMessage):
    user_input = msg.message
    
    if user_input.lower() == "clear":
        await session.clear_session()
        await trigger_event({"type": "session_cleared"})
        return {"response": "Session cleared."}
    
    # Trim & Summarize
    await trim_and_summarize(
        session=session,
        system_prompt=SYSTEM_PROMPT,
        azure_client=azure_client,
        summary_deployment=AZURE_OPENAI_SUMMARY_DEPLOYMENT,
        event_callback=trigger_event
    )
    
    # Run Agent
    try:
        await trigger_event({"type": "agent_running"})
        result = await Runner.run(
            agent,
            user_input,
            session=session,
        )
        
        # Optionally send updated full session state
        items = await session.get_items()
        
        await trigger_event({
            "type": "agent_response", 
            "response": result.final_output,
            "session_items": items
        })
        return {"response": result.final_output}
    except Exception as e:
        log.exception("Agent run failed")
        await trigger_event({"type": "error", "message": str(e)})
        return {"error": str(e)}

@app.get("/api/session")
async def get_session():
    items = await session.get_items()
    return {"items": items}
    
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            # We don't really expect client to send much here except maybe ping
    except WebSocketDisconnect:
        manager.disconnect(websocket)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=True)