Blame for src/techdocs_graphrag/extraction.py

import logging

from pydantic import ValidationError

from techdocs_graphrag.models import Chunk, GraphExtraction
from techdocs_graphrag.ports import LanguageModel
from techdocs_graphrag.prompts import EXTRACTION_SYSTEM_PROMPT

logger = logging.getLogger(__name__)


class GraphExtractor:
    def __init__(self, llm: LanguageModel) -> None:
        self._llm = llm

    def extract(self, chunk: Chunk) -> GraphExtraction:
        payload = self._llm.complete(EXTRACTION_SYSTEM_PROMPT, chunk.text, json_mode=True)
        try:
            return GraphExtraction.model_validate_json(payload)
        except ValidationError as error:
            logger.warning(
                "Graph extraction rejected for chunk %s (%s errors)", chunk.id, error.error_count()
            )
            return GraphExtraction()