""" 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