agentic_rag/services/knowledge/graph/index.py

"""Knowledge Graph Index — entity extraction, relationship mapping, and graph storage.

Inspired by RAG-Anything / LightRAG:
- Multi-modal entities from text + images + tables + equations
- Cross-modal relationship inference
- Hierarchical structure preservation ("belongs_to" chains)
- Weighted relationship scoring
"""

import asyncio
import hashlib
import uuid
from typing import Any, Callable, Optional

from agentic_rag.services.knowledge.content_list import ContentItem, ContentList, ContentType


class Entity:
    """A node in the knowledge graph."""

    def __init__(self, entity_id: str, name: str, entity_type: str,
                 content_item: Optional[ContentItem] = None, properties: dict = None):
        self.id = entity_id or uuid.uuid4().hex
        self.name = name
        self.type = entity_type  # "text_chunk", "image", "table", "equation", "concept"
        self.content_item = content_item
        self.properties = properties or {}
        self.embedding: Optional[list[float]] = None

    def __repr__(self):
        return f"Entity({self.name[:30]}, type={self.type})"


class Relation:
    """An edge in the knowledge graph."""

    def __init__(self, source_id: str, target_id: str, relation_type: str,
                 weight: float = 1.0, metadata: dict = None):
        self.source_id = source_id
        self.target_id = target_id
        self.type = relation_type  # "belongs_to", "references", "nearby", "describes"
        self.weight = weight
        self.metadata = metadata or {}

    def __repr__(self):
        return f"Relation({self.type}: {self.source_id[:8]} -> {self.target_id[:8]})"


class KnowledgeGraph:
    """In-memory knowledge graph with entities and relations.

    For production, this would be backed by a graph database (Neo4j, ArangoDB)
    or LightRAG's graph storage.
    """

    def __init__(self):
        self._entities: dict[str, Entity] = {}
        self._relations: list[Relation] = []
        # Indexes
        self._by_type: dict[str, list[str]] = {}   # type -> [entity_ids]
        self._by_page: dict[int, list[str]] = {}   # page_idx -> [entity_ids]
        self._adjacency: dict[str, list[Relation]] = {}  # entity_id -> [relations]

    # ── Entity CRUD ──────────────────────────────

    def add_entity(self, entity: Entity) -> str:
        self._entities[entity.id] = entity
        self._by_type.setdefault(entity.type, []).append(entity.id)
        if entity.content_item:
            page = entity.content_item.page_idx
            self._by_page.setdefault(page, []).append(entity.id)
        return entity.id

    def get_entity(self, entity_id: str) -> Optional[Entity]:
        return self._entities.get(entity_id)

    def get_entities_by_type(self, entity_type: str) -> list[Entity]:
        ids = self._by_type.get(entity_type, [])
        return [self._entities[eid] for eid in ids if eid in self._entities]

    # ── Relation CRUD ────────────────────────────

    def add_relation(self, source_id: str, target_id: str,
                     relation_type: str = "references", weight: float = 1.0,
                     metadata: dict = None) -> Relation | None:
        # Deduplicate: skip if an identical (source, target, type) edge exists
        for existing in self._adjacency.get(source_id, []):
            if (existing.target_id == target_id
                    and existing.source_id == source_id
                    and existing.type == relation_type):
                return None  # already exists, skip
        rel = Relation(source_id, target_id, relation_type, weight, metadata)
        self._relations.append(rel)
        self._adjacency.setdefault(source_id, []).append(rel)
        # Inverse edge for bidirectional traversal
        inv_rel = Relation(
            target_id, source_id, f"inverse_{relation_type}", weight, metadata)
        self._adjacency.setdefault(target_id, []).append(inv_rel)
        return rel

    def get_neighbors(self, entity_id: str, relation_type: str = "") -> list[Entity]:
        """Get neighboring entities, optionally filtered by relation type."""
        rels = self._adjacency.get(entity_id, [])
        if relation_type:
            rels = [r for r in rels if r.type == relation_type]
        neighbors = []
        for rel in rels:
            neighbor_id = rel.target_id if rel.source_id == entity_id else rel.source_id
            if neighbor_id in self._entities:
                neighbors.append(self._entities[neighbor_id])
        return neighbors

    def traverse(self, entity_id: str, max_depth: int = 2,
                 relation_types: list[str] = None) -> list[Entity]:
        """BFS traversal from an entity."""
        visited = set()
        queue = [(entity_id, 0)]
        result = []
        while queue:
            eid, depth = queue.pop(0)
            if eid in visited or depth > max_depth:
                continue
            visited.add(eid)
            if eid in self._entities:
                result.append(self._entities[eid])
            for rel in self._adjacency.get(eid, []):
                neighbor = rel.target_id if rel.source_id == eid else rel.source_id
                if (relation_types is None or rel.type in relation_types):
                    queue.append((neighbor, depth + 1))
        return result

    # ── Properties ───────────────────────────────

    @property
    def entity_count(self) -> int:
        return len(self._entities)

    @property
    def relation_count(self) -> int:
        return len(self._relations)

    def stats(self) -> dict:
        return {
            "entities": self.entity_count,
            "relations": self.relation_count,
            "by_type": {t: len(ids) for t, ids in self._by_type.items()},
            "by_page": {p: len(ids) for p, ids in sorted(self._by_page.items())},
        }

    # ── Serialization ────────────────────────────

    def to_dict(self) -> dict:
        """Serialize graph for storage."""
        return {
            "entities": {
                eid: {
                    "id": e.id, "name": e.name, "type": e.type,
                    "properties": e.properties,
                    "content_type": e.content_item.type.value if e.content_item else None,
                    "content_text": e.content_item.to_searchable_text()[:500] if e.content_item else "",
                    # Persist enough ContentItem fields to reconstruct graph edges
                    "_item_text": e.content_item.text if e.content_item else "",
                    "_item_page_idx": e.content_item.page_idx if e.content_item else 0,
                }
                for eid, e in self._entities.items()
            },
            "relations": [
                {"source": r.source_id, "target": r.target_id,
                 "type": r.type, "weight": r.weight}
                for r in self._relations
            ],
        }

    def save_json(self, path: str) -> None:
        """Persist graph to a JSON file."""
        import json
        data = self.to_dict()
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

    @classmethod
    def load_json(cls, path: str) -> "KnowledgeGraph":
        """Restore graph from a JSON file. Returns empty graph if file missing."""
        import json
        from pathlib import Path
        graph = cls()
        if not Path(path).exists():
            return graph
        try:
            data = json.loads(Path(path).read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            return graph

        # Rebuild entities as lightweight stubs (no ContentItem — only used
        # for graph-traversal retrieval, not for full-text reconstruction).
        from agentic_rag.services.knowledge.content_list import ContentItem, ContentType
        for eid, ent in data.get("entities", {}).items():
            ctype_str = ent.get("content_type", "text")
            try:
                ctype = ContentType(ctype_str)
            except ValueError:
                ctype = ContentType.TEXT
            item = ContentItem(
                type=ctype,
                text=ent.get("_item_text", "")[:500],
                page_idx=ent.get("_item_page_idx", 0),
            )
            entity = Entity(
                entity_id=ent["id"],
                name=ent["name"],
                entity_type=ent["type"],
                content_item=item,
                properties=ent.get("properties", {}),
            )
            graph._entities[entity.id] = entity
            graph._by_type.setdefault(entity.type, []).append(entity.id)
            if entity.content_item:
                page = entity.content_item.page_idx
                graph._by_page.setdefault(page, []).append(entity.id)

        for rel in data.get("relations", []):
            graph.add_relation(
                rel["source"], rel["target"], rel["type"],
                weight=rel.get("weight", 1.0),
            )

        return graph