A simple example of RAG

import ollama
import chromadb

# Sample documents
documents = [
    "Llamas are members of the camelid family meaning they're pretty closely related to vicuñas and camels",
    "Llamas were first domesticated and used as pack animals 4,000 to 5,000 years ago in the Peruvian highlands",
    "Llamas can grow as much as 6 feet tall though the average llama between 5 feet 6 inches and 5 feet 9 inches tall",
]

# Set up ChromaDB
client = chromadb.Client()
# ✅ Fix: Use get_or_create_collection to avoid errors on re-runs
collection = client.get_or_create_collection(name="llama_docs")

# ✅ Fix: Embed documents one at a time
try:
    embeddings = []
    for doc in documents:
        response = ollama.embeddings(model="all-minilm", prompt=doc)
        embeddings.append(response['embedding'])  # Note: 'embedding' (singular)
    
    print(f"Generated {len(embeddings)} embeddings")
    
    collection.add(
        ids=[str(i) for i in range(len(documents))],
        embeddings=embeddings,
        documents=documents
    )
    print("Documents added to collection")
except Exception as e:
    print(f"Error embedding documents: {e}")
    exit()

# Query
#query = "What animals are llamas related to?"
query = "size of llamas?"
try:
    query_response = ollama.embeddings(model="all-minilm", prompt=query)
    print(f"Query embedding generated")
    
    results = collection.query(
        query_embeddings=[query_response['embedding']],
        n_results=2
    )
    relevant_doc = results['documents'][0][0]
    print(f"Relevant document: {results}")
except Exception as e:
    print(f"Error querying: {e}")
    exit()

# Generate response with gemma
prompt = f"Using this data: {relevant_doc}. Respond to: {query}"
try:
    # ✅ Fix: Check your available models (might be gemma2:12b or gemma:7b)
    output = ollama.generate(model="gemma3:12b", prompt=prompt)
    print(f"\nGenerated response:\n{output['response']}")
except Exception as e:
    print(f"Error generating response: {e}")
    # Try alternative model name
    print("Trying alternative model name...")
    try:
        output = ollama.generate(model="llama2-uncensored:7b", prompt=prompt)
        print(f"\nGenerated response:\n{output['response']}")
    except Exception as e2:
        print(f"Error with alternative model: {e2}")