resolvekit 0.0.1__py3-none-any.whl

resolvekit/data/membership_repository.py

@@ -0,0 +1,107 @@
+"""Repository for group membership queries."""
+
+from datetime import date
+
+from sqlalchemy import text
+from sqlmodel import Session
+
+from resolvekit.data.db_manager import DatabaseManager
+
+
+class MembershipRepository:
+    """Repository for querying group memberships with temporal support."""
+
+    def __init__(self, db_manager: DatabaseManager):
+        """
+        Initialize repository.
+
+        Args:
+            db_manager: Database manager instance
+        """
+        self.db = db_manager
+
+    def check_memberships_batch(
+        self, entity_dcids: list[str], group_dcid: str, as_of: date
+    ) -> set[str]:
+        """
+        Check which entities are members of a group at a given date.
+
+        Uses a single batch query for efficiency. Respects overlay databases -
+        if any source (base or overlay) indicates membership at the given date,
+        the entity is considered a member.
+
+        Args:
+            entity_dcids: List of entity DCIDs to check
+            group_dcid: Group DCID to check membership in
+            as_of: Date to check membership at
+
+        Returns:
+            Set of entity DCIDs that are members of the group at the given date
+
+        Raises:
+            RuntimeError: If database is not connected
+        """
+        if not self.db.engine:
+            raise RuntimeError("Database not connected. Call connect() first.")
+
+        if not entity_dcids:
+            return set()
+
+        # Build parameterized query with IN clause
+        placeholders = ", ".join(f":dcid_{i}" for i in range(len(entity_dcids)))
+
+        # Build union of all membership sources (main + overlays)
+        # Collect all rows first, then deduplicate by precedence
+        all_sources = [
+            f"""
+            SELECT entity_dcid, group_dcid, valid_from, valid_until, 0 AS precedence
+            FROM main.memberships
+            WHERE entity_dcid IN ({placeholders})
+              AND group_dcid = :group_dcid
+            """
+        ]
+
+        for schema_name, precedence in self.db.overlays:
+            all_sources.append(f"""
+            SELECT entity_dcid, group_dcid, valid_from, valid_until, {precedence} AS precedence
+            FROM {schema_name}.memberships
+            WHERE entity_dcid IN ({placeholders})
+              AND group_dcid = :group_dcid
+            """)
+
+        # Deduplicate by (entity_dcid, group_dcid) keeping highest precedence,
+        # THEN apply temporal filter. This ensures overlay rows override base rows.
+        query = text(f"""
+            WITH best_memberships AS (
+                SELECT entity_dcid, valid_from, valid_until
+                FROM (
+                    SELECT entity_dcid, group_dcid, valid_from, valid_until,
+                           ROW_NUMBER() OVER (
+                               PARTITION BY entity_dcid, group_dcid
+                               ORDER BY precedence DESC
+                           ) as rn
+                    FROM (
+                        {" UNION ALL ".join(all_sources)}
+                    ) all_memberships
+                ) ranked
+                WHERE rn = 1
+            )
+            SELECT DISTINCT entity_dcid
+            FROM best_memberships
+            WHERE valid_from <= :as_of
+              AND (valid_until IS NULL OR valid_until > :as_of)
+        """)
+
+        # Build parameters dict
+        params = {
+            "group_dcid": group_dcid,
+            "as_of": as_of.isoformat(),
+        }
+        for i, dcid in enumerate(entity_dcids):
+            params[f"dcid_{i}"] = dcid
+
+        with Session(self.db.engine) as session:
+            result = session.execute(query, params)
+            members = {row[0] for row in result.fetchall()}
+
+        return members

resolvekit/data/query_builder.py

@@ -0,0 +1,177 @@
+"""Query builder for overlay precedence."""
+
+from datetime import date
+from typing import Any
+
+
+class QueryBuilder:
+    """
+    Builds SQL queries with overlay precedence support.
+
+    Constructs UNION queries across main + overlay databases,
+    with deduplication keeping highest precedence.
+    """
+
+    def __init__(self, db_manager: Any):
+        """
+        Initialize query builder.
+
+        Args:
+            db_manager: DatabaseManager instance
+        """
+        self.db = db_manager
+        self.overlays = db_manager.overlays
+
+    def build_union_query(
+        self,
+        table: str,
+        columns: list[str],
+        where_clause: str,
+        params: dict[str, Any],
+        unique_key: str,
+        as_of: date | None = None,
+    ) -> tuple[str, dict[str, Any]]:
+        """
+        Build UNION query across all databases with deduplication.
+
+        Args:
+            table: Table name
+            columns: Column names to select
+            where_clause: WHERE clause (using named parameters)
+            params: Query parameters
+            unique_key: Column for deduplication (e.g., "dcid", "alias_uid")
+            as_of: Optional date for temporal filtering
+
+        Returns:
+            Tuple of (SQL query, parameters dict)
+        """
+        queries = []
+        col_list = ", ".join(columns)
+
+        # Add temporal filtering to where clause if as_of is provided
+        full_where = where_clause
+        if as_of is not None:
+            full_where += " AND (valid_from IS NULL OR valid_from <= :as_of)"
+            full_where += " AND (valid_until IS NULL OR valid_until >= :as_of)"
+            params = {**params, "as_of": as_of.isoformat()}
+
+        # Main database (precedence 0)
+        queries.append(f"""
+            SELECT {col_list}, 0 AS precedence, 'main' AS source_db
+            FROM main.{table}
+            WHERE {full_where}
+        """)
+
+        # Overlay databases with their precedence
+        for schema_name, precedence in self.overlays:
+            queries.append(f"""
+                SELECT {col_list}, {precedence} AS precedence, '{schema_name}' AS source_db
+                FROM {schema_name}.{table}
+                WHERE {full_where}
+            """)
+
+        # Union all results
+        union_query = " UNION ALL ".join(queries)
+
+        # Deduplicate by unique_key keeping highest precedence using window functions
+        final_query = f"""
+            WITH all_results AS ({union_query}),
+            ranked AS (
+                SELECT *, ROW_NUMBER() OVER (PARTITION BY {unique_key} ORDER BY precedence DESC) as rn
+                FROM all_results
+            )
+            SELECT * FROM ranked WHERE rn = 1
+            ORDER BY precedence DESC
+        """
+
+        return final_query, params
+
+    def build_code_lookup_union(
+        self,
+        code_system: str,
+        code_values: str | list[str],
+        entity_columns: list[str],
+        include_code_value: bool = False,
+    ) -> tuple[str, dict[str, Any]]:
+        """
+        Build UNION query for code lookups across all databases.
+
+        Args:
+            code_system: Code system to search
+            code_values: Single code value or list of code values
+            entity_columns: Entity column names to select
+            include_code_value: Whether to include code_value in results (for batch lookups)
+
+        Returns:
+            Tuple of (SQL query, parameters dict)
+        """
+        # Build column list
+        entity_cols = ", ".join(f"e.{col}" for col in entity_columns)
+
+        # Build CTE columns and WHERE clause
+        is_batch = isinstance(code_values, list)
+        params: dict[str, Any] = {"code_system": code_system}
+
+        if is_batch:
+            placeholders = ", ".join(f":code_{i}" for i in range(len(code_values)))
+            params.update({f"code_{i}": code for i, code in enumerate(code_values)})
+            where_clause = (
+                f"code_system = :code_system AND code_value IN ({placeholders})"
+            )
+            cte_columns = (
+                "entity_dcid, code_value" if include_code_value else "entity_dcid"
+            )
+        else:
+            params["code_value"] = code_values
+            where_clause = "code_system = :code_system AND code_value = :code_value"
+            cte_columns = (
+                "entity_dcid, code_value" if include_code_value else "entity_dcid"
+            )
+
+        # Select clause for results
+        result_columns = (
+            f"{entity_cols}, c.code_value" if include_code_value else entity_cols
+        )
+
+        # Build query for main database
+        sql = f"""
+            WITH code_lookup AS (
+                SELECT {cte_columns}
+                FROM main.codes
+                WHERE {where_clause}
+            )
+            SELECT {result_columns}, 0 AS precedence, 'main' AS source_db
+            FROM main.entities e
+            INNER JOIN code_lookup c ON e.dcid = c.entity_dcid
+        """
+
+        # Add overlays
+        for schema_name, precedence in self.overlays:
+            sql += f"""
+                UNION ALL
+                SELECT {result_columns}, {precedence} AS precedence, '{schema_name}' AS source_db
+                FROM {schema_name}.entities e
+                INNER JOIN (
+                    SELECT {cte_columns}
+                    FROM {schema_name}.codes
+                    WHERE {where_clause}
+                ) c ON e.dcid = c.entity_dcid
+            """
+
+        return sql, params
+
+    def escape_fts_query(self, query: str) -> str:
+        """
+        Escape FTS5 special characters.
+
+        Args:
+            query: Raw query string
+
+        Returns:
+            Escaped query string
+        """
+        # Escape double quotes
+        escaped = query.replace('"', '""')
+
+        # Wrap in quotes for phrase search
+        return f'"{escaped}"'

resolvekit/data/schema.py

@@ -0,0 +1,122 @@
+"""Database schema definitions using SQLModel.
+
+Schema is automatically generated from SQLModel definitions:
+- Table definitions in types.py (EntityRow, AliasRow, CodeRow, MembershipRow)
+- Constraints, indexes, and foreign keys defined in SQLModel
+- FTS5 virtual table requires raw SQL (SQLAlchemy doesn't support FTS5)
+"""
+
+from sqlalchemy import Index, text
+from sqlalchemy.engine import Engine
+from sqlmodel import SQLModel
+
+from resolvekit.types import AliasRow, CodeRow, EntityRow, MembershipRow
+
+# Schema version
+SCHEMA_VERSION = 1
+
+# ==============================================================================
+# Indexes (additional indexes beyond those defined in SQLModel)
+# ==============================================================================
+
+# Note: Primary indexes are already defined in the SQLModel classes
+# These are additional composite or coverage indexes
+
+indexes = [
+    Index("idx_entities_type", EntityRow.__table__.c.entity_type),
+    Index("idx_entities_parent", EntityRow.__table__.c.parent_dcid),
+    Index(
+        "idx_entities_valid",
+        EntityRow.__table__.c.valid_from,
+        EntityRow.__table__.c.valid_until,
+    ),
+    Index("idx_aliases_dcid", AliasRow.__table__.c.entity_dcid),
+    # idx_aliases_norm already defined in SQLModel via index=True
+    Index("idx_aliases_type", AliasRow.__table__.c.alias_type),
+    Index(
+        "idx_codes_lookup",
+        CodeRow.__table__.c.code_system,
+        CodeRow.__table__.c.code_value,
+    ),
+    Index(
+        "idx_memberships_entity",
+        MembershipRow.__table__.c.entity_dcid,
+        MembershipRow.__table__.c.valid_from,
+        MembershipRow.__table__.c.valid_until,
+    ),
+    Index("idx_memberships_group", MembershipRow.__table__.c.group_dcid),
+]
+
+# ==============================================================================
+# FTS Virtual Table (requires raw SQL)
+# ==============================================================================
+
+# FTS5 virtual tables must be created with raw SQL
+# (SQLAlchemy doesn't have native FTS5 support)
+ALIASES_FTS_SQL = """
+CREATE VIRTUAL TABLE IF NOT EXISTS aliases_fts USING fts5(
+    alias_norm,
+    content='aliases',
+    content_rowid='alias_id',
+    tokenize = "unicode61 remove_diacritics 2 tokenchars '.-'",
+    prefix='2,3'
+)
+"""
+
+
+def create_schema(engine: Engine) -> None:
+    """
+    Create database schema with all tables and indexes.
+
+    Uses SQLModel metadata to create tables, then raw SQL for FTS5.
+
+    Args:
+        engine: SQLAlchemy engine
+    """
+    # Create all tables via SQLModel metadata
+    SQLModel.metadata.create_all(engine)
+
+    # Create additional indexes
+    for index in indexes:
+        index.create(engine, checkfirst=True)
+
+    # Create FTS virtual table (requires raw SQL)
+    with engine.connect() as conn:
+        conn.execute(text(ALIASES_FTS_SQL))
+
+        # Create schema_version table manually (not a SQLModel)
+        conn.execute(
+            text("""
+            CREATE TABLE IF NOT EXISTS schema_version (
+                version INTEGER NOT NULL,
+                updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
+            )
+        """)
+        )
+
+        # Insert schema version if not exists
+        result = conn.execute(text("SELECT COUNT(*) FROM schema_version"))
+        if result.fetchone()[0] == 0:
+            conn.execute(
+                text("INSERT INTO schema_version (version) VALUES (:version)"),
+                {"version": SCHEMA_VERSION},
+            )
+        conn.commit()
+
+
+def get_schema_version(engine: Engine) -> int:
+    """
+    Get current schema version.
+
+    Args:
+        engine: SQLAlchemy engine
+
+    Returns:
+        Schema version number
+    """
+    with engine.connect() as conn:
+        result = conn.execute(
+            text("SELECT version FROM schema_version ORDER BY updated_at DESC LIMIT 1")
+        )
+        row = result.fetchone()
+        return row[0] if row else 0

resolvekit/disambiguation/README.md

@@ -0,0 +1,72 @@
+# Disambiguation Module
+
+## Purpose
+
+The disambiguation module handles ambiguous queries where multiple entities are plausible matches, using context and semantic understanding to select the most likely entity.
+
+## Components
+
+### Core Components
+
+1. **Ambiguity Detector** (`detector.py`)
+   - Computes score margin between top candidates
+   - Checks ambiguity registry for known ambiguous terms
+   - Determines when to invoke semantic disambiguation
+
+2. **Ambiguity Registry** (`registry.py`)
+   - SQLite table of known ambiguous surface forms
+   - Maps ambiguous terms to likely entity types and notes
+   - Examples: "Georgia" → {country, state}, "Congo" → {COD, COG}
+
+3. **Micro-Semantic Sidecar** (`semantic_sidecar.py`)
+   - HNSW ANN index for semantic similarity search
+   - Only stores embeddings for curated ambiguous aliases
+   - Uses PQ/4-bit quantization for memory efficiency
+   - Optional component (requires `resolver[semantic]`)
+
+4. **Context Analyzer** (`context.py`)
+   - Extracts disambiguation hints from context
+   - Handles co-mentioned entities, parent regions, coordinates
+   - Computes context-candidate compatibility scores
+
+5. **Default Heuristics** (`heuristics.py`)
+   - Non-semantic disambiguation rules
+   - Examples: prefer countries over subdivisions by default
+   - Global prominence scores, population-based ranking
+
+### Key Files
+
+- `ambiguity_model.py`: Data models for ambiguity metadata
+- `encoder.py`: Optional tiny encoder (MiniLM-class, ~80MB)
+- `sidecar_builder.py`: Tools to build ambiguity sidecar from registry
+
+## Ambiguity Detection Logic
+
+```python
+# Compute margin between top 2 candidates
+margin = score(top1) - score(top2)
+
+# Trigger semantic if:
+is_ambiguous = (
+    margin < threshold  # Small margin (e.g., < 0.15)
+    or query in ambiguity_registry  # Known ambiguous term
+)
+
+if is_ambiguous and semantic_available():
+    # Re-rank using semantic similarity
+    rerank_with_semantics(candidates, query_embedding)
+else:
+    # Use lexicon-only heuristics
+    apply_default_heuristics(candidates)
+```
+
+## Design Principles
+
+1. **Selective invocation**: Only use expensive semantic search when needed
+2. **Explainable defaults**: Clear rules for non-semantic disambiguation
+3. **Graceful degradation**: Works without semantic components
+4. **User override**: Allow explicit disambiguation via parameters
+
+## Implementation Priority
+
+**Phase E** - Ambiguity subsystem

resolvekit/extraction/README.md

@@ -0,0 +1,204 @@
+# Extraction Module
+
+## Purpose
+
+The extraction module identifies and extracts entities from unstructured text (sentences, paragraphs, documents), resolving them to canonical entities.
+
+## Components
+
+### Core Components
+
+1. **Text Extractor** (`extractor.py`)
+   - Main extraction pipeline
+   - Orchestrates NER and resolution
+   - Returns entities with spans and confidence
+
+2. **Dictionary Matcher** (`dictionary_matcher.py`)
+   - Aho-Corasick or FlashText for fast dictionary matching
+   - Matches known entity names and aliases in text
+   - Efficient for large dictionaries
+
+3. **NER Assistant** (`ner.py`)
+   - Optional spaCy integration for named entity recognition
+   - Helps identify entity boundaries
+   - Filters false positives (common words)
+
+4. **Context Extractor** (`context.py`)
+   - Extract surrounding context for disambiguation
+   - Co-mentioned entities, geographic hints
+   - Sentence/paragraph boundaries
+
+5. **Span Manager** (`spans.py`)
+   - Handle overlapping and nested entities
+   - Deduplication (same entity mentioned multiple times)
+   - Link acronyms to full forms
+
+### Filters and Validators
+
+- `stoplist.py`: Common words that aren't entities ("March", "Reading")
+- `pos_filter.py`: Part-of-speech filters to reduce false positives
+- `confidence_thresholds.py`: Extraction-specific confidence tuning
+
+## Extraction Pipeline
+
+```python
+from resolvekit.extraction import TextExtractor
+
+extractor = TextExtractor(min_confidence=0.8)
+
+text = """
+The agreement was signed by representatives from France, Germany,
+and Côte d'Ivoire. Regional coordination will be managed through
+the West African Economic and Monetary Union (UEMOA) headquarters
+in Burkina Faso.
+"""
+
+# Extract entities
+entities = extractor.extract(text)
+
+# Results
+for entity in entities:
+    print(f"{entity.text} [{entity.span}]")
+    print(f"  → {entity.dcid} ({entity.type})")
+    print(f"  Confidence: {entity.confidence}")
+```
+
+## Extracted Entity Structure
+
+```python
+@dataclass
+class ExtractedEntity:
+    """Entity extracted from text."""
+
+    text: str                   # Original text span
+    span: tuple[int, int]       # Character offsets (start, end)
+    dcid: str                   # Resolved entity DCID
+    canonical_name: str         # Canonical entity name
+    entity_type: str            # Entity type (country, org, etc.)
+    confidence: float           # Resolution confidence
+    context: str | None         # Surrounding context
+    method: str                 # Extraction method (dictionary, ner, etc.)
+```
+
+## Extraction Modes
+
+### 1. Dictionary-First (Fast)
+
+```python
+# Use only dictionary matching
+extractor = TextExtractor(mode="dictionary")
+entities = extractor.extract(text)
+```
+
+- Fast: O(n) where n = text length
+- High precision for known names
+- May miss creative mentions
+
+### 2. NER-Assisted (Accurate)
+
+```python
+# Use spaCy NER + dictionary
+extractor = TextExtractor(mode="ner")
+entities = extractor.extract(text)
+```
+
+- Better entity boundary detection
+- Catches variations not in dictionary
+- Slower, requires spaCy model
+
+### 3. Hybrid (Recommended)
+
+```python
+# Dictionary first, NER for gaps
+extractor = TextExtractor(mode="hybrid")
+entities = extractor.extract(text)
+```
+
+- Best of both: fast + accurate
+- Dictionary for common entities
+- NER for rare/variant mentions
+
+## Context-Aware Disambiguation
+
+```python
+text = "Georgia joined the UN in 1992."
+
+# Without context: ambiguous
+entities = extractor.extract(text, use_context=False)
+# → Georgia (country) with confidence 0.6
+
+# With context: disambiguated
+entities = extractor.extract(text, use_context=True)
+# → Georgia (country) with confidence 0.95 (UN membership hint)
+```
+
+Context signals:
+- Co-mentioned entities (nearby countries suggest geography)
+- Temporal hints (dates in EU/UN membership context)
+- Acronym expansions (UEMOA → West African Union)
+- Document topic (academic paper vs. news article)
+
+## Handling Special Cases
+
+### Overlapping/Nested Entities
+
+```text
+"Paris, France"
+       ^^^^^^  → France (country)
+ ^^^^^        → Paris (city)
+```
+
+Strategy: Return both with relationship noted
+
+### Acronyms and Full Forms
+
+```text
+"West African Economic and Monetary Union (UEMOA)"
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  → org/UEMOA (full form)
+                                           ^^^^^^ → org/UEMOA (acronym)
+```
+
+Strategy: Link acronym to full form, return single entity
+
+### False Positives
+
+```text
+"We will meet in March in Reading."
+                 ^^^^^  → Month, not country
+                       ^^^^^^^ → City, not event
+```
+
+Filtering:
+- Stoplist (common words)
+- POS tagging (proper nouns only)
+- Capitalization checks
+- Context validation
+
+## Performance
+
+### Dictionary Size vs. Speed
+
+- 10K aliases: ~1ms per sentence
+- 100K aliases: ~5ms per sentence
+- 1M aliases: ~50ms per sentence
+
+Use Aho-Corasick for efficient multi-pattern matching.
+
+### Document Length
+
+- Sentences (<100 words): <10ms
+- Paragraphs (<500 words): <50ms
+- Full documents (>5000 words): <1s
+
+Batch processing recommended for large corpora.
+
+## Design Principles
+
+1. **Precision over recall**: Avoid false positives
+2. **Context-aware**: Use surrounding text for disambiguation
+3. **Configurable**: Adjustable confidence thresholds
+4. **Efficient**: Handle large documents without timeouts
+
+## Implementation Priority
+
+**Phase F** - Entity extraction add-on