如何为Session ID分配FastAPI Uvicorn线程？

问题说明 #

• 本地部署的Web应用集成需维持对话上下文的AI模型
• 多次发送聊天POST请求时，请求可能被Uvicorn线程池中的不同线程处理，此前因错误使用线程本地存储导致上下文分散在各线程内存
• 要求：不能用Cookie或文件存储上下文，需通过Session Token关联对话上下文
• 此前尝试自定义线程管理无效，原因是自定义线程与Uvicorn线程池不属于同一体系

当前代码分析 #

你提供的代码已采用全局字典conversations存储会话上下文，这是正确的方向，但存在线程安全隐患：多线程同时读写字典可能引发竞争条件，导致数据不一致。此外核心误解在于试图绑定Session到特定线程——Uvicorn会自动调度线程处理请求，无法手动指定线程分配，这种思路本身不可行。

正确解决方案 #

放弃“按Session Token分配线程”的思路，改用进程内共享的线程安全会话存储，通过Session ID关联每个请求与对应的对话上下文：

• 使用全局字典存储会话，并用异步锁保护读写操作，确保多线程/异步环境下的数据一致性
• 所有请求通过Session ID从共享字典中获取对应上下文，无需绑定线程

改进后的代码实现 #

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from dotenv import load_dotenv
from openai import AzureOpenAI
import os
import asyncio
from typing import Dict, List
from uuid import UUID, uuid4

load_dotenv()
router = APIRouter()

class Message(BaseModel):
    """对话消息模型"""
    role: str  # 'user' 或 'assistant'
    content: str

class Prompt(BaseModel):
    """请求消息模型"""
    session_id: UUID
    prompt: str

class NewSessionResponse(BaseModel):
    """新建会话响应模型"""
    session_id: UUID
    
class ResetRequest(BaseModel):
    """重置会话请求模型"""
    session_id: UUID

# 存储会话上下文的全局字典
conversations: Dict[UUID, List[Message]] = {}
# 异步锁，保护字典的读写操作
conversations_lock = asyncio.Lock()

# 全局初始化Azure OpenAI客户端，避免重复创建
api_key = os.getenv("key")
api_url = os.getenv("endpoint url")
if not api_key or not api_url:
    raise RuntimeError('未找到API密钥或端点，请检查配置')

client = AzureOpenAI(
    api_key=api_key,
    azure_endpoint=api_url,
    api_version="2023-05-15"  # 根据实际使用版本调整
)

@router.post("/ai/chat")
async def chat(prompt: Prompt):
    """处理AI聊天请求，维持对话上下文"""
    session_id = prompt.session_id
    
    # 加锁保护会话字典的读写
    async with conversations_lock:
        # 初始化会话（如果不存在）
        if session_id not in conversations:
            conversations[session_id] = []
        # 添加用户消息到上下文
        conversations[session_id].append(Message(role="user", content=prompt.prompt))
        # 转换为API要求的格式
        context = [{'role': msg.role, 'content': msg.content} for msg in conversations[session_id]]
    
    # 用线程池包装同步调用，避免阻塞FastAPI事件循环
    response = await asyncio.to_thread(
        client.chat.completions.create,
        model="gpt-35-turbo-4k-0613",
        messages=context
    )
    
    model_response = response.choices[0].message.content
    
    # 加锁保存助手响应
    async with conversations_lock:
        conversations[session_id].append(Message(role='assistant', content=model_response))
    
    return {"response": model_response}

@router.post("/ai/new_session", response_model=NewSessionResponse)
async def new_session():
    """创建新的对话会话"""
    session_id = uuid4()
    async with conversations_lock:
        conversations[session_id] = []
    return NewSessionResponse(session_id=session_id)

@router.post("/ai/reset_session")
async def reset_session(request: ResetRequest):
    """重置指定会话的上下文"""
    async with conversations_lock:
        if request.session_id in conversations:
            conversations[request.session_id] = []
            return {"detail": "会话已重置"}
        raise HTTPException(status_code=404, detail='会话不存在')

关键改进点 #

• 添加asyncio.Lock确保多异步任务下字典操作的线程安全
• 将Azure OpenAI客户端初始化移到全局，避免每次请求重复创建连接
• 用asyncio.to_thread包装同步的API调用，避免阻塞FastAPI事件循环
• 补充会话重置接口，完善功能逻辑
• 优化代码注释与配置检查逻辑

额外说明 #

• 若应用采用多进程部署（如Uvicorn启动多个worker），全局字典无法跨进程共享，此时需用Redis等分布式存储替代；单进程部署下全局字典完全够用

内容的提问来源于stack exchange，提问作者Javier Martín Pizarro