Agent Zero
362 lines
12 KiB
Python
362 lines
12 KiB
Python
"""Knowledge extraction — LLM-based relationship and entity extraction (H-EXT, H-ENT, H-AUTO, H-CONF).
|
|
|
|
Analyzes texts from knowledge sources (wiki, DMS, mail, communication)
|
|
and extracts:
|
|
- Named entities (persons, companies, projects) — H-ENT
|
|
- Relationships between entities — H-EXT
|
|
- Auto-creates relationships in GraphRAG — H-AUTO
|
|
- Confidence scores with low-confidence → review queue — H-CONF
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ExtractedEntity:
|
|
"""A named entity extracted from text (H-ENT)."""
|
|
|
|
name: str
|
|
entity_type: str # person, company, project, location, date, other
|
|
mentions: list[int] = field(default_factory=list) # character positions
|
|
confidence: float = 1.0
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class ExtractedRelationship:
|
|
"""A relationship between two entities extracted from text (H-EXT)."""
|
|
|
|
source_entity: str
|
|
source_type: str
|
|
target_entity: str
|
|
target_type: str
|
|
relationship_type: str # works_for, related_to, has_email, etc.
|
|
confidence: float = 0.0
|
|
evidence: str = "" # Text snippet that supports this relationship
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class ExtractionResult:
|
|
"""Result of a knowledge extraction run."""
|
|
|
|
entities: list[ExtractedEntity] = field(default_factory=list)
|
|
relationships: list[ExtractedRelationship] = field(default_factory=list)
|
|
source_type: str = ""
|
|
source_id: str = ""
|
|
tenant_id: str = ""
|
|
overall_confidence: float = 0.0
|
|
|
|
|
|
# ─── Extraction prompt ───────────────────────────────────────────────────────
|
|
|
|
_EXTRACTION_SYSTEM_PROMPT = """You are a knowledge extraction assistant. Analyze the given text and extract:
|
|
|
|
1. Named entities (persons, companies, projects, locations, dates)
|
|
2. Relationships between entities (e.g., "works_for", "related_to", "has_email", "part_of")
|
|
|
|
Return JSON in this format:
|
|
{
|
|
"entities": [
|
|
{"name": "John Doe", "type": "person", "confidence": 0.95}
|
|
],
|
|
"relationships": [
|
|
{
|
|
"source": "John Doe",
|
|
"source_type": "person",
|
|
"target": "Acme Corp",
|
|
"target_type": "company",
|
|
"type": "works_for",
|
|
"confidence": 0.9,
|
|
"evidence": "John Doe works at Acme Corp"
|
|
}
|
|
]
|
|
}
|
|
|
|
Only extract explicitly stated facts. Do not infer or hallucinate.
|
|
If no entities or relationships are found, return empty arrays."""
|
|
|
|
|
|
async def extract_knowledge(
|
|
text: str,
|
|
tenant_id: uuid.UUID,
|
|
source_type: str = "",
|
|
source_id: str = "",
|
|
llm_model: str | None = None,
|
|
) -> ExtractionResult:
|
|
"""Extract entities and relationships from text using LLM (H-EXT, H-ENT).
|
|
|
|
Args:
|
|
text: The text to analyze.
|
|
tenant_id: Tenant ID for multi-tenancy.
|
|
source_type: Source type (wiki, dms, mail, communication).
|
|
source_id: Source entity ID.
|
|
llm_model: Optional LLM model override.
|
|
|
|
Returns:
|
|
ExtractionResult with entities and relationships.
|
|
"""
|
|
if not text or len(text.strip()) < 10:
|
|
return ExtractionResult(
|
|
source_type=source_type,
|
|
source_id=source_id,
|
|
tenant_id=str(tenant_id),
|
|
)
|
|
|
|
try:
|
|
from app.ai.llm_client import llm_complete
|
|
|
|
messages = [
|
|
{"role": "system", "content": _EXTRACTION_SYSTEM_PROMPT},
|
|
{"role": "user", "content": f"Analyze this text:\n\n{text[:8000]}"},
|
|
]
|
|
|
|
response = await llm_complete(
|
|
model=llm_model or "ollama/deepseek-v4-flash",
|
|
messages=messages,
|
|
temperature=0.1,
|
|
max_tokens=2000,
|
|
)
|
|
|
|
# Parse LLM response
|
|
import json
|
|
|
|
raw = response.get("content", "")
|
|
# Try to extract JSON from response
|
|
try:
|
|
data = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
# Try to find JSON in the response
|
|
start = raw.find("{")
|
|
end = raw.rfind("}") + 1
|
|
if start >= 0 and end > start:
|
|
data = json.loads(raw[start:end])
|
|
else:
|
|
logger.warning("Failed to parse extraction response as JSON")
|
|
return ExtractionResult(
|
|
source_type=source_type,
|
|
source_id=source_id,
|
|
tenant_id=str(tenant_id),
|
|
)
|
|
|
|
# Build ExtractionResult
|
|
entities: list[ExtractedEntity] = []
|
|
for ent in data.get("entities", []):
|
|
entities.append(ExtractedEntity(
|
|
name=ent.get("name", ""),
|
|
entity_type=ent.get("type", "other"),
|
|
confidence=ent.get("confidence", 0.5),
|
|
metadata=ent.get("metadata", {}),
|
|
))
|
|
|
|
relationships: list[ExtractedRelationship] = []
|
|
for rel in data.get("relationships", []):
|
|
relationships.append(ExtractedRelationship(
|
|
source_entity=rel.get("source", ""),
|
|
source_type=rel.get("source_type", "other"),
|
|
target_entity=rel.get("target", ""),
|
|
target_type=rel.get("target_type", "other"),
|
|
relationship_type=rel.get("type", "related_to"),
|
|
confidence=rel.get("confidence", 0.5),
|
|
evidence=rel.get("evidence", ""),
|
|
metadata=rel.get("metadata", {}),
|
|
))
|
|
|
|
# Calculate overall confidence
|
|
all_confidences = [e.confidence for e in entities] + [r.confidence for r in relationships]
|
|
overall = sum(all_confidences) / len(all_confidences) if all_confidences else 0.0
|
|
|
|
return ExtractionResult(
|
|
entities=entities,
|
|
relationships=relationships,
|
|
source_type=source_type,
|
|
source_id=source_id,
|
|
tenant_id=str(tenant_id),
|
|
overall_confidence=overall,
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.warning("Knowledge extraction failed: %s", e)
|
|
return ExtractionResult(
|
|
source_type=source_type,
|
|
source_id=source_id,
|
|
tenant_id=str(tenant_id),
|
|
)
|
|
|
|
|
|
# ─── Confidence scoring (H-CONF) ─────────────────────────────────────────────
|
|
|
|
LOW_CONFIDENCE_THRESHOLD = 0.6
|
|
|
|
|
|
def is_low_confidence(confidence: float) -> bool:
|
|
"""Check if a confidence score is below the review threshold (H-CONF)."""
|
|
return confidence < LOW_CONFIDENCE_THRESHOLD
|
|
|
|
|
|
def filter_high_confidence(
|
|
relationships: list[ExtractedRelationship],
|
|
threshold: float = LOW_CONFIDENCE_THRESHOLD,
|
|
) -> tuple[list[ExtractedRelationship], list[ExtractedRelationship]]:
|
|
"""Split relationships into high-confidence and low-confidence (review queue).
|
|
|
|
Returns:
|
|
Tuple of (high_confidence, low_confidence) lists.
|
|
"""
|
|
high = [r for r in relationships if r.confidence >= threshold]
|
|
low = [r for r in relationships if r.confidence < threshold]
|
|
return high, low
|
|
|
|
|
|
# ─── Auto-relationship creation in GraphRAG (H-AUTO) ─────────────────────────
|
|
|
|
|
|
async def auto_create_relationships(
|
|
db: AsyncSession,
|
|
tenant_id: uuid.UUID,
|
|
extraction: ExtractionResult,
|
|
min_confidence: float = LOW_CONFIDENCE_THRESHOLD,
|
|
) -> dict[str, Any]:
|
|
"""Auto-create extracted relationships in GraphRAG (H-AUTO).
|
|
|
|
Only creates relationships with confidence >= min_confidence.
|
|
Low-confidence relationships are returned for the review queue.
|
|
|
|
Returns:
|
|
Dict with ``created``, ``skipped_low_confidence``, ``errors`` counts.
|
|
"""
|
|
from app.plugins.builtins.graph_rag.models import EntityRelationship
|
|
|
|
created = 0
|
|
skipped = 0
|
|
errors = 0
|
|
|
|
high_conf, low_conf = filter_high_confidence(extraction.relationships, min_confidence)
|
|
|
|
for rel in high_conf:
|
|
try:
|
|
# Try to resolve entity names to actual entity IDs
|
|
# For now, store as typed relationships with name-based references
|
|
source_id = await _resolve_entity_id(db, tenant_id, rel.source_entity, rel.source_type)
|
|
target_id = await _resolve_entity_id(db, tenant_id, rel.target_entity, rel.target_type)
|
|
|
|
if source_id is None or target_id is None:
|
|
skipped += 1
|
|
continue
|
|
|
|
# Check if relationship already exists
|
|
existing = await db.execute(
|
|
select(EntityRelationship).where(
|
|
EntityRelationship.tenant_id == tenant_id,
|
|
EntityRelationship.source_type == rel.source_type,
|
|
EntityRelationship.source_id == source_id,
|
|
EntityRelationship.target_type == rel.target_type,
|
|
EntityRelationship.target_id == target_id,
|
|
EntityRelationship.relationship_type == rel.relationship_type,
|
|
)
|
|
)
|
|
if existing.scalar_one_or_none() is not None:
|
|
skipped += 1
|
|
continue
|
|
|
|
# Create new relationship
|
|
er = EntityRelationship(
|
|
tenant_id=tenant_id,
|
|
source_type=rel.source_type,
|
|
source_id=source_id,
|
|
target_type=rel.target_type,
|
|
target_id=target_id,
|
|
relationship_type=rel.relationship_type,
|
|
confidence=rel.confidence,
|
|
metadata={
|
|
"evidence": rel.evidence,
|
|
"source_type": extraction.source_type,
|
|
"source_id": extraction.source_id,
|
|
"auto_extracted": True,
|
|
},
|
|
)
|
|
db.add(er)
|
|
created += 1
|
|
except Exception as e:
|
|
logger.warning("Failed to auto-create relationship: %s", e)
|
|
errors += 1
|
|
|
|
await db.flush()
|
|
|
|
return {
|
|
"created": created,
|
|
"skipped_low_confidence": len(low_conf),
|
|
"skipped_existing": skipped,
|
|
"errors": errors,
|
|
"review_queue": [
|
|
{
|
|
"source": r.source_entity,
|
|
"target": r.target_entity,
|
|
"type": r.relationship_type,
|
|
"confidence": r.confidence,
|
|
"evidence": r.evidence,
|
|
}
|
|
for r in low_conf
|
|
],
|
|
}
|
|
|
|
|
|
async def _resolve_entity_id(
|
|
db: AsyncSession,
|
|
tenant_id: uuid.UUID,
|
|
name: str,
|
|
entity_type: str,
|
|
) -> uuid.UUID | None:
|
|
"""Try to resolve an entity name to an actual entity ID.
|
|
|
|
Searches contacts, companies, etc. by name.
|
|
Returns None if no match found.
|
|
"""
|
|
try:
|
|
if entity_type == "person":
|
|
from app.models.contact import Contact
|
|
result = await db.execute(
|
|
select(Contact.id).where(
|
|
Contact.tenant_id == tenant_id,
|
|
Contact.deleted_at.is_(None),
|
|
Contact.name.ilike(f"%{name}%"),
|
|
).limit(1)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
elif entity_type == "company":
|
|
from app.models.contact import Contact
|
|
result = await db.execute(
|
|
select(Contact.id).where(
|
|
Contact.tenant_id == tenant_id,
|
|
Contact.deleted_at.is_(None),
|
|
Contact.is_company.is_(True),
|
|
Contact.name.ilike(f"%{name}%"),
|
|
).limit(1)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
__all__ = [
|
|
"ExtractedEntity",
|
|
"ExtractedRelationship",
|
|
"ExtractionResult",
|
|
"extract_knowledge",
|
|
"is_low_confidence",
|
|
"filter_high_confidence",
|
|
"auto_create_relationships",
|
|
"LOW_CONFIDENCE_THRESHOLD",
|
|
]
|