161 lines
6.9 KiB
Python
161 lines
6.9 KiB
Python
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",
|
||
}
|