src/components/source_tracker.py - master - alsoalgo/CTF

src/components/source_tracker.py
93 строки3 KB
alsoalgo
some improvements
14 дек 2025, 22:23
14 дек 2025, 22:2337d2ec0
"""
Трекер источников данных - отслеживает, откуда именно в документе были взяты данные
"""
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, asdict
import json


@dataclass
class SourceLocation:
    """
    Информация о местоположении данных в документе
    """
    file_path: str
    file_name: str
    page_number: Optional[int] = None  # Номер страницы (для PDF)
    coordinates: Optional[Dict[str, float]] = None  # Координаты (x, y, width, height)
    context: Optional[str] = None  # Окружающий текст
    extraction_method: str = "standard"  # standard, ocr, table_extraction
    confidence: float = 1.0  # Уверенность в извлечении (0.0-1.0)
    
    def to_dict(self) -> Dict:
        """Преобразует в словарь для сериализации"""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict) -> 'SourceLocation':
        """Создает из словаря"""
        return cls(**data)


class SourceTracker:
    """
    Класс для отслеживания источников извлеченных данных
    """
    
    def __init__(self):
        self.sources: Dict[str, SourceLocation] = {}  # key -> SourceLocation
    
    def add_source(
        self,
        key: str,
        file_path: str,
        page_number: Optional[int] = None,
        coordinates: Optional[Dict[str, float]] = None,
        context: Optional[str] = None,
        extraction_method: str = "standard",
        confidence: float = 1.0
    ):
        """
        Добавляет источник данных.
        
        :param key: Уникальный ключ (например, "supplier" или "row_0_col_1")
        :param file_path: Путь к файлу
        :param page_number: Номер страницы
        :param coordinates: Координаты в документе
        :param context: Окружающий текст
        :param extraction_method: Метод извлечения
        :param confidence: Уверенность
        """
        import os
        self.sources[key] = SourceLocation(
            file_path=file_path,
            file_name=os.path.basename(file_path),
            page_number=page_number,
            coordinates=coordinates,
            context=context,
            extraction_method=extraction_method,
            confidence=confidence
        )
    
    def get_source(self, key: str) -> Optional[SourceLocation]:
        """Получает источник по ключу"""
        return self.sources.get(key)
    
    def get_all_sources(self) -> Dict[str, SourceLocation]:
        """Возвращает все источники"""
        return self.sources.copy()
    
    def to_dict(self) -> Dict:
        """Преобразует в словарь для сериализации"""
        return {
            key: source.to_dict()
            for key, source in self.sources.items()
        }
    
    @classmethod
    def from_dict(cls, data: Dict) -> 'SourceTracker':
        """Создает из словаря"""
        tracker = cls()
        for key, source_data in data.items():
            tracker.sources[key] = SourceLocation.from_dict(source_data)
        return tracker