llm_service.py 3.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. import aiohttp
  2. import asyncio
  3. import json
  4. import os
  5. from sqlalchemy.ext.asyncio import AsyncSession
  6. from sqlalchemy import select
  7. from app.core.config import settings
  8. from app.core.biz_exception import NotFoundError
  9. from app.schemas.llm import ParseUserInputsPayload, ParseUserInputsOut
  10. from app.models.schema import VasSchema
  11. # --- 配置区 ---
  12. # 请换成你新生成的 API Key
  13. API_KEY = settings.openai_api_key
  14. # API_URL = "https://api.openai.com/v1/chat/completions"
  15. API_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
  16. class LlmService:
  17. async def handle_parse(db: AsyncSession, payload: ParseUserInputsPayload):
  18. stmt = select(VasSchema).where(VasSchema.id == payload.schema_id)
  19. obj = (await db.execute(stmt)).scalar_one_or_none()
  20. if not obj:
  21. raise NotFoundError("Schema not exist")
  22. parsed_obj = await LlmService.parse_data_async(payload.input_raw_str, obj.schema_json)
  23. out = ParseUserInputsOut(parsed_obj=parsed_obj)
  24. return out
  25. @staticmethod
  26. async def parse_data_async(user_text: str, json_schema: dict):
  27. """
  28. [异步版本] 调用 LLM 解析数据
  29. """
  30. headers = {
  31. "Authorization": f"Bearer {API_KEY}",
  32. "Content-Type": "application/json"
  33. }
  34. # 构造 Prompt
  35. system_instruction = "You are a specialized data extraction API. Output valid JSON only."
  36. user_prompt = f"""
  37. Extract data from the text strictly based on the provided JSON Schema.
  38. [JSON Schema]
  39. {json.dumps(json_schema)}
  40. [User Text]
  41. {user_text}
  42. """
  43. payload = {
  44. # "model": "gpt-4o", # 或 gpt-3.5-turbo
  45. "model": "qwen-plus",
  46. "messages": [
  47. {"role": "system", "content": system_instruction},
  48. {"role": "user", "content": user_prompt}
  49. ],
  50. "temperature": 0,
  51. "response_format": {"type": "json_object"} # 强制 JSON 模式
  52. }
  53. async with aiohttp.ClientSession() as session:
  54. async with session.post(API_URL, headers=headers, json=payload, timeout=30) as response:
  55. # 1. 检查 HTTP 状态码
  56. if response.status != 200:
  57. error_text = await response.text()
  58. return {"error": f"HTTP {response.status}", "detail": error_text}
  59. # 2. 获取响应体
  60. result = await response.json()
  61. # 3. 提取并解析内容
  62. content_str = result['choices'][0]['message']['content']
  63. return json.loads(content_str)
  64. # --- 测试运行 ---
  65. if __name__ == "__main__":
  66. # 定义你的 Schema
  67. my_schema = {
  68. "type": "object",
  69. "properties": {
  70. "full_name": {"type": "string"},
  71. "budget": {"type": "integer"},
  72. "items": {"type": "array", "items": {"type": "string"}}
  73. },
  74. "required": ["full_name", "budget"]
  75. }
  76. # 模拟用户输入
  77. user_input = "我是张三,打算花2000块钱买个耳机和键盘。"
  78. print("正在解析...")
  79. result = LlmService.parse_data_async(user_input, my_schema)
  80. print(json.dumps(result, ensure_ascii=False, indent=2))