backend/graph.py

# backend/graph.py

import json
from pathlib import Path

import chromadb
import networkx as nx
from sentence_transformers import SentenceTransformer

from backend.entity_resolver import resolve_entity
from backend.relation_extractor import extract_relations

RELATION_TYPE_BY_LABEL = {
    "ORG": "MENTIONS_ORG",
    "PERSON": "MENTIONS_PERSON",
    "PRODUCT": "MENTIONS_PRODUCT",
    "MONEY": "MENTIONS_MONEY",
    "DATE": "MENTIONS_DATE",
    "GPE": "MENTIONS_LOCATION",
}


class FinancialGraph:
    def __init__(self):
        self.G = nx.Graph()

        # Vector DB
        self.embedder = SentenceTransformer("all-MiniLM-L6-v2")
        self.chroma = chromadb.PersistentClient(path="data/chroma")
        self.collection = self.chroma.get_or_create_collection("finsight")

    def add_document(self, parsed: dict, use_llm_fallback: bool = True):
        company = parsed["company"]
        year = str(parsed["year"])
        filing_id = f"{company}_{year}"

        self.G.add_node(company, type="company")

        self.G.add_node(
            filing_id,
            type="filing",
            company=company,
            year=year,
            sector=parsed.get("sector", "GENERAL"),
            metrics=parsed["metrics"]
        )
        self.G.add_edge(company, filing_id, rel="FILED")

        # chunk_id -> list of (mention_text, resolved_eid) for ORG entities
        # found in that chunk's text, used for Phase 2 relation extraction below
        chunk_org_mentions = {}

        for chunk in parsed["chunks"]:
            cid = chunk["chunk_id"]

            self.G.add_node(
                cid,
                type="chunk",
                text=chunk["text"],
                company=company,
                year=year,
                page=chunk.get("page")
            )

            self.G.add_edge(filing_id, cid, rel="CONTAINS")
            chunk_org_mentions[cid] = []

        # NOTE: existing_orgs holds CANONICAL names (e.g. "Apple"), not raw
        # mention text, so fuzzy matching in resolve_entity compares like
        # with like. Previously this pulled data.get("text") (raw mentions),
        # which made fuzzy-match quality depend on whichever variant got
        # stored first — fixed here.
        existing_orgs = [
            data.get("canonical_id", "").replace("ORG_", "")
            for node, data in self.G.nodes(data=True)
            if data.get("type") == "entity" and data.get("label") == "ORG"
        ]

        for ent in parsed["entities"]:
            eid = resolve_entity(ent["text"], ent["label"], existing_orgs)
            rel_type = RELATION_TYPE_BY_LABEL.get(ent["label"], "MENTIONS")

            if not self.G.has_node(eid):
                self.G.add_node(
                    eid,
                    type="entity",
                    label=ent["label"],
                    text=ent["text"],
                    canonical_id=eid
                )
                if ent["label"] == "ORG":
                    existing_orgs.append(eid.replace("ORG_", ""))

            if not self.G.has_edge(company, eid):
                self.G.add_edge(company, eid, rel=rel_type)
            else:
                # update the edge label if a more specific relationship exists
                existing_rel = self.G.edges[company, eid].get("rel")
                if existing_rel == "MENTIONS" and rel_type != "MENTIONS":
                    self.G.edges[company, eid]["rel"] = rel_type

            # track which chunk(s) this ORG mention actually appears in,
            # so Phase 2 relation extraction runs on the right chunk text
            if ent["label"] == "ORG":
                for chunk in parsed["chunks"]:
                    if ent["text"] in chunk["text"]:
                        chunk_org_mentions[chunk["chunk_id"]].append(
                            (ent["text"], eid)
                        )

        # ===== Phase 2: typed relationship extraction between ORG entities
        # co-occurring in the same chunk (e.g. SUPPLIER_TO, COMPETITOR_OF) =====
        for chunk in parsed["chunks"]:
            cid = chunk["chunk_id"]
            mentions = chunk_org_mentions.get(cid, [])

            if len(mentions) < 2:
                continue

            mention_texts = [m[0] for m in mentions]
            mention_to_eid = dict(mentions)

            relations = extract_relations(
                chunk["text"],
                mention_texts,
                use_llm_fallback=use_llm_fallback
            )

            for rel in relations:
                eid_a = mention_to_eid.get(rel["org_a"])
                eid_b = mention_to_eid.get(rel["org_b"])

                if not eid_a or not eid_b or eid_a == eid_b:
                    continue

                edge_attrs = {
                    "rel": rel["relation"],
                    "source": rel["source"],
                    "chunk_id": cid,
                    # nx.Graph is undirected — it does NOT preserve which
                    # node was passed first to add_edge(). Relation types
                    # like SUPPLIER_TO/SUBSIDIARY_OF/ACQUIRED are directional
                    # in meaning, so the direction has to be stored explicitly
                    # as data, not inferred from tuple order.
                    "from": eid_a,
                    "to": eid_b
                }
                if rel["source"] == "llm":
                    edge_attrs["confidence"] = rel.get("confidence", "low")

                if self.G.has_edge(eid_a, eid_b):
                    # don't overwrite a pattern-sourced edge with a lower
                    # confidence LLM-sourced guess for the same pair
                    existing_source = self.G.edges[eid_a, eid_b].get("source")
                    if existing_source == "pattern" and rel["source"] == "llm":
                        continue
                    self.G.edges[eid_a, eid_b].update(edge_attrs)
                else:
                    self.G.add_edge(eid_a, eid_b, **edge_attrs)

        # ===== ChromaDB indexing =====
        texts = [c["text"] for c in parsed["chunks"]]
        ids = [c["chunk_id"] for c in parsed["chunks"]]
        metadatas = [
            {
                "company": company,
                "year": year,
                # Chroma metadata can't hold None — 0 means "page unknown"
                "page": c.get("page") or 0
            }
            for c in parsed["chunks"]
        ]

        if texts:
            embeddings = self.embedder.encode(texts).tolist()

            batch_size = 100

            for i in range(0, len(texts), batch_size):
                self.collection.upsert(
                    documents=texts[i:i + batch_size],
                    ids=ids[i:i + batch_size],
                    embeddings=embeddings[i:i + batch_size],
                    metadatas=metadatas[i:i + batch_size]
                )

        print(
            f"Graph: {self.G.number_of_nodes()} nodes, "
            f"{self.G.number_of_edges()} edges"
        )

    def get_company_metrics(self, company: str) -> dict:
        metrics_by_year = {}

        for node, data in self.G.nodes(data=True):
            if data.get("type") != "filing":
                continue

            if data.get("company") != company:
                continue

            year = str(data.get("year", ""))

            if not year.isdigit():
                continue

            if int(year) < 2020 or int(year) > 2030:
                continue

            metrics_by_year[year] = data.get("metrics", {})

        return metrics_by_year

    def get_filing_sector(self, company: str, year: str) -> str | None:
        """Returns the detected sector for a specific company+year filing,
        or None if no such filing exists. Used by /red_flags and
        /recommendation endpoints to route to the correct sector-specific
        rule set without the caller needing to re-detect it."""
        filing_id = f"{company}_{year}"
        if not self.G.has_node(filing_id):
            return None
        return self.G.nodes[filing_id].get("sector")

    def get_relevant_chunks(
        self,
        question: str,
        company: str = None,
        top_k: int = 3
    ) -> list:
        """Returns list of {"text", "company", "year", "page"} dicts —
        page provenance travels with every retrieved chunk so answers
        can cite where they came from. page is None when unknown (chunks
        indexed before pages were tracked)."""

        try:
            q_embedding = self.embedder.encode([question]).tolist()[0]

            where = {"company": company} if company else None

            results = self.collection.query(
                query_embeddings=[q_embedding],
                n_results=top_k,
                where=where
            )

            if results and results.get("documents") and results["documents"][0]:
                docs = results["documents"][0]
                metas = (results.get("metadatas") or [[]])[0]
                out = []
                for i, doc in enumerate(docs):
                    meta = metas[i] if i < len(metas) else {}
                    out.append({
                        "text": doc,
                        "company": meta.get("company"),
                        "year": meta.get("year"),
                        "page": meta.get("page") or None
                    })
                return out

        except Exception as e:
            print("Chroma query failed:", e)

        # ===== Fallback keyword search =====
        stopwords = {
            "what", "was", "the", "in", "of", "a", "an",
            "is", "are", "how", "did", "does", "we",
            "if", "that", "you", "as", "based", "on",
            "which", "has", "for"
        }

        keywords = [
            w.lower().strip("?.,")
            for w in question.split()
            if w.lower() not in stopwords and len(w) > 2
        ]

        chunk_scores = {}

        for node, data in self.G.nodes(data=True):
            if data.get("type") != "chunk":
                continue

            if (
                company
                and data.get("company", "").lower()
                != company.lower()
            ):
                continue

            text = data.get("text", "").lower()

            score = sum(
                1 for kw in keywords
                if kw in text
            )

            if score > 0:
                chunk_scores[node] = score

        top = sorted(
            chunk_scores,
            key=chunk_scores.get,
            reverse=True
        )[:top_k]

        return [
            {
                "text": self.G.nodes[n]["text"],
                "company": self.G.nodes[n].get("company"),
                "year": self.G.nodes[n].get("year"),
                "page": self.G.nodes[n].get("page")
            }
            for n in top
        ]

    def compare_companies(
        self,
        companies: list,
        metric: str
    ) -> dict:

        result = {}

        for company in companies:
            metrics = self.get_company_metrics(company)

            result[company] = {
                year: data.get(metric)
                for year, data in metrics.items()
                if data.get(metric) is not None
            }

        return result

    def get_company_graph(self, company: str) -> dict:
        """Returns this company's subgraph as JSON-safe data for the
        /graph/{company} demo endpoint — nodes the company connects to
        (filings, entities) plus all edges among them, with typed
        entity-relation edges separated out for visibility."""
        if not self.G.has_node(company):
            return {"company": company, "found": False}

        node_ids = {company}
        node_ids.update(self.G.neighbors(company))

        # also pull in neighbors-of-neighbors one hop further, so
        # entity-to-entity relation edges (e.g. Foxconn-Apple) show up
        # even though Foxconn isn't directly linked to the company node
        for n in list(node_ids):
            node_ids.update(self.G.neighbors(n))

        nodes = []
        for n in node_ids:
            data = dict(self.G.nodes[n])
            data.pop("text", None)  # skip full chunk text, too noisy for this view
            data["id"] = n
            nodes.append(data)

        edges = []
        relations = []
        seen = set()

        for n in node_ids:
            for neighbor in self.G.neighbors(n):
                if neighbor not in node_ids:
                    continue
                edge_key = frozenset([n, neighbor])
                if edge_key in seen:
                    continue
                seen.add(edge_key)

                edge_data = dict(self.G.edges[n, neighbor])
                rel = edge_data.get("rel")

                if rel in (
                    "SUBSIDIARY_OF", "COMPETITOR_OF", "SUPPLIER_TO",
                    "PARTNERED_WITH", "ACQUIRED", "BOARD_OVERLAP_WITH"
                ):
                    relations.append({
                        "from": edge_data.get("from", n),
                        "to": edge_data.get("to", neighbor),
                        "relation": rel,
                        "source": edge_data.get("source"),
                        "confidence": edge_data.get("confidence")
                    })
                else:
                    edges.append({
                        "from": n, "to": neighbor, "rel": rel
                    })

        return {
            "company": company,
            "found": True,
            "node_count": len(nodes),
            "nodes": nodes,
            "structural_edges": edges,
            "typed_relations": relations
        }

    def save(self, path: str):
        data = nx.node_link_data(self.G)

        with open(path, "w") as f:
            json.dump(data, f)

        print(f"Graph saved to {path}")

    def load(self, path: str):
        with open(path) as f:
            data = json.load(f)

        self.G = nx.node_link_graph(data)

        print(
            f"Graph loaded: "
            f"{self.G.number_of_nodes()} nodes"
        )


if __name__ == "__main__":
    fg = FinancialGraph()

    doc1 = {
        "company": "Apple", "year": "2022", "file": "test1.pdf",
        "char_count": 100, "chunk_count": 1,
        "metrics": {"revenue": 394300000000.0},
        "entities": [{"text": "Apple Inc.", "label": "ORG"}],
        "chunks": [{"chunk_id": "Apple_2022_0000", "company": "Apple", "year": "2022",
                    "text": "Apple Inc. reported revenue of $394.3 billion."}]
    }

    doc2 = {
        "company": "Apple", "year": "2023", "file": "test2.pdf",
        "char_count": 100, "chunk_count": 1,
        "metrics": {"revenue": 383300000000.0},
        "entities": [{"text": "AAPL", "label": "ORG"}, {"text": "Apple", "label": "ORG"}],
        "chunks": [{"chunk_id": "Apple_2023_0000", "company": "Apple", "year": "2023",
                    "text": "AAPL reported revenue for fiscal 2023."}]
    }

    doc3 = {
        "company": "Apple", "year": "2024", "file": "test3.pdf",
        "char_count": 150, "chunk_count": 1,
        "metrics": {"revenue": 391000000000.0},
        "entities": [
            {"text": "Apple", "label": "ORG"},
            {"text": "Foxconn", "label": "ORG"}
        ],
        "chunks": [{"chunk_id": "Apple_2024_0000", "company": "Apple", "year": "2024",
                    "text": "Foxconn is a major supplier to Apple for iPhone assembly."}]
    }

    fg.add_document(doc1, use_llm_fallback=False)
    fg.add_document(doc2, use_llm_fallback=False)
    fg.add_document(doc3, use_llm_fallback=False)

    org_nodes = [n for n, d in fg.G.nodes(data=True)
                 if d.get("type") == "entity" and d.get("label") == "ORG"]
    print("ORG entity nodes:", org_nodes)
    print("Total nodes:", fg.G.number_of_nodes())
    print("Total edges:", fg.G.number_of_edges())

    print("\nEntity-to-entity typed relations:")
    for u, v, data in fg.G.edges(data=True):
        if data.get("rel") not in (
            "FILED", "CONTAINS", "MENTIONS", "MENTIONS_ORG",
            "MENTIONS_LOCATION", "MENTIONS_PERSON", "MENTIONS_PRODUCT",
            "MENTIONS_MONEY", "MENTIONS_DATE"
        ):
            from_node = data.get("from", u)
            to_node = data.get("to", v)
            print(f"  {from_node} --[{data.get('rel')}]--> {to_node}  (source={data.get('source')})")