spb/uqo-research-db

1# ============================================================2# SPID — knowledge graph builder3#4# Derives nodes (Company / Project / Technology / Location /5# Partner) and edges (owns / uses / located_in / partners_with)6# from the resolved projects, then exports to parquet + Cypher.7# ============================================================8from __future__ import annotations910from pathlib import Path1112from spid.config import get_logger13from spid.store import Store1415log = get_logger("graph")1617_BUILD = """18DELETE FROM kg_nodes;19DELETE FROM kg_edges;2021-- Company nodes (one per cik; a cik may carry varying name/sector across projects)22INSERT INTO kg_nodes23SELECT 'C:' || cik, 'Company', any_value(coalesce(company_name, ticker, cik)),24       any_value(to_json({'cik': cik, 'ticker': ticker, 'sector': sector}))25FROM projects WHERE cik IS NOT NULL26GROUP BY cik;2728-- Project nodes29INSERT INTO kg_nodes30SELECT 'P:' || project_id, 'Project', project_name,31       to_json({'type': project_type, 'status': status,32                'amount_usd': total_amount_usd, 'n_filings': n_filings,33                'first_seen': first_seen, 'last_seen': last_seen})34FROM projects;3536-- Technology nodes37INSERT INTO kg_nodes38SELECT DISTINCT 'T:' || t, 'Technology', t, NULL39FROM (SELECT unnest(technologies) AS t FROM projects) WHERE t IS NOT NULL AND t <> '';4041-- Location nodes42INSERT INTO kg_nodes43SELECT DISTINCT 'L:' || canonical_location, 'Location', canonical_location, NULL44FROM projects WHERE canonical_location IS NOT NULL AND canonical_location <> '';4546-- Partner nodes (from mentions; id is case-folded so collapse case variants to one row)47INSERT INTO kg_nodes48SELECT 'PR:' || lower(p), 'Partner', any_value(p), NULL49FROM (SELECT unnest(partners) AS p FROM project_mentions)50WHERE p IS NOT NULL AND p <> ''51GROUP BY lower(p);5253-- Edges: Company -owns-> Project54INSERT INTO kg_edges55SELECT 'C:' || cik, 'owns', 'P:' || project_id,56       to_json({'amount_usd': total_amount_usd})57FROM projects WHERE cik IS NOT NULL;5859-- Edges: Project -uses-> Technology60INSERT INTO kg_edges61SELECT 'P:' || project_id, 'uses', 'T:' || t, NULL62FROM (SELECT project_id, unnest(technologies) AS t FROM projects)63WHERE t IS NOT NULL AND t <> '';6465-- Edges: Project -located_in-> Location66INSERT INTO kg_edges67SELECT 'P:' || project_id, 'located_in', 'L:' || canonical_location, NULL68FROM projects WHERE canonical_location IS NOT NULL AND canonical_location <> '';6970-- Edges: Project -requires_capital (self-annotation kept as company edge prop already)71"""727374def build_graph(store: Store) -> dict:75    store.con.execute(_BUILD)76    n_nodes = store.count("kg_nodes")77    n_edges = store.count("kg_edges")78    by_type = dict(store.con.execute(79        "SELECT node_type, COUNT(*) FROM kg_nodes GROUP BY 1 ORDER BY 2 DESC").fetchall())80    log.info("Knowledge graph: %d nodes, %d edges | %s", n_nodes, n_edges, by_type)81    return {"nodes": n_nodes, "edges": n_edges, "nodes_by_type": by_type}828384def export(store: Store, out_dir: Path) -> dict:85    out_dir.mkdir(parents=True, exist_ok=True)86    np = out_dir / "kg_nodes.parquet"87    ep = out_dir / "kg_edges.parquet"88    store.con.execute(f"COPY kg_nodes TO '{np}' (FORMAT PARQUET)")89    store.con.execute(f"COPY kg_edges TO '{ep}' (FORMAT PARQUET)")9091    # Neo4j-style Cypher dump (capped for portability)92    cypher = out_dir / "kg_import.cypher"93    rows_n = store.con.execute(94        "SELECT node_id, node_type, label FROM kg_nodes").fetchall()95    rows_e = store.con.execute(96        "SELECT src, rel, dst FROM kg_edges").fetchall()97    with cypher.open("w") as fh:98        fh.write("// SPID knowledge graph — Neo4j import\n")99        for nid, ntype, label in rows_n:100            safe = (label or "").replace("'", "\\'")[:120]101            fh.write(f"MERGE (n:{ntype} {{id:'{nid}'}}) SET n.label='{safe}';\n")102        for src, rel, dst in rows_e:103            fh.write(f"MATCH (a {{id:'{src}'}}),(b {{id:'{dst}'}}) "104                     f"MERGE (a)-[:{rel.upper()}]->(b);\n")105    log.info("Exported KG: %s, %s, %s", np.name, ep.name, cypher.name)106    return {"nodes_parquet": str(np), "edges_parquet": str(ep), "cypher": str(cypher)}107