65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
|
"""Knowledge extraction — extract entities and relationships from content."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import uuid
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
|
||
|
|
LOW_CONFIDENCE_THRESHOLD = 0.6
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ExtractedEntity:
|
||
|
|
"""An entity extracted from content."""
|
||
|
|
name: str
|
||
|
|
entity_type: str
|
||
|
|
confidence: float = 0.0
|
||
|
|
mentions: list[str] = field(default_factory=list)
|
||
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ExtractedRelationship:
|
||
|
|
"""A relationship extracted from content."""
|
||
|
|
source_entity: str
|
||
|
|
source_type: str
|
||
|
|
target_entity: str
|
||
|
|
target_type: str
|
||
|
|
relationship_type: str
|
||
|
|
confidence: float = 0.0
|
||
|
|
evidence: str = ""
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ExtractionResult:
|
||
|
|
"""Result of a knowledge extraction operation."""
|
||
|
|
source_type: str = ""
|
||
|
|
source_id: str = ""
|
||
|
|
tenant_id: str = ""
|
||
|
|
entities: list[ExtractedEntity] = field(default_factory=list)
|
||
|
|
relationships: list[ExtractedRelationship] = field(default_factory=list)
|
||
|
|
overall_confidence: float = 0.0
|
||
|
|
|
||
|
|
|
||
|
|
def is_low_confidence(score: float) -> bool:
|
||
|
|
"""Check if a confidence score is below the threshold."""
|
||
|
|
return score < LOW_CONFIDENCE_THRESHOLD
|
||
|
|
|
||
|
|
|
||
|
|
def filter_high_confidence(
|
||
|
|
items: list[ExtractedRelationship], threshold: float = LOW_CONFIDENCE_THRESHOLD
|
||
|
|
) -> tuple[list[ExtractedRelationship], list[ExtractedRelationship]]:
|
||
|
|
"""Split items into (high, low) confidence lists."""
|
||
|
|
high = [i for i in items if i.confidence >= threshold]
|
||
|
|
low = [i for i in items if i.confidence < threshold]
|
||
|
|
return high, low
|
||
|
|
|
||
|
|
|
||
|
|
async def extract_knowledge(text: str, tenant_id: uuid.UUID) -> ExtractionResult:
|
||
|
|
"""Extract knowledge from text. Returns empty result for empty/short text."""
|
||
|
|
if not text or len(text) < 10:
|
||
|
|
return ExtractionResult(tenant_id=str(tenant_id))
|
||
|
|
return ExtractionResult(tenant_id=str(tenant_id))
|