고급 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"엔티티: {node.text}")
            context_parts.append(f"유형: {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
)