Spaces:
Sleeping
Sleeping
| 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() | |