# 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')})")