query.py

import requests
import chromadb

# Configuration (must match index.py)
OLLAMA_API_BASE = "http://localhost:11434"
EMBEDDING_MODEL = "granite-embedding:30m"
GENERATION_MODEL = "gemma3:4b"  # Your chat model
CHROMA_DB_PATH = "./chromadb_storage"

def get_embedding(text: str) -> list:
    """Get embeddings from Ollama"""
    try:
        response = requests.post(
            f"{OLLAMA_API_BASE}/api/embeddings",
            json={"model": EMBEDDING_MODEL, "prompt": text}
        )
        response.raise_for_status()
        return response.json()["embedding"]
    except Exception as e:
        print(f"Error getting embedding: {e}")
        raise

def retrieve_context(query: str, n_results: int = 3) -> list:
    """
    Search ChromaDB for relevant chunks
    Returns list of relevant text chunks
    """
    # Get query embedding
    query_embedding = get_embedding(query)
    
    # Search ChromaDB
    client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
    collection = client.get_or_create_collection(name="documents")
    
    # Query the collection
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=n_results
    )
    
    # Extract the text chunks
    if results and results['documents']:
        return results['documents'][0]  # Returns list of matching chunks
    return []

def generate_answer(question: str, context_chunks: list) -> str:
    """
    Generate answer using Ollama with retrieved context
    """
    # Combine context chunks
    context = "\n\n".join(context_chunks)
    
    # Create prompt
    prompt = f"""Based on the following context, answer the question. If the answer cannot be found in the context, say "I don't have enough information to answer that."

Context:
{context}

Question: {question}

Answer:"""
    
    # Call Ollama generation API
    try:
        response = requests.post(
            f"{OLLAMA_API_BASE}/api/generate",
            json={
                "model": GENERATION_MODEL,
                "prompt": prompt,
                "stream": False
            }
        )
        response.raise_for_status()
        return response.json()["response"]
    except Exception as e:
        return f"Error generating answer: {e}"

def ask_question(question: str) -> str:
    """
    Main RAG pipeline:
    1. Embed the question
    2. Retrieve relevant chunks from ChromaDB
    3. Generate answer using retrieved context
    """
    # Step 1 & 2: Retrieve relevant context
    print("  → Searching for relevant information...")
    context_chunks = retrieve_context(question, n_results=3)
    
    if not context_chunks:
        return "No relevant information found in the indexed documents. Please index some documents first."
    
    print(f"  → Found {len(context_chunks)} relevant chunks")
    
    # Step 3: Generate answer
    print("  → Generating answer...")
    answer = generate_answer(question, context_chunks)
    
    return answer

# For standalone testing
if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python query.py <your question>")
        sys.exit(1)
    
    question = " ".join(sys.argv[1:])
    print(f"\nQuestion: {question}\n")
    answer = ask_question(question)
    print(f"Answer: {answer}\n")