如何为AI智能体添加RAG以实现上下文理解

class ContextManager: 
    def __init__(self, collection_name: str, persist_dir: str = "context_db", chunk_size: int = 1000, chunk_overlap: int = 200): 
        self.persist_dir = persist_dir 
        self.collection_name = collection_name 
        self.chunk_size = chunk_size 
        self.chunk_overlap = chunk_overlap 
        # Initialize ChromaDB 
        self.client = chromadb.PersistentClient(path=persist_dir) 
        self.collection = self.client.get_or_create_collection(
            name=collection_name, metadata={"hnsw:space": "cosine"}
        )

def index_document(self, 
                   pdf_path: str, 
                   metadata: Optional[Union[Dict[str, Any], DocumentMetadata]] = None) -> bool:
    """
    Index a PDF document into the vector store.
    """
    try:
        # Convert metadata to DocumentMetadata if it's a dict
        if isinstance(metadata, dict):
            metadata = DocumentMetadata.from_dict(metadata)
        elif metadata is None:
            metadata = DocumentMetadata(source=os.path.basename(pdf_path))
        
        # Extract text from PDF
        text = self._extract_text_from_pdf(pdf_path)
        
        # Split text into chunks
        chunks = self._text_splitter.split_text(text)
        
        # Generate unique IDs and prepare metadata
        ids = [self._generate_document_id(chunk, metadata.to_dict()) for chunk in chunks]
        metadatas = [metadata.to_dict() for _ in chunks]
        
        # Add to ChromaDB
        self.collection.add(
            documents=chunks,
            ids=ids,
            metadatas=metadatas
        )
        
        return True
    except Exception as e:
        logger.error(f"Error indexing document {pdf_path}: {str(e)}")
        return False

def query(self, 
         query: str, 
         num_results: int = 3,
         filter_metadata: Optional[Dict[str, Any]] = None) -> str:
    """
    Query the context and return relevant information.
    """
    try:
        self._current_query = query
        
        # Prepare query parameters
        query_params = {
            "query_texts": [query],
            "n_results": num_results
        }
        if filter_metadata:
            query_params["where"] = filter_metadata
        
        # Execute query
        results = self.collection.query(**query_params)
        
        # Format the context with metadata
        context_parts = []
        for i, (doc, metadata) in enumerate(zip(
            results['documents'][0], 
            results['metadatas'][0]
        ), 1):
            source = metadata.get('source', 'Unknown source')
            context_parts.append(
                f"Relevant Context {i} (from {source}):\n{doc}\n"
            )
        
        self._current_context = "\n".join(context_parts)
        return self._current_context
        
    except Exception as e:
        logger.error(f"Error executing query: {str(e)}")
        self._current_context = ""
        return ""

# 1. Initialize context and index the PDF
print("\nStep 1: Initializing context and indexing document...")
context = ContextManager.initialize(
    collection_name="simple_docs",
    persist_dir="context_db"
)
 
pdf_path = "quantum_computing.pdf"
metadata = DocumentMetadata(
    source=os.path.basename(pdf_path),
    doc_type="pdf",
    author="Demo Author",
    created_at=datetime.now(),
    tags="technical,quantum,computing"
)
 
success = context.index_document(pdf_path, metadata)
 
# 2. Create and configure the agent with context
print("\nStep 2: Creating agent with context...")
agent = Agent("rag_agent", context=context)
 
# Set agent persona and instruction
agent.persona = """I am a helpful AI assistant that provides accurate information 
based on the given context. I analyze documents and explain complex topics in a 
clear and understandable way."""
 
agent.instruction = """When explaining concepts from the document:
1. Focus on key principles and fundamentals
2. Use clear and precise language
3. Provide relevant examples where applicable"""
 
# 3. Query and get response using context
print("\nStep 3: Setting context and executing task...")
context_query = "What are the main principles of quantum computing?"
agent.set_context_query(context_query)
 
agent.task = """Based on the provided context:
1. Identify and explain the key principles of quantum computing"""
 
response = agent.execute()