PMFlow/backend/app/services/llm_service/mixin_task_helpers.py

from __future__ import annotations

from .shared import *

class LLMTaskHelpersMixin:
    @staticmethod
    def _extract_embedding_vector(payload: Any) -> list[float]:
        if not isinstance(payload, dict):
            return []
        if isinstance(payload.get("embedding"), list):
            return [float(item) for item in payload.get("embedding") or []]
        embeddings = payload.get("embeddings")
        if isinstance(embeddings, list) and embeddings and isinstance(embeddings[0], list):
            return [float(item) for item in embeddings[0]]
        data = payload.get("data")
        if isinstance(data, list) and data and isinstance(data[0], dict):
            embedding = data[0].get("embedding")
            if isinstance(embedding, list):
                return [float(item) for item in embedding]
        return []

    async def _request_openai_embedding(self, *, endpoint: str, headers: dict[str, str], payload: dict[str, Any]) -> list[float]:
        async with httpx.AsyncClient(timeout=self.request_timeout, trust_env=False) as client:
            resp = await client.post(endpoint, headers=headers, json=payload)
            resp.raise_for_status()
            return self._extract_embedding_vector(resp.json())

    async def _request_ollama_embedding(self, *, host: str, model: str, text: str) -> list[float]:
        base_url = str(host or os.getenv("OLLAMA_HOST", "http://127.0.0.1:11434")).rstrip("/")
        headers = {"Content-Type": "application/json"}
        candidates = [
            (f"{base_url}/api/embed", {"model": model, "input": text}),
            (f"{base_url}/api/embeddings", {"model": model, "prompt": text}),
        ]
        for endpoint, payload in candidates:
            try:
                async with httpx.AsyncClient(timeout=self.request_timeout, trust_env=False) as client:
                    resp = await client.post(endpoint, headers=headers, json=payload)
                    resp.raise_for_status()
                    vector = self._extract_embedding_vector(resp.json())
                if vector:
                    return vector
            except Exception:
                continue
        return []

    async def generate_task_breakdown_async(self, description: str, **kwargs: Any) -> list[dict]:
        if not description.strip():
            return []
        prompt = (
            "请将以下需求拆解为 5-8 条执行任务，返回 JSON 数组，"
            "每项格式: {\"task_name\":\"\",\"type\":\"backend|frontend|database|ai|devops\","
            "\"estimated_hours\":4,\"priority\":\"high|medium|low\"}。\n"
            f"需求：{description[:1200]}"
        )
        text = await self.generate(prompt, function_name="generate_task_breakdown", **kwargs)
        rows = self._parse_json_array(text)
        if rows:
            return rows
        return [
            {
                "task_name": "需求分析",
                "type": "ai",
                "priority": "medium",
                "estimated_hours": 4,
                "description": description[:200],
            }
        ]

    async def get_embedding_async(self, text: str, **kwargs: Any) -> list[float]:
        clean_text = str(text or "").strip()
        if not clean_text:
            return []
        db = kwargs.get("db") or kwargs.get("db_session")
        tenant_id = int(kwargs.get("tenant_id") or 1)
        if db is None:
            async with AsyncSessionLocal() as session:
                return await self.get_embedding_async(clean_text, db=session, tenant_id=tenant_id)
        engine: dict[str, Any] = {}
        try:
            engine = await self._resolve_active_engine(db, tenant_id=tenant_id)
            provider = str(engine.get("source_provider") or engine.get("provider") or "").strip().lower()
            if provider == "openai":
                endpoint = f"{str(engine.get('base_url') or '').rstrip('/')}/embeddings"
                headers = {
                    "Authorization": f"Bearer {str(engine.get('api_key') or '').strip()}",
                    "Content-Type": "application/json",
                }
                payload = {
                    "model": str(engine.get("model") or "text-embedding-3-small"),
                    "input": clean_text,
                }
                vector = await self._request_openai_embedding(endpoint=endpoint, headers=headers, payload=payload)
                if vector:
                    return vector
            if provider == "ollama":
                vector = await self._request_ollama_embedding(
                    host=str(engine.get("base_url") or os.getenv("OLLAMA_HOST", "http://127.0.0.1:11434")),
                    model=str(engine.get("model") or os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text:latest")),
                    text=clean_text,
                )
                if vector:
                    return vector
        except Exception:
            engine = dict(engine) if isinstance(engine, dict) else {}
        vector = await self._request_ollama_embedding(
            host=str(os.getenv("OLLAMA_HOST", "http://127.0.0.1:11434")),
            model=str(os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text:latest")),
            text=clean_text,
        )
        if vector:
            return vector
        return []

    async def extract_context_insights_async(self, text: str, **kwargs: Any) -> dict:
        prompt = (
            "请提炼以下技术文档上下文的关键结论，返回 JSON："
            "{\"summary\":\"\",\"keywords\":[\"\",\"\"],\"risks\":[\"\",\"\"]}\n"
            f"上下文：{text[:2000]}"
        )
        raw = await self.generate(prompt, function_name="extract_context_insights", **kwargs)
        parsed = self._parse_json_object(raw)
        if parsed:
            return parsed
        return {"summary": text[:120], "keywords": []}

    async def find_relevant_context(self, **_: Any) -> list[dict]:
        return []

    async def evaluate_with_voting_async(self, signal_content: str, impact_score: float, **_: Any) -> dict:
        is_strategic = impact_score >= 0.8
        return {
            "is_strategic": is_strategic,
            "confidence": min(1.0, max(0.2, impact_score)),
            "decision_proposal": {
                "summary": signal_content[:120],
                "action": "follow_up" if is_strategic else "observe",
            },
            "tasks": [],
            "_voting": {
                "enabled": False,
                "models_used": [],
                "strategic_votes": 0,
                "total_votes": 0,
                "unanimous": True,
            },
        }

    async def audit_code_change_async(
        self,
        task_name: str,
        commit_message: str,
        code_diff: str,
        **_: Any,
    ) -> dict:
        passed = bool(task_name and commit_message and code_diff is not None)
        return {
            "passed": passed,
            "reason": "basic_stub_check_passed" if passed else "invalid_input",
        }