""" AI对话API路由 提供LLM对话的RESTful API端点，支持联网搜索功能 需求: 2.1, 3.1, 3.2, 5.5, 6.4, 10.5 流式搜索需求: 9.1, 9.2, 9.3, 9.4, 9.5 """ from typing import List, AsyncGenerator from fastapi import APIRouter, Depends, Request, HTTPException from fastapi.responses import StreamingResponse from sqlalchemy.orm import Session from app.database import get_db, SessionLocal from app.services.llm_service import LLMService from app.services.system_config_manager import get_config_bool from app.schemas.model_schema import ApiResponse, ModelResponse from app.schemas.llm_schema import ( ChatRequest, ChatResponse, EnhancedChatRequest, EnhancedChatResponse ) from app.models.user import User from app.middleware import get_current_user_from_request router = APIRouter(prefix="/api/llm", tags=["AI对话"]) async def _stream_with_db(generator: AsyncGenerator) -> AsyncGenerator[str, None]: """ 包装流式生成器，确保 db 连接在流完全结束后才关闭。 生成器本身负责持有 db 引用，此函数只负责透传数据。 """ async for chunk in generator: yield chunk @router.post("/chat") async def chat( request: EnhancedChatRequest, req: Request, db: Session = Depends(get_db), current_user: User = Depends(get_current_user_from_request) ): """ 统一对话API端点 支持流式和非流式输出，自动检测搜索选项 需要用户认证，使用用户的apikey调用百炼平台 需要余额检查，余额不足时返回402错误 向后兼容：如果没有提供search_options，则使用普通对话模式 """ api_key = current_user.apikey if not api_key: return ApiResponse(code=403, message="未配置API密钥，请在用户设置中配置apikey", data=None) # 优先使用模型自带的 api_key（爬虫同步的），没有才 fallback 到用户自己配置的 apikey from app.services.crypto_utils import get_effective_api_key api_key = get_effective_api_key(db, request.model, api_key) search_enabled = ( hasattr(request, 'search_options') and request.search_options and request.search_options.enable_search ) if search_enabled and not get_config_bool("enable_search", True): return ApiResponse(code=403, message="系统暂未开放联网搜索功能", data=None) try: if request.stream: # 流式请求：手动管理 db 生命周期，确保流结束后才关闭连接 stream_db = SessionLocal() async def stream_and_close(): try: service = LLMService(stream_db, api_key=api_key, user_id=str(current_user.id)) if search_enabled: gen = service.chat_stream_with_search(request, conversation_id=request.conversation_id) else: gen = service.chat_stream(request, conversation_id=request.conversation_id) async for chunk in gen: yield chunk finally: stream_db.close() return StreamingResponse( stream_and_close(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no" } ) else: service = LLMService(db, api_key=api_key, user_id=str(current_user.id)) if search_enabled: data = service.chat_with_search(request, conversation_id=request.conversation_id) else: data = service.chat(request, conversation_id=request.conversation_id) return ApiResponse(code=200, message="success", data=data) except HTTPException: raise except Exception as e: return ApiResponse(code=500, message=f"对话服务异常: {str(e)}", data=None) @router.post("/chat/search") async def chat_with_search( request: EnhancedChatRequest, req: Request, db: Session = Depends(get_db), current_user: User = Depends(get_current_user_from_request) ): """ 支持搜索的对话API端点 支持流式和非流式输出，集成联网搜索功能 需要用户认证，使用用户的apikey调用百炼平台 需求: 5.5, 6.4, 9.1, 9.2, 9.3, 9.4, 9.5 """ if not get_config_bool("enable_search", True): return ApiResponse(code=403, message="系统暂未开放联网搜索功能", data=None) api_key = current_user.apikey if not api_key: return ApiResponse(code=403, message="未配置API密钥，请在用户设置中配置apikey", data=None) # 优先使用模型自带的 api_key（爬虫同步的），没有才 fallback 到用户自己配置的 apikey from app.services.crypto_utils import get_effective_api_key api_key = get_effective_api_key(db, request.model, api_key) try: if request.stream: stream_db = SessionLocal() async def stream_and_close(): try: service = LLMService(stream_db, api_key=api_key, user_id=str(current_user.id)) async for chunk in service.chat_stream_with_search(request, conversation_id=request.conversation_id): yield chunk finally: stream_db.close() return StreamingResponse( stream_and_close(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no" } ) else: service = LLMService(db, api_key=api_key, user_id=str(current_user.id)) data = service.chat_with_search(request, conversation_id=request.conversation_id) return ApiResponse(code=200, message="success", data=data) except HTTPException: raise except Exception as e: return ApiResponse(code=500, message=f"搜索增强对话服务异常: {str(e)}", data=None) @router.get("/search/models", response_model=ApiResponse[List[str]]) def get_search_supported_models( db: Session = Depends(get_db), current_user: User = Depends(get_current_user_from_request) ): """ 获取支持搜索功能的模型列表 返回支持联网搜索的模型名称列表 需求: 10.1, 10.4 """ api_key = current_user.apikey if not api_key: return ApiResponse(code=403, message="未配置API密钥，请在用户设置中配置apikey", data=None) service = LLMService(db, api_key=api_key, user_id=str(current_user.id)) data = service.get_search_supported_models() return ApiResponse(code=200, message="success", data=data) @router.get("/search/check/{model}") def check_search_support( model: str, db: Session = Depends(get_db), current_user: User = Depends(get_current_user_from_request) ): """ 检查指定模型是否支持搜索功能 Args: model: 模型名称 Returns: 是否支持搜索功能的布尔值 需求: 10.1, 10.4 """ api_key = current_user.apikey if not api_key: return ApiResponse(code=403, message="未配置API密钥，请在用户设置中配置apikey", data=None) service = LLMService(db, api_key=api_key, user_id=str(current_user.id)) is_supported = service.is_search_supported(model) return ApiResponse( code=200, message="success", data={"model": model, "search_supported": is_supported} ) @router.get("/models", response_model=ApiResponse[List[ModelResponse]]) def get_llm_models(db: Session = Depends(get_db)): """ 获取所有可用的LLM模型列表 返回type=0的语言模型 """ service = LLMService(db) data = service.get_llm_models() return ApiResponse(code=200, message="success", data=data)