123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
"""Enrichment agent orchestrating LLM calls for entity detection.
|
|
|
|
Calls the LLM with extracted text, parses the structured response,
|
|
filters entities by confidence threshold, and handles parse errors with retry.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
|
|
from ingestion.config import IngestionConfig
|
|
from ingestion.integrations.llm import LLMClient, LLMError
|
|
|
|
from .models import EnrichmentResult, Entity
|
|
from .prompts import STRICT_RETRY_PROMPT, build_system_prompt, build_user_prompt
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class EnrichmentError(Exception):
|
|
"""Raised when enrichment fails after retries."""
|
|
|
|
pass
|
|
|
|
|
|
class EnrichmentAgent:
|
|
"""Orchestrates LLM-based document enrichment."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: IngestionConfig,
|
|
llm_client: LLMClient | None = None,
|
|
known_projects: list[str] | None = None,
|
|
known_people: list[str] | None = None,
|
|
):
|
|
self.config = config
|
|
self.llm = llm_client or LLMClient(config)
|
|
self.confidence_threshold = config.confidence_threshold
|
|
self.system_prompt = build_system_prompt(known_projects, known_people)
|
|
|
|
def enrich(self, text: str) -> EnrichmentResult:
|
|
"""Enrich extracted text with LLM-detected entities and metadata.
|
|
|
|
Args:
|
|
text: The extracted document text to analyze.
|
|
|
|
Returns:
|
|
EnrichmentResult with entities filtered by confidence threshold.
|
|
|
|
Raises:
|
|
EnrichmentError: If LLM call and retry both fail.
|
|
"""
|
|
messages = [
|
|
{"role": "system", "content": self.system_prompt},
|
|
{"role": "user", "content": build_user_prompt(text)},
|
|
]
|
|
|
|
# First attempt
|
|
try:
|
|
response_text = self.llm.complete(messages)
|
|
result = self._parse_response(response_text)
|
|
except (json.JSONDecodeError, ValueError) as parse_err:
|
|
logger.warning(
|
|
"First LLM response parse failed (%s), retrying with strict prompt",
|
|
parse_err,
|
|
)
|
|
# Retry with stricter prompt
|
|
result = self._retry_with_strict_prompt(messages, response_text)
|
|
except LLMError as e:
|
|
raise EnrichmentError(f"LLM call failed: {e}") from e
|
|
|
|
# Filter entities by confidence threshold
|
|
result.entities = self._filter_by_confidence(result.entities)
|
|
|
|
return result
|
|
|
|
def _retry_with_strict_prompt(
|
|
self, original_messages: list[dict[str, str]], failed_response: str
|
|
) -> EnrichmentResult:
|
|
"""Retry LLM call with a stricter prompt after parse failure."""
|
|
retry_messages = original_messages + [
|
|
{"role": "assistant", "content": failed_response},
|
|
{"role": "user", "content": STRICT_RETRY_PROMPT},
|
|
]
|
|
|
|
try:
|
|
response_text = self.llm.complete(retry_messages)
|
|
return self._parse_response(response_text)
|
|
except (json.JSONDecodeError, ValueError, LLMError) as e:
|
|
logger.error("Retry also failed (%s), using defaults", e)
|
|
return self._default_result()
|
|
|
|
def _parse_response(self, response_text: str) -> EnrichmentResult:
|
|
"""Parse LLM response text into an EnrichmentResult.
|
|
|
|
Handles responses that may be wrapped in markdown code fences.
|
|
"""
|
|
cleaned = response_text.strip()
|
|
|
|
# Strip markdown code fences if present
|
|
if cleaned.startswith("```"):
|
|
lines = cleaned.split("\n")
|
|
# Remove first line (```json) and last line (```)
|
|
lines = [l for l in lines[1:] if not l.strip() == "```"]
|
|
cleaned = "\n".join(lines)
|
|
|
|
data = json.loads(cleaned)
|
|
return EnrichmentResult.model_validate(data)
|
|
|
|
def _filter_by_confidence(self, entities: list[Entity]) -> list[Entity]:
|
|
"""Keep only entities meeting the confidence threshold."""
|
|
return [e for e in entities if e.confidence >= self.confidence_threshold]
|
|
|
|
def _default_result(self) -> EnrichmentResult:
|
|
"""Return a minimal default result when all parsing fails."""
|
|
return EnrichmentResult(
|
|
title="Untitled Note",
|
|
category="inbox",
|
|
tags=[],
|
|
entities=[],
|
|
action_items=[],
|
|
summary=None,
|
|
)
|