groq_rag.py

import os

import chromadb
from dotenv import load_dotenv
from groq import Groq
from sentence_transformers import SentenceTransformer


# --------------------------------------------------
# Configuration
# --------------------------------------------------

load_dotenv()

GROQ_API_KEY = os.getenv("GROQ_API_KEY")

if not GROQ_API_KEY:
    raise ValueError("GROQ_API_KEY is not set")

groq_client = Groq(api_key=GROQ_API_KEY)


# --------------------------------------------------
# Embedding model
# --------------------------------------------------

embedding_model = SentenceTransformer("all-MiniLM-L6-v2")


# --------------------------------------------------
# ChromaDB
# --------------------------------------------------

chroma_client = chromadb.PersistentClient(path="./chroma")

collection = chroma_client.get_or_create_collection(
    name="company_docs"
)


# --------------------------------------------------
# Documents
# --------------------------------------------------

documents = [
    """
    Refund Policy

    Customers can request a refund within 7 days of receiving
    their order. The product must be unused and in its original
    condition. Refunds are processed within 5 business days after
    approval.
    """,

    """
    Shipping Policy

    Standard shipping usually takes 3 to 5 business days.
    Express shipping takes 1 to 2 business days.
    Shipping is available throughout Bangladesh.
    """
] 


# --------------------------------------------------
# Add documents to ChromaDB
# --------------------------------------------------

embeddings = embedding_model.encode(documents).tolist()
print( f"Adding documents to ChromaDB...{len(documents)} documents added.")
print(f"Embeddings shape: {len(embeddings)} x {len(embeddings[0])} x {len(embeddings[1])})")
# print(f"Sample embedding for first document: {embeddings[0]}")

collection.upsert(
    ids=["refund-policy", "shipping-policy"],
    documents=documents,
    embeddings=embeddings
)


# --------------------------------------------------
# Retrieval
# --------------------------------------------------

def retrieve_documents(query: str, top_k: int = 3):
    query_embedding = embedding_model.encode([query]).tolist()
    print(f"Query embedding shape: {len(query_embedding)} x {len(query_embedding[0])}")
    # print(f"Sample query embedding: {query_embedding[0]}")

    results = collection.query(
        query_embeddings=query_embedding,
        n_results=top_k
    )

    return results["documents"][0]


# --------------------------------------------------
# Generate answer using Groq
# --------------------------------------------------

def generate_answer(query: str, context: list[str]):
    # print(f"Retrieved context: {context}")
    
    for c in context:
        print(f"Context snippet: {c[:]}...\n\n")  # Print first 100 characters of each context snippet

    context_text = "\n\n".join(context)

    prompt = f"""
You are a helpful assistant.

Answer the user's question using ONLY the provided context.

If the answer cannot be found in the context, say:
"I don't know based on the provided information."

Context:
----------------
{context_text}
----------------

Question:
{query}

Answer:
"""

    response = groq_client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ],
        temperature=0
    )

    return response.choices[0].message.content


# --------------------------------------------------
# RAG pipeline
# --------------------------------------------------

def ask(query: str):

    # 1. Retrieve relevant documents
    context = retrieve_documents(query)

    # 2. Generate answer using retrieved context
    answer = generate_answer(query, context)

    return answer


# --------------------------------------------------
# Interactive chat
# --------------------------------------------------

if __name__ == "__main__":

    print("RAG chatbot started.")
    print("Type 'exit' to quit.\n")

    while True:

        query = input("You: ")

        if query.lower() == "exit":
            break

        answer = ask(query)

        print(f"\nAssistant: {answer}\n")