# -*- coding: utf-8 -*- import re import asyncio import aiofiles from pathlib import Path from typing import Any, AsyncGenerator, List from concurrent.futures import ThreadPoolExecutor from fastapi import UploadFile from app.core.exceptions import CustomException from app.core.logger import log from app.core.database import async_db_session from app.config.setting import settings from app.api.v1.module_system.auth.schema import AuthSchema from .tools.ai_util import AIClient # 延迟导入向量化工具，避免影响路由加载 # from .tools.document_processor import DocumentProcessor # from .tools.embedding_util import EmbeddingUtil from .schema import ( McpCreateSchema, McpUpdateSchema, McpOutSchema, ChatQuerySchema, McpQueryParam, AIProviderCreateSchema, AIProviderUpdateSchema, AIProviderOutSchema, AIProviderQueryParam, EmbeddingConfigCreateSchema, EmbeddingConfigUpdateSchema, EmbeddingConfigOutSchema, EmbeddingConfigQueryParam, KnowledgeBaseCreateSchema, KnowledgeBaseUpdateSchema, KnowledgeBaseOutSchema, KnowledgeBaseQueryParam, AIModelConfigCreateSchema, AIModelConfigUpdateSchema, AIModelConfigOutSchema, AIModelConfigQueryParam, AIModelTrainingMessageCreateSchema, AIModelTrainingMessageOutSchema, AIModelTrainingChatSchema, AI_MODEL_TYPES ) from .crud import McpCRUD, AIProviderCRUD, EmbeddingConfigCRUD, KnowledgeBaseCRUD, AIModelConfigCRUD, AIModelTrainingMessageCRUD from .model import EmbeddingConfigModel, AIModelConfigModel class McpService: """MCP服务层""" @classmethod async def detail_service(cls, auth: AuthSchema, id: int) -> dict[str, Any]: """ 获取MCP服务器详情 参数: - auth (AuthSchema): 认证信息模型 - id (int): MCP服务器ID 返回: - dict[str, Any]: MCP服务器详情字典 """ obj = await McpCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='MCP 服务器不存在') return McpOutSchema.model_validate(obj).model_dump() @classmethod async def list_service(cls, auth: AuthSchema, search: McpQueryParam | None = None, order_by: list[dict[str, str]] | None = None) -> list[dict[str, Any]]: """ 列表查询MCP服务器 参数: - auth (AuthSchema): 认证信息模型 - search (McpQueryParam | None): 查询参数模型 - order_by (list[dict[str, str]] | None): 排序参数列表 返回: - list[dict[str, Any]]: MCP服务器详情字典列表 """ search_dict = search.__dict__ if search else None obj_list = await McpCRUD(auth).get_list_crud(search=search_dict, order_by=order_by) return [McpOutSchema.model_validate(obj).model_dump() for obj in obj_list] @classmethod async def create_service(cls, auth: AuthSchema, data: McpCreateSchema) -> dict[str, Any]: """ 创建MCP服务器 参数: - auth (AuthSchema): 认证信息模型 - data (McpCreateSchema): 创建MCP服务器模型 返回: - dict[str, Any]: 创建的MCP服务器详情字典 """ obj = await McpCRUD(auth).get_by_name_crud(name=data.name) if obj: raise CustomException(msg='创建失败，MCP 服务器已存在') obj = await McpCRUD(auth).create_crud(data=data) return McpOutSchema.model_validate(obj).model_dump() @classmethod async def update_service(cls, auth: AuthSchema, id: int, data: McpUpdateSchema) -> dict[str, Any]: """ 更新MCP服务器 参数: - auth (AuthSchema): 认证信息模型 - id (int): MCP服务器ID - data (McpUpdateSchema): 更新MCP服务器模型 返回: - dict[str, Any]: 更新的MCP服务器详情字典 """ obj = await McpCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='更新失败，该数据不存在') exist_obj = await McpCRUD(auth).get_by_name_crud(name=data.name) if exist_obj and exist_obj.id != id: raise CustomException(msg='更新失败，MCP 服务器名称重复') obj = await McpCRUD(auth).update_crud(id=id, data=data) return McpOutSchema.model_validate(obj).model_dump() @classmethod async def delete_service(cls, auth: AuthSchema, ids: list[int]) -> None: """ 批量删除MCP服务器 参数: - auth (AuthSchema): 认证信息模型 - ids (list[int]): MCP服务器ID列表 返回: - None """ if len(ids) < 1: raise CustomException(msg='删除失败，删除对象不能为空') for id in ids: obj = await McpCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='删除失败，该数据不存在') await McpCRUD(auth).delete_crud(ids=ids) @classmethod async def chat_query(cls, query: ChatQuerySchema) -> AsyncGenerator[str, Any]: """ 处理聊天查询 参数: - query (ChatQuerySchema): 聊天查询模型 返回: - AsyncGenerator[str, None]: 异步生成器,每次返回一个聊天响应 """ # 创建MCP客户端实例 mcp_client = AIClient() try: # 处理消息 async for response in mcp_client.process(query.message): yield response finally: # 确保关闭客户端连接 await mcp_client.close() class AIProviderService: """AI供应商服务层""" @classmethod async def detail_service(cls, auth: AuthSchema, id: int) -> dict[str, Any]: obj = await AIProviderCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='AI供应商不存在') return AIProviderOutSchema.model_validate(obj).model_dump() @classmethod async def list_service(cls, auth: AuthSchema, search: AIProviderQueryParam | None = None, order_by: list[dict[str, str]] | None = None) -> list[dict[str, Any]]: search_dict = search.__dict__ if search else None obj_list = await AIProviderCRUD(auth).get_list_crud(search=search_dict, order_by=order_by) return [AIProviderOutSchema.model_validate(obj).model_dump() for obj in obj_list] @classmethod async def create_service(cls, auth: AuthSchema, data: AIProviderCreateSchema) -> dict[str, Any]: obj = await AIProviderCRUD(auth).get_by_name_crud(name=data.name) if obj: raise CustomException(msg='创建失败，AI供应商名称已存在') # 如果设置为默认，先清除其他默认 if data.is_default == 1: await AIProviderCRUD(auth).clear_default_crud() obj = await AIProviderCRUD(auth).create_crud(data=data) return AIProviderOutSchema.model_validate(obj).model_dump() @classmethod async def update_service(cls, auth: AuthSchema, id: int, data: AIProviderUpdateSchema) -> dict[str, Any]: obj = await AIProviderCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='更新失败，该数据不存在') exist_obj = await AIProviderCRUD(auth).get_by_name_crud(name=data.name) if exist_obj and exist_obj.id != id: raise CustomException(msg='更新失败，AI供应商名称重复') # 如果设置为默认，先清除其他默认 if data.is_default == 1: await AIProviderCRUD(auth).clear_default_crud() obj = await AIProviderCRUD(auth).update_crud(id=id, data=data) return AIProviderOutSchema.model_validate(obj).model_dump() @classmethod async def delete_service(cls, auth: AuthSchema, ids: list[int]) -> None: if len(ids) < 1: raise CustomException(msg='删除失败，删除对象不能为空') for id in ids: obj = await AIProviderCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='删除失败，该数据不存在') await AIProviderCRUD(auth).delete_crud(ids=ids) class EmbeddingConfigService: """向量化配置服务层""" @classmethod async def detail_service(cls, auth: AuthSchema, id: int) -> dict[str, Any]: obj = await EmbeddingConfigCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='向量化配置不存在') return EmbeddingConfigOutSchema.model_validate(obj).model_dump() @classmethod async def list_service(cls, auth: AuthSchema, search: EmbeddingConfigQueryParam | None = None, order_by: list[dict[str, str]] | None = None) -> list[dict[str, Any]]: search_dict = search.__dict__ if search else None obj_list = await EmbeddingConfigCRUD(auth).get_list_crud(search=search_dict, order_by=order_by) return [EmbeddingConfigOutSchema.model_validate(obj).model_dump() for obj in obj_list] @classmethod async def create_service(cls, auth: AuthSchema, data: EmbeddingConfigCreateSchema) -> dict[str, Any]: obj = await EmbeddingConfigCRUD(auth).get_by_name_crud(name=data.name) if obj: raise CustomException(msg='创建失败，向量化配置名称已存在') # 远程模式验证 if data.embedding_type == 1: if not data.base_url or not data.api_key: raise CustomException(msg='远程模式必须提供接口地址和API Key') # 如果设置为默认，先清除其他默认 if data.is_default == 1: await EmbeddingConfigCRUD(auth).clear_default_crud() obj = await EmbeddingConfigCRUD(auth).create_crud(data=data) return EmbeddingConfigOutSchema.model_validate(obj).model_dump() @classmethod async def update_service(cls, auth: AuthSchema, id: int, data: EmbeddingConfigUpdateSchema) -> dict[str, Any]: obj = await EmbeddingConfigCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='更新失败，该数据不存在') exist_obj = await EmbeddingConfigCRUD(auth).get_by_name_crud(name=data.name) if exist_obj and exist_obj.id != id: raise CustomException(msg='更新失败，向量化配置名称重复') # 远程模式验证 if data.embedding_type == 1: if not data.base_url or not data.api_key: raise CustomException(msg='远程模式必须提供接口地址和API Key') # 如果设置为默认，先清除其他默认 if data.is_default == 1: await EmbeddingConfigCRUD(auth).clear_default_crud() obj = await EmbeddingConfigCRUD(auth).update_crud(id=id, data=data) return EmbeddingConfigOutSchema.model_validate(obj).model_dump() @classmethod async def delete_service(cls, auth: AuthSchema, ids: list[int]) -> None: if len(ids) < 1: raise CustomException(msg='删除失败，删除对象不能为空') for id in ids: obj = await EmbeddingConfigCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='删除失败，该数据不存在') await EmbeddingConfigCRUD(auth).delete_crud(ids=ids) class KnowledgeBaseService: """知识库服务层""" @staticmethod def _generate_collection_name(name: str) -> str: """根据知识库名称生成ChromaDB集合名称""" import time # 移除特殊字符，只保留字母数字下划线 clean_name = re.sub(r'[^a-zA-Z0-9_]', '_', name) timestamp = int(time.time()) return f"kb_{clean_name}_{timestamp}" @classmethod async def detail_service(cls, auth: AuthSchema, id: int) -> dict[str, Any]: obj = await KnowledgeBaseCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='知识库不存在') return KnowledgeBaseOutSchema.model_validate(obj).model_dump() @classmethod async def list_service(cls, auth: AuthSchema, search: KnowledgeBaseQueryParam | None = None, order_by: list[dict[str, str]] | None = None) -> list[dict[str, Any]]: search_dict = search.__dict__ if search else None obj_list = await KnowledgeBaseCRUD(auth).get_list_crud(search=search_dict, order_by=order_by) return [KnowledgeBaseOutSchema.model_validate(obj).model_dump() for obj in obj_list] @classmethod async def create_service(cls, auth: AuthSchema, data: KnowledgeBaseCreateSchema, files: List[UploadFile] | None = None) -> dict[str, Any]: obj = await KnowledgeBaseCRUD(auth).get_by_name_crud(name=data.name) if obj: raise CustomException(msg='创建失败，知识库名称已存在') # 验证向量化配置是否存在 embedding_config = None if data.embedding_config_id: embedding_config = await EmbeddingConfigCRUD(auth).get_by_id_crud(id=data.embedding_config_id) if not embedding_config: raise CustomException(msg='向量化配置不存在') # 生成集合名称 collection_name = cls._generate_collection_name(data.name) # 计算文档数量 document_count = len(files) if files else 0 # 创建知识库记录，初始状态为处理中(如果有文件)或已完成(如果没有文件) initial_status = 1 if files and len(files) > 0 else 2 # 1=处理中, 2=已完成 # 保存文件到磁盘 file_paths = [] file_contents = [] if files and len(files) > 0: # 创建知识库专用目录 kb_dir = settings.UPLOAD_FILE_PATH / "knowledge_base" / collection_name kb_dir.mkdir(parents=True, exist_ok=True) for file in files: if not file.filename: continue # 读取文件内容 content = await file.read() # 保存文件到磁盘 file_path = kb_dir / file.filename async with aiofiles.open(file_path, 'wb') as f: await f.write(content) file_paths.append(str(file_path)) file_contents.append({ 'filename': file.filename, 'content': content, 'content_type': file.content_type }) # 创建知识库记录 create_data = { 'name': data.name, 'embedding_config_id': data.embedding_config_id, 'collection_name': collection_name, 'description': data.description, 'document_count': document_count, 'vector_count': 0, 'kb_status': initial_status, 'file_paths': file_paths if file_paths else None } obj = await KnowledgeBaseCRUD(auth).create_crud(data=create_data) knowledge_base_id = obj.id # 如果有文件上传，在后台进行向量化处理 if file_contents: # 提取 embedding 配置信息 embedding_info = None if embedding_config: embedding_info = { 'embedding_type': embedding_config.embedding_type, 'model_name': embedding_config.model_name, 'base_url': embedding_config.base_url, 'api_key': embedding_config.api_key, } # 启动后台向量化任务 asyncio.create_task( cls._process_vectorization( knowledge_base_id=knowledge_base_id, collection_name=collection_name, file_contents=file_contents, embedding_info=embedding_info ) ) return KnowledgeBaseOutSchema.model_validate(obj).model_dump() @classmethod async def _process_vectorization( cls, knowledge_base_id: int, collection_name: str, file_contents: List[dict], embedding_info: dict | None ) -> None: """ 后台向量化处理任务 参数: - knowledge_base_id: 知识库ID - collection_name: ChromaDB集合名称 - file_contents: 文件内容列表 [{'filename': str, 'content': bytes, 'content_type': str}] - embedding_info: Embedding配置信息 """ from langchain.schema import Document try: # 1. 处理文档 all_documents = [] for file_info in file_contents: filename = file_info['filename'] content = file_info['content'] if not filename: continue # 根据文件类型处理 ext = filename.rsplit('.', 1)[-1].lower() if '.' in filename else '' try: if ext in ['txt', 'md']: # 文本文件 text = None for encoding in ['utf-8', 'gbk', 'gb2312', 'latin-1']: try: text = content.decode(encoding) break except UnicodeDecodeError: continue if text: all_documents.append(Document( page_content=text, metadata={"source": filename, "type": ext} )) elif ext == 'pdf': # PDF 文件 try: import pypdf from io import BytesIO pdf_file = BytesIO(content) reader = pypdf.PdfReader(pdf_file) for page_num, page in enumerate(reader.pages): text = page.extract_text() if text and text.strip(): all_documents.append(Document( page_content=text, metadata={ "source": filename, "type": "pdf", "page": page_num + 1 } )) except Exception as e: log.error(f"处理 PDF 文件失败: {filename}, 错误: {e}") elif ext == 'docx': # DOCX 文件 try: from docx import Document as DocxDocument from io import BytesIO docx_file = BytesIO(content) doc = DocxDocument(docx_file) paragraphs = [] for para in doc.paragraphs: if para.text.strip(): paragraphs.append(para.text) text = '\n'.join(paragraphs) if text.strip(): all_documents.append(Document( page_content=text, metadata={"source": filename, "type": "docx"} )) except Exception as e: log.error(f"处理 DOCX 文件失败: {filename}, 错误: {e}") except Exception as e: log.error(f"处理文件失败: {filename}, 错误: {e}") if not all_documents: # 没有有效文档，更新状态为失败向量化处理失败 await cls._update_kb_status( knowledge_base_id=knowledge_base_id, kb_status=3, # 失败 error_message="没有有效的文档内容" ) return # 2. 创建 Embedding 工具（延迟导入） from .tools.embedding_util import EmbeddingUtil if embedding_info: embedding_util = EmbeddingUtil( embedding_type=embedding_info['embedding_type'], model_name=embedding_info['model_name'], base_url=embedding_info['base_url'], api_key=embedding_info['api_key'] ) else: # 使用默认配置 embedding_util = EmbeddingUtil( embedding_type=1, # 远程模式 model_name='text-embedding-3-small', base_url='https://api.openai.com/v1', api_key='your-api-key' # 需要配置 ) # 3. 在线程池中执行向量化（因为 ChromaDB 操作是同步的） loop = asyncio.get_event_loop() with ThreadPoolExecutor(max_workers=1) as executor: vector_count = await loop.run_in_executor( executor, embedding_util.add_documents, collection_name, all_documents ) # 4. 更新知识库状态为已完成 await cls._update_kb_status( knowledge_base_id=knowledge_base_id, kb_status=2, # 已完成 vector_count=vector_count ) except Exception as e: log.error(f"向量化处理失败，知识库ID: {knowledge_base_id}, 错误: {e}") # 更新状态为失败，截断错误信息到50字 error_msg = str(e)[:50] if len(str(e)) > 50 else str(e) await cls._update_kb_status( knowledge_base_id=knowledge_base_id, kb_status=3, # 失败 error_message=error_msg ) @classmethod async def _update_kb_status( cls, knowledge_base_id: int, kb_status: int, vector_count: int = 0, error_message: str | None = None ) -> None: """ 更新知识库状态 参数: - knowledge_base_id: 知识库ID - kb_status: 状态 (0=待处理, 1=处理中, 2=已完成, 3=失败) - vector_count: 向量数量 - error_message: 错误信息 """ from .model import KnowledgeBaseModel from sqlalchemy import update try: async with async_db_session() as session: async with session.begin(): update_data = { 'kb_status': kb_status, 'vector_count': vector_count, } if error_message: update_data['error_message'] = error_message stmt = ( update(KnowledgeBaseModel) .where(KnowledgeBaseModel.id == knowledge_base_id) .values(**update_data) ) await session.execute(stmt) await session.commit() except Exception as e: log.error(f"更新知识库状态失败，ID: {knowledge_base_id}, 错误: {e}") @classmethod async def update_service(cls, auth: AuthSchema, id: int, data: KnowledgeBaseUpdateSchema) -> dict[str, Any]: obj = await KnowledgeBaseCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='更新失败，该数据不存在') if data.name: exist_obj = await KnowledgeBaseCRUD(auth).get_by_name_crud(name=data.name) if exist_obj and exist_obj.id != id: raise CustomException(msg='更新失败，知识库名称重复') # 验证向量化配置是否存在 if data.embedding_config_id: embedding_config = await EmbeddingConfigCRUD(auth).get_by_id_crud(id=data.embedding_config_id) if not embedding_config: raise CustomException(msg='向量化配置不存在') obj = await KnowledgeBaseCRUD(auth).update_crud(id=id, data=data) return KnowledgeBaseOutSchema.model_validate(obj).model_dump() @classmethod async def delete_service(cls, auth: AuthSchema, ids: list[int]) -> None: if len(ids) < 1: raise CustomException(msg='删除失败，删除对象不能为空') for id in ids: obj = await KnowledgeBaseCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='删除失败，该数据不存在') # 删除物理文件 if obj.file_paths: for file_path in obj.file_paths: try: path = Path(file_path) if path.exists(): path.unlink() except Exception as e: log.warning(f"删除文件失败: {file_path}, 错误: {e}") # 删除知识库目录（如果为空） try: kb_dir = settings.UPLOAD_FILE_PATH / "knowledge_base" / obj.collection_name if kb_dir.exists() and not any(kb_dir.iterdir()): kb_dir.rmdir() except Exception as e: log.warning(f"删除知识库目录失败: {obj.collection_name}, 错误: {e}") # 删除 ChromaDB 中的集合（延迟导入） if obj.collection_name: try: from .tools.embedding_util import EmbeddingUtil # 获取 embedding 配置 embedding_info = None if obj.embedding_config_id: embedding_config = await EmbeddingConfigCRUD(auth).get_by_id_crud(id=obj.embedding_config_id) if embedding_config: embedding_info = { 'embedding_type': embedding_config.embedding_type, 'model_name': embedding_config.model_name, 'base_url': embedding_config.base_url, 'api_key': embedding_config.api_key, } # 创建 EmbeddingUtil 并删除集合 if embedding_info: embedding_util = EmbeddingUtil( embedding_type=embedding_info['embedding_type'], model_name=embedding_info['model_name'], base_url=embedding_info['base_url'], api_key=embedding_info['api_key'] ) else: embedding_util = EmbeddingUtil() embedding_util.delete_collection(obj.collection_name) except Exception as e: log.warning(f"删除 ChromaDB 集合失败: {obj.collection_name}, 错误: {e}") await KnowledgeBaseCRUD(auth).delete_crud(ids=ids) @classmethod async def retry_service(cls, auth: AuthSchema, id: int) -> dict[str, Any]: """ 重试向量化 参数: - auth (AuthSchema): 认证信息模型 - id (int): 知识库ID 返回: - dict[str, Any]: 知识库详情字典 """ obj = await KnowledgeBaseCRUD(auth).get_by_id_crud(id=id) if not obj: raise CustomException(msg='知识库不存在') if obj.kb_status != 3: raise CustomException(msg='只有处理失败的知识库才能重试') if not obj.file_paths or len(obj.file_paths) == 0: raise CustomException(msg='没有找到已保存的文件，请删除后重新创建知识库') # 从磁盘读取文件内容 file_contents = [] for file_path in obj.file_paths: path = Path(file_path) if not path.exists(): log.warning(f"文件不存在: {file_path}") continue async with aiofiles.open(path, 'rb') as f: content = await f.read() file_contents.append({ 'filename': path.name, 'content': content, 'content_type': None }) if not file_contents: raise CustomException(msg='所有文件均已丢失，请删除后重新创建知识库') # 获取 embedding 配置 embedding_info = None if obj.embedding_config_id: embedding_config = await EmbeddingConfigCRUD(auth).get_by_id_crud(id=obj.embedding_config_id) if embedding_config: embedding_info = { 'embedding_type': embedding_config.embedding_type, 'model_name': embedding_config.model_name, 'base_url': embedding_config.base_url, 'api_key': embedding_config.api_key, } # 更新状态为处理中，清空错误信息 await cls._update_kb_status( knowledge_base_id=id, kb_status=1, # 处理中 vector_count=0, error_message='' ) # 启动后台向量化任务 asyncio.create_task( cls._process_vectorization( knowledge_base_id=id, collection_name=obj.collection_name, file_contents=file_contents, embedding_info=embedding_info ) ) # 返回更新后的知识库信息 obj = await KnowledgeBaseCRUD(auth).get_by_id_crud(id=id) return KnowledgeBaseOutSchema.model_validate(obj).model_dump() class AIModelConfigService: """AI模型配置服务层""" @classmethod async def detail_service(cls, auth: AuthSchema, model_type: str) -> dict[str, Any]: """获取模型配置详情，如果不存在则创建默认配置""" if model_type not in AI_MODEL_TYPES: raise CustomException(msg=f'无效的模型类型: {model_type}') obj = await AIModelConfigCRUD(auth).get_by_model_type_crud(model_type=model_type) if not obj: # 创建默认配置 create_data = { 'model_type': model_type, 'model_name': None, 'provider_id': None, 'system_prompt': None, 'temperature': 1.0, 'knowledge_base_ids': None } obj = await AIModelConfigCRUD(auth).create_crud(data=create_data) return AIModelConfigOutSchema.model_validate(obj).model_dump() @classmethod async def list_service(cls, auth: AuthSchema, search: AIModelConfigQueryParam | None = None, order_by: list[dict[str, str]] | None = None) -> list[dict[str, Any]]: """查询模型配置列表""" search_dict = search.__dict__ if search else None obj_list = await AIModelConfigCRUD(auth).get_list_crud(search=search_dict, order_by=order_by) return [AIModelConfigOutSchema.model_validate(obj).model_dump() for obj in obj_list] @classmethod async def update_service(cls, auth: AuthSchema, model_type: str, data: AIModelConfigUpdateSchema) -> dict[str, Any]: """更新模型配置""" if model_type not in AI_MODEL_TYPES: raise CustomException(msg=f'无效的模型类型: {model_type}') obj = await AIModelConfigCRUD(auth).get_by_model_type_crud(model_type=model_type) if not obj: # 如果不存在，先创建 create_data = { 'model_type': model_type, 'model_name': data.model_name, 'provider_id': data.provider_id, 'system_prompt': data.system_prompt, 'temperature': data.temperature if data.temperature is not None else 1.0, 'knowledge_base_ids': data.knowledge_base_ids } obj = await AIModelConfigCRUD(auth).create_crud(data=create_data) else: # 更新现有配置 update_data = {} if data.model_name is not None: update_data['model_name'] = data.model_name if data.provider_id is not None: update_data['provider_id'] = data.provider_id if data.system_prompt is not None: update_data['system_prompt'] = data.system_prompt if data.temperature is not None: update_data['temperature'] = data.temperature if data.knowledge_base_ids is not None: update_data['knowledge_base_ids'] = data.knowledge_base_ids if update_data: obj = await AIModelConfigCRUD(auth).update_crud(id=obj.id, data=update_data) return AIModelConfigOutSchema.model_validate(obj).model_dump() @classmethod async def get_available_models(cls, auth: AuthSchema, provider_id: int) -> list[dict[str, Any]]: """根据供应商配置的baseurl和key远程获取可用模型列表""" import httpx provider = await AIProviderCRUD(auth).get_by_id_crud(id=provider_id) if not provider: raise CustomException(msg='AI供应商不存在') # 构建请求URL，兼容OpenAI标准API base_url = provider.base_url.rstrip('/') # 确保base_url包含/v1路径 if not base_url.endswith('/v1'): base_url = f"{base_url}/v1" models_url = f"{base_url}/models" headers = { 'Authorization': f'Bearer {provider.api_key}', 'Content-Type': 'application/json' } try: async with httpx.AsyncClient(timeout=30.0, verify=False) as client: response = await client.get(models_url, headers=headers) response.raise_for_status() data = response.json() # 解析OpenAI标准格式的models响应 models_data = data.get('data', []) models = [] for model in models_data: model_id = model.get('id', '') # 使用id作为name，如果没有单独的name字段 model_name = model.get('name', model_id) if model_id: models.append({'id': model_id, 'name': model_name}) # 按模型id排序 models.sort(key=lambda x: x['id']) return models except httpx.HTTPStatusError as e: log.error(f"获取模型列表失败, HTTP错误: {e.response.status_code}, {e.response.text}") if e.response.status_code == 401: raise CustomException(msg='API Key无效或已过期') elif e.response.status_code == 403: raise CustomException(msg='无权限访问模型列表') else: raise CustomException(msg=f'获取模型列表失败: HTTP {e.response.status_code}') except httpx.RequestError as e: log.error(f"获取模型列表失败, 网络错误: {str(e)}") raise CustomException(msg=f'网络请求失败: {str(e)}') except Exception as e: log.error(f"获取模型列表失败: {str(e)}") raise CustomException(msg=f'获取模型列表失败: {str(e)}') class AIModelTrainingService: """AI模型训练对话服务层""" @classmethod async def get_messages_service(cls, auth: AuthSchema, model_type: str) -> list[dict[str, Any]]: """获取某个模型的训练对话记录""" if model_type not in AI_MODEL_TYPES: raise CustomException(msg=f'无效的模型类型: {model_type}') # 获取模型配置 config = await AIModelConfigCRUD(auth).get_by_model_type_crud(model_type=model_type) if not config: return [] messages = await AIModelTrainingMessageCRUD(auth).get_list_by_config_crud(model_config_id=config.id) return [AIModelTrainingMessageOutSchema.model_validate(msg).model_dump() for msg in messages] @classmethod async def delete_message_service(cls, auth: AuthSchema, message_id: int) -> None: """删除单条训练对话记录""" msg = await AIModelTrainingMessageCRUD(auth).get_by_id_crud(id=message_id) if not msg: raise CustomException(msg='对话记录不存在') await AIModelTrainingMessageCRUD(auth).delete_crud(ids=[message_id]) @classmethod async def clear_messages_service(cls, auth: AuthSchema, model_type: str) -> None: """清空某个模型的所有训练对话记录""" if model_type not in AI_MODEL_TYPES: raise CustomException(msg=f'无效的模型类型: {model_type}') config = await AIModelConfigCRUD(auth).get_by_model_type_crud(model_type=model_type) if config: await AIModelTrainingMessageCRUD(auth).delete_by_config_crud(model_config_id=config.id) @classmethod async def chat_stream(cls, auth: AuthSchema, chat_data: AIModelTrainingChatSchema) -> AsyncGenerator[str, Any]: """ 训练对话流式输出 注意：由于 StreamingResponse 的特性，依赖注入的数据库会话在流式响应期间会被关闭， 因此需要使用独立的数据库会话来保存消息。 重要：yield 不能在 async with session.begin() 块内使用，否则会导致生成器无法正确执行 """ from app.core.database import async_db_session model_type = chat_data.model_type if model_type not in AI_MODEL_TYPES: yield f'无效的模型类型: {model_type}' return # 用于存储错误消息 error_message = None # 用于存储AI调用所需的配置信息 config_id = None config_system_prompt = None config_model_name = None config_temperature = None provider_api_key = None provider_base_url = None messages = [] # 1. 使用独立的数据库会话处理所有数据库操作（不在此块内yield） async with async_db_session() as session: async with session.begin(): stream_auth = AuthSchema(db=session, user=auth.user, check_data_scope=False) # 如果配置有变动，先保存配置 if chat_data.config_changed and chat_data.config_data: await AIModelConfigService.update_service(auth=stream_auth, model_type=model_type, data=chat_data.config_data) # 获取模型配置 config = await AIModelConfigCRUD(stream_auth).get_by_model_type_crud(model_type=model_type) if not config: error_message = '模型配置不存在，请先配置模型' else: # 保存用户消息 user_msg = await AIModelTrainingMessageCRUD(stream_auth).create_crud(data={ 'model_config_id': config.id, 'role': 'user', 'content': chat_data.message }) # 获取供应商配置 provider = None if config.provider_id: provider = await AIProviderCRUD(stream_auth).get_by_id_crud(id=config.provider_id) if not provider: providers = await AIProviderCRUD(stream_auth).get_list_crud(search={'is_default': 1}) if providers: provider = providers[0] if not provider: error_message = '未配置AI供应商，请先在AI配置中添加供应商' else: # 获取历史对话记录 history_messages = await AIModelTrainingMessageCRUD(stream_auth).get_list_by_config_crud(model_config_id=config.id) # 提取需要的信息 config_id = config.id config_system_prompt = config.system_prompt config_model_name = config.model_name config_temperature = config.temperature provider_api_key = provider.api_key provider_base_url = provider.base_url # 构建消息列表 if config_system_prompt: messages.append({"role": "system", "content": config_system_prompt}) for msg in history_messages: if msg.id != user_msg.id: messages.append({"role": msg.role, "content": msg.content}) messages.append({"role": "user", "content": chat_data.message}) # 提交事务 await session.commit() # 2. 如果有错误，yield错误消息并返回 if error_message: log.warning(f"[训练对话] 配置错误: {error_message}") yield error_message return # 检查配置是否完整 if not provider_api_key or not provider_base_url: log.error("[训练对话] 供应商配置不完整") yield "供应商配置不完整，请检查API Key和Base URL" return # 3. 创建AI客户端并流式输出 import httpx from openai import AsyncOpenAI # 确保base_url包含/v1路径 if provider_base_url and not provider_base_url.rstrip('/').endswith('/v1'): provider_base_url = provider_base_url.rstrip('/') + '/v1' log.info(f"[训练对话] 开始调用AI，模型: {config_model_name}, 消息数: {len(messages)}") log.info(f"[训练对话] API配置: base_url={provider_base_url}") http_client = httpx.AsyncClient( timeout=60.0, follow_redirects=True, verify=False # 禁用 SSL 验证 ) client = AsyncOpenAI( api_key=provider_api_key, base_url=provider_base_url, http_client=http_client ) full_response = "" is_cancelled = False response = None try: response = await client.chat.completions.create( model=config_model_name or 'gpt-3.5-turbo', messages=messages, temperature=config_temperature, stream=True ) log.info(f"[训练对话] AI响应开始流式输出") async for chunk in response: # 处理不同的响应格式 content = None if hasattr(chunk, 'choices') and chunk.choices: delta = chunk.choices[0].delta if hasattr(delta, 'content') and delta.content: content = delta.content if content: full_response += content yield content log.info(f"[训练对话] AI响应完成，总长度: {len(full_response)}") except asyncio.CancelledError: # 客户端断开连接，停止AI输出 is_cancelled = True log.info(f"[训练对话] 客户端断开连接，停止AI输出，已生成长度: {len(full_response)}") # 关闭AI流式响应 if response: try: await response.close() except Exception: pass raise # 重新抛出以便 FastAPI 知道请求已取消 except GeneratorExit: # 生成器被关闭，通常是客户端断开 is_cancelled = True log.info(f"[训练对话] 生成器关闭，停止AI输出，已生成长度: {len(full_response)}") if response: try: await response.close() except Exception: pass raise except Exception as e: log.error(f"AI训练对话失败: {str(e)}") import traceback log.error(f"AI训练对话异常详情: {traceback.format_exc()}") error_msg = f"对话失败: {str(e)}" full_response = error_msg yield error_msg finally: await http_client.aclose() # 4. 保存AI响应 if full_response and config_id: async with async_db_session() as session: async with session.begin(): save_auth = AuthSchema(db=session, user=auth.user, check_data_scope=False) await AIModelTrainingMessageCRUD(save_auth).create_crud(data={ 'model_config_id': config_id, 'role': 'assistant', 'content': full_response }) await session.commit() class AIModelTestService: """起名测试服务层（用于小程序/外部调用，仅读不写）""" @classmethod async def test_naming(cls, model_type: str, text: str) -> str: """ 起名测试接口 - 读取模型配置和训练信息 - 内部流式调用AI - 组合响应后一次性返回 - 仅读数据，不保存对话记录 参数: model_type: 模型类型(enterprise_naming/personal_naming等) text: 用户输入的文本 返回: AI的完整响应文本 """ from app.core.database import async_db_session from app.api.v1.module_application.ai.schema import AI_MODEL_TYPES if model_type not in AI_MODEL_TYPES: raise CustomException(msg=f'无效的模型类型: {model_type}') # 用于存储AI调用所需的配置信息 config_system_prompt = None config_model_name = None config_temperature = None provider_api_key = None provider_base_url = None messages = [] # 1. 读取模型配置和训练记录（只读，不保存任何数据） async with async_db_session() as session: async with session.begin(): from app.api.v1.module_system.auth.schema import AuthSchema read_auth = AuthSchema(db=session, user=None, check_data_scope=False) # 获取模型配置 config = await AIModelConfigCRUD(read_auth).get_by_model_type_crud(model_type=model_type) if not config: raise CustomException(msg='模型配置不存在，请先配置模型') # 获取供应商配置 provider = None if config.provider_id: provider = await AIProviderCRUD(read_auth).get_by_id_crud(id=config.provider_id) if not provider: providers = await AIProviderCRUD(read_auth).get_list_crud(search={'is_default': 1}) if providers: provider = providers[0] if not provider: raise CustomException(msg='未配置AI供应商，请先在AI配置中添加供应商') # 获取历史训练对话记录（作为上下文） history_messages = await AIModelTrainingMessageCRUD(read_auth).get_list_by_config_crud(model_config_id=config.id) # 提取配置信息 config_system_prompt = config.system_prompt config_model_name = config.model_name config_temperature = config.temperature provider_api_key = provider.api_key provider_base_url = provider.base_url # 构建消息列表 if config_system_prompt: messages.append({"role": "system", "content": config_system_prompt}) # 添加历史训练对话作为上下文 for msg in history_messages: messages.append({"role": msg.role, "content": msg.content}) # 添加当前用户输入 messages.append({"role": "user", "content": text}) # 检查配置是否完整 if not provider_api_key or not provider_base_url: raise CustomException(msg="供应商配置不完整，请检查API Key和Base URL") # 2. 调用AI（内部流式，收集后一次性返回） import httpx import asyncio from openai import AsyncOpenAI # 确保base_url包含/v1路径 if provider_base_url and not provider_base_url.rstrip('/').endswith('/v1'): provider_base_url = provider_base_url.rstrip('/') + '/v1' log.info(f"[起名测试] 开始调用AI，模型类型: {model_type}, 模型: {config_model_name}, 消息数: {len(messages)}") def _extract_status_code(err: Exception) -> int | None: if hasattr(err, "status_code") and isinstance(getattr(err, "status_code"), int): return getattr(err, "status_code") response = getattr(err, "response", None) if response is not None and hasattr(response, "status_code"): return response.status_code return None def _extract_error_message(err: Exception) -> str: body = getattr(err, "body", None) if isinstance(body, dict) and "error" in body: inner = body.get("error") or {} msg = inner.get("message") if isinstance(msg, str) and msg.strip(): return msg.strip() return str(err) retry_times = 3 base_sleep = 0.8 last_err: Exception | None = None for attempt in range(1, retry_times + 1): http_client = httpx.AsyncClient( timeout=120.0, follow_redirects=True, verify=False ) client = AsyncOpenAI( api_key=provider_api_key, base_url=provider_base_url, http_client=http_client ) full_response = "" try: response = await client.chat.completions.create( model=config_model_name or 'gpt-3.5-turbo', messages=messages, temperature=config_temperature, stream=True ) async for chunk in response: if hasattr(chunk, 'choices') and chunk.choices: delta = chunk.choices[0].delta if hasattr(delta, 'content') and delta.content: full_response += delta.content log.info(f"[起名测试] AI响应完成，总长度: {len(full_response)}") return full_response except Exception as e: last_err = e status_code = _extract_status_code(e) err_msg = _extract_error_message(e) transient = status_code in {408, 429, 500, 502, 503, 504} or isinstance(e, (httpx.TimeoutException, httpx.NetworkError)) if not transient or attempt >= retry_times: log.error(f"[起名测试] AI调用失败(终止): {str(e)}") import traceback log.error(f"[起名测试] 异常详情: {traceback.format_exc()}") raise CustomException(msg=f"服务暂时不可用，请稍后重试。({err_msg})", status_code=503) sleep_s = base_sleep * (2 ** (attempt - 1)) log.warning(f"[起名测试] AI调用失败(第{attempt}次/{retry_times}次)，{sleep_s:.2f}s后重试: {err_msg}") await asyncio.sleep(sleep_s) finally: await http_client.aclose() raise CustomException(msg=f"服务暂时不可用，请稍后重试。({_extract_error_message(last_err) if last_err else ''})", status_code=503)