Agent Zero
242 lines
9.6 KiB
Python
242 lines
9.6 KiB
Python
"""Knowledge extraction services — LLM-based entity/relationship extraction."""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.ai.llm_client import llm_complete
|
|
from app.plugins.builtins.knowledge.models import KnowledgeExtraction
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
EXTRACTION_PROMPT = """You are a knowledge extraction assistant for a CRM system.
|
|
Analyze the following text and extract entities and relationships.
|
|
|
|
Return JSON with this structure:
|
|
{
|
|
"entities": [
|
|
{"type": "person|company|project|topic", "name": "...", "description": "..."}
|
|
],
|
|
"relationships": [
|
|
{"source": "entity_name", "target": "entity_name", "type": "works_for|related_to|part_of|mentions"}
|
|
],
|
|
"confidence": 0.0-1.0
|
|
}
|
|
|
|
Text to analyze:
|
|
"""
|
|
|
|
async def extract_knowledge(
|
|
db: AsyncSession,
|
|
tenant_id: uuid.UUID,
|
|
source_type: str,
|
|
source_id: uuid.UUID,
|
|
source_title: str | None,
|
|
source_text: str,
|
|
user_id: uuid.UUID | None = None,
|
|
model: str = "openai/gpt-4o-mini",
|
|
) -> dict[str, Any]:
|
|
"""Extract entities and relationships from text using LLM."""
|
|
messages = [
|
|
{"role": "system", "content": "You are a knowledge extraction assistant. Return only valid JSON."},
|
|
{"role": "user", "content": EXTRACTION_PROMPT + source_text[:4000]},
|
|
]
|
|
response = await llm_complete(
|
|
model=model,
|
|
messages=messages,
|
|
temperature=0.2,
|
|
max_tokens=2000,
|
|
tenant_id=tenant_id,
|
|
db=db,
|
|
)
|
|
import json
|
|
try:
|
|
result = json.loads(response.get("content", "{}"))
|
|
except (json.JSONDecodeError, TypeError):
|
|
result = {"entities": [], "relationships": [], "confidence": 0.0}
|
|
extraction = KnowledgeExtraction(
|
|
tenant_id=tenant_id,
|
|
source_type=source_type,
|
|
source_id=source_id,
|
|
source_title=source_title,
|
|
extracted_entities=result.get("entities", []),
|
|
extracted_relationships=result.get("relationships", []),
|
|
confidence=float(result.get("confidence", 0.0)),
|
|
status="auto_created" if result.get("confidence", 0.0) >= 0.8 else "pending",
|
|
llm_model=model,
|
|
llm_cost_usd=response.get("cost_usd", 0.0),
|
|
created_by=user_id,
|
|
)
|
|
db.add(extraction)
|
|
await db.flush()
|
|
# Auto-create relationships in GraphRAG if confidence >= 0.8
|
|
if extraction.confidence >= 0.8 and extraction.extracted_relationships:
|
|
try:
|
|
from app.plugins.builtins.graph_rag.services import create_relationship
|
|
for rel in extraction.extracted_relationships:
|
|
# Only create if both source and target have IDs (resolved entities)
|
|
if rel.get("source_id") and rel.get("target_id"):
|
|
await create_relationship(
|
|
db=db,
|
|
tenant_id=tenant_id,
|
|
source_type=rel.get("source_type", "topic"),
|
|
source_id=uuid.UUID(rel["source_id"]),
|
|
target_type=rel.get("target_type", "topic"),
|
|
target_id=uuid.UUID(rel["target_id"]),
|
|
relationship_type=rel.get("type", "related_to"),
|
|
metadata={"extraction_id": str(extraction.id), "confidence": extraction.confidence},
|
|
owner_id=user_id,
|
|
)
|
|
except Exception as e:
|
|
logger.warning("Failed to auto-create relationships: %s", e)
|
|
await db.commit()
|
|
return {
|
|
"id": str(extraction.id),
|
|
"entities": extraction.extracted_entities,
|
|
"relationships": extraction.extracted_relationships,
|
|
"confidence": extraction.confidence,
|
|
"status": extraction.status,
|
|
}
|
|
|
|
async def ask_knowledge(
|
|
db: AsyncSession,
|
|
tenant_id: uuid.UUID,
|
|
question: str,
|
|
model: str = "openai/gpt-4o-mini",
|
|
) -> dict[str, Any]:
|
|
"""Answer a knowledge question using wiki articles + graph_rag as context."""
|
|
# Search wiki articles for context
|
|
from app.plugins.builtins.unified_search.provider_registry import get_search_registry
|
|
registry = get_search_registry()
|
|
evidence: list[dict[str, Any]] = []
|
|
context_parts = []
|
|
wiki_provider = registry.get("wiki_article")
|
|
if wiki_provider:
|
|
results = await wiki_provider._search_fts_filtered(
|
|
db=db, tsquery=question, tenant_id=tenant_id, limit=5, visible_ids=None
|
|
)
|
|
for r in results:
|
|
snippet = (r.get("content", "") or "")[:300]
|
|
title = r.get("title", "")
|
|
context_parts.append(f"Title: {title}\nContent: {snippet}")
|
|
evidence.append({
|
|
"id": str(r.get("id", "")),
|
|
"source_type": "wiki_article",
|
|
"source_id": str(r.get("id", "")),
|
|
"title": title,
|
|
"snippet": snippet,
|
|
"score": float(r.get("rank", 0.0)),
|
|
"url": f"/wiki?article={r.get('slug', '')}",
|
|
})
|
|
# Search graph_rag for relationships
|
|
graph_provider = registry.get("graph_relationship")
|
|
if graph_provider:
|
|
try:
|
|
results = await graph_provider._search_fts_filtered(
|
|
db=db, tsquery=question, tenant_id=tenant_id, limit=5, visible_ids=None
|
|
)
|
|
for r in results:
|
|
title = f"{r.get('source_type', '')} → {r.get('relationship_type', '')} → {r.get('target_type', '')}"
|
|
snippet = str(r.get("meta", ""))
|
|
context_parts.append(f"Relationship: {title}")
|
|
evidence.append({
|
|
"id": str(r.get("id", "")),
|
|
"source_type": "graph_relationship",
|
|
"source_id": str(r.get("id", "")),
|
|
"title": title,
|
|
"snippet": snippet,
|
|
"score": float(r.get("rank", 0.0)),
|
|
"url": None,
|
|
})
|
|
except Exception:
|
|
logger.warning("GraphRAG search failed in ask_knowledge", exc_info=True)
|
|
context = "\n\n".join(context_parts) if context_parts else "No knowledge base content found."
|
|
messages = [
|
|
{"role": "system", "content": f"You are a knowledge assistant. Answer based on this context:\n\n{context}"},
|
|
{"role": "user", "content": question},
|
|
]
|
|
response = await llm_complete(
|
|
model=model, messages=messages, temperature=0.3, max_tokens=1000,
|
|
tenant_id=tenant_id, db=db,
|
|
)
|
|
return {
|
|
"answer": response.get("content", ""),
|
|
"evidence": evidence,
|
|
"sources_used": list(set(e["source_type"] for e in evidence)),
|
|
"cost_usd": response.get("cost_usd", 0.0),
|
|
}
|
|
|
|
async def get_review_queue(
|
|
db: AsyncSession,
|
|
tenant_id: uuid.UUID,
|
|
page: int = 1,
|
|
page_size: int = 20,
|
|
) -> dict[str, Any]:
|
|
"""Get pending knowledge extractions for review."""
|
|
q = select(KnowledgeExtraction).where(
|
|
KnowledgeExtraction.tenant_id == tenant_id,
|
|
KnowledgeExtraction.status == "pending",
|
|
).order_by(KnowledgeExtraction.created_at.desc())
|
|
from sqlalchemy import func
|
|
count_q = select(func.count()).select_from(q.subquery())
|
|
total = (await db.execute(count_q)).scalar() or 0
|
|
offset = (page - 1) * page_size
|
|
result = await db.execute(q.offset(offset).limit(page_size))
|
|
items = [
|
|
{
|
|
"id": str(e.id), "source_type": e.source_type, "source_id": str(e.source_id),
|
|
"source_title": e.source_title, "entities": e.extracted_entities,
|
|
"relationships": e.extracted_relationships, "confidence": e.confidence,
|
|
"status": e.status, "created_at": e.created_at.isoformat() if e.created_at else None,
|
|
}
|
|
for e in result.scalars().all()
|
|
]
|
|
return {"items": items, "total": total, "page": page, "page_size": page_size}
|
|
|
|
async def review_extraction(
|
|
db: AsyncSession,
|
|
tenant_id: uuid.UUID,
|
|
extraction_id: uuid.UUID,
|
|
approved: bool,
|
|
user_id: uuid.UUID,
|
|
notes: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Approve or reject a knowledge extraction."""
|
|
from datetime import datetime, timezone
|
|
result = await db.execute(
|
|
select(KnowledgeExtraction).where(
|
|
KnowledgeExtraction.tenant_id == tenant_id,
|
|
KnowledgeExtraction.id == extraction_id,
|
|
)
|
|
)
|
|
extraction = result.scalar_one_or_none()
|
|
if not extraction:
|
|
return {"error": "Extraction not found"}
|
|
extraction.status = "approved" if approved else "rejected"
|
|
extraction.reviewed_by = user_id
|
|
extraction.reviewed_at = datetime.now(timezone.utc)
|
|
extraction.review_notes = notes
|
|
if approved and extraction.confidence < 0.8 and extraction.extracted_relationships:
|
|
try:
|
|
from app.plugins.builtins.graph_rag.services import create_relationship
|
|
for rel in extraction.extracted_relationships:
|
|
if rel.get("source_id") and rel.get("target_id"):
|
|
await create_relationship(
|
|
db=db, tenant_id=tenant_id,
|
|
source_type=rel.get("source_type", "topic"),
|
|
source_id=uuid.UUID(rel["source_id"]),
|
|
target_type=rel.get("target_type", "topic"),
|
|
target_id=uuid.UUID(rel["target_id"]),
|
|
relationship_type=rel.get("type", "related_to"),
|
|
metadata={"extraction_id": str(extraction.id), "reviewed": True},
|
|
owner_id=user_id,
|
|
)
|
|
except Exception as e:
|
|
logger.warning("Failed to create relationships after approval: %s", e)
|
|
await db.commit()
|
|
return {"id": str(extraction.id), "status": extraction.status}
|