Spaces:

Sanjam19
/

FinSight

Sleeping

Deploy FinSight demo (single-container Docker Space)

d4f8959 verified 4 months ago

16.4 kB

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

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