corp-extractor 0.9.3__py3-none-any.whl → 0.9.4__py3-none-any.whl

statement_extractor/database/models.py

@@ -1,5 +1,11 @@
 """
 Pydantic models for organization/entity database records.
+
+v2 Schema Changes:
+- Added SimplifiedLocationType enum for location categorization
+- Added SourceTypeEnum for normalized source references
+- Added RoleRecord and LocationRecord models for new tables
+- Models support both TEXT-based v1 and FK-based v2 schemas
 """
 
 from enum import Enum
@@ -8,7 +14,36 @@ from typing import Any, Literal, Optional
 from pydantic import BaseModel, Field
 
 
-SourceType = Literal["gleif", "sec_edgar", "companies_house", "wikipedia"]
+# Legacy source types for backward compatibility with v1 schema
+SourceType = Literal["gleif", "sec_edgar", "companies_house", "wikipedia", "wikidata"]
+
+
+class SourceTypeEnum(str, Enum):
+    """
+    Data source enumeration for v2 normalized schema.
+
+    Used as foreign key reference to source_types table.
+    """
+    GLEIF = "gleif"               # id=1: GLEIF LEI data
+    SEC_EDGAR = "sec_edgar"       # id=2: SEC EDGAR filings
+    COMPANIES_HOUSE = "companies_house"  # id=3: UK Companies House
+    WIKIDATA = "wikidata"         # id=4: Wikidata/Wikipedia
+
+
+class SimplifiedLocationType(str, Enum):
+    """
+    Simplified location type categories for querying.
+
+    Maps detailed Wikidata location types to canonical categories.
+    Used for filtering searches (e.g., "find all cities").
+    """
+    CONTINENT = "continent"       # id=1: Continents (Q5107)
+    COUNTRY = "country"           # id=2: Countries and sovereign states
+    SUBDIVISION = "subdivision"   # id=3: States, provinces, regions, counties
+    CITY = "city"                 # id=4: Cities, towns, municipalities, communes
+    DISTRICT = "district"         # id=5: Districts, boroughs, neighborhoods
+    HISTORIC = "historic"         # id=6: Former countries, historic territories
+    OTHER = "other"               # id=7: Unclassified locations
 
 
 class EntityType(str, Enum):
@@ -105,9 +140,98 @@ class CompanyRecord(BaseModel):
         }
 
 
+# Person sources (same as org sources but without GLEIF)
 PersonSourceType = Literal["wikidata", "sec_edgar", "companies_house"]
 
 
+# =============================================================================
+# ROLE RECORD MODEL (v2)
+# =============================================================================
+
+
+class RoleRecord(BaseModel):
+    """
+    A role/job title record for the roles table.
+
+    Used for normalizing job titles across sources and enabling role-based search.
+    Supports canonicalization to group equivalent roles (e.g., CEO, Chief Executive).
+    """
+    name: str = Field(..., description="Role/title name (e.g., 'Chief Executive Officer')")
+    source: SourceType = Field(default="wikidata", description="Data source")
+    source_id: Optional[str] = Field(default=None, description="Source identifier (e.g., Q484876 for CEO)")
+    qid: Optional[int] = Field(default=None, description="Wikidata QID as integer (e.g., 484876)")
+    record: dict[str, Any] = Field(default_factory=dict, description="Original record from source")
+
+    @property
+    def canonical_id(self) -> str:
+        """Generate canonical ID in format source:source_id."""
+        if self.source_id:
+            return f"{self.source}:{self.source_id}"
+        return f"{self.source}:{self.name}"
+
+    def model_dump_for_db(self) -> dict[str, Any]:
+        """Convert to dict suitable for database storage."""
+        return {
+            "name": self.name,
+            "source": self.source,
+            "source_id": self.source_id or "",
+            "qid": self.qid,
+            "record": self.record,
+        }
+
+
+# =============================================================================
+# LOCATION RECORD MODEL (v2)
+# =============================================================================
+
+
+class LocationRecord(BaseModel):
+    """
+    A location/place record for the locations table.
+
+    Used for storing geopolitical entities (countries, states, cities) with
+    hierarchical relationships and type classification.
+    """
+    name: str = Field(..., description="Location name (e.g., 'United States', 'California')")
+    source: SourceType = Field(default="wikidata", description="Data source")
+    source_id: Optional[str] = Field(default=None, description="Source identifier (e.g., 'US', 'Q30')")
+    qid: Optional[int] = Field(default=None, description="Wikidata QID as integer (e.g., 30 for USA)")
+    location_type: str = Field(default="country", description="Detailed location type (e.g., 'us_state', 'city')")
+    simplified_type: SimplifiedLocationType = Field(
+        default=SimplifiedLocationType.COUNTRY,
+        description="Simplified type for filtering"
+    )
+    parent_ids: list[int] = Field(
+        default_factory=list,
+        description="Parent location IDs in hierarchy (e.g., [country_id, state_id])"
+    )
+    from_date: Optional[str] = Field(default=None, description="Start date (ISO format YYYY-MM-DD)")
+    to_date: Optional[str] = Field(default=None, description="End date (ISO format YYYY-MM-DD)")
+    record: dict[str, Any] = Field(default_factory=dict, description="Original record from source")
+
+    @property
+    def canonical_id(self) -> str:
+        """Generate canonical ID in format source:source_id."""
+        if self.source_id:
+            return f"{self.source}:{self.source_id}"
+        return f"{self.source}:{self.name}"
+
+    def model_dump_for_db(self) -> dict[str, Any]:
+        """Convert to dict suitable for database storage."""
+        import json
+        return {
+            "name": self.name,
+            "source": self.source,
+            "source_id": self.source_id or "",
+            "qid": self.qid,
+            "location_type": self.location_type,
+            "parent_ids": json.dumps(self.parent_ids),
+            "from_date": self.from_date or "",
+            "to_date": self.to_date or "",
+            "record": self.record,
+        }
+
+
 class PersonRecord(BaseModel):
     """
     A person record for the embedding database.

statement_extractor/database/schema_v2.py

@@ -0,0 +1,409 @@
+"""
+Database schema v2 with normalized foreign key references.
+
+This module contains DDL statements for the normalized entity database schema
+that replaces TEXT-based enum storage with INTEGER FK references to lookup tables.
+
+Changes from v1:
+- Enum tables: source_types, people_types, organization_types, location_types
+- New tables: roles, locations, simplified_location_types
+- organizations_v2: source_id FK, entity_type_id FK, region_id FK (to locations)
+- people_v2: source_id FK, person_type_id FK, country_id FK, known_for_role_id FK
+- qid_labels: qid stored as INTEGER (Q prefix stripped)
+- Human-readable views with JOINs
+"""
+
+# =============================================================================
+# ENUM LOOKUP TABLES
+# =============================================================================
+
+CREATE_SOURCE_TYPES = """
+CREATE TABLE IF NOT EXISTS source_types (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL UNIQUE
+);
+"""
+
+CREATE_PEOPLE_TYPES = """
+CREATE TABLE IF NOT EXISTS people_types (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL UNIQUE
+);
+"""
+
+CREATE_ORGANIZATION_TYPES = """
+CREATE TABLE IF NOT EXISTS organization_types (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL UNIQUE
+);
+"""
+
+CREATE_SIMPLIFIED_LOCATION_TYPES = """
+CREATE TABLE IF NOT EXISTS simplified_location_types (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL UNIQUE
+);
+"""
+
+CREATE_LOCATION_TYPES = """
+CREATE TABLE IF NOT EXISTS location_types (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL UNIQUE,
+    qid INTEGER,
+    simplified_id INTEGER NOT NULL,
+    FOREIGN KEY (simplified_id) REFERENCES simplified_location_types(id)
+);
+"""
+
+# =============================================================================
+# ROLES TABLE
+# =============================================================================
+
+CREATE_ROLES = """
+CREATE TABLE IF NOT EXISTS roles (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    qid INTEGER,
+    name TEXT NOT NULL,
+    name_normalized TEXT NOT NULL,
+    source_id INTEGER NOT NULL,
+    source_identifier TEXT,
+    record TEXT NOT NULL DEFAULT '{}',
+    canon_id INTEGER DEFAULT NULL,
+    canon_size INTEGER DEFAULT 1,
+    FOREIGN KEY (source_id) REFERENCES source_types(id),
+    UNIQUE(name_normalized, source_id)
+);
+"""
+
+CREATE_ROLES_INDEXES = """
+CREATE INDEX IF NOT EXISTS idx_roles_name ON roles(name);
+CREATE INDEX IF NOT EXISTS idx_roles_name_normalized ON roles(name_normalized);
+CREATE INDEX IF NOT EXISTS idx_roles_qid ON roles(qid);
+CREATE INDEX IF NOT EXISTS idx_roles_source_id ON roles(source_id);
+CREATE INDEX IF NOT EXISTS idx_roles_canon_id ON roles(canon_id);
+"""
+
+# =============================================================================
+# LOCATIONS TABLE
+# =============================================================================
+
+CREATE_LOCATIONS = """
+CREATE TABLE IF NOT EXISTS locations (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    qid INTEGER,
+    name TEXT NOT NULL,
+    name_normalized TEXT NOT NULL,
+    source_id INTEGER NOT NULL,
+    source_identifier TEXT,
+    parent_ids TEXT,
+    location_type_id INTEGER NOT NULL,
+    record TEXT NOT NULL DEFAULT '{}',
+    from_date TEXT DEFAULT NULL,
+    to_date TEXT DEFAULT NULL,
+    canon_id INTEGER DEFAULT NULL,
+    canon_size INTEGER DEFAULT 1,
+    FOREIGN KEY (source_id) REFERENCES source_types(id),
+    FOREIGN KEY (location_type_id) REFERENCES location_types(id),
+    UNIQUE(source_identifier, source_id)
+);
+"""
+
+CREATE_LOCATIONS_INDEXES = """
+CREATE INDEX IF NOT EXISTS idx_locations_name ON locations(name);
+CREATE INDEX IF NOT EXISTS idx_locations_name_normalized ON locations(name_normalized);
+CREATE INDEX IF NOT EXISTS idx_locations_qid ON locations(qid);
+CREATE INDEX IF NOT EXISTS idx_locations_source_id ON locations(source_id);
+CREATE INDEX IF NOT EXISTS idx_locations_location_type_id ON locations(location_type_id);
+CREATE INDEX IF NOT EXISTS idx_locations_canon_id ON locations(canon_id);
+"""
+
+# =============================================================================
+# ORGANIZATIONS V2 TABLE
+# =============================================================================
+
+CREATE_ORGANIZATIONS_V2 = """
+CREATE TABLE IF NOT EXISTS organizations (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    qid INTEGER,
+    name TEXT NOT NULL,
+    name_normalized TEXT NOT NULL,
+    source_id INTEGER NOT NULL,
+    source_identifier TEXT NOT NULL,
+    region_id INTEGER,
+    entity_type_id INTEGER NOT NULL DEFAULT 17,
+    from_date TEXT DEFAULT NULL,
+    to_date TEXT DEFAULT NULL,
+    record TEXT NOT NULL DEFAULT '{}',
+    canon_id INTEGER DEFAULT NULL,
+    canon_size INTEGER DEFAULT 1,
+    FOREIGN KEY (source_id) REFERENCES source_types(id),
+    FOREIGN KEY (region_id) REFERENCES locations(id),
+    FOREIGN KEY (entity_type_id) REFERENCES organization_types(id),
+    UNIQUE(source_identifier, source_id)
+);
+"""
+
+CREATE_ORGANIZATIONS_V2_INDEXES = """
+CREATE INDEX IF NOT EXISTS idx_orgs_name ON organizations(name);
+CREATE INDEX IF NOT EXISTS idx_orgs_name_normalized ON organizations(name_normalized);
+CREATE INDEX IF NOT EXISTS idx_orgs_qid ON organizations(qid);
+CREATE INDEX IF NOT EXISTS idx_orgs_source_id ON organizations(source_id);
+CREATE INDEX IF NOT EXISTS idx_orgs_source_identifier ON organizations(source_identifier);
+CREATE INDEX IF NOT EXISTS idx_orgs_region_id ON organizations(region_id);
+CREATE INDEX IF NOT EXISTS idx_orgs_entity_type_id ON organizations(entity_type_id);
+CREATE INDEX IF NOT EXISTS idx_orgs_canon_id ON organizations(canon_id);
+CREATE UNIQUE INDEX IF NOT EXISTS idx_orgs_name_region_source ON organizations(name, region_id, source_id);
+"""
+
+# =============================================================================
+# PEOPLE V2 TABLE
+# =============================================================================
+
+CREATE_PEOPLE_V2 = """
+CREATE TABLE IF NOT EXISTS people (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    qid INTEGER,
+    name TEXT NOT NULL,
+    name_normalized TEXT NOT NULL,
+    source_id INTEGER NOT NULL,
+    source_identifier TEXT NOT NULL,
+    country_id INTEGER,
+    person_type_id INTEGER NOT NULL DEFAULT 15,
+    known_for_role_id INTEGER,
+    known_for_org TEXT NOT NULL DEFAULT '',
+    known_for_org_id INTEGER,
+    from_date TEXT DEFAULT NULL,
+    to_date TEXT DEFAULT NULL,
+    birth_date TEXT DEFAULT NULL,
+    death_date TEXT DEFAULT NULL,
+    record TEXT NOT NULL DEFAULT '{}',
+    canon_id INTEGER DEFAULT NULL,
+    canon_size INTEGER DEFAULT 1,
+    FOREIGN KEY (source_id) REFERENCES source_types(id),
+    FOREIGN KEY (country_id) REFERENCES locations(id),
+    FOREIGN KEY (person_type_id) REFERENCES people_types(id),
+    FOREIGN KEY (known_for_role_id) REFERENCES roles(id),
+    FOREIGN KEY (known_for_org_id) REFERENCES organizations(id),
+    UNIQUE(source_identifier, source_id, known_for_role_id, known_for_org_id)
+);
+"""
+
+CREATE_PEOPLE_V2_INDEXES = """
+CREATE INDEX IF NOT EXISTS idx_people_name ON people(name);
+CREATE INDEX IF NOT EXISTS idx_people_name_normalized ON people(name_normalized);
+CREATE INDEX IF NOT EXISTS idx_people_qid ON people(qid);
+CREATE INDEX IF NOT EXISTS idx_people_source_id ON people(source_id);
+CREATE INDEX IF NOT EXISTS idx_people_source_identifier ON people(source_identifier);
+CREATE INDEX IF NOT EXISTS idx_people_country_id ON people(country_id);
+CREATE INDEX IF NOT EXISTS idx_people_person_type_id ON people(person_type_id);
+CREATE INDEX IF NOT EXISTS idx_people_known_for_role_id ON people(known_for_role_id);
+CREATE INDEX IF NOT EXISTS idx_people_known_for_org_id ON people(known_for_org_id);
+CREATE INDEX IF NOT EXISTS idx_people_canon_id ON people(canon_id);
+"""
+
+# =============================================================================
+# QID LABELS TABLE (V2 - INTEGER QID)
+# =============================================================================
+
+CREATE_QID_LABELS_V2 = """
+CREATE TABLE IF NOT EXISTS qid_labels (
+    qid INTEGER PRIMARY KEY,
+    label TEXT NOT NULL
+);
+"""
+
+# =============================================================================
+# EMBEDDING VIRTUAL TABLES
+# =============================================================================
+
+def get_create_organization_embeddings(embedding_dim: int = 768) -> str:
+    """Get DDL for organization embeddings virtual table."""
+    return f"""
+CREATE VIRTUAL TABLE IF NOT EXISTS organization_embeddings USING vec0(
+    org_id INTEGER PRIMARY KEY,
+    embedding float[{embedding_dim}]
+);
+"""
+
+
+def get_create_person_embeddings(embedding_dim: int = 768) -> str:
+    """Get DDL for person embeddings virtual table."""
+    return f"""
+CREATE VIRTUAL TABLE IF NOT EXISTS person_embeddings USING vec0(
+    person_id INTEGER PRIMARY KEY,
+    embedding float[{embedding_dim}]
+);
+"""
+
+
+def get_create_organization_embeddings_scalar(embedding_dim: int = 768) -> str:
+    """Get DDL for organization scalar (int8) embeddings virtual table."""
+    return f"""
+CREATE VIRTUAL TABLE IF NOT EXISTS organization_embeddings_scalar USING vec0(
+    org_id INTEGER PRIMARY KEY,
+    embedding int8[{embedding_dim}]
+);
+"""
+
+
+def get_create_person_embeddings_scalar(embedding_dim: int = 768) -> str:
+    """Get DDL for person scalar (int8) embeddings virtual table."""
+    return f"""
+CREATE VIRTUAL TABLE IF NOT EXISTS person_embeddings_scalar USING vec0(
+    person_id INTEGER PRIMARY KEY,
+    embedding int8[{embedding_dim}]
+);
+"""
+
+# =============================================================================
+# HUMAN-READABLE VIEWS
+# =============================================================================
+
+CREATE_ORGANIZATIONS_VIEW = """
+CREATE VIEW IF NOT EXISTS organizations_view AS
+SELECT
+    o.id,
+    o.qid,
+    o.name,
+    o.name_normalized,
+    s.name as source,
+    o.source_identifier,
+    l.name as region,
+    slt.name as region_type,
+    ot.name as entity_type,
+    o.from_date,
+    o.to_date,
+    o.canon_id,
+    o.canon_size
+FROM organizations o
+JOIN source_types s ON o.source_id = s.id
+LEFT JOIN locations l ON o.region_id = l.id
+LEFT JOIN location_types lt ON l.location_type_id = lt.id
+LEFT JOIN simplified_location_types slt ON lt.simplified_id = slt.id
+JOIN organization_types ot ON o.entity_type_id = ot.id;
+"""
+
+CREATE_PEOPLE_VIEW = """
+CREATE VIEW IF NOT EXISTS people_view AS
+SELECT
+    p.id,
+    p.qid,
+    p.name,
+    p.name_normalized,
+    s.name as source,
+    p.source_identifier,
+    l.name as country,
+    pt.name as person_type,
+    r.name as known_for_role,
+    p.known_for_org,
+    p.known_for_org_id,
+    p.from_date,
+    p.to_date,
+    p.birth_date,
+    p.death_date,
+    p.canon_id,
+    p.canon_size
+FROM people p
+JOIN source_types s ON p.source_id = s.id
+LEFT JOIN locations l ON p.country_id = l.id
+JOIN people_types pt ON p.person_type_id = pt.id
+LEFT JOIN roles r ON p.known_for_role_id = r.id;
+"""
+
+CREATE_ROLES_VIEW = """
+CREATE VIEW IF NOT EXISTS roles_view AS
+SELECT
+    r.id,
+    r.qid,
+    r.name,
+    r.name_normalized,
+    s.name as source,
+    r.source_identifier,
+    r.canon_id,
+    r.canon_size
+FROM roles r
+JOIN source_types s ON r.source_id = s.id;
+"""
+
+CREATE_LOCATIONS_VIEW = """
+CREATE VIEW IF NOT EXISTS locations_view AS
+SELECT
+    l.id,
+    l.qid,
+    l.name,
+    l.name_normalized,
+    s.name as source,
+    l.source_identifier,
+    l.parent_ids,
+    lt.name as location_type,
+    slt.name as simplified_type,
+    l.from_date,
+    l.to_date,
+    l.canon_id,
+    l.canon_size
+FROM locations l
+JOIN source_types s ON l.source_id = s.id
+JOIN location_types lt ON l.location_type_id = lt.id
+JOIN simplified_location_types slt ON lt.simplified_id = slt.id;
+"""
+
+# =============================================================================
+# ALL DDL STATEMENTS IN ORDER
+# =============================================================================
+
+ALL_DDL_STATEMENTS = [
+    # Enum tables first (no dependencies)
+    CREATE_SOURCE_TYPES,
+    CREATE_PEOPLE_TYPES,
+    CREATE_ORGANIZATION_TYPES,
+    CREATE_SIMPLIFIED_LOCATION_TYPES,
+    CREATE_LOCATION_TYPES,
+    # New entity tables
+    CREATE_ROLES,
+    CREATE_ROLES_INDEXES,
+    CREATE_LOCATIONS,
+    CREATE_LOCATIONS_INDEXES,
+    # Main entity tables
+    CREATE_ORGANIZATIONS_V2,
+    CREATE_ORGANIZATIONS_V2_INDEXES,
+    CREATE_PEOPLE_V2,
+    CREATE_PEOPLE_V2_INDEXES,
+    # Reference tables
+    CREATE_QID_LABELS_V2,
+]
+
+VIEW_DDL_STATEMENTS = [
+    CREATE_ORGANIZATIONS_VIEW,
+    CREATE_PEOPLE_VIEW,
+    CREATE_ROLES_VIEW,
+    CREATE_LOCATIONS_VIEW,
+]
+
+
+def create_all_tables(conn, embedding_dim: int = 768) -> None:
+    """
+    Create all v2 schema tables.
+
+    Args:
+        conn: SQLite connection
+        embedding_dim: Dimension for embedding vectors
+    """
+    for ddl in ALL_DDL_STATEMENTS:
+        for statement in ddl.strip().split(";"):
+            statement = statement.strip()
+            if statement:
+                conn.execute(statement)
+
+    # Create embedding virtual tables (float32)
+    conn.execute(get_create_organization_embeddings(embedding_dim))
+    conn.execute(get_create_person_embeddings(embedding_dim))
+
+    # Create scalar embedding virtual tables (int8) for 75% storage reduction
+    conn.execute(get_create_organization_embeddings_scalar(embedding_dim))
+    conn.execute(get_create_person_embeddings_scalar(embedding_dim))
+
+    # Create views
+    for ddl in VIEW_DDL_STATEMENTS:
+        conn.execute(ddl)
+
+    conn.commit()