Long Context 처리를 위한 효율적인 방법: Recursive Abstractive Processing for Tree-Organized Retrieval (RAPTOR)

import openai
import textwrap
from typing import List, Dict, Any, Optional
import numpy as np

# OpenAI API 설정
openai.api_key = "your-api-key"

class RAPNode:
    def __init__(self, content: str, level: int, parent=None):
        self.content = content
        self.summary = None
        self.level = level
        self.parent = parent
        self.children = []
        self.embedding = None
    
    def add_child(self, child):
        self.children.append(child)
        child.parent = self
    
    def generate_summary(self):
        if not self.children:  # 리프 노드인 경우
            self.summary = self._summarize_content(self.content)
        else:  # 내부 노드인 경우
            child_summaries = [child.summary for child in self.children]
            combined_summary = "\n".join(child_summaries)
            self.summary = self._summarize_content(combined_summary)
        return self.summary
    
    def _summarize_content(self, text: str) -> str:
        """OpenAI API를 사용하여 텍스트 요약"""
        try:
            response = openai.ChatCompletion.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "다음 텍스트의 핵심 내용을 200단어 이내로 요약해주세요."},
                    {"role": "user", "content": text}
                ],
                max_tokens=300,
                temperature=0.3
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            print(f"요약 생성 중 오류 발생: {e}")
            return "요약 생성 실패"
    
    def generate_embedding(self):
        """텍스트에 대한 임베딩 생성"""
        if self.summary:
            try:
                response = openai.Embedding.create(
                    model="text-embedding-ada-002",
                    input=self.summary
                )
                self.embedding = response['data'][0]['embedding']
            except Exception as e:
                print(f"임베딩 생성 중 오류 발생: {e}")
                self.embedding = None
        return self.embedding

class RAP:
    def __init__(self, max_chunk_size: int = 2000):
        self.max_chunk_size = max_chunk_size
        self.root = None
    
    def process_document(self, document: str):
        """문서를 처리하여 RAP 트리 구조 생성"""
        # 문서를 청크로 분할
        chunks = self._split_into_chunks(document)
        
        # 리프 노드 생성
        leaf_nodes = [RAPNode(chunk, level=0) for chunk in chunks]
        
        # 트리 구조 구축
        self.root = self._build_tree(leaf_nodes)
        
        # 요약 생성 (상향식)
        self._generate_summaries(self.root)
        
        # 임베딩 생성
        self._generate_embeddings(self.root)
        
        return self.root
    
    def _split_into_chunks(self, document: str) -> List[str]:
        """문서를 관리 가능한 크기의 청크로 분할"""
        return textwrap.wrap(document, self.max_chunk_size, break_long_words=False, replace_whitespace=False)
    
    def _build_tree(self, nodes: List[RAPNode], level: int = 0, max_children: int = 4) -> RAPNode:
        """노드 리스트로부터 계층적 트리 구조 구축"""
        if len(nodes) <= max_children:
            # 상위 노드 생성
            parent = RAPNode("", level + 1)
            for node in nodes:
                parent.add_child(node)
            return parent
        
        # 노드를 그룹으로 분할
        groups = [nodes[i:i+max_children] for i in range(0, len(nodes), max_children)]
        
        # 각 그룹에 대한 부모 노드 생성
        parent_nodes = []
        for group in groups:
            parent = RAPNode("", level + 1)
            for node in group:
                parent.add_child(node)
            parent_nodes.append(parent)
        
        # 재귀적으로 트리 구축 계속
        return self._build_tree(parent_nodes, level + 1, max_children)
    
    def _generate_summaries(self, node: RAPNode):
        """트리의 모든 노드에 대한 요약 생성 (상향식)"""
        for child in node.children:
            self._generate_summaries(child)
        
        node.generate_summary()
    
    def _generate_embeddings(self, node: RAPNode):
        """트리의 모든 노드에 대한 임베딩 생성"""
        node.generate_embedding()
        for child in node.children:
            self._generate_embeddings(child)
    
    def query(self, query_text: str, top_k: int = 3) -> List[str]:
        """쿼리와 가장 관련성 높은 청크 검색"""
        # 쿼리 임베딩 생성
        try:
            response = openai.Embedding.create(
                model="text-embedding-ada-002",
                input=query_text
            )
            query_embedding = response['data'][0]['embedding']
        except Exception as e:
            print(f"쿼리 임베딩 생성 중 오류 발생: {e}")
            return []
        
        # 트리를 순회하며 관련 노드 찾기
        relevant_chunks = self._search_tree(self.root, query_embedding, top_k)
        
        return [node.content for node in relevant_chunks]
    
    def _search_tree(self, node: RAPNode, query_embedding: List[float], top_k: int) -> List[RAPNode]:
        """트리를 순회하며 쿼리와 가장 관련성 높은 노드 검색"""
        if not node.children:  # 리프 노드인 경우
            return [node]
        
        # 자식 노드들과 쿼리 간의 유사도 계산
        similarities = []
        for child in node.children:
            if child.embedding:
                similarity = self._calculate_similarity(query_embedding, child.embedding)
                similarities.append((child, similarity))
        
        # 유사도에 따라 정렬
        similarities.sort(key=lambda x: x[1], reverse=True)
        
        # 상위 k개의 관련 자식 노드 선택
        top_children = [item[0] for item in similarities[:top_k]]
        
        # 선택된 자식 노드들에 대해 재귀적으로 검색
        results = []
        for child in top_children:
            results.extend(self._search_tree(child, query_embedding, top_k))
        
        return results[:top_k]
    
    def _calculate_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
        """두 임베딩 간의 코사인 유사도 계산"""
        embedding1 = np.array(embedding1)
        embedding2 = np.array(embedding2)
        
        return np.dot(embedding1, embedding2) / (np.linalg.norm(embedding1) * np.linalg.norm(embedding2))

# 사용 예시
def main():
    # 예시 문서
    document = """
    [여기에 긴 문서 내용이 들어갑니다...]
    """
    
    # RAP 인스턴스 생성 및 문서 처리
    rap = RAP(max_chunk_size=1500)
    rap.process_document(document)
    
    # 쿼리 실행
    query = "인공지능의 윤리적 문제는 무엇인가요?"
    relevant_chunks = rap.query(query, top_k=2)
    
    print("관련 청크:")
    for i, chunk in enumerate(relevant_chunks):
        print(f"청크 {i+1}:\n{chunk}\n")

if __name__ == "__main__":
    main()