109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
import json
|
||
import os
|
||
from typing import Any, Dict
|
||
|
||
from openai import AsyncOpenAI
|
||
|
||
|
||
_client: AsyncOpenAI | None = None
|
||
|
||
|
||
def get_ai_client() -> AsyncOpenAI:
|
||
"""
|
||
Create (or reuse) a singleton AsyncOpenAI client.
|
||
|
||
The client is configured via:
|
||
- AI_API_KEY / OPENAI_API_KEY
|
||
- AI_BASE_URL (optional, defaults to official OpenAI endpoint)
|
||
- AI_MODEL (optional, defaults to gpt-4.1-mini or a similar capable model)
|
||
"""
|
||
global _client
|
||
if _client is not None:
|
||
return _client
|
||
|
||
api_key = os.getenv("AI_API_KEY") or os.getenv("OPENAI_API_KEY")
|
||
if not api_key:
|
||
raise RuntimeError("AI_API_KEY or OPENAI_API_KEY must be set in environment.")
|
||
|
||
base_url = os.getenv("AI_BASE_URL") # can point to OpenAI, DeepSeek, Qwen, etc.
|
||
|
||
_client = AsyncOpenAI(
|
||
api_key=api_key,
|
||
base_url=base_url or None,
|
||
)
|
||
return _client
|
||
|
||
|
||
def _build_requirement_prompt(raw_text: str) -> str:
|
||
"""
|
||
Build a clear system/user prompt for requirement analysis.
|
||
The model must output valid JSON only.
|
||
"""
|
||
return (
|
||
"你是一名资深的系统架构师，请阅读以下来自客户的原始需求文本，"
|
||
"提炼出清晰的交付方案，并严格按照指定 JSON 结构输出。\n\n"
|
||
"【要求】\n"
|
||
"1. 按功能模块拆分需求。\n"
|
||
"2. 每个模块给出简要说明和技术实现思路。\n"
|
||
"3. 估算建议工时（以人天或人小时为单位，使用数字）。\n"
|
||
"4. 可以根据你的经验给出每个模块的单价与小计金额，并给出总金额，"
|
||
"方便后续生成报价单。\n\n"
|
||
"【返回格式】请只返回 JSON，不要包含任何额外说明文字：\n"
|
||
"{\n"
|
||
' "modules": [\n'
|
||
" {\n"
|
||
' "name": "模块名称",\n'
|
||
' "description": "模块说明（可以为 Markdown 格式）",\n'
|
||
' "technical_approach": "技术实现思路（Markdown 格式）",\n'
|
||
' "estimated_hours": 16,\n'
|
||
' "unit_price": 800,\n'
|
||
' "subtotal": 12800\n'
|
||
" }\n"
|
||
" ],\n"
|
||
' "total_estimated_hours": 40,\n'
|
||
' "total_amount": 32000,\n'
|
||
' "notes": "整体方案备注（可选，Markdown 格式）"\n'
|
||
"}\n\n"
|
||
f"【客户原始需求】\n{raw_text}"
|
||
)
|
||
|
||
|
||
async def analyze_requirement(raw_text: str) -> Dict[str, Any]:
|
||
"""
|
||
Call the AI model to analyze customer requirements.
|
||
|
||
Returns a Python dict matching the JSON structure described
|
||
in `_build_requirement_prompt`.
|
||
"""
|
||
client = get_ai_client()
|
||
model = os.getenv("AI_MODEL", "gpt-4.1-mini")
|
||
|
||
prompt = _build_requirement_prompt(raw_text)
|
||
|
||
completion = await client.chat.completions.create(
|
||
model=model,
|
||
response_format={"type": "json_object"},
|
||
messages=[
|
||
{
|
||
"role": "system",
|
||
"content": (
|
||
"你是一名严谨的系统架构师，只能输出有效的 JSON，不要输出任何解释文字。"
|
||
),
|
||
},
|
||
{
|
||
"role": "user",
|
||
"content": prompt,
|
||
},
|
||
],
|
||
temperature=0.2,
|
||
)
|
||
|
||
content = completion.choices[0].message.content or "{}"
|
||
try:
|
||
data: Dict[str, Any] = json.loads(content)
|
||
except json.JSONDecodeError as exc:
|
||
raise RuntimeError(f"AI 返回的内容不是合法 JSON：{content}") from exc
|
||
|
||
return data
|
||
|