leocrm/app/plugins/builtins/knowledge/services.py

"""Knowledge extraction services — LLM-based entity/relationship extraction."""
from __future__ import annotations
import logging
import uuid
from typing import Any
from sqlalchemy import select, update
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}