lib/evidence_graph.py

"""
lib/evidence_graph.py — V6 Evidence Graph

Builds a lightweight graph structure from extracted evidence.
Nodes are evidence pieces, companies, roles, and skills.
Edges connect related nodes (temporal, causal, skill-to-evidence).

The graph is used by the reasoning engine to produce richer,
more connected explanations. Not a full graph database —
just an adjacency list for traversal.

Graph structure:
  Node types: evidence, company, role, skill, metric, time_period
  Edge types: achieved_at, used_skill, demonstrates_ownership,
              shows_impact, during_period, at_company
"""
from __future__ import annotations
from dataclasses import dataclass, field
from lib.evidence import Evidence, extract_all_evidence
from lib import schema


@dataclass
class GraphNode:
    """A node in the evidence graph."""
    node_id: str
    node_type: str        # evidence, company, role, skill, metric, time_period
    label: str
    score: float = 0.0
    metadata: dict = field(default_factory=dict)


@dataclass
class GraphEdge:
    """A directed edge in the evidence graph."""
    source: str           # node_id
    target: str           # node_id
    edge_type: str        # achieved_at, used_skill, etc.
    weight: float = 1.0


@dataclass
class EvidenceGraph:
    """Lightweight evidence graph for a candidate."""
    nodes: dict[str, GraphNode] = field(default_factory=dict)
    edges: list[GraphEdge] = field(default_factory=list)
    _adj: dict[str, list[str]] = field(default_factory=dict)

    def add_node(self, node: GraphNode) -> None:
        self.nodes[node.node_id] = node
        if node.node_id not in self._adj:
            self._adj[node.node_id] = []

    def add_edge(self, edge: GraphEdge) -> None:
        self.edges.append(edge)
        if edge.source in self._adj:
            self._adj[edge.source].append(edge.target)
        else:
            self._adj[edge.source] = [edge.target]

    def neighbors(self, node_id: str) -> list[str]:
        return self._adj.get(node_id, [])

    def get_nodes_by_type(self, node_type: str) -> list[GraphNode]:
        return [n for n in self.nodes.values() if n.node_type == node_type]

    def get_edges_by_type(self, edge_type: str) -> list[GraphEdge]:
        return [e for e in self.edges if e.edge_type == edge_type]

    def strongest_evidence(self, n: int = 3) -> list[GraphNode]:
        """Get the N strongest evidence nodes."""
        ev_nodes = self.get_nodes_by_type("evidence")
        ev_nodes.sort(key=lambda x: x.score, reverse=True)
        return ev_nodes[:n]

    def connected_companies(self) -> list[str]:
        """Get all companies connected to evidence nodes."""
        companies = set()
        for edge in self.get_edges_by_type("achieved_at"):
            companies.add(edge.target)
        return list(companies)

    def skills_in_evidence(self) -> list[str]:
        """Get all skills connected to evidence nodes."""
        skills = set()
        for edge in self.get_edges_by_type("used_skill"):
            skills.add(edge.target)
        return list(skills)

    def evidence_for_company(self, company_id: str) -> list[GraphNode]:
        """Get all evidence achieved at a specific company."""
        ev_ids = [e.source for e in self.edges
                  if e.edge_type == "achieved_at" and e.target == company_id]
        return [self.nodes[eid] for eid in ev_ids if eid in self.nodes]

    def impact_chain(self, evidence_id: str) -> list[GraphNode]:
        """Follow the impact chain from an evidence node."""
        chain = []
        visited = set()
        current = evidence_id
        while current and current not in visited:
            visited.add(current)
            if current in self.nodes:
                chain.append(self.nodes[current])
            # Follow to metric, company, skill
            neighbors = self.neighbors(current)
            if neighbors:
                # Prefer metric/skill edges over company
                metric_edges = [n for n in neighbors
                               if any(e.source == current and e.target == n
                                      and e.edge_type in ("shows_metric", "used_skill")
                                      for e in self.edges)]
                current = metric_edges[0] if metric_edges else neighbors[0]
            else:
                break
        return chain


def build_graph(candidate: dict) -> EvidenceGraph:
    """
    Build an evidence graph from a candidate's extracted evidence.
    """
    graph = EvidenceGraph()
    all_evidence = extract_all_evidence(candidate)

    # Build company and role nodes first
    ch = schema.career_history(cand=candidate)
    company_nodes = {}
    for i, role in enumerate(ch):
        company = role.get("company", "")
        title = role.get("title", "")
        if company:
            cid = f"company_{i}_{company.lower().replace(' ', '_')}"
            graph.add_node(GraphNode(
                node_id=cid, node_type="company",
                label=company,
                metadata={"title": title, "role_index": i},
            ))
            company_nodes[i] = cid

    # Add evidence nodes and edges
    for ev in all_evidence:
        ev_id = f"evidence_{ev.type}_{ev.metric or ev.ownership or 'x'}".replace(" ", "_")[:80]
        # Ensure unique ID
        base_id = ev_id
        counter = 0
        while ev_id in graph.nodes:
            counter += 1
            ev_id = f"{base_id}_{counter}"

        graph.add_node(GraphNode(
            node_id=ev_id,
            node_type="evidence",
            label=ev.context,
            score=ev.score,
            metadata={
                "type": ev.type,
                "metric": ev.metric,
                "ownership": ev.ownership,
                "company": ev.company,
                "year_range": ev.year_range,
                "domain": ev.domain,
            },
        ))

        # Edge to company
        if ev.company:
            for ci, cid in company_nodes.items():
                company_name = ch[ci].get("company", "")
                if company_name.lower() == ev.company.lower():
                    graph.add_edge(GraphEdge(
                        source=ev_id, target=cid,
                        edge_type="achieved_at", weight=ev.score / 20.0,
                    ))
                    break

        # Edge to skill/domain
        if ev.domain:
            skill_id = f"skill_{ev.domain}"
            if skill_id not in graph.nodes:
                graph.add_node(GraphNode(
                    node_id=skill_id, node_type="skill",
                    label=ev.domain,
                ))
            graph.add_edge(GraphEdge(
                source=ev_id, target=skill_id,
                edge_type="used_skill", weight=0.8,
            ))

        # Edge for metric
        if ev.metric:
            metric_id = f"metric_{ev.metric.replace(' ', '_')[:40]}".replace("/", "_")
            if metric_id not in graph.nodes:
                graph.add_node(GraphNode(
                    node_id=metric_id, node_type="metric",
                    label=ev.metric,
                ))
            graph.add_edge(GraphEdge(
                source=ev_id, target=metric_id,
                edge_type="shows_metric", weight=ev.score / 20.0,
            ))

        # Edge for ownership
        if ev.ownership:
            own_id = f"ownership_{ev.ownership.replace(' ', '_')}"
            if own_id not in graph.nodes:
                graph.add_node(GraphNode(
                    node_id=own_id, node_type="ownership",
                    label=ev.ownership,
                    metadata={"weight": ev.ownership},
                ))
            graph.add_edge(GraphEdge(
                source=ev_id, target=own_id,
                edge_type="demonstrates_ownership", weight=0.7,
            ))

    return graph