Vector Search RAG

import openai
from sentence_transformers import SentenceTransformer
import numpy as np
from typing import List, Dict

class DocumentEmbedder:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)
        self.openai_client = openai.OpenAI()
    
    def generate_embeddings(self, documents: List[str]) -> np.ndarray:
        """Generate embeddings for a list of documents"""
        return self.model.encode(documents, show_progress_bar=True)
    
    def generate_openai_embeddings(self, documents: List[str]) -> List[List[float]]:
        """Generate embeddings using OpenAI's API"""
        embeddings = []
        for doc in documents:
            response = self.openai_client.embeddings.create(
                model="text-embedding-ada-002",
                input=doc
            )
            embeddings.append(response.data[0].embedding)
        return embeddings
    
    def chunk_documents(self, documents: List[str], chunk_size: int = 512) -> List[str]:
        """Split documents into smaller chunks for better embedding"""
        chunks = []
        for doc in documents:
            words = doc.split()
            for i in range(0, len(words), chunk_size):
                chunk = " ".join(words[i:i + chunk_size])
                chunks.append(chunk)
        return chunks

import chromadb
from chromadb.config import Settings
import numpy as np
from typing import List, Dict, Optional

class VectorDatabase:
    def __init__(self, persist_directory: str = "./chroma_db"):
        self.client = chromadb.PersistentClient(
            path=persist_directory,
            settings=Settings(anonymized_telemetry=False)
        )
        self.collection = self.client.get_or_create_collection(
            name="documents",
            metadata={"hnsw:space": "cosine"}
        )
    
    def add_documents(self, documents: List[str], metadata: List[Dict], ids: List[str]):
        """Add documents to the vector database"""
        # Generate embeddings for documents
        embedder = DocumentEmbedder()
        embeddings = embedder.generate_embeddings(documents)
        
        # Add to collection
        self.collection.add(
            embeddings=embeddings.tolist(),
            documents=documents,
            metadatas=metadata,
            ids=ids
        )
    
    def search(self, query: str, n_results: int = 5) -> Dict:
        """Search for similar documents"""
        embedder = DocumentEmbedder()
        query_embedding = embedder.generate_embeddings([query])
        
        results = self.collection.query(
            query_embeddings=query_embedding.tolist(),
            n_results=n_results,
            include=["documents", "metadatas", "distances"]
        )
        
        return results
    
    def hybrid_search(self, query: str, n_results: int = 5, 
                     semantic_weight: float = 0.7) -> Dict:
        """Combine semantic and keyword search"""
        # Semantic search
        semantic_results = self.search(query, n_results)
        
        # Keyword search (simple implementation)
        keyword_results = self.keyword_search(query, n_results)
        
        # Combine and rerank results
        combined_results = self.combine_results(
            semantic_results, keyword_results, semantic_weight
        )
        
        return combined_results

from typing import List, Dict
import openai
from dataclasses import dataclass

@dataclass
class SearchResult:
    document: str
    metadata: Dict
    similarity_score: float
    source: str

class VectorSearchRAG:
    def __init__(self, vector_db: VectorDatabase, openai_api_key: str):
        self.vector_db = vector_db
        self.openai_client = openai.OpenAI(api_key=openai_api_key)
    
    def retrieve_context(self, query: str, n_results: int = 5) -> List[SearchResult]:
        """Retrieve relevant context for a query"""
        search_results = self.vector_db.search(query, n_results)
        
        results = []
        for i in range(len(search_results['documents'][0])):
            result = SearchResult(
                document=search_results['documents'][0][i],
                metadata=search_results['metadatas'][0][i],
                similarity_score=1 - search_results['distances'][0][i],
                source=search_results['metadatas'][0][i].get('source', 'unknown')
            )
            results.append(result)
        
        return results
    
    def generate_response(self, query: str, context: List[SearchResult]) -> str:
        """Generate response using retrieved context"""
        # Prepare context for the LLM
        context_text = "\n\n".join([
            f"Document {i+1} (Score: {result.similarity_score:.3f}):\n{result.document}"
            for i, result in enumerate(context)
        ])
        
        prompt = f"""Based on the following context, answer the user's question. 
        If the context doesn't contain enough information to answer the question, 
        say so. Always cite the source documents when possible.

        Context:
        {context_text}

        Question: {query}

        Answer:"""
        
        response = self.openai_client.chat.completions.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": "You are a helpful assistant that answers questions based on provided context."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=500,
            temperature=0.3
        )
        
        return response.choices[0].message.content
    
    def query(self, question: str, n_results: int = 5) -> Dict:
        """Complete RAG pipeline: retrieve context and generate response"""
        # Retrieve relevant context
        context = self.retrieve_context(question, n_results)
        
        # Generate response
        answer = self.generate_response(question, context)
        
        return {
            "question": question,
            "answer": answer,
            "context": context,
            "sources": [result.source for result in context]
        }