feat(H): H-EXT/H-ENT/H-AUTO/H-CONF — LLM knowledge extraction (entities, relationships, auto-create in GraphRAG, confidence scoring with review queue), 31 tests passing
This commit is contained in:
Agent Zero
@@ -0,0 +1,361 @@
|
||||
"""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",
|
||||
]
|
||||
Reference in New Issue
Block a user