4장: LlamaIndex -- 데이터 프레임워크와 워크플로우

from llama_index.core import SimpleDirectoryReader
from llama_index.readers.web import SimpleWebPageReader
from llama_index.readers.database import DatabaseReader
 
# 로컬 파일 (PDF, DOCX, TXT 등 자동 감지)
documents = SimpleDirectoryReader(
    input_dir="./data",
    recursive=True,
    required_exts=[".pdf", ".md", ".txt"],
).load_data()
 
# 웹 페이지
web_docs = SimpleWebPageReader(html_to_text=True).load_data(
    urls=["https://docs.example.com/guide"]
)
 
# 데이터베이스
db_docs = DatabaseReader(
    sql_database=sql_db,
).load_data(
    query="SELECT title, content FROM articles WHERE published = true"
)

from llama_index.core.node_parser import (
    SentenceSplitter,
    SemanticSplitterNodeParser,
    HierarchicalNodeParser,
)
from llama_index.core.ingestion import IngestionPipeline
 
# 문장 기반 분할 (가장 기본적)
sentence_splitter = SentenceSplitter(
    chunk_size=1024,
    chunk_overlap=200,
)
 
# 의미 기반 분할 (임베딩 유사도로 경계 결정)
semantic_splitter = SemanticSplitterNodeParser(
    buffer_size=1,
    breakpoint_percentile_threshold=95,
    embed_model=embed_model,
)
 
# 계층적 분할 (부모-자식 관계 유지)
hierarchical_splitter = HierarchicalNodeParser.from_defaults(
    chunk_sizes=[2048, 512, 128],
)
 
# 인제스천 파이프라인으로 조합
pipeline = IngestionPipeline(
    transformations=[
        sentence_splitter,
        embed_model,
    ]
)
nodes = pipeline.run(documents=documents)

from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
 
# Chroma 벡터 스토어 연결
chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_or_create_collection("docs")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
 
storage_context = StorageContext.from_defaults(
    vector_store=vector_store
)
 
# 인덱스 생성
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    show_progress=True,
)

from llama_index.core import TreeIndex
 
tree_index = TreeIndex.from_documents(
    documents,
    num_children=10,  # 각 노드의 자식 수
    build_tree_from_text=True,
)

from llama_index.core import KeywordTableIndex
 
keyword_index = KeywordTableIndex.from_documents(
    documents,
    max_keywords_per_node=10,
)

from llama_index.core import KnowledgeGraphIndex
 
kg_index = KnowledgeGraphIndex.from_documents(
    documents,
    max_triplets_per_chunk=10,
    include_embeddings=True,
)

# 기본 쿼리 엔진
query_engine = index.as_query_engine(
    similarity_top_k=5,
    response_mode="compact",
)
 
response = query_engine.query("LlamaIndex의 핵심 개념은 무엇인가요?")
print(response.response)
print(f"소스 노드: {len(response.source_nodes)}")

from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector
from llama_index.core.tools import QueryEngineTool
 
# 각 인덱스를 도구로 래핑
vector_tool = QueryEngineTool.from_defaults(
    query_engine=vector_index.as_query_engine(),
    description="특정 기술 세부사항을 찾을 때 사용",
)
summary_tool = QueryEngineTool.from_defaults(
    query_engine=tree_index.as_query_engine(),
    description="전체적인 내용 요약이 필요할 때 사용",
)
 
# 라우터로 조합
router_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=[vector_tool, summary_tool],
)
 
response = router_engine.query("이 프로젝트의 전반적인 아키텍처를 설명해줘")

from llama_index.core.workflow import (
    Workflow,
    Event,
    StartEvent,
    StopEvent,
    step,
)
 
# 커스텀 이벤트 정의
class RetrievalEvent(Event):
    nodes: list
    query: str
 
class SynthesisEvent(Event):
    response: str
 
# 워크플로우 정의
class RAGWorkflow(Workflow):
    @step
    async def retrieve(self, ev: StartEvent) -> RetrievalEvent:
        """문서 검색 단계"""
        query = ev.query
        nodes = await self.index.aretrieve(query)
        return RetrievalEvent(nodes=nodes, query=query)
 
    @step
    async def synthesize(self, ev: RetrievalEvent) -> StopEvent:
        """응답 합성 단계"""
        context = "\n".join(n.text for n in ev.nodes)
        response = await self.llm.acomplete(
            f"컨텍스트: {context}\n\n질문: {ev.query}"
        )
        return StopEvent(result=str(response))
 
# 실행
workflow = RAGWorkflow()
result = await workflow.run(query="LlamaIndex Workflows란?")

class ParallelRAGWorkflow(Workflow):
    @step
    async def route(self, ev: StartEvent) -> RetrievalEvent:
        """질문을 여러 검색 경로로 분기"""
        # 동시에 두 이벤트 발행
        self.send_event(
            RetrievalEvent(source="vector", query=ev.query)
        )
        self.send_event(
            RetrievalEvent(source="keyword", query=ev.query)
        )
        return None  # 이 스텝은 StopEvent를 반환하지 않음
 
    @step
    async def retrieve(self, ev: RetrievalEvent) -> ResultEvent:
        """각 소스에서 검색"""
        if ev.source == "vector":
            nodes = await self.vector_index.aretrieve(ev.query)
        else:
            nodes = await self.keyword_index.aretrieve(ev.query)
        return ResultEvent(nodes=nodes, source=ev.source)
 
    @step(num_workers=2)  # 2개의 결과를 모두 수집
    async def combine(self, ev: ResultEvent) -> StopEvent:
        """결과 합성"""
        results = self.collect_events(ev, [ResultEvent, ResultEvent])
        if results is None:
            return None  # 아직 모든 결과가 도착하지 않음
 
        all_nodes = []
        for r in results:
            all_nodes.extend(r.nodes)
 
        response = await self.synthesize(all_nodes)
        return StopEvent(result=response)

from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.core.tools import FunctionTool
 
# 도구 정의
def search_web(query: str) -> str:
    """웹 검색을 수행합니다."""
    return f"검색 결과: {query}에 대한 정보..."
 
def calculate(expression: str) -> str:
    """수학 계산을 수행합니다."""
    return str(eval(expression))
 
# 에이전트 워크플로우 생성
agent = AgentWorkflow.from_tools_or_functions(
    tools_or_functions=[
        FunctionTool.from_defaults(fn=search_web),
        FunctionTool.from_defaults(fn=calculate),
    ],
    llm=llm,
    system_prompt="당신은 유능한 연구 어시스턴트입니다.",
)
 
# 실행
response = await agent.run("서울의 인구와 면적을 검색하고, 인구밀도를 계산해줘")