| """ |
| DocMap — lightweight structural index built from the Graph. |
| |
| The agent reads this first for every question (~300 tokens). |
| Answers structural questions with zero retrieval: |
| - what sections exist? |
| - which are most important? |
| - what does section X depend on? |
| - what type of paper is this? |
| |
| Built once at ingestion, cached to disk, never rebuilt unless graph changes. |
| """ |
|
|
| from __future__ import annotations |
| from dataclasses import dataclass, field |
| from typing import Optional |
| import json |
|
|
| from graph.math_graph import Graph |
|
|
|
|
| @dataclass |
| class SectionEntry: |
| section_id: str |
| title: str |
| importance: int |
| math_object_count: int |
| proof_count: int |
| depends_on: list[str] |
| referenced_by: list[str] |
|
|
|
|
| @dataclass |
| class DocMap: |
| paper_id: str |
| paper_title: str |
| sections: list[SectionEntry] |
| top_nodes: list[dict] |
| total_nodes: int |
| total_edges: int |
| dominant_type: str |
|
|
| |
|
|
| def section_by_id(self, section_id: str) -> Optional[SectionEntry]: |
| for s in self.sections: |
| if s.section_id == section_id: |
| return s |
| return None |
|
|
| def most_important_sections(self, n: int = 3) -> list[SectionEntry]: |
| return sorted(self.sections, key=lambda s: s.importance, reverse=True)[:n] |
|
|
| def to_prompt_block(self) -> str: |
| """Compact text for agent prompt injection. ~300 tokens.""" |
| is_prose = "prose" in self.dominant_type |
| lines = [ |
| f'Paper: "{self.paper_title}"', |
| f"Sections ({len(self.sections)} total, " |
| f"{self.total_edges} cross-references, {self.dominant_type}):", |
| ] |
| for s in self.sections: |
| dep_str = (f" → depends on: {', '.join(s.depends_on)}" |
| if s.depends_on else "") |
| if is_prose: |
| lines.append( |
| f" [{s.section_id}] {s.title} " |
| f"(importance={s.importance}, concepts={s.math_object_count})" |
| f"{dep_str}" |
| ) |
| else: |
| lines.append( |
| f" [{s.section_id}] {s.title} " |
| f"(importance={s.importance}, " |
| f"math_objects={s.math_object_count}, " |
| f"proofs={s.proof_count})" |
| f"{dep_str}" |
| ) |
| if self.top_nodes: |
| lines.append("\nMost referenced nodes:") |
| for n in self.top_nodes[:5]: |
| pt = f", proof_type={n['proof_type']}" if n.get("proof_type") else "" |
| lines.append( |
| f" {n['label']} " |
| f"({n['node_type']}, in-degree={n['in_degree']}{pt})" |
| ) |
| return "\n".join(lines) |
|
|
| |
|
|
| def to_dict(self) -> dict: |
| return { |
| "paper_id": self.paper_id, |
| "paper_title": self.paper_title, |
| "sections": [s.__dict__ for s in self.sections], |
| "top_nodes": self.top_nodes, |
| "total_nodes": self.total_nodes, |
| "total_edges": self.total_edges, |
| "dominant_type": self.dominant_type, |
| } |
|
|
| @classmethod |
| def from_dict(cls, data: dict) -> "DocMap": |
| return cls( |
| paper_id=data["paper_id"], |
| paper_title=data["paper_title"], |
| sections=[SectionEntry(**s) for s in data["sections"]], |
| top_nodes=data["top_nodes"], |
| total_nodes=data["total_nodes"], |
| total_edges=data["total_edges"], |
| dominant_type=data["dominant_type"], |
| ) |
|
|
| def to_json(self) -> str: |
| return json.dumps(self.to_dict(), ensure_ascii=False, indent=2) |
|
|
| @classmethod |
| def from_json(cls, s: str) -> "DocMap": |
| return cls.from_dict(json.loads(s)) |
|
|
|
|
| |
|
|
| def build_english_doc_map(doc, section_data: list[dict]) -> "DocMap": |
| """ |
| Build DocMap for non-math papers from section keyword data. |
| section_data already excludes bibliography sections. |
| importance = keyword richness; math_object_count = named concept count. |
| """ |
| sections = [] |
| for s in section_data: |
| sections.append(SectionEntry( |
| section_id=s["section_id"], |
| title=s["title"], |
| importance=len(s.get("keywords", [])), |
| math_object_count=len(s.get("named_phrases", [])), |
| proof_count=0, |
| depends_on=[], |
| referenced_by=[], |
| )) |
| return DocMap( |
| paper_id=doc.paper_id, |
| paper_title=doc.title, |
| sections=sections, |
| top_nodes=[], |
| total_nodes=len(sections), |
| total_edges=0, |
| dominant_type="prose-heavy", |
| ) |
|
|
|
|
| def build_doc_map(graph: Graph, paper_id: str, paper_title: str) -> DocMap: |
| """Build DocMap from a finalized Graph. No LLM calls.""" |
|
|
| |
| sec_depends_on: dict[str, set[str]] = {} |
| sec_referenced_by: dict[str, set[str]] = {} |
| for node in graph.nodes.values(): |
| if node.node_type == "section": |
| sec_depends_on.setdefault(node.section_id, set()) |
| sec_referenced_by.setdefault(node.section_id, set()) |
|
|
| for edge in graph.edges: |
| fn = graph.nodes.get(edge.from_id) |
| tn = graph.nodes.get(edge.to_id) |
| if not fn or not tn: |
| continue |
| fs, ts = fn.section_id, tn.section_id |
| if fs != ts: |
| sec_depends_on.setdefault(fs, set()).add(ts) |
| sec_referenced_by.setdefault(ts, set()).add(fs) |
|
|
| |
| math_counts: dict[str, int] = {} |
| proof_counts: dict[str, int] = {} |
| sec_importance: dict[str, int] = {} |
| for node in graph.nodes.values(): |
| sid = node.section_id |
| sec_importance[sid] = sec_importance.get(sid, 0) + node.in_degree |
| if node.node_type == "proof": |
| proof_counts[sid] = proof_counts.get(sid, 0) + 1 |
| elif node.node_type != "section": |
| math_counts[sid] = math_counts.get(sid, 0) + 1 |
|
|
| |
| sections = [] |
| for sn in graph.section_nodes(): |
| sid = sn.section_id |
| sections.append(SectionEntry( |
| section_id=sid, |
| title=sn.label, |
| importance=sec_importance.get(sid, 0), |
| math_object_count=math_counts.get(sid, 0), |
| proof_count=proof_counts.get(sid, 0), |
| depends_on=sorted(sec_depends_on.get(sid, set())), |
| referenced_by=sorted(sec_referenced_by.get(sid, set())), |
| )) |
|
|
| |
| top_nodes = [ |
| {"node_id": n.node_id, "label": n.label, "node_type": n.node_type, |
| "section_id": n.section_id, "in_degree": n.in_degree, |
| "proof_type": n.proof_type} |
| for n in graph.nodes_by_importance() |
| if n.node_type != "section" and n.in_degree > 0 |
| ][:10] |
|
|
| |
| type_counts: dict[str, int] = {} |
| for node in graph.nodes.values(): |
| if node.node_type not in ("section", "proof"): |
| type_counts[node.node_type] = type_counts.get(node.node_type, 0) + 1 |
| dominant = max(type_counts, key=lambda k: type_counts[k]) if type_counts else "mixed" |
| dominant_type = f"{dominant}-heavy" |
|
|
| return DocMap( |
| paper_id=paper_id, |
| paper_title=paper_title, |
| sections=sections, |
| top_nodes=top_nodes, |
| total_nodes=len(graph.nodes), |
| total_edges=len(graph.edges), |
| dominant_type=dominant_type, |
| ) |
|
|
