monish563/NU-KIOSK-API

1"""Lightweight in-memory catalog for loading CSV data into named tables."""2 3from __future__ import annotations4 5from dataclasses import dataclass, field6from pathlib import Path7from typing import Any, Callable, Dict, Iterable, List, Optional8import logging9 10from .sources import DataSource, default_sources11from .utils import (12    canonicalize_name,13    centers_for_faculty,14    extract_leadership_names,15    extract_advisor_names,16    generate_name_variants,17    tokenize_name,18)19 20 21@dataclass22class EntityRecords:23    """Container for a structured dataset plus helpful indices."""24 25    name: str26    records: List[Dict[str, Any]]27    key_field: Optional[str] = None28    origin: Optional[str] = None29    _index: Dict[str, Dict[str, Any]] = field(default_factory=dict, init=False, repr=False)30 31    def build_index(self, normalizer: Optional[Callable[[str], str]] = None) -> None:32        if self.key_field is None:33            return34        normalizer = normalizer or canonicalize_name35        for row in self.records:36            key = row.get(self.key_field)37            if not key:38                continue39            self._index[normalizer(str(key))] = row40 41    def get_by_key(self, value: str, normalizer: Optional[Callable[[str], str]] = None) -> Optional[Dict[str, Any]]:42        if not value or self.key_field is None:43            return None44        normalizer = normalizer or canonicalize_name45        return self._index.get(normalizer(value))46 47 48class DataCatalog:49    """50    Registry that mirrors Satyrn's Ring layer.51 52    Each entity (faculty, staff, events...) is stored with metadata so53    analysis blueprints can retrieve and join information deterministically.54    """55 56    def __init__(self) -> None:57        self.entities: Dict[str, EntityRecords] = {}58        self.metadata: Dict[str, Dict[str, Any]] = {}59        self.relationships: Dict[str, RelationshipDefinition] = {}60 61    def register_entity(62        self,63        name: str,64        records: Iterable[Dict[str, Any]],65        *,66        key_field: Optional[str] = None,67        origin: Optional[str] = None,68        normalizer: Optional[Callable[[str], str]] = None,69    ) -> None:70        dataset = EntityRecords(71            name=name,72            records=list(records),73            key_field=key_field,74            origin=origin,75        )76        dataset.build_index(normalizer)77        self.entities[name] = dataset78 79    def get(self, name: str) -> EntityRecords:80        if name not in self.entities:81            raise KeyError(f"Unknown entity '{name}'")82        return self.entities[name]83 84    def try_get(self, name: str) -> Optional[EntityRecords]:85        return self.entities.get(name)86 87    # Relationship handling -------------------------------------------------88    def register_relationship(self, relationship: RelationshipDefinition) -> None:89        self.relationships[relationship.name] = relationship90 91    def resolve_relationship(self, name: str, source_row: Dict[str, Any]) -> List[Dict[str, Any]]:92        relationship = self.relationships.get(name)93        if not relationship:94            raise KeyError(f"Relationship '{name}' not registered.")95 96        source_value = source_row.get(relationship.from_field)97        if not source_value:98            return []99 100        target_entity = self.try_get(relationship.to_entity)101        if not target_entity:102            return []103 104        candidates = relationship.variant_generator(source_value) if relationship.variant_generator else [source_value]105        matches: List[Dict[str, Any]] = []106        for candidate in candidates:107            target_row = target_entity.get_by_key(candidate, relationship.normalizer)108            if target_row:109                matches.append(target_row)110 111        if matches:112            return matches113 114        if relationship.allow_token_subset:115            source_tokens = tokenize_name(source_value)116            if source_tokens:117                for row in target_entity.records:118                    dest_tokens = tokenize_name(row.get(relationship.to_field, ""))119                    if source_tokens.issubset(dest_tokens):120                        matches.append(row)121            if matches:122                return matches123 124        if relationship.fallback_matcher:125            extra = relationship.fallback_matcher(source_row, target_entity.records)126            if extra:127                matches.extend(extra)128 129        return matches130 131 132@dataclass133class RelationshipDefinition:134    name: str135    from_entity: str136    to_entity: str137    from_field: str138    to_field: str139    normalizer: Callable[[str], str] = canonicalize_name140    variant_generator: Optional[Callable[[str], Iterable[str]]] = lambda value: generate_name_variants(value)141    allow_token_subset: bool = False142    fallback_matcher: Optional[Callable[[Dict[str, Any], List[Dict[str, Any]]], List[Dict[str, Any]]]] = None143 144 145def load_default_catalog(146    base_dir: Path | str = "Archive",147    *,148    sources: Optional[Iterable[DataSource]] = None,149) -> DataCatalog:150    """151    Load the project CSVs into a catalog with sensible defaults.152 153    This mirrors the datasets referenced in the exploratory notebook,154    giving downstream code a single call to bootstrap the backend.155 156    A custom ``sources`` iterable can be supplied to extend or override157    the default data sources (e.g., to add TA office hours).158    """159 160    base_path = Path(base_dir)161    catalog = DataCatalog()162    logger = logging.getLogger(__name__)163 164    source_list = list(sources) if sources is not None else default_sources(base_path, name_normalizer=canonicalize_name)165 166    for source in source_list:167        result = source.load()168        for entity in result.entities:169            if not entity.records:170                continue171            catalog.register_entity(172                entity.name,173                entity.records,174                key_field=entity.key_field,175                origin=entity.origin,176                normalizer=entity.normalizer,177            )178            # Log which file provided this entity so we can trace provenance179            try:180                origin_txt = entity.origin or "(unknown)"181            except Exception:182                origin_txt = "(unknown)"183            logger.info("Loaded entity '%s' with %d records from %s", entity.name, len(entity.records), origin_txt)184        for key, value in result.metadata.items():185            existing = catalog.metadata.get(key)186            if isinstance(existing, dict) and isinstance(value, dict):187                existing.update(value)188            else:189                catalog.metadata[key] = value190 191    # Register common relationships between entities to mirror Satyrn's Ring metadata.192    if catalog.try_get("faculty") and catalog.try_get("faculty_offices"):193        catalog.register_relationship(194            RelationshipDefinition(195                name="faculty_to_office",196                from_entity="faculty",197                to_entity="faculty_offices",198                from_field="Name",199                to_field="Assignee Name",200                allow_token_subset=True,201            )202        )203 204    if catalog.try_get("centers") and catalog.try_get("faculty"):205        catalog.register_relationship(206            RelationshipDefinition(207                name="center_to_faculty_leads",208                from_entity="centers",209                to_entity="faculty",210                from_field="Leadership",211                to_field="Name",212                normalizer=canonicalize_name,213                variant_generator=lambda value: extract_leadership_names(value),214                allow_token_subset=False,215            )216        )217        catalog.register_relationship(218            RelationshipDefinition(219                name="faculty_to_centers",220                from_entity="faculty",221                to_entity="centers",222                from_field="Name",223                to_field="Leadership",224                normalizer=canonicalize_name,225                variant_generator=None,226                allow_token_subset=False,227                fallback_matcher=centers_for_faculty,228            )229        )230 231    if catalog.try_get("students") and catalog.try_get("mudd_seating"):232        catalog.register_relationship(233            RelationshipDefinition(234                name="student_to_mudd_seat",235                from_entity="students",236                to_entity="mudd_seating",237                from_field="Name",238                to_field="Student/Visitor",239                allow_token_subset=True,240            )241        )242 243    if catalog.try_get("students") and catalog.try_get("faculty"):244        catalog.register_relationship(245            RelationshipDefinition(246                name="student_to_advisors",247                from_entity="students",248                to_entity="faculty",249                from_field="Advisor(s)",250                to_field="Name",251                normalizer=canonicalize_name,252                variant_generator=lambda value: extract_advisor_names(value),253                allow_token_subset=False,254            )255        )256        catalog.register_relationship(257            RelationshipDefinition(258                name="faculty_to_advisees",259                from_entity="faculty",260                to_entity="students",261                from_field="Name",262                to_field="Advisor(s)",263                normalizer=canonicalize_name,264                variant_generator=None,265                allow_token_subset=False,266                fallback_matcher=lambda source, students: [267                    row for row in students if canonicalize_name(source.get("Name", "")) in {268                        canonicalize_name(name) for name in extract_advisor_names(row.get("Advisor(s)"))269                    }270                ],271            )272        )273 274    return catalog275