LlamaIndex 完整教程

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.vector_stores.qdrant import QdrantVectorStore
import qdrant_client

class EnterpriseKnowledgeBase:
    def __init__(self, collection_name="company_docs"):
        # 全局配置
        Settings.llm = OpenAI(model="gpt-4", temperature=0)
        Settings.embed_model = OpenAIEmbedding()
        Settings.chunk_size = 512
        Settings.chunk_overlap = 50
        
        # 向量数据库
        self.client = qdrant_client.QdrantClient(path="./qdrant_data")
        self.vector_store = QdrantVectorStore(
            client=self.client,
            collection_name=collection_name
        )
        
    def index_documents(self, directory):
        """索引文档目录"""
        # 加载文档
        documents = SimpleDirectoryReader(
            input_dir=directory,
            recursive=True,
            required_exts=[".pdf", ".docx", ".txt", ".md"]
        ).load_data()
        
        print(f"加载 {len(documents)} 个文档")
        
        # 创建索引
        from llama_index.core import StorageContext
        
        storage_context = StorageContext.from_defaults(
            vector_store=self.vector_store
        )
        
        self.index = VectorStoreIndex.from_documents(
            documents,
            storage_context=storage_context,
            show_progress=True
        )
        
        print("索引创建完成")
        return self.index
    
    def create_query_engine(self):
        """创建查询引擎"""
        from llama_index.postprocessor.cohere_rerank import CohereRerank
        
        query_engine = self.index.as_query_engine(
            similarity_top_k=10,
            node_postprocessors=[
                CohereRerank(top_n=3)  # 重排序
            ],
            response_mode="compact",
            streaming=True
        )
        
        return query_engine
    
    def create_chat_engine(self):
        """创建对话引擎"""
        from llama_index.core.memory import ChatMemoryBuffer
        
        chat_engine = self.index.as_chat_engine(
            chat_mode="condense_plus_context",
            memory=ChatMemoryBuffer.from_defaults(token_limit=3000),
            system_prompt="""你是公司的智能助手，负责回答员工的问题。
            
【规则】
1. 必须基于公司文档回答
2. 标注信息来源
3. 信息不足时明确说明
4. 保持专业和友好的语气"""
        )
        
        return chat_engine

# 使用示例
kb = EnterpriseKnowledgeBase()
kb.index_documents("./company_docs")

# 问答模式
query_engine = kb.create_query_engine()
response = query_engine.query("公司的休假政策是什么？")

# 对话模式
chat_engine = kb.create_chat_engine()
response = chat_engine.chat("你好，请问年假有多少天？")