591 lines
21 KiB
Python
591 lines
21 KiB
Python
"""
|
||
PMFlow AI 核心能力封装
|
||
|
||
技术栈：LangChain 0.3+、OpenAI SDK、Redis
|
||
- 封装 GPT-4o mini 调用（流式/非流式）
|
||
- Prompt 集中在 config/ai_prompts
|
||
- AI 调用日志写入 ai_call_logs
|
||
- Redis 结果缓存（1 小时 TTL）
|
||
- 指数退避重试
|
||
"""
|
||
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import time
|
||
from datetime import datetime
|
||
from typing import Any, AsyncGenerator, Optional, Union
|
||
from sqlalchemy import and_, select
|
||
|
||
# LangChain & OpenAI
|
||
try:
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
|
||
from langchain_core.language_models.chat_models import BaseChatModel
|
||
LANGCHAIN_AVAILABLE = True
|
||
except ImportError:
|
||
LANGCHAIN_AVAILABLE = False
|
||
BaseChatModel = object
|
||
|
||
# Redis
|
||
try:
|
||
import redis.asyncio as redis
|
||
REDIS_AVAILABLE = True
|
||
except ImportError:
|
||
REDIS_AVAILABLE = False
|
||
|
||
# Token 定价 (GPT-4o-mini, 2024)
|
||
INPUT_PRICE_PER_1K = 0.00015
|
||
OUTPUT_PRICE_PER_1K = 0.0006
|
||
CACHE_TTL_SECONDS = 3600 # 1 小时
|
||
PROMPT_PROFILE_CACHE_TTL_SECONDS = 60
|
||
|
||
from ..config.ai_prompts import build_default_role_prompt_bundle
|
||
from ..database import AsyncSessionLocal
|
||
from ..models import Integration
|
||
|
||
|
||
def _t2s(v: Union[str, int]) -> str:
|
||
"""统一转为 str"""
|
||
return str(v) if v is not None else ""
|
||
|
||
|
||
def _hash_input(fn: str, *args, **kwargs) -> str:
|
||
"""生成缓存 key"""
|
||
payload = json.dumps({"fn": fn, "args": args, "kwargs": kwargs}, sort_keys=True, default=str)
|
||
return hashlib.sha256(payload.encode()).hexdigest()
|
||
|
||
|
||
def _estimate_cost(prompt_tokens: int, completion_tokens: int) -> float:
|
||
return (prompt_tokens * INPUT_PRICE_PER_1K + completion_tokens * OUTPUT_PRICE_PER_1K) / 1000
|
||
|
||
|
||
def _tenant_to_int(tenant_id: Union[str, int]) -> int:
|
||
try:
|
||
return int(tenant_id)
|
||
except (TypeError, ValueError):
|
||
return 1
|
||
|
||
|
||
class AIService:
|
||
"""PMFlow AI 服务"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_key: Optional[str] = None,
|
||
model: str = "gpt-4o-mini",
|
||
redis_url: Optional[str] = None,
|
||
max_retries: int = 3,
|
||
):
|
||
self.api_key = api_key or os.getenv("OPENAI_API_KEY", "")
|
||
self.model = model
|
||
self.redis_url = redis_url or os.getenv("REDIS_URL")
|
||
self.max_retries = max_retries
|
||
self._llm: Optional[BaseChatModel] = None
|
||
self._redis: Optional[Any] = None
|
||
self._prompt_profile_cache: dict[int, tuple[float, tuple[str, str]]] = {}
|
||
|
||
def _get_llm(self) -> BaseChatModel:
|
||
if not LANGCHAIN_AVAILABLE:
|
||
raise RuntimeError("请安装 langchain-openai: pip install langchain-openai langchain-core")
|
||
if not self.api_key:
|
||
raise RuntimeError("请设置 OPENAI_API_KEY 环境变量")
|
||
if self._llm is None:
|
||
self._llm = ChatOpenAI(
|
||
model=self.model,
|
||
api_key=self.api_key,
|
||
temperature=0.3,
|
||
)
|
||
return self._llm
|
||
|
||
async def _get_redis(self):
|
||
if not REDIS_AVAILABLE or not self.redis_url:
|
||
return None
|
||
if self._redis is None:
|
||
self._redis = redis.from_url(self.redis_url, decode_responses=True)
|
||
return self._redis
|
||
|
||
async def _cache_get(self, key: str) -> Optional[str]:
|
||
r = await self._get_redis()
|
||
if r is None:
|
||
return None
|
||
try:
|
||
return await r.get(key)
|
||
except Exception:
|
||
return None
|
||
|
||
async def _cache_set(self, key: str, value: str, ttl: int = CACHE_TTL_SECONDS):
|
||
r = await self._get_redis()
|
||
if r is None:
|
||
return
|
||
try:
|
||
await r.setex(key, ttl, value)
|
||
except Exception:
|
||
pass
|
||
|
||
async def _resolve_active_system_prompt(
|
||
self,
|
||
*,
|
||
tenant_id: Union[str, int],
|
||
db=None,
|
||
) -> tuple[str, str]:
|
||
tenant_int = _tenant_to_int(tenant_id)
|
||
now = time.time()
|
||
cached = self._prompt_profile_cache.get(tenant_int)
|
||
if cached and cached[0] > now:
|
||
return cached[1]
|
||
|
||
role_name = ""
|
||
system_prompt = ""
|
||
try:
|
||
async def _load_profile(session) -> tuple[str, str]:
|
||
row = (
|
||
await session.execute(
|
||
select(Integration).where(
|
||
and_(
|
||
Integration.tenant_id == int(tenant_int),
|
||
Integration.type == "ai_prompt_profiles",
|
||
)
|
||
)
|
||
)
|
||
).scalar_one_or_none()
|
||
bundle = row.config if row is not None and isinstance(row.config, dict) else {}
|
||
active_key = str(bundle.get("active_role_key") or "").strip()
|
||
roles = bundle.get("roles") if isinstance(bundle.get("roles"), list) else []
|
||
selected = None
|
||
for item in roles:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
if str(item.get("role_key") or "").strip() == active_key:
|
||
selected = item
|
||
break
|
||
if selected is None:
|
||
for item in roles:
|
||
if isinstance(item, dict) and bool(item.get("enabled", True)):
|
||
selected = item
|
||
break
|
||
if not isinstance(selected, dict):
|
||
return ("", "")
|
||
return (
|
||
str(selected.get("role_name") or "").strip(),
|
||
str(selected.get("system_prompt") or "").strip(),
|
||
)
|
||
|
||
if db is not None:
|
||
role_name, system_prompt = await _load_profile(db)
|
||
else:
|
||
async with AsyncSessionLocal() as session:
|
||
role_name, system_prompt = await _load_profile(session)
|
||
except Exception:
|
||
role_name = ""
|
||
system_prompt = ""
|
||
|
||
if len(system_prompt) < 20:
|
||
fallback = build_default_role_prompt_bundle()
|
||
active_key = str(fallback.get("active_role_key") or "").strip()
|
||
roles = fallback.get("roles") if isinstance(fallback.get("roles"), list) else []
|
||
selected = None
|
||
for row in roles:
|
||
if not isinstance(row, dict):
|
||
continue
|
||
if str(row.get("role_key") or "").strip() == active_key:
|
||
selected = row
|
||
break
|
||
if selected is None and roles and isinstance(roles[0], dict):
|
||
selected = roles[0]
|
||
if isinstance(selected, dict):
|
||
role_name = str(selected.get("role_name") or role_name or "默认角色").strip()
|
||
system_prompt = str(selected.get("system_prompt") or "").strip()
|
||
|
||
result = (role_name, system_prompt)
|
||
self._prompt_profile_cache[tenant_int] = (now + PROMPT_PROFILE_CACHE_TTL_SECONDS, result)
|
||
return result
|
||
|
||
async def _log_call(
|
||
self,
|
||
db,
|
||
tenant_id: str,
|
||
function_name: str,
|
||
input_hash: str,
|
||
status: str,
|
||
prompt_tokens: int = 0,
|
||
completion_tokens: int = 0,
|
||
latency_ms: Optional[int] = None,
|
||
error_message: Optional[str] = None,
|
||
):
|
||
try:
|
||
from ..models import AICallLog
|
||
|
||
total = prompt_tokens + completion_tokens
|
||
cost = _estimate_cost(prompt_tokens, completion_tokens)
|
||
log = AICallLog(
|
||
tenant_id=_t2s(tenant_id),
|
||
function_name=function_name,
|
||
input_hash=input_hash,
|
||
status=status,
|
||
prompt_tokens=prompt_tokens,
|
||
completion_tokens=completion_tokens,
|
||
total_tokens=total,
|
||
cost_usd=cost,
|
||
latency_ms=latency_ms,
|
||
error_message=error_message,
|
||
)
|
||
db.add(log)
|
||
await db.commit()
|
||
except Exception:
|
||
pass
|
||
|
||
async def _invoke_with_retry(
|
||
self,
|
||
prompt: str,
|
||
db=None,
|
||
tenant_id: str = "default",
|
||
function_name: str = "unknown",
|
||
use_cache: bool = True,
|
||
cache_key: Optional[str] = None,
|
||
project_id: Optional[int] = None,
|
||
) -> str:
|
||
"""带缓存、重试、日志的 LLM 调用"""
|
||
input_hash = cache_key or _hash_input(function_name, prompt)
|
||
|
||
if use_cache and db is None:
|
||
cached = await self._cache_get(input_hash)
|
||
if cached:
|
||
return cached
|
||
|
||
# 预算前置检查（仅在有 db 会话时启用）
|
||
if db is not None:
|
||
try:
|
||
from .budget_enforcement_service import BudgetEnforcementService
|
||
|
||
tenant_int = _tenant_to_int(tenant_id)
|
||
# 预估：按 prompt 长度粗略估算，保守给定 completion 预算
|
||
estimated_prompt_tokens = max(len(prompt) // 4, 1)
|
||
estimated_completion_tokens = 1200
|
||
estimated_cost = _estimate_cost(estimated_prompt_tokens, estimated_completion_tokens)
|
||
|
||
budget_service = BudgetEnforcementService(db)
|
||
budget_check = await budget_service.check_budget_before_ai_call(
|
||
tenant_id=tenant_int,
|
||
estimated_cost_usd=estimated_cost,
|
||
project_id=project_id,
|
||
)
|
||
if not budget_check.get("allowed", True):
|
||
raise RuntimeError("预算超限，AI 调用被阻止（hard limit）")
|
||
except RuntimeError:
|
||
raise
|
||
except Exception:
|
||
# 预算系统异常不阻断主流程，避免影响现有可用性
|
||
pass
|
||
|
||
llm = self._get_llm()
|
||
_, system_prompt = await self._resolve_active_system_prompt(
|
||
tenant_id=tenant_id,
|
||
db=db,
|
||
)
|
||
last_err = None
|
||
for attempt in range(self.max_retries):
|
||
try:
|
||
t0 = time.perf_counter()
|
||
messages = [HumanMessage(content=prompt)]
|
||
if system_prompt:
|
||
messages = [SystemMessage(content=system_prompt), *messages]
|
||
resp = await llm.ainvoke(messages)
|
||
latency_ms = int((time.perf_counter() - t0) * 1000)
|
||
text = resp.content if hasattr(resp, "content") else str(resp)
|
||
pt = getattr(resp, "response_metadata", None) or {}
|
||
usage = (pt or {}).get("usage", {}) or (getattr(resp, "response_metadata", None) or {}).get("usage", {})
|
||
prompt_tokens = usage.get("input_tokens", 0) or usage.get("prompt_tokens", 0)
|
||
completion_tokens = usage.get("output_tokens", 0) or usage.get("completion_tokens", 0)
|
||
|
||
if db:
|
||
await self._log_call(
|
||
db, tenant_id, function_name, input_hash, "success",
|
||
prompt_tokens, completion_tokens, latency_ms,
|
||
)
|
||
# 按实际 token 成本累计预算
|
||
try:
|
||
from .budget_enforcement_service import BudgetEnforcementService
|
||
|
||
actual_cost = _estimate_cost(prompt_tokens, completion_tokens)
|
||
await BudgetEnforcementService(db).update_budget_usage_after_ai_call(
|
||
tenant_id=_tenant_to_int(tenant_id),
|
||
actual_cost_usd=actual_cost,
|
||
project_id=project_id,
|
||
)
|
||
except Exception:
|
||
pass
|
||
if use_cache:
|
||
await self._cache_set(input_hash, text)
|
||
return text
|
||
except Exception as e:
|
||
last_err = e
|
||
wait = 2 ** attempt
|
||
time.sleep(wait)
|
||
|
||
if db:
|
||
await self._log_call(
|
||
db, tenant_id, function_name, input_hash, "failed",
|
||
error_message=str(last_err),
|
||
)
|
||
raise last_err
|
||
|
||
def _parse_json(self, text: str) -> dict:
|
||
text = text.strip()
|
||
for p in ("```json", "```"):
|
||
if text.startswith(p):
|
||
text = text[len(p):]
|
||
if text.endswith("```"):
|
||
text = text[:-3]
|
||
return json.loads(text.strip())
|
||
|
||
# ========================================================================
|
||
# 核心函数
|
||
# ========================================================================
|
||
|
||
async def generate_charter_from_signal(
|
||
self,
|
||
signal_content: str,
|
||
tenant_id: Union[str, int] = "default",
|
||
db=None,
|
||
) -> dict:
|
||
"""从市场/模型信号生成需求宪章草案"""
|
||
from ..config.ai_prompts import PROMPT_GENERATE_CHARTER
|
||
|
||
tid = _t2s(tenant_id)
|
||
prompt = PROMPT_GENERATE_CHARTER.format(signal_content=signal_content)
|
||
text = await self._invoke_with_retry(
|
||
prompt, db=db, tenant_id=tid, function_name="generate_charter_from_signal",
|
||
)
|
||
return self._parse_json(text)
|
||
|
||
async def decompose_tasks_from_charter(
|
||
self,
|
||
charter_content: str,
|
||
project_id: Union[str, int],
|
||
tenant_id: Union[str, int] = "default",
|
||
db=None,
|
||
) -> dict:
|
||
"""从宪章拆解任务，自动识别所需技能标签，返回带技能要求的任务列表"""
|
||
from ..config.ai_prompts import PROMPT_DECOMPOSE_TASKS
|
||
|
||
tid = _t2s(tenant_id)
|
||
prompt = PROMPT_DECOMPOSE_TASKS.format(
|
||
charter_content=charter_content,
|
||
project_id=_t2s(project_id),
|
||
)
|
||
text = await self._invoke_with_retry(
|
||
prompt, db=db, tenant_id=tid, function_name="decompose_tasks_from_charter",
|
||
)
|
||
return self._parse_json(text)
|
||
|
||
async def recommend_assignees_for_task(
|
||
self,
|
||
task_content: str,
|
||
required_skill_tag_ids: list,
|
||
project_id: Union[str, int],
|
||
tenant_id: Union[str, int] = "default",
|
||
db=None,
|
||
candidates_json: Optional[str] = None,
|
||
) -> dict:
|
||
"""基于人员能力图谱、负载、熟练度，推荐 Top3 人员"""
|
||
from ..config.ai_prompts import PROMPT_RECOMMEND_ASSIGNEES
|
||
|
||
tid = _t2s(tenant_id)
|
||
candidates = candidates_json or "[]"
|
||
prompt = PROMPT_RECOMMEND_ASSIGNEES.format(
|
||
task_content=task_content,
|
||
required_skill_tag_ids=json.dumps(required_skill_tag_ids),
|
||
candidates_json=candidates,
|
||
)
|
||
text = await self._invoke_with_retry(
|
||
prompt, db=db, tenant_id=tid, function_name="recommend_assignees_for_task",
|
||
)
|
||
return self._parse_json(text)
|
||
|
||
async def update_user_skill_scores(
|
||
self,
|
||
task_id: Union[str, int],
|
||
tenant_id: Union[str, int] = "default",
|
||
db=None,
|
||
quality_score: float = 0.0,
|
||
efficiency_score: float = 0.0,
|
||
assignee_ids: Optional[list] = None,
|
||
skill_tag_ids: Optional[list] = None,
|
||
) -> dict:
|
||
"""任务完成后，根据质量/效率评分更新相关人员技能分"""
|
||
from sqlalchemy import select
|
||
from ..models import UserSkillBinding
|
||
from uuid import UUID
|
||
|
||
def to_uuid(v):
|
||
if v is None:
|
||
return None
|
||
if isinstance(v, UUID):
|
||
return v
|
||
s = str(v)
|
||
try:
|
||
return UUID(s) if len(s) == 36 else None
|
||
except ValueError:
|
||
return None
|
||
|
||
if db is None:
|
||
return {"updated_count": 0, "updates": []}
|
||
|
||
tid = _t2s(tenant_id)
|
||
updates = []
|
||
if assignee_ids and skill_tag_ids:
|
||
for uid in assignee_ids:
|
||
for sid in skill_tag_ids:
|
||
uid_u = to_uuid(uid)
|
||
sid_u = to_uuid(sid)
|
||
if not uid_u or not sid_u:
|
||
continue
|
||
try:
|
||
r = await db.execute(
|
||
select(UserSkillBinding).where(
|
||
UserSkillBinding.tenant_id == tid,
|
||
UserSkillBinding.user_id == uid_u,
|
||
UserSkillBinding.skill_tag_id == sid_u,
|
||
)
|
||
)
|
||
b = r.scalar_one_or_none()
|
||
if b:
|
||
b.quality_score = min(100, max(0, (b.quality_score or 0) + quality_score))
|
||
b.efficiency_score = min(100, max(0, (b.efficiency_score or 0) + efficiency_score))
|
||
b.last_used_at = datetime.utcnow()
|
||
await db.commit()
|
||
updates.append({
|
||
"user_id": str(uid),
|
||
"skill_tag_id": str(sid),
|
||
"quality_score_delta": quality_score,
|
||
"efficiency_score_delta": efficiency_score,
|
||
"last_used_at": b.last_used_at.isoformat(),
|
||
})
|
||
except Exception:
|
||
pass
|
||
|
||
return {"updated_count": len(updates), "updates": updates}
|
||
|
||
async def calculate_alignment_score(
|
||
self,
|
||
user_id: Union[str, int],
|
||
truth_node_id: Union[str, int],
|
||
tenant_id: Union[str, int] = "default",
|
||
db=None,
|
||
truth_content: str = "",
|
||
user_context: str = "",
|
||
) -> dict:
|
||
"""计算用户对真理节点的认知对齐度（0-100）"""
|
||
from ..config.ai_prompts import PROMPT_ALIGNMENT_SCORE
|
||
from sqlalchemy import select
|
||
from ..models import TruthNode
|
||
|
||
tid = _t2s(tenant_id)
|
||
if db and not truth_content:
|
||
r = await db.execute(
|
||
select(TruthNode).where(
|
||
TruthNode.id == truth_node_id,
|
||
TruthNode.tenant_id == tid,
|
||
)
|
||
)
|
||
tn = r.scalar_one_or_none()
|
||
if tn:
|
||
truth_content = tn.content or tn.title or ""
|
||
|
||
prompt = PROMPT_ALIGNMENT_SCORE.format(
|
||
user_id=_t2s(user_id),
|
||
truth_node_id=_t2s(truth_node_id),
|
||
truth_content=truth_content or "(无内容)",
|
||
user_context=user_context or "(无)",
|
||
)
|
||
text = await self._invoke_with_retry(
|
||
prompt, db=db, tenant_id=tid, function_name="calculate_alignment_score",
|
||
)
|
||
data = self._parse_json(text)
|
||
data["user_id"] = _t2s(user_id)
|
||
data["truth_node_id"] = _t2s(truth_node_id)
|
||
return data
|
||
|
||
async def detect_cognitive_conflicts(
|
||
self,
|
||
truth_node_id: Union[str, int],
|
||
tenant_id: Union[str, int] = "default",
|
||
db=None,
|
||
truth_content: str = "",
|
||
user_viewpoints: str = "",
|
||
) -> dict:
|
||
"""识别真理节点下的隐性认知冲突"""
|
||
from ..config.ai_prompts import PROMPT_COGNITIVE_CONFLICTS
|
||
from sqlalchemy import select
|
||
from ..models import TruthNode
|
||
|
||
tid = _t2s(tenant_id)
|
||
if db and not truth_content:
|
||
r = await db.execute(
|
||
select(TruthNode).where(
|
||
TruthNode.id == truth_node_id,
|
||
TruthNode.tenant_id == tid,
|
||
)
|
||
)
|
||
tn = r.scalar_one_or_none()
|
||
if tn:
|
||
truth_content = tn.content or tn.title or ""
|
||
|
||
prompt = PROMPT_COGNITIVE_CONFLICTS.format(
|
||
truth_node_id=_t2s(truth_node_id),
|
||
truth_content=truth_content or "(无内容)",
|
||
user_viewpoints=user_viewpoints or "(无)",
|
||
)
|
||
text = await self._invoke_with_retry(
|
||
prompt, db=db, tenant_id=tid, function_name="detect_cognitive_conflicts",
|
||
)
|
||
data = self._parse_json(text)
|
||
data["truth_node_id"] = _t2s(truth_node_id)
|
||
data["total_count"] = len(data.get("conflicts", []))
|
||
return data
|
||
|
||
async def evaluate_model_requirement(
|
||
self,
|
||
charter_content: str,
|
||
model_version_id: Union[str, int],
|
||
tenant_id: Union[str, int] = "default",
|
||
db=None,
|
||
model_info: str = "",
|
||
) -> dict:
|
||
"""评估需求与模型的匹配度"""
|
||
from ..config.ai_prompts import PROMPT_EVALUATE_MODEL_REQUIREMENT
|
||
|
||
tid = _t2s(tenant_id)
|
||
prompt = PROMPT_EVALUATE_MODEL_REQUIREMENT.format(
|
||
charter_content=charter_content,
|
||
model_version_id=_t2s(model_version_id),
|
||
model_info=model_info or "(未提供)",
|
||
)
|
||
text = await self._invoke_with_retry(
|
||
prompt, db=db, tenant_id=tid, function_name="evaluate_model_requirement",
|
||
)
|
||
return self._parse_json(text)
|
||
|
||
# ========================================================================
|
||
# 流式调用
|
||
# ========================================================================
|
||
|
||
async def stream_chat(self, prompt: str) -> AsyncGenerator[str, None]:
|
||
"""流式生成"""
|
||
llm = self._get_llm()
|
||
async for chunk in llm.astream([HumanMessage(content=prompt)]):
|
||
if hasattr(chunk, "content") and chunk.content:
|
||
yield chunk.content
|
||
|
||
|
||
# 单例
|
||
_ai_service: Optional[AIService] = None
|
||
|
||
|
||
def get_ai_service() -> AIService:
|
||
global _ai_service
|
||
if _ai_service is None:
|
||
_ai_service = AIService()
|
||
return _ai_service
|