src/ingest.py

"""Fetch company data from accelerator sources and generate embeddings."""

import argparse
import json
from pathlib import Path

import numpy as np

from src.embeddings import Embedder, embeddings_filename
from src.sources import ALL_SOURCES

DATA_DIR = Path(__file__).resolve().parent.parent / "data"


FIELD_NAMES = ["problem", "solution", "industry", "customer"]


def build_search_text(c: dict) -> str:
    parts = []
    if c.get("one_liner"):
        parts.append(c["one_liner"])
    if c.get("long_description"):
        parts.append(c["long_description"])
    if c.get("tags"):
        parts.append("Tags: " + ", ".join(c["tags"]))
    if c.get("industries"):
        parts.append("Industries: " + ", ".join(c["industries"]))
    return ". ".join(parts) if parts else c.get("name", "")


def build_field_texts(c: dict) -> list[str]:
    return [
        c.get("_problem") or c.get("one_liner") or c.get("name", ""),
        c.get("_solution") or c.get("one_liner") or "",
        c.get("_industry") or " ".join(c.get("industries") or c.get("tags") or []),
        c.get("_customer") or "",
    ]


def main():
    parser = argparse.ArgumentParser(description="Ingest accelerator company data")
    parser.add_argument(
        "--source",
        choices=[*ALL_SOURCES.keys(), "all"],
        default="all",
        help="Which source to ingest (default: all)",
    )
    parser.add_argument(
        "--embed-only",
        action="store_true",
        help="Skip fetching; re-embed existing data/companies.json only",
    )
    parser.add_argument(
        "--multi",
        action="store_true",
        help="Create per-field embeddings (problem/solution/industry/customer) instead of single blob",
    )
    args = parser.parse_args()

    DATA_DIR.mkdir(parents=True, exist_ok=True)

    if args.embed_only:
        companies_path = DATA_DIR / "companies.json"
        if not companies_path.exists():
            print(f"No {companies_path} found. Run a full ingest first.")
            return
        with open(companies_path) as f:
            all_companies = json.load(f)
        print(f"Loaded {len(all_companies)} companies from {companies_path}")
    else:
        sources_to_run = ALL_SOURCES if args.source == "all" else {args.source: ALL_SOURCES[args.source]}

        all_companies: list[dict] = []
        for name, module in sources_to_run.items():
            print(f"\n{'='*60}")
            print(f"Ingesting: {name}")
            print(f"{'='*60}")
            try:
                companies = module.fetch()
                all_companies.extend(companies)
                print(f"  -> {len(companies)} companies from {name}")
            except Exception as e:
                print(f"  !! Error fetching {name}: {e}")

        if not all_companies:
            print("No companies fetched. Exiting.")
            return

        print(f"\nTotal companies across all sources: {len(all_companies)}")

    embedder = Embedder()

    if args.multi:
        print(f"Generating MULTI-FIELD embeddings for {len(all_companies)} companies "
              f"(provider: {embedder.provider}, dim: {embedder.dim}, fields: {len(FIELD_NAMES)}) ...")
        per_field: list[list[str]] = [[] for _ in FIELD_NAMES]
        for c in all_companies:
            fields = build_field_texts(c)
            for k, text in enumerate(fields):
                per_field[k].append(text if text else " ")
        field_matrices = []
        for k, field_name in enumerate(FIELD_NAMES):
            print(f"\n  [{field_name}] Embedding {len(per_field[k])} texts ...")
            mat = embedder.embed_batch(per_field[k])
            field_matrices.append(mat)
        embeddings = np.stack(field_matrices, axis=1)  # (N, 4, dim)
        print(f"  -> Multi-field embedding shape: {embeddings.shape}")
    else:
        texts = [build_search_text(c) for c in all_companies]
        print(f"Generating embeddings for {len(texts)} companies (provider: {embedder.provider}, dim: {embedder.dim}) ...")
        embeddings = embedder.embed_batch(texts)
        print(f"  -> Embedding matrix shape: {embeddings.shape}")

    companies_path = DATA_DIR / "companies.json"
    embeddings_path = DATA_DIR / embeddings_filename(embedder.provider, multi=args.multi)

    if not args.embed_only:
        with open(companies_path, "w") as f:
            json.dump(all_companies, f, separators=(",", ":"))
        print(f"Saved {companies_path} ({companies_path.stat().st_size / 1024:.0f} KB)")

    np.save(embeddings_path, embeddings)
    print(f"Saved {embeddings_path} ({embeddings_path.stat().st_size / 1024 / 1024:.1f} MB)")

    if not args.embed_only:
        by_source = {}
        for c in all_companies:
            by_source.setdefault(c.get("source", "unknown"), 0)
            by_source[c["source"]] += 1
        print(f"\nBreakdown: {by_source}")
    print("Done!")


if __name__ == "__main__":
    main()