Ollama + FastAPI 私有大模型部署

本文将指导你如何使用 FastAPI 搭建一个支持流式响应的 Ollama API 代理服务,让你能够在局域网内安全地调用私有大模型,实现类似 ChatGPT 的对话体验。

本文基于上一篇关于 Ollama 环境搭建的文章进行扩展,我们将重点关注如何使用 FastAPI 构建一个安全、高效的 API 服务。如果你已经完成了 Ollama 的安装和配置,那么就让我们开始吧!

1- 开发环境准备

1.1- 安装依赖

FastAPI 依赖 uvicorn 作为 ASGI 服务器。在命令行中运行:

pip install fastapi uvicorn

2- FastAPI 接口开发

2.1- 创建 API 服务

以下是完整的 API 服务代码,包含了安全验证、流式响应等核心功能:

from fastapi import FastAPI, HTTPException, Header, Body, Request
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, ValidationError
from typing import Optional
import requests
import logging
import json

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# 配置 CORS 中间件(解决 OPTIONS 405 错误)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["POST", "GET", "OPTIONS"],  # 允许 OPTIONS 预检请求
    allow_headers=["Authorization", "Content-Type"],
)

# 自定义错误处理
@app.exception_handler(ValidationError)
async def validation_exception_handler(request, exc):
    return {
        "status": "error",
        "message": "请求参数格式错误",
        "details": exc.errors(),
    }

# ollama服务地址
OLLAMA_URL = "http://localhost:11434/api/generate"
# API密钥
SECRET_KEY = "your_secret_token"
# 限制可允许访问的IP地址
ALLOWED_IP_RANGE = "192.168.50."

class Message(BaseModel):
    id: str = "default_id"
    role: str = "user"
    content: Optional[str] = None
    class Config:
        extra = "allow"

class ChatRequest(BaseModel):
    messages: list[Message] = None
    prompt: str = None
    model: str = "qwq:32b"
    temperature: float = 0.7
    top_p: float = 0.9
    stream: bool = True

@app.post("/generate/chat/completions")
async def generate(
    request: Request,
    chat_input: ChatRequest = Body(...),
    authorization: str = Header(..., description="Bearer Token")
):
    try:
        logger.info(f"Received request body: {await request.json()}")
        # 验证 Token
        token = authorization.split(" ")[1] if "Bearer " in authorization else None
        if token != SECRET_KEY:
            raise HTTPException(status_code=401, detail="无效的 Token") 

        # 检查 IP
        client_ip = request.client.host
        if not client_ip.startswith(ALLOWED_IP_RANGE):
            raise HTTPException(status_code=403, detail="IP 不允许")

        # 提取 prompt
        if chat_input.messages:
            user_messages = [
                msg for msg in chat_input.messages 
                if msg.role == "user" and msg.content and msg.content.strip()
            ]
            if not user_messages:
                raise HTTPException(status_code=400, detail="未提供有效用户消息")
            prompt = user_messages[-1].content.strip()
        elif chat_input.prompt and chat_input.prompt.strip():
            prompt = chat_input.prompt.strip()
        else:
            raise HTTPException(status_code=400, detail="未提供有效提示")

        payload = {
            "model": chat_input.model,
            "prompt": prompt,
            "temperature": chat_input.temperature,
            "top_p": chat_input.top_p,
        }

        # 流式响应适配
        if chat_input.stream:
            response = requests.post(
                OLLAMA_URL,
                json=payload,
                stream=True,
                timeout=120
            )
            def stream_response():
                for line in response.iter_lines():
                    if line:
                        data = json.loads(line)
                        chunk = {
                            "choices": [
                                {
                                    "delta": {
                                        "content": data.get("response", "")
                                    }
                                }
                            ]
                        }
                        yield f"data: {json.dumps(chunk)}\n\n"
            return StreamingResponse(
                stream_response(),
                media_type="text/event-stream"
            )
        else:
            response = requests.post(OLLAMA_URL, json=payload, timeout=30)
            response.raise_for_status()
            ollama_data = response.json()
            chatbox_response = {
                "choices": [
                    {
                        "message": {
                            "content": ollama_data.get("response", "")
                        }
                    }
                ]
            }
            return chatbox_response  # 确保返回格式包含 choices 字段

    except ValidationError as ve:
        logger.error(f"验证错误: {ve}")
        raise HTTPException(status_code=422, detail="无效的请求体") 
    except requests.exceptions.RequestException as re:
        logger.error(f"Ollama 请求失败: {re}")
        raise HTTPException(status_code=500, detail="Ollama 服务异常")
    except Exception as e:
        logger.error(f"未知错误: {str(e)}")
        raise HTTPException(status_code=500, detail="服务器内部错误")

# 处理根路径的友好提示(解决 GET / 404)
@app.get("/")
async def root():
    return {"message": "Ollama Proxy API. 使用 POST 请求访问 /generate/chat/completions"}

# 处理 GET /generate 的友好提示(解决 GET /generate 404)
@app.get("/generate")
async def generate_docs():
    return {"message": "请使用 POST 请求访问 /generate/chat/completions"}

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

2.2- 配置防火墙

添加一条入站规则,允许外部流量通过 TCP 协议 访问本地的 8000 端口

netsh advfirewall firewall add rule name="Open Port 8000" dir=in action=allow protocol=TCP localport=8000

2.3- 启动服务

你可以选择以下任一方式启动 FastAPI 服务:

  • 直接运行 Python 文件
  • 通过 uvicorn 命令启动(推荐)
uvicorn main:app --host 0.0.0.0 --port 8000 --reload

3- 客户端配置

3.1- Chatbox 客户端配置示例

以下是使用 Chatbox 客户端配置的示例:

Chatbox 配置示例

3.2- 验证服务

完成以上步骤后,局域网中的其他用户就可以通过配置好的客户端安全地访问你的私有大模型服务了。请确保:

  1. 客户端 IP 地址在允许范围内
  2. 使用正确的 API 密钥
  3. 选择合适的模型名称

至此,你已经成功搭建了一个安全可靠的私有大模型服务。这个服务不仅支持流式响应,还包含了必要的安全验证机制,可以放心地在局域网中使用。