高度なRAG：LongRAG、Self-RAG、GraphRAGを解説

from typing import List, Dict
import numpy as np

class LongRAGRetriever:
    def __init__(self, model, chunk_size=8000, overlap=1600):
        self.model = model
        self.chunk_size = chunk_size
        self.overlap = overlap
        self.doc_embeddings = []
        self.documents = []
    
    def create_long_chunks(self, text: str) -> List[str]:
        """Create overlapping large chunks"""
        chunks = []
        start = 0
        while start < len(text):
            end = start + self.chunk_size
            chunk = text[start:end]
            chunks.append(chunk)
            start += (self.chunk_size - self.overlap)
        return chunks
    
    def index_document(self, doc: str, metadata: Dict):
        """Index document with hierarchical embedding"""
        # Embed entire document
        doc_embedding = self.model.embed(doc)
        
        # Create large chunks with overlap
        chunks = self.create_long_chunks(doc)
        chunk_embeddings = [self.model.embed(c) for c in chunks]
        
        self.doc_embeddings.append({
            'doc_id': len(self.documents),
            'doc_embedding': doc_embedding,
            'chunk_embeddings': chunk_embeddings,
            'chunks': chunks,
            'full_text': doc,
            'metadata': metadata
        })
        self.documents.append(doc)
    
    def retrieve(self, query: str, top_k: int = 3) -> List[Dict]:
        """Retrieve relevant long-form content"""
        query_embedding = self.model.embed(query)
        
        # Score at document level first
        doc_scores = [
            np.dot(query_embedding, doc['doc_embedding'])
            for doc in self.doc_embeddings
        ]
        
        # Get top documents
        top_doc_indices = np.argsort(doc_scores)[-top_k:][::-1]
        
        results = []
        for idx in top_doc_indices:
            doc_data = self.doc_embeddings[idx]
            
            # For each document, find best chunks
            chunk_scores = [
                np.dot(query_embedding, emb)
                for emb in doc_data['chunk_embeddings']
            ]
            best_chunk_idx = np.argmax(chunk_scores)
            
            # Return extended context around best chunk
            context_chunks = self._get_extended_context(
                doc_data['chunks'], 
                best_chunk_idx
            )
            
            results.append({
                'text': ''.join(context_chunks),
                'score': doc_scores[idx],
                'metadata': doc_data['metadata']
            })
        
        return results
    
    def _get_extended_context(self, chunks: List[str], 
                             center_idx: int) -> List[str]:
        """Get extended context around relevant chunk"""
        start = max(0, center_idx - 1)
        end = min(len(chunks), center_idx + 2)
        return chunks[start:end]

class SelfRAGSystem:
    def __init__(self, retriever, generator, critic):
        self.retriever = retriever
        self.generator = generator
        self.critic = critic
    
    def generate_with_reflection(self, query: str, 
                                 max_iterations: int = 3):
        """Generate answer with self-reflection"""
        
        # Phase 1: Decide if retrieval needed
        retrieve_decision = self.critic.should_retrieve(query)
        
        if not retrieve_decision:
            # Direct generation without retrieval
            return self.generator.generate(query)
        
        # Phase 2: Retrieve and evaluate relevance
        retrieved_docs = self.retriever.retrieve(query)
        relevant_docs = []
        
        for doc in retrieved_docs:
            relevance_score = self.critic.assess_relevance(
                query, doc
            )
            if relevance_score > 0.7:  # Threshold
                relevant_docs.append(doc)
        
        if not relevant_docs:
            # Fallback to generation without retrieval
            return self.generator.generate(query)
        
        # Phase 3: Generate and verify support
        best_answer = None
        best_score = -1
        
        for _ in range(max_iterations):
            # Generate candidate answer
            answer = self.generator.generate(
                query, context=relevant_docs
            )
            
            # Evaluate support and quality
            support_score = self.critic.check_support(
                answer, relevant_docs
            )
            quality_score = self.critic.assess_quality(answer)
            
            total_score = 0.6 * support_score + 0.4 * quality_score
            
            if total_score > best_score:
                best_score = total_score
                best_answer = answer
            
            # Early stopping if high quality achieved
            if total_score > 0.9:
                break
        
        return {
            'answer': best_answer,
            'confidence': best_score,
            'sources': relevant_docs,
            'reflections': {
                'retrieved': retrieve_decision,
                'relevance': len(relevant_docs),
                'support': support_score
            }
        }

class GraphRAGBuilder:
    def __init__(self, entity_extractor, relation_extractor):
        self.entity_extractor = entity_extractor
        self.relation_extractor = relation_extractor
        self.graph = NetworkGraph()
    
    def build_graph(self, documents: List[str]):
        """Build knowledge graph from documents"""
        for doc in documents:
            # Extract entities
            entities = self.entity_extractor.extract(doc)
            
            # Add entities as nodes
            for entity in entities:
                self.graph.add_node(
                    entity['text'],
                    entity_type=entity['type'],
                    context=entity['surrounding_text']
                )
            
            # Extract relationships
            relations = self.relation_extractor.extract(
                doc, entities
            )
            
            # Add relationships as edges
            for rel in relations:
                self.graph.add_edge(
                    rel['source'],
                    rel['target'],
                    relation_type=rel['type'],
                    confidence=rel['score'],
                    evidence=rel['text_span']
                )
    
    def enrich_graph(self):
        """Add derived relationships and metadata"""
        # Compute node importance (PageRank, etc.)
        self.graph.compute_centrality()
        
        # Identify communities/clusters
        self.graph.detect_communities()
        
        # Add temporal ordering if timestamps available
        self.graph.add_temporal_edges()

class GraphRAGRetriever:
    def __init__(self, graph, embedder):
        self.graph = graph
        self.embedder = embedder
    
    def retrieve_subgraph(self, query: str, 
                         max_hops: int = 2,
                         max_nodes: int = 50):
        """Retrieve relevant subgraph for query"""
        
        # Identify seed entities in query
        query_entities = self.entity_extractor.extract(query)
        
        # Find matching nodes in graph
        seed_nodes = []
        for entity in query_entities:
            matches = self.graph.find_similar_nodes(
                entity['text'],
                similarity_threshold=0.85
            )
            seed_nodes.extend(matches)
        
        # Expand subgraph via traversal
        subgraph = self.graph.create_subgraph()
        visited = set()
        
        for seed in seed_nodes:
            self._expand_from_node(
                seed, 
                subgraph, 
                visited,
                current_hop=0,
                max_hops=max_hops
            )
        
        # Rank nodes by relevance
        ranked_nodes = self._rank_subgraph_nodes(
            subgraph, query
        )
        
        # Extract and format context
        context = self._format_graph_context(
            ranked_nodes[:max_nodes],
            subgraph
        )
        
        return context
    
    def _expand_from_node(self, node, subgraph, visited,
                         current_hop, max_hops):
        """Recursively expand subgraph"""
        if current_hop >= max_hops or node in visited:
            return
        
        visited.add(node)
        subgraph.add_node(node)
        
        # Get neighbors
        neighbors = self.graph.get_neighbors(node)
        
        for neighbor, edge_data in neighbors:
            # Add edge to subgraph
            subgraph.add_edge(node, neighbor, edge_data)
            
            # Recursively expand
            self._expand_from_node(
                neighbor,
                subgraph,
                visited,
                current_hop + 1,
                max_hops
            )
    
    def _format_graph_context(self, nodes, subgraph):
        """Convert subgraph to textual context"""
        context_parts = []
        
        for node in nodes:
            # Add node context
            context_parts.append(f"Entity: {node.text}")
            context_parts.append(f"Type: {node.entity_type}")
            
            # Add relationship information
            edges = subgraph.get_edges(node)
            for edge in edges:
                context_parts.append(
                    f"- {edge.relation_type} -> {edge.target.text}"
                )
        
        return "\n".join(context_parts)

# Install dependencies
# pip install sentence-transformers faiss-cpu langchain

from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

# Basic setup for experimenting with long chunks
model = SentenceTransformer('all-MiniLM-L6-v2')

documents = [
    # Your long-form documents here
]

# Create embeddings
embeddings = model.encode(documents)

# Build FAISS index
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings.astype('float32'))

# Query
query = "Your question here"
query_embedding = model.encode([query])
distances, indices = index.search(
    query_embedding.astype('float32'), k=3
)