RAG技术：让大模型更智能、更准确

import chromadb
from openai import OpenAI
import tiktoken

class RAGSystem:
    def __init__(self, api_key):
        self.client = OpenAI(api_key=api_key)
        self.chroma = chromadb.Client()
        self.collection = self.chroma.create_collection("knowledge_base")
        self.tokenizer = tiktoken.encoding_for_model("gpt-4")
        
    def add_documents(self, documents):
        """添加文档到知识库"""
        for doc in documents:
            chunks = self.split_document(doc)
            embeddings = self.create_embeddings(chunks)
            
            self.collection.add(
                embeddings=embeddings,
                documents=chunks,
                ids=[f"doc_{i}" for i in range(len(chunks))],
                metadatas=[{"source": doc.source} for _ in chunks]
            )
    
    def split_document(self, doc, chunk_size=500, overlap=50):
        """智能文档分割"""
        tokens = self.tokenizer.encode(doc.text)
        chunks = []
        
        for i in range(0, len(tokens), chunk_size - overlap):
            chunk_tokens = tokens[i:i + chunk_size]
            chunk_text = self.tokenizer.decode(chunk_tokens)
            chunks.append(chunk_text)
            
        return chunks
    
    def create_embeddings(self, texts):
        """创建文本向量"""
        response = self.client.embeddings.create(
            model="text-embedding-ada-002",
            input=texts
        )
        return [data.embedding for data in response.data]
    
    def retrieve(self, query, top_k=5):
        """检索相关文档"""
        query_embedding = self.create_embeddings([query])[0]
        
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k
        )
        
        return results['documents'][0]
    
    def generate_answer(self, query):
        """生成RAG答案"""
        # 检索相关上下文
        contexts = self.retrieve(query)
        
        # 构建prompt
        context_str = "\n\n".join(contexts)
        prompt = f"""基于以下参考信息回答用户问题。

参考信息：
{context_str}

用户问题：{query}

请提供准确、详细的回答。如果参考信息不足，请明确说明。
"""
        
        # 调用LLM生成答案
        response = self.client.chat.completions.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": "你是一个基于知识库的问答助手"},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,
            max_tokens=1000
        )
        
        return {
            "answer": response.choices[0].message.content,
            "contexts": contexts,
            "usage": response.usage
        }

# 使用示例
rag = RAGSystem(api_key="your-api-key")

# 添加文档
rag.add_documents([
    {"text": "公司成立于2020年...", "source": "company_intro.pdf"},
    {"text": "产品特性包括...", "source": "product_manual.pdf"}
])

# 问答
result = rag.generate_answer("公司什么时候成立的？")
print(f"答案：{result['answer']}")
print(f"参考来源：{result['contexts']}")