Simple RAG with chatlas

rag-chatlas.py

import numpy as np
from chatlas import ChatAnthropic
from sentence_transformers import SentenceTransformer

# Initialize the SentenceTransformer model and chat client
embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
chat = ChatAnthropic(
    system_prompt="""
    You are a helpful AI assistant. Using the provided context, 
    answer the user's question. If you cannot answer the question based on the 
    context, say so.
"""
)

# Our list of documents (one document per list element)
documents = [
    "The Python programming language was created by Guido van Rossum.",
    "Python is known for its simple, readable syntax.",
    "Python supports multiple programming paradigms.",
]

# Compute embeddings for each document
embeddings = [
    embed_model.encode([doc])[0] for doc in documents
]

# Compute embedding for the user query
user_query = "Who created Python?"
query_embedding = embed_model.encode([user_query])[0]

# Calculate cosine similarity between the query and each document
similarities = np.dot(embeddings, query_embedding) / (
    np.linalg.norm(embeddings, axis=1) * np.linalg.norm(query_embedding)
)

# Get the top-k most similar documents
top_k = 3
top_indices = np.argsort(similarities)[-top_k:][::-1]

# Give the model both the context and the question
chat.chat(
    f"Context:\n{documents[top_indices[0]]}\n\nQuestion: {user_query}"
)