Retrieval-Augmented Generation (RAG)

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
from baml_client import b

class VectorStore:
    def __init__(self, vectorizer, tfidf_matrix, documents):
        self.vectorizer = vectorizer
        self.tfidf_matrix = tfidf_matrix
        self.documents = documents

    @classmethod
    def from_documents(cls, documents: list[str]) -> "VectorStore":
        vectorizer = TfidfVectorizer()
        tfidf_matrix = vectorizer.fit_transform(documents)
        return cls(vectorizer, tfidf_matrix, documents)

    def retrieve_with_scores(self, query: str, k: int = 2) -> list[dict]:
        query_vector = self.vectorizer.transform([query])
        similarities = cosine_similarity(query_vector, self.tfidf_matrix).flatten()
        top_k_indices = np.argsort(similarities)[-k:][::-1]
        return [
            {"document": self.documents[i], "relevance": float(similarities[i])}
            for i in top_k_indices
        ]

    def retrieve_context(self, query: str, k: int = 2) -> str:
        documents = self.retrieve_with_scores(query, k)
        return "\n".join([item["document"] for item in documents])

# Example usage
if __name__ == "__main__":
    documents = [
        "SpaceX is an American spacecraft manufacturer founded by Elon Musk in 2002.",
        "Fiji is a country in the South Pacific known for its beaches and coral reefs.",
        "Dunkirk is a 2017 war film depicting the Dunkirk evacuation of World War II.",
        "BoundaryML makes BAML, the best way to get structured outputs with LLMs."
    ]

    vector_store = VectorStore.from_documents(documents)

    questions = [
        "What is BAML?",
        "Which aircraft was featured in Dunkirk?",
        "When was SpaceX founded?",
        "Where is Fiji located?",
        "What is the capital of Fiji?"
    ]

    for question in questions:
        context = vector_store.retrieve_context(question)
        response = b.RAG(question, context)
        print(f"Q: {response.question}")
        print(f"A: {response.answer}")
        print("-" * 40)

from baml_client import b

def rag_with_citations(question: str, context: str):
    response = b.RAGWithCitations(question, context)
    
    print(f"Question: {response.question}")
    print(f"Answer: {response.answer}")
    print("\nCitations:")
    for i, citation in enumerate(response.citations, 1):
        print(f"  [{i}] {citation}")
    
    return response

# Example usage
context = """
SpaceX is an American spacecraft manufacturer founded by Elon Musk in 2002.
The company has revolutionized space travel with reusable rockets.
"""

result = rag_with_citations(
    "When was SpaceX founded and by whom?",
    context
)

from pinecone import Pinecone, ServerlessSpec
from sentence_transformers import SentenceTransformer
from baml_client import b

class PineconeStore:
    def __init__(self, api_key: str, index_name: str):
        self.pc = Pinecone(api_key=api_key)
        self.index_name = index_name
        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
        
        # Create index if it doesn't exist
        if index_name not in self.pc.list_indexes().names():
            self.pc.create_index(
                name=index_name,
                dimension=384,  # all-MiniLM-L6-v2 dimension
                metric='cosine',
                spec=ServerlessSpec(
                    cloud='aws',
                    region='us-east-1'
                )
            )
        self.index = self.pc.Index(index_name)

    def add_documents(self, documents: list[str], ids: list[str] = None):
        if ids is None:
            ids = [str(i) for i in range(len(documents))]
        
        # Create embeddings
        embeddings = self.encoder.encode(documents)
        
        # Create vector records
        vectors = [
            (id, emb.tolist(), {"text": doc}) 
            for id, emb, doc in zip(ids, embeddings, documents)
        ]
        
        # Upsert to Pinecone
        self.index.upsert(vectors=vectors)

    def retrieve_context(self, query: str, k: int = 2) -> str:
        # Create query embedding
        query_embedding = self.encoder.encode(query).tolist()
        
        # Query Pinecone
        results = self.index.query(
            vector=query_embedding,
            top_k=k,
            include_metadata=True
        )
        
        # Extract document texts
        contexts = [match.metadata["text"] for match in results.matches]
        return "\n".join(contexts)

# Example usage
if __name__ == "__main__":
    # Initialize Pinecone store
    vector_store = PineconeStore(
        api_key="YOUR_API_KEY",
        index_name="baml-rag-demo"
    )
    
    # Sample documents
    documents = [
        "SpaceX is an American spacecraft manufacturer founded by Elon Musk in 2002.",
        "Fiji is a country in the South Pacific known for its beaches.",
        "BoundaryML makes BAML, the best way to get structured outputs with LLMs."
    ]
    
    # Add documents to Pinecone
    vector_store.add_documents(documents)
    
    # Query using BAML
    questions = [
        "What is BAML?",
        "When was SpaceX founded?",
    ]

    for question in questions:
        context = vector_store.retrieve_context(question)
        response = b.RAGWithCitations(question, context)
        print(f"Q: {response.question}")
        print(f"A: {response.answer}")
        print(f"Citations: {response.citations}")
        print("-" * 40)

from baml_client import b

def multi_query_rag(question: str, vector_store):
    # Generate multiple query variations
    queries = b.GenerateQueries(question)
    
    print(f"Original: {question}")
    print(f"Variations: {queries}\n")
    
    # Retrieve context for each query
    all_contexts = []
    for query in queries:
        context = vector_store.retrieve_context(query, k=2)
        all_contexts.append(context)
    
    # Combine and deduplicate contexts
    combined_context = "\n".join(set(all_contexts))
    
    # Get final answer with citations
    response = b.RAGWithCitations(question, combined_context)
    
    print(f"Answer: {response.answer}")
    print(f"Citations: {response.citations}")
    
    return response

# Example usage
if __name__ == "__main__":
    vector_store = VectorStore.from_documents([...])
    
    result = multi_query_rag(
        "How does SpaceX technology work?",
        vector_store
    )

from sklearn.feature_extraction.text import TfidfVectorizer
from sentence_transformers import SentenceTransformer
import numpy as np

class HybridVectorStore:
    def __init__(self, documents: list[str]):
        self.documents = documents
        
        # Keyword search (TF-IDF)
        self.tfidf_vectorizer = TfidfVectorizer()
        self.tfidf_matrix = self.tfidf_vectorizer.fit_transform(documents)
        
        # Semantic search (embeddings)
        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
        self.embeddings = self.encoder.encode(documents)

    def retrieve_context(self, query: str, k: int = 3, alpha: float = 0.5) -> str:
        # Keyword search scores
        query_tfidf = self.tfidf_vectorizer.transform([query])
        keyword_scores = cosine_similarity(query_tfidf, self.tfidf_matrix).flatten()
        
        # Semantic search scores
        query_embedding = self.encoder.encode([query])
        semantic_scores = cosine_similarity(query_embedding, self.embeddings).flatten()
        
        # Combine scores (alpha controls the balance)
        hybrid_scores = alpha * keyword_scores + (1 - alpha) * semantic_scores
        
        # Get top k documents
        top_k_indices = np.argsort(hybrid_scores)[-k:][::-1]
        contexts = [self.documents[i] for i in top_k_indices]
        
        return "\n".join(contexts)

# Example usage
documents = [...]
hybrid_store = HybridVectorStore(documents)

question = "What is machine learning?"
context = hybrid_store.retrieve_context(question, k=3, alpha=0.6)
response = b.RAG(question, context)