query.py

import argparse
import time
from rich.console import Console
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings

# Optional: use a local model if installed (like Ollama or GPT4All)
try:
    from langchain_community.llms import Ollama
    HAS_OLLAMA = True
except ImportError:
    HAS_OLLAMA = False

console = Console()

def run_query(question):
    start = time.time()

    # Load local embeddings + vectorstore
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    vectordb = Chroma(persist_directory="vectorstore", embedding_function=embeddings)

    # Retrieve top 3 relevant documents
    docs = vectordb.similarity_search(question, k=3)

    # Combine retrieved context
    context = "\n\n".join([d.page_content for d in docs])

    if HAS_OLLAMA:
        console.print("[bold green]Using local Ollama model for answer generation...[/bold green]")
        llm = Ollama(model="llama3")  # change to any local Ollama model you have
        answer = llm.invoke(f"Context:\n{context}\n\nQuestion: {question}\nAnswer:")
    else:
        console.print("[yellow]No local LLM found. Showing retrieved context only.[/yellow]")
        answer = "⚠️ No local LLM installed. Below are the most relevant retrieved snippets:\n\n" + context[:1200]

    end = time.time()

    console.rule(f"[bold cyan]Question[/bold cyan]")
    console.print(question, style="bold white")
    console.rule("[bold cyan]Answer[/bold cyan]")
    console.print(answer, style="green")
    console.rule("[bold cyan]Sources[/bold cyan]")
    for i, doc in enumerate(docs, start=1):
        src = doc.metadata.get("source")
        console.print(f"{i}. {src}")

    console.print(f"\n[blue]Total Latency: {end - start:.2f}s[/blue]")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--question", required=True, help="Question to ask the local RAG system")
    args = parser.parse_args()
    run_query(args.question)