spb/uqo-research-db
Public
1# ============================================================2# SPID — DuckDB store3#4# Source corpus (filings / sections / companies) is exposed as5# read-only views over the existing parquet. SPID writes its6# own tables: mentions, projects, timeline, graph, embeddings.7# ============================================================8from __future__ import annotations910from pathlib import Path1112import duckdb1314from spid.config import Settings, get_settings15from spid.extract.base import ProjectMention1617SCHEMA = """18CREATE TABLE IF NOT EXISTS project_mentions (19 mention_id VARCHAR PRIMARY KEY,20 accession_number VARCHAR,21 cik VARCHAR,22 section_id VARCHAR,23 section_name VARCHAR,24 project_name VARCHAR,25 project_type VARCHAR,26 description VARCHAR,27 objective VARCHAR,28 amount_usd DOUBLE,29 amount_raw VARCHAR,30 locations VARCHAR[],31 technologies VARCHAR[],32 partners VARCHAR[],33 suppliers VARCHAR[],34 benefits VARCHAR[],35 risks VARCHAR[],36 years INTEGER[],37 status VARCHAR,38 confidence DOUBLE,39 backend VARCHAR,40 -- denormalised filing context (filled on load)41 ticker VARCHAR,42 company_name VARCHAR,43 form_type VARCHAR,44 filing_date DATE,45 fiscal_year INTEGER,46 sector VARCHAR47);4849CREATE TABLE IF NOT EXISTS projects (50 project_id VARCHAR PRIMARY KEY,51 cik VARCHAR,52 ticker VARCHAR,53 company_name VARCHAR,54 sector VARCHAR,55 project_type VARCHAR,56 project_name VARCHAR,57 canonical_location VARCHAR,58 technologies VARCHAR[],59 status VARCHAR,60 total_amount_usd DOUBLE,61 n_mentions INTEGER,62 n_filings INTEGER,63 first_seen DATE,64 last_seen DATE,65 first_year INTEGER,66 last_year INTEGER,67 avg_confidence DOUBLE,68 description VARCHAR69);7071CREATE TABLE IF NOT EXISTS project_timeline (72 project_id VARCHAR,73 accession_number VARCHAR,74 filing_date DATE,75 form_type VARCHAR,76 status VARCHAR,77 amount_usd DOUBLE,78 confidence DOUBLE,79 snippet VARCHAR80);8182CREATE TABLE IF NOT EXISTS kg_nodes (83 node_id VARCHAR PRIMARY KEY,84 node_type VARCHAR, -- Company | Project | Technology | Location | Partner85 label VARCHAR,86 props JSON87);8889CREATE TABLE IF NOT EXISTS kg_edges (90 src VARCHAR,91 rel VARCHAR, -- owns | uses | located_in | partners_with | requires_capital92 dst VARCHAR,93 props JSON94);9596CREATE TABLE IF NOT EXISTS project_embeddings (97 project_id VARCHAR PRIMARY KEY,98 model VARCHAR,99 dim INTEGER,100 vector FLOAT[]101);102"""103104105class Store:106 def __init__(self, settings: Settings | None = None, read_only: bool = False):107 self.s = settings or get_settings()108 self.s.db_path.parent.mkdir(parents=True, exist_ok=True)109 self.con = duckdb.connect(str(self.s.db_path), read_only=read_only)110 if not read_only:111 self._init()112113 def _init(self) -> None:114 self.con.execute(115 "CREATE TABLE IF NOT EXISTS spid_sections ("116 "section_id VARCHAR PRIMARY KEY, accession_number VARCHAR, section_name VARCHAR,"117 "item_number VARCHAR, text VARCHAR, word_count INTEGER)")118 self._register_sources()119 self.con.execute(SCHEMA)120121 def _register_sources(self) -> None:122 c = self.con123 c.execute(124 "CREATE TABLE IF NOT EXISTS extracted_sections ("125 "section_id VARCHAR PRIMARY KEY, accession_number VARCHAR, backend VARCHAR,"126 "n_mentions INTEGER, extracted_at TIMESTAMP DEFAULT current_timestamp)")127 c.execute(f"CREATE OR REPLACE VIEW src_filings AS SELECT * FROM read_parquet('{self.s.filings_parquet}')")128 c.execute(f"CREATE OR REPLACE VIEW src_sections AS SELECT * FROM read_parquet('{self.s.sections_parquet}')")129 c.execute(f"CREATE OR REPLACE VIEW src_companies AS SELECT * FROM read_parquet('{self.s.companies_parquet}')")130 # unified, enriched section feed: pre-parsed sections (all forms since 2026-10-27)131 # UNION ingested 10-K items (legacy S&P 500 run); a 10-K already ingested in132 # spid_sections is not read a second time from the parsed corpus.133 scope = self.s.section_scope134 scope_sql = ""135 if scope:136 parts = []137 for form, items in scope.items():138 q = ",".join("'" + i + "'" for i in items)139 parts.append(f"(form_type = '{form}' AND section_name IN ({q}))")140 others = ",".join("'" + f + "'" for f in scope)141 scope_sql = " WHERE " + " OR ".join(parts) + f" OR form_type NOT IN ({others})"142 c.execute(f"""143 CREATE OR REPLACE VIEW v_sections AS144 WITH parsed AS (145 SELECT s.section_id, s.accession_number, s.section_name, s.item_number, s.text, s.word_count,146 f.form_type147 FROM src_sections s JOIN src_filings f USING (accession_number)148 WHERE NOT (f.form_type = '10-K'149 AND s.accession_number IN (SELECT DISTINCT accession_number FROM spid_sections))150 ),151 all_sec AS (152 SELECT section_id, accession_number, section_name, item_number, text, word_count153 FROM parsed{scope_sql}154 UNION ALL155 SELECT section_id, accession_number, section_name, item_number, text, word_count156 FROM spid_sections157 )158 SELECT s.section_id, s.accession_number, s.section_name, s.item_number,159 s.text, s.word_count,160 f.cik, f.ticker, f.company_name AS filing_company, f.form_type,161 CAST(f.filing_date AS DATE) AS filing_date, f.fiscal_year,162 co.company_name AS company_name, co.sector, co.sub_industry, co.sic_code, co.state163 FROM all_sec s164 JOIN src_filings f USING (accession_number)165 LEFT JOIN src_companies co ON f.cik = co.cik166 """)167168 # ---- read side ----------------------------------------------------169 def companies(self) -> list[tuple]:170 return self.con.execute(171 "SELECT DISTINCT cik, ticker FROM src_filings ORDER BY cik"172 ).fetchall()173174 def sections_for(self, *, ciks: list[str] | None = None, forms: list[str] | None = None,175 limit: int | None = None, min_words: int = 30, skip_done: bool = True) -> list[dict]:176 where = [f"word_count >= {int(min_words)}"]177 if skip_done:178 where.append("section_id NOT IN (SELECT section_id FROM extracted_sections)")179 if ciks:180 quoted = ",".join("'" + c.replace("'", "") + "'" for c in ciks)181 where.append(f"cik IN ({quoted})")182 if forms:183 quoted = ",".join("'" + f.replace("'", "") + "'" for f in forms)184 where.append(f"form_type IN ({quoted})")185 sql = "SELECT * FROM v_sections WHERE " + " AND ".join(where)186 sql += " ORDER BY filing_date"187 if limit:188 sql += f" LIMIT {int(limit)}"189 res = self.con.execute(sql)190 table = res.to_arrow_table() if hasattr(res, "to_arrow_table") else res.fetch_arrow_table() # DuckDB >= 2.0 / < 2.0191 return table.to_pylist()192193 # ---- write side ---------------------------------------------------194 def upsert_mentions(self, mentions: list[ProjectMention], ctx: dict) -> int:195 """ctx maps accession_number -> filing context dict."""196 if not mentions:197 return 0198 rows = []199 for m in mentions:200 cx = ctx.get(m.accession_number, {})201 rows.append((202 m.mention_id, m.accession_number, m.cik, m.section_id, m.section_name,203 m.project_name, m.project_type, m.description, m.objective,204 m.amount_usd, m.amount_raw, m.locations, m.technologies, m.partners,205 m.suppliers, m.benefits, m.risks, m.years, m.status, m.confidence, m.backend,206 cx.get("ticker"), cx.get("company_name") or cx.get("filing_company"),207 cx.get("form_type"), cx.get("filing_date"), cx.get("fiscal_year"), cx.get("sector"),208 ))209 self.con.executemany(210 "INSERT OR REPLACE INTO project_mentions VALUES "211 "(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", rows)212 return len(rows)213214 def mark_extracted(self, rows: list[tuple]) -> None:215 """rows: (section_id, accession_number, backend, n_mentions)."""216 if rows:217 self.con.executemany(218 "INSERT OR REPLACE INTO extracted_sections (section_id, accession_number, backend, n_mentions) "219 "VALUES (?,?,?,?)", rows)220221 def backfill_extracted(self, backend: str = "llm") -> int:222 """Record every section currently visible in v_sections as already extracted.223224 Used once, before the Russell 3000 corpus is merged, so the legacy S&P 500 run225 (whose sections were processed without a ledger) is not re-billed.226 """227 self.con.execute(f"""228 INSERT OR REPLACE INTO extracted_sections (section_id, accession_number, backend, n_mentions)229 SELECT v.section_id, v.accession_number, '{backend}',230 (SELECT COUNT(*) FROM project_mentions m WHERE m.section_id = v.section_id)231 FROM v_sections v232 """)233 return self.count("extracted_sections")234235 def count(self, table: str) -> int:236 return self.con.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]237238 def close(self) -> None:239 self.con.close()240