高度な 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]:
        """重複する大きなチャンクを作成"""
        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):
        """階層的埋め込みでドキュメントをインデックス登録"""
        # 全文を埋め込み
        doc_embedding = self.model.embed(doc)
        
        # 重複を含む大きなチャンクを作成
        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]:
        """関連する長文形式のコンテンツを検索"""
        query_embedding = self.model.embed(query)
        
        # まずドキュメントレベルでスコアリング
        doc_scores = [
            np.dot(query_embedding, doc['doc_embedding'])
            for doc in self.doc_embeddings
        ]
        
        # 上位のドキュメントを取得
        top_doc_indices = np.argsort(doc_scores)[-top_k:][::-1]
        
        results = []
        for idx in top_doc_indices:
            doc_data = self.doc_embeddings[idx]
            
            # 各ドキュメントについて最良のチャンクを見つける
            chunk_scores = [
                np.dot(query_embedding, emb)
                for emb in doc_data['chunk_embeddings']
            ]
            best_chunk_idx = np.argmax(chunk_scores)
            
            # 最良のチャンク周辺の拡張コンテキストを返す
            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]:
        """関連するチャンク周辺の拡張コンテキストを取得"""
        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):
        """自己反省を伴う回答生成"""
        
        # フェーズ 1: 検索が必要かどうかを決定
        retrieve_decision = self.critic.should_retrieve(query)
        
        if not retrieve_decision:
            # 検索なしで直接生成
            return self.generator.generate(query)
        
        # フェーズ 2: 検索と関連性の評価
        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:  # しきい値
                relevant_docs.append(doc)
        
        if not relevant_docs:
            # 検索なしでの生成にフォールバック
            return self.generator.generate(query)
        
        # フェーズ 3: 生成とサポートの確認
        best_answer = None
        best_score = -1
        
        for _ in range(max_iterations):
            # 候補回答を生成
            answer = self.generator.generate(
                query, context=relevant_docs
            )
            
            # サポートと品質を評価
            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
            
            # 高品質が達成されたら早期停止
            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]):
        """ドキュメントから知識グラフを構築"""
        for doc in documents:
            # エンティティを抽出
            entities = self.entity_extractor.extract(doc)
            
            # エンティティをノードとして追加
            for entity in entities:
                self.graph.add_node(
                    entity['text'],
                    entity_type=entity['type'],
                    context=entity['surrounding_text']
                )
            
            # 関係性を抽出
            relations = self.relation_extractor.extract(
                doc, entities
            )
            
            # 関係性をエッジとして追加
            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):
        """派生関係性とメタデータを追加"""
        # ノードの重要性を計算（PageRank など）
        self.graph.compute_centrality()
        
        # コミュニティ/クラスターを特定
        self.graph.detect_communities()
        
        # タイムスタンプがあれば時間的順序を追加
        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):
        """クエリに関連するサブグラフを検索"""
        
        # クエリ内のシードエンティティを特定
        query_entities = self.entity_extractor.extract(query)
        
        # グラフ内の一致するノードを検索
        seed_nodes = []
        for entity in query_entities:
            matches = self.graph.find_similar_nodes(
                entity['text'],
                similarity_threshold=0.85
            )
            seed_nodes.extend(matches)
        
        # トラバーサルによってサブグラフを拡張
        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
            )
        
        # 関連性でノードをランク付け
        ranked_nodes = self._rank_subgraph_nodes(
            subgraph, query
        )
        
        # コンテキストを抽出してフォーマット
        context = self._format_graph_context(
            ranked_nodes[:max_nodes],
            subgraph
        )
        
        return context
    
    def _expand_from_node(self, node, subgraph, visited,
                         current_hop, max_hops):
        """再帰的にサブグラフを拡張"""
        if current_hop >= max_hops or node in visited:
            return
        
        visited.add(node)
        subgraph.add_node(node)
        
        # 近隣ノードを取得
        neighbors = self.graph.get_neighbors(node)
        
        for neighbor, edge_data in neighbors:
            # サブグラフにエッジを追加
            subgraph.add_edge(node, neighbor, edge_data)
            
            # 再帰的に拡張
            self._expand_from_node(
                neighbor,
                subgraph,
                visited,
                current_hop + 1,
                max_hops
            )
    
    def _format_graph_context(self, nodes, subgraph):
        """サブグラフをテキストコンテキストに変換"""
        context_parts = []
        
        for node in nodes:
            # ノードコンテキストを追加
            context_parts.append(f"Entity: {node.text}")
            context_parts.append(f"Type: {node.entity_type}")
            
            # 関係性情報を追加
            edges = subgraph.get_edges(node)
            for edge in edges:
                context_parts.append(
                    f"- {edge.relation_type} -> {edge.target.text}"
                )
        
        return "\n".join(context_parts)

# 依存関係のインストール
# pip install sentence-transformers faiss-cpu langchain

from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

# 大きなチャンクを実験するための基本的なセットアップ
model = SentenceTransformer('all-MiniLM-L6-v2')

documents = [
    # ここに長文形式のドキュメントを配置
]

# 埋め込みを作成
embeddings = model.encode(documents)

# FAISS インデックスを構築
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings.astype('float32'))

# クエリ
query = "ここに質問を入力"
query_embedding = model.encode([query])
distances, indices = index.search(
    query_embedding.astype('float32'), k=3
)