1 line
20 KiB
JSON
1 line
20 KiB
JSON
{"repo": "apconw/aix-db", "files": {"services/aimodel_service.py": {"9d2e15e7436e98b0c221ab1df5b94a0e45787647": {"t": 1785674028.3543525, "n": 18242, "h": "29d807d61e7c2357", "c": "import json\nimport logging\nimport httpx\nfrom datetime import datetime\nfrom typing import List, Optional\n\nfrom sqlalchemy import desc\n\nfrom common.exception import MyException\nfrom constants.code_enum import SysCodeEnum\nfrom model.db_connection_pool import get_db_pool\nfrom model.db_models import TAiModel\nfrom model.serializers import model_to_dict\n\nlogger = logging.getLogger(__name__)\npool = get_db_pool()\n\n\ndef _build_openai_headers(api_key: str) -> dict:\n headers = {\"Content-Type\": \"application/json\"}\n if api_key:\n headers[\"Authorization\"] = f\"Bearer {api_key}\"\n return headers\n\n\ndef _http_error_detail(error: httpx.HTTPStatusError) -> str:\n detail = error.response.text.strip()\n if not detail:\n return f\"HTTP {error.response.status_code}\"\n return detail[:500]\n\n\ndef _has_chat_choices(response: httpx.Response) -> bool:\n try:\n choices = response.json().get(\"choices\")\n except (ValueError, AttributeError):\n return False\n return isinstance(choices, list) and bool(choices)\n\n\nasync def _check_embedding_model_legacy(\n supplier: int,\n api_domain: str,\n api_key: str,\n) -> dict:\n domain = api_domain.rstrip(\"/\")\n if supplier == 3:\n if domain.endswith(\"/v1\"):\n domain = domain[:-3]\n url = f\"{domain}/api/tags\"\n headers = None\n timeout = 5\n else:\n url = f\"{domain}/models\"\n headers = _build_openai_headers(api_key)\n timeout = 10\n\n async with httpx.AsyncClient() as client:\n request_kwargs = {\"timeout\": timeout}\n if headers is not None:\n request_kwargs[\"headers\"] = headers\n response = await client.get(url, **request_kwargs)\n response.raise_for_status()\n\n return {\"success\": True, \"message\": \"连接成功\"}\n\n\nasync def _check_rerank_model(\n supplier: int,\n api_domain: str,\n api_key: str,\n base_model: str,\n) -> dict:\n documents = [\n \"文本排序模型根据相关性对候选文本进行排序\",\n \"量子计算是计算科学的前沿领域\",\n ]\n is_dashscope = supplier == 6 or \"aliyuncs\" in api_domain\n if is_dashscope:\n payload = {\n \"model\": base_model,\n \"input\": {\n \"query\": \"什么是文本排序模型\",\n \"documents\": documents,\n },\n \"parameters\": {\n \"return_documents\": False,\n \"top_n\": len(documents),\n },\n }\n else:\n payload = {\n \"query\": \"什么是文本排序模型\",\n \"documents\": documents,\n }\n\n async with httpx.AsyncClient() as client:\n response = await client.post(\n api_domain.rstrip(\"/\"),\n headers=_build_openai_headers(api_key),\n json=payload,\n timeout=30,\n )\n response.raise_for_status()\n\n try:\n response_data = response.json()\n except ValueError:\n return {\"success\": False, \"message\": \"Rerank 测试失败: 响应不是有效 JSON\"}\n\n if is_dashscope:\n results = (response_data.get(\"output\") or {}).get(\"results\")\n else:\n results = response_data.get(\"results\") if isinstance(response_data, dict) else response_data\n if not isinstance(results, list) or not results:\n return {\"success\": False, \"message\": \"Rerank 测试失败: 响应中未返回排序结果\"}\n\n return {\"success\": True, \"message\": \"连接成功\"}\n\n\nasync def _check_openai_compatible_model(\n api_domain: str,\n api_key: str,\n base_model: str,\n) -> dict:\n domain = api_domain.rstrip(\"/\")\n headers = _build_openai_headers(api_key)\n capabilities = {\n \"models\": False,\n \"chat\": False,\n \"tool_calling\": False,\n }\n\n async with httpx.AsyncClient() as client:\n try:\n response = await client.get(\n f\"{domain}/models\",\n headers=headers,\n timeout=10,\n )\n response.raise_for_status()\n capabilities[\"models\"] = True\n except httpx.HTTPStatusError:\n # 部分兼容服务未实现 /models，但不影响按已配置模型名称调用。\n pass\n\n chat_body = {\n \"model\": base_model,\n \"messages\": [{\"role\": \"user\", \"content\": \"请回复 ok\"}],\n \"stream\": False,\n }\n try:\n response = await client.post(\n f\"{domain}/chat/completions\",\n headers=headers,\n json=chat_body,\n timeout=30,\n )\n response.raise_for_status()\n if not _has_chat_choices(response):\n return {\n \"success\": False,\n \"message\": \"基础对话测试失败: 模型响应格式异常，未返回 choices\",\n \"capabilities\": capabilities,\n }\n capabilities[\"chat\"] = True\n except httpx.HTTPStatusError as error:\n return {\n \"success\": False,\n \"message\": f\"基础对话测试失败: {_http_error_detail(error)}\",\n \"capabilities\": capabilities,\n }\n\n tool_body = {\n **chat_body,\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": \"请判断是否需要调用 connection_test 工具\",\n }\n ],\n \"tools\": [\n {\n \"type\": \"function\",\n \"function\": {\n \"name\": \"connection_test\",\n \"description\": \"用于测试模型服务是否支持 Agent 工具调用\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {},\n },\n },\n }\n ],\n \"tool_choice\": \"auto\",\n }\n try:\n response = await client.post(\n f\"{domain}/chat/completions\",\n headers=headers,\n json=tool_body,\n timeout=30,\n )\n response.raise_for_status()\n if not _has_chat_choices(response):\n return {\n \"success\": False,\n \"message\": \"基础对话可用，但 Agent 工具调用测试响应格式异常\",\n \"capabilities\": capabilities,\n }\n capabilities[\"tool_calling\"] = True\n except httpx.HTTPStatusError as error:\n detail = _http_error_detail(error)\n if \"--enable-auto-tool-choice\" in detail:\n message = (\n \"基础对话可用，但 Agent 工具调用不可用: \"\n f\"{detail}。请在模型服务启用自动工具调用。\"\n )\n else:\n message = f\"基础对话可用，但 Agent 工具调用测试失败: {detail}\"\n return {\n \"success\": False,\n \"message\": message,\n \"capabilities\": capabilities,\n }\n\n message = \"连接成功，基础对话和 Agent 工具调用均可用\"\n if not capabilities[\"models\"]:\n message += \"；模型列表接口不可用\"\n\n return {\n \"success\": True,\n \"message\": message,\n \"capabilities\": capabilities,\n }\n\n\nasync def query_model_list(keyword: str = None, model_type: int = None) -> List[dict]:\n with pool.get_session() as session:\n query = session.query(TAiModel)\n if keyword:\n query = query.filter(TAiModel.name.like(f\"%{keyword}%\"))\n \n if model_type:\n query = query.filter(TAiModel.model_type == model_type)\n \n # Order by default_model desc, then name\n models = query.order_by(desc(TAiModel.default_model), TAiModel.name).all()\n \n result = []\n for model in models:\n m_dict = model_to_dict(model)\n result.append(m_dict)\n return result\n\nasync def get_model_detail(model_id: int) -> dict:\n with pool.get_session() as session:\n model = session.query(TAiModel).filter(TAiModel.id == model_id).first()\n if not model:\n raise MyException(SysCodeEnum.PARAM_ERROR, \"Model not found\")\n \n data = model_to_dict(model)\n if data.get('config'):\n try:\n data['config_list'] = json.loads(data['config'])\n except:\n data['config_list'] = []\n else:\n data['config_list'] = []\n return data\n\nasync def add_model(data: dict) -> bool:\n with pool.get_session() as session:\n # Check if default\n model_type = data.get('model_type', 1)\n \n # Only LLM (type 1) can be default\n is_default = False\n if model_type == 1:\n count = session.query(TAiModel).filter(\n TAiModel.model_type == 1\n ).count()\n is_default = (count == 0) # First LLM is default\n\n config_list = data.get('config_list', [])\n config_str = json.dumps(config_list)\n\n # 处理 api_key：空字符串转换为 None\n api_key = data.get('api_key')\n if api_key is not None and api_key.strip() == '':\n api_key = None\n\n new_model = TAiModel(\n name=data['name'],\n base_model=data['base_model'],\n model_type=data.get('model_type', 1), # Default to 1 (LLM)\n supplier=data.get('supplier', 1),\n protocol=data.get('protocol', 1),\n api_domain=data['api_domain'],\n api_key=api_key,\n config=config_str,\n default_model=is_default,\n status=1,\n create_time=int(datetime.now().timestamp())\n )\n session.add(new_model)\n session.commit()\n return True\n\nasync def update_model(model_id: int, data: dict) -> bool:\n with pool.get_session() as session:\n model = session.query(TAiModel).filter(TAiModel.id == model_id).first()\n if not model:\n raise MyException(SysCodeEnum.PARAM_ERROR, \"Model not found\")\n \n # 更新所有可修改的字段\n if 'name' in data:\n model.name = data['name']\n if 'base_model' in data:\n model.base_model = data['base_model']\n if 'supplier' in data:\n model.supplier = data['supplier']\n if 'model_type' in data:\n model.model_type = data['model_type']\n if 'protocol' in data:\n model.protocol = data['protocol']\n if 'api_domain' in data:\n model.api_domain = data['api_domain']\n if 'api_key' in data:\n # 处理 api_key：空字符串转换为 None\n api_key = data['api_key']\n if api_key is not None and api_key.strip() == '':\n api_key = None\n model.api_key = api_key\n \n if 'config_list' in data:\n model.config = json.dumps(data['config_list'])\n \n session.commit()\n return True\n\nasync def delete_model(model_id: int) -> bool:\n with pool.get_session() as session:\n model = session.query(TAiModel).filter(TAiModel.id == model_id).first()\n if not model:\n raise MyException(SysCodeEnum.PARAM_ERROR, \"Model not found\")\n \n if model.default_model:\n raise MyException(SysCodeEnum.PARAM_ERROR, \"Cannot delete default model\")\n \n session.delete(model)\n session.commit()\n return True\n\nasync def set_default_model(model_id: int) -> bool:\n with pool.get_session() as session:\n model = session.query(TAiModel).filter(TAiModel.id == model_id).first()\n if not model:\n raise MyException(SysCodeEnum.PARAM_ERROR, \"Model not found\")\n \n if model.model_type != 1:\n raise MyException(SysCodeEnum.PARAM_ERROR, \"Only LLM can be set as default\")\n\n if model.default_model:\n return True\n \n # Unset previous default for LLM\n session.query(TAiModel).filter(\n TAiModel.default_model == True,\n TAiModel.model_type == 1\n ).update({TAiModel.default_model: False})\n \n model.default_model = True\n session.commit()\n return True\n\nasync def get_default_model() -> Optional[dict]:\n \"\"\"\n 查询默认模型\n :return: 默认模型信息，如果不存在返回None\n \"\"\"\n with pool.get_session() as session:\n model = session.query(TAiModel).filter(\n TAiModel.default_model == True,\n TAiModel.model_type == 1\n ).first()\n \n if not model:\n return None\n \n return model_to_dict(model)\n\nasync def check_llm_status(data: dict) -> dict:\n \"\"\"\n 测试模型连接状态\n :param data: 模型配置数据\n :return: 测试结果\n \"\"\"\n supplier = data.get('supplier', 1)\n model_type = data.get('model_type', 1)\n api_key = data.get('api_key') or ''\n api_domain = data.get('api_domain', '')\n base_model = data.get('base_model', '')\n \n if not api_domain:\n return {\"success\": False, \"message\": \"API 域名不能为空\"}\n \n try:\n if model_type == 2:\n return await _check_embedding_model_legacy(supplier, api_domain, api_key)\n if model_type == 3:\n if not base_model:\n return {\"success\": False, \"message\": \"基础模型不能为空\"}\n return await _check_rerank_model(\n supplier,\n api_domain,\n api_key,\n base_model,\n )\n\n if supplier == 3:\n domain = api_domain.rstrip(\"/\")\n if domain.endswith('/v1'):\n domain = domain[:-3]\n async with httpx.AsyncClient() as client:\n resp = await client.get(f\"{domain}/api/tags\", timeout=5)\n resp.raise_for_status()\n return {\n \"success\": True,\n \"message\": \"连接成功\",\n \"capabilities\": {\"models\": True},\n }\n\n if not base_model:\n return {\"success\": False, \"message\": \"基础模型不能为空\"}\n\n return await _check_openai_compatible_model(api_domain, api_key, base_model)\n \n except httpx.HTTPStatusError as e:\n if e.response.status_code == 401:\n return {\"success\": False, \"message\": \"API Key 无效或未授权\"}\n elif e.response.status_code == 404:\n return {\"success\": False, \"message\": \"API 端点不存在，请检查 API 域名\"}\n else:\n return {\"success\": False, \"message\": f\"连接失败: HTTP {e.response.status_code}\"}\n except httpx.TimeoutException:\n return {\"success\": False, \"message\": \"连接超时，请检查网络或 API 域名\"}\n except httpx.ConnectError:\n return {\"success\": False, \"message\": \"无法连接到服务器，请检查 API 域名\"}\n except Exception as e:\n logger.error(f\"测试模型连接失败: {e}\")\n return {\"success\": False, \"message\": f\"连接失败: {str(e)}\"}\n\nasync def fetch_base_models(supplier: int, api_key: str = None, api_domain: str = None) -> List[str]:\n try:\n # OpenAI\n if supplier == 1:\n if not api_key:\n return []\n domain = api_domain or \"https://api.openai.com/v1\"\n url = f\"{domain}/models\"\n headers = {\"Authorization\": f\"Bearer {api_key}\"}\n async with httpx.AsyncClient() as client:\n resp = await client.get(url, headers=headers, timeout=10)\n resp.raise_for_status()\n data = resp.json()\n models = [m['id'] for m in data.get('data', []) if 'gpt' in m['id']]\n return sorted(models)\n \n # Ollama\n elif supplier == 3:\n domain = api_domain or \"http://localhost:11434\"\n # Ollama API structure: GET /api/tags\n if domain.endswith('/v1'):\n domain = domain[:-3] # Strip /v1 if present\n \n url = f\"{domain}/api/tags\"\n async with httpx.AsyncClient() as client:\n resp = await client.get(url, timeout=5)\n resp.raise_for_status()\n data = resp.json()\n models = [m['name'] for m in data.get('models', [])]\n return sorted(models)\n \n # vLLM\n elif supplier == 4:\n if not api_domain:\n return []\n url = f\"{api_domain}/models\"\n headers = {}\n if api_key:\n headers[\"Authorization\"] = f\"Bearer {api_key}\"\n async with httpx.AsyncClient() as client:\n resp = await client.get(url, headers=headers, timeout=5)\n resp.raise_for_status()\n data = resp.json()\n models = [m['id'] for m in data.get('data', [])]\n return sorted(models)\n \n # DeepSeek\n elif supplier == 5:\n # Similar to OpenAI\n if not api_key:\n return []\n domain = api_domain or \"https://api.deepseek.com\"\n url = f\"{domain}/models\"\n headers = {\"Authorization\": f\"Bearer {api_key}\"}\n async with httpx.AsyncClient() as client:\n resp = await client.get(url, headers=headers, timeout=10)\n resp.raise_for_status()\n data = resp.json()\n models = [m['id'] for m in data.get('data', [])]\n return sorted(models)\n \n # MiniMax\n elif supplier == 10:\n if not api_key:\n return []\n domain = api_domain or \"https://api.minimaxi.com/v1\"\n url = f\"{domain}/models\"\n headers = {\"Authorization\": f\"Bearer {api_key}\"}\n \n try:\n async with httpx.AsyncClient() as client:\n resp = await client.get(url, headers=headers, timeout=10)\n resp.raise_for_status()\n data = resp.json()\n # 根据OpenAI兼容API格式提取模型列表\n if 'data' in data:\n models = [m['id'] for m in data.get('data', [])]\n else:\n # 如果没有data字段，尝试其他格式\n models = [m['id'] for m in data] if isinstance(data, list) else []\n return sorted(models)\n except:\n # 如果获取模型列表失败，返回一些常见的MiniMax模型名称\n return [\"MiniMax-M2.1\", \"abab6.5s-chat\", \"abab5.5-chat\"]\n \n # Fallback or other providers: Return empty list or hardcoded common ones?\n # For now return empty, frontend can allow manual entry\n return []\n\n except Exception as e:\n logger.error(f\"Failed to fetch models for supplier {supplier}: {e}\")\n # Return empty list on error\n return []\n"}}}} |