lambdagent/agentexample/qaagent67/scripts/graph_engine.py

#!/usr/bin/env python3
"""
Entity Relation Graph Engine
Loads relations from wiki compilation, provides N-hop traversal
for query expansion and relation-based retrieval.

Lambda: Tool("graph_walk", λ(entities, n_hops). BFS → subgraph)
"""
import json
import os
import re
from collections import defaultdict, deque
from pathlib import Path


class EntityGraph:
    """
    Adjacency-list graph built from wiki relations.json.

    Each relation triple: {subject, relation, object, description, source}
    Example: 海事局 --隶属--> 交通运输部
    """

    def __init__(self):
        # entity → [(relation_type, target, source_doc, description)]
        self.adjacency = defaultdict(list)
        # entity → [(relation_type, source_entity, source_doc)]
        self.reverse = defaultdict(list)
        # entity metadata
        self.entity_meta = {}
        # all unique relation types
        self.relation_types = set()
        # stats
        self.num_entities = 0
        self.num_relations = 0

    def load(self, relations_path=None, progress_path=None, wiki_dir=None):
        """Load graph from relations.json or compile_progress.json"""
        loaded = False

        # Try relations.json first
        if relations_path and Path(relations_path).exists():
            with open(relations_path, 'r', encoding='utf-8') as f:
                relations = json.load(f)
            self._ingest_relations(relations)
            loaded = True

        # Try .compile_progress.json
        if progress_path and Path(progress_path).exists():
            with open(progress_path, 'r', encoding='utf-8') as f:
                progress = json.load(f)
            relations = progress.get('relations', [])
            if relations:
                self._ingest_relations(relations)
                loaded = True
            # Also load entity metadata
            for ename, einfo in progress.get('entities', {}).items():
                if ename not in self.entity_meta:
                    self.entity_meta[ename] = {
                        'type': einfo.get('type', 'unknown'),
                        'refs_count': len(einfo.get('refs', [])),
                    }

        # Supplement from wiki entity pages
        if wiki_dir:
            entities_dir = Path(wiki_dir) / 'entities'
            if entities_dir.exists():
                for f in entities_dir.glob('*.md'):
                    ename = f.stem
                    if ename not in self.entity_meta:
                        self.entity_meta[ename] = {'type': 'unknown', 'refs_count': 0}
                    self.num_entities = max(self.num_entities, len(self.entity_meta))

        if loaded:
            self.num_entities = len(set(
                list(self.adjacency.keys()) + list(self.reverse.keys()) + list(self.entity_meta.keys())
            ))

        return loaded

    def _ingest_relations(self, relations):
        """Add relation triples to graph"""
        for r in relations:
            subj = r.get('subject', '').strip()
            rel = r.get('relation', '').strip()
            obj = r.get('object', '').strip()
            desc = r.get('description', '')
            src = r.get('source', '')

            if not subj or not rel or not obj:
                continue

            self.adjacency[subj].append((rel, obj, src, desc))
            self.reverse[obj].append((rel, subj, src))
            self.relation_types.add(rel)
            self.num_relations += 1

    def find_entities(self, text, max_results=10):
        """Fuzzy match text against known entity names"""
        text_lower = text.lower()
        all_entities = set(self.adjacency.keys()) | set(self.reverse.keys()) | set(self.entity_meta.keys())

        matches = []
        for ename in all_entities:
            ename_lower = ename.lower()
            # Exact substring match
            if ename_lower in text_lower or text_lower in ename_lower:
                # Score: longer match = better
                score = len(ename) / max(len(text), 1)
                matches.append((ename, score))

        # Sort by score descending, deduplicate
        matches.sort(key=lambda x: -x[1])
        seen = set()
        result = []
        for name, score in matches:
            if name not in seen:
                seen.add(name)
                result.append(name)
            if len(result) >= max_results:
                break
        return result

    def walk(self, seed_entities, max_hops=2, max_nodes=20):
        """
        BFS N-hop traversal from seed entities.
        Returns subgraph with entities, relations, and paths.
        """
        visited = set()
        queue = deque()
        hop_counts = {}
        paths = {}
        found_relations = []

        # Initialize with seeds
        for e in seed_entities:
            if e in self.adjacency or e in self.reverse or e in self.entity_meta:
                queue.append((e, 0, []))
                visited.add(e)
                hop_counts[e] = 0
                paths[e] = []

        while queue and len(visited) < max_nodes:
            entity, hop, path = queue.popleft()

            if hop >= max_hops:
                continue

            # Forward edges
            for rel, target, src, desc in self.adjacency.get(entity, []):
                found_relations.append({
                    'subject': entity,
                    'relation': rel,
                    'object': target,
                    'source': src,
                })
                if target not in visited and len(visited) < max_nodes:
                    visited.add(target)
                    hop_counts[target] = hop + 1
                    new_path = path + [f"{entity} --{rel}--> {target}"]
                    paths[target] = new_path
                    queue.append((target, hop + 1, new_path))

            # Reverse edges
            for rel, source_entity, src in self.reverse.get(entity, []):
                found_relations.append({
                    'subject': source_entity,
                    'relation': rel,
                    'object': entity,
                    'source': src,
                })
                if source_entity not in visited and len(visited) < max_nodes:
                    visited.add(source_entity)
                    hop_counts[source_entity] = hop + 1
                    new_path = path + [f"{source_entity} --{rel}--> {entity}"]
                    paths[source_entity] = new_path
                    queue.append((source_entity, hop + 1, new_path))

        # Deduplicate relations
        seen_rels = set()
        unique_rels = []
        for r in found_relations:
            key = (r['subject'], r['relation'], r['object'])
            if key not in seen_rels:
                seen_rels.add(key)
                unique_rels.append(r)

        return {
            'entities': list(visited),
            'relations': unique_rels,
            'paths': paths,
            'hop_counts': hop_counts,
        }

    def expand_query(self, query, seed_entities=None, max_hops=2):
        """
        Expand a query using graph-discovered related entities.
        Returns expanded query string + relation context.
        """
        if seed_entities is None:
            seed_entities = self.find_entities(query, max_results=5)

        if not seed_entities:
            return {'expanded_query': query, 'entities': [], 'relations': [], 'relation_context': ''}

        walked = self.walk(seed_entities, max_hops=max_hops)

        # Build relation context for LLM prompt
        relation_lines = []
        for r in walked['relations'][:15]:
            relation_lines.append(f"- {r['subject']} --{r['relation']}--> {r['object']}")
        relation_context = '\n'.join(relation_lines)

        # Expand query with discovered entity names
        new_terms = [e for e in walked['entities'] if e not in seed_entities][:10]
        expanded_query = query
        if new_terms:
            expanded_query += '\n\n[Related: ' + ', '.join(new_terms) + ']'

        return {
            'expanded_query': expanded_query,
            'entities': walked['entities'],
            'relations': walked['relations'],
            'relation_context': relation_context,
            'seed_entities': seed_entities,
        }

    def shortest_path(self, entity_a, entity_b, max_depth=5):
        """Find shortest relation path between two entities"""
        if entity_a == entity_b:
            return [entity_a]

        visited = {entity_a}
        queue = deque([(entity_a, [entity_a])])

        while queue:
            current, path = queue.popleft()
            if len(path) > max_depth:
                break

            for rel, target, _, _ in self.adjacency.get(current, []):
                if target == entity_b:
                    return path + [f"--{rel}-->", target]
                if target not in visited:
                    visited.add(target)
                    queue.append((target, path + [f"--{rel}-->", target]))

            for rel, source, _ in self.reverse.get(current, []):
                if source == entity_b:
                    return path + [f"<--{rel}--", source]
                if source not in visited:
                    visited.add(source)
                    queue.append((source, path + [f"<--{rel}--", source]))

        return None  # No path found

    def stats(self):
        return {
            'entities': self.num_entities,
            'relations': self.num_relations,
            'relation_types': len(self.relation_types),
            'top_relation_types': sorted(
                [(rt, sum(1 for edges in self.adjacency.values() for e in edges if e[0] == rt))
                 for rt in self.relation_types],
                key=lambda x: -x[1]
            )[:10],
        }


# ── Singleton loader ──
_graph_cache = None

def get_graph(wiki_dir=None, config_path=None):
    """Load or return cached EntityGraph"""
    global _graph_cache
    if _graph_cache is not None:
        return _graph_cache

    graph = EntityGraph()

    # Resolve paths from config
    if config_path:
        import yaml
        with open(config_path) as f:
            cfg = yaml.safe_load(f)
        wiki_cfg = cfg.get('wiki', cfg.get('retrieval', {}).get('skills', {}).get('wiki', {}))
        wiki_dir = wiki_cfg.get('dir', wiki_dir)

    if wiki_dir is None:
        wiki_dir = os.environ.get('WIKI_DIR', '/app/agentexample/qaagent67wiki/wiki')

    wiki_path = Path(wiki_dir)
    relations_file = wiki_path / 'relations.json'
    progress_file = wiki_path / '.compile_progress.json'

    graph.load(
        relations_path=str(relations_file) if relations_file.exists() else None,
        progress_path=str(progress_file) if progress_file.exists() else None,
        wiki_dir=str(wiki_path),
    )

    _graph_cache = graph
    return graph


if __name__ == '__main__':
    import sys
    wiki_dir = sys.argv[1] if len(sys.argv) > 1 else '/app/agentexample/qaagent67wiki/wiki'

    graph = get_graph(wiki_dir=wiki_dir)
    stats = graph.stats()
    print(f"Graph loaded: {stats['entities']} entities, {stats['relations']} relations")
    print(f"Relation types: {stats['relation_types']}")
    for rt, count in stats['top_relation_types'][:5]:
        print(f"  {rt}: {count}")

    # Test query expansion
    query = sys.argv[2] if len(sys.argv) > 2 else "海上交通安全法与水污染防治法"
    print(f"\nExpanding: {query}")
    result = graph.expand_query(query)
    print(f"Seed entities: {result['seed_entities']}")
    print(f"All entities ({len(result['entities'])}): {result['entities'][:10]}")
    print(f"Relations ({len(result['relations'])}):")
    for r in result['relations'][:5]:
        print(f"  {r['subject']} --{r['relation']}--> {r['object']}")