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