# query.py

import argparse
from langchain.vectorstores import Chroma
from langchain.prompts import PromptTemplate
from langchain.llms import HuggingFaceHub
from embedding_function import get_embedding
import os

# Set your Hugging Face API token
HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")

# Path to the Chroma database
CHROMA_PATH = "chroma"

# Prompt template for the LLM
PROMPT_TEMPLATE = """
You are a helpful assistant. Use the following context to answer the question concisely and accurately.

Context:
{context}

Question:
{question}

Answer:
"""

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("query_text", type=str, help="The query text.")
    args = parser.parse_args()
    response = query_rag(args.query_text)
    print("\n💬 Response:")
    print(response)

def query_rag(query_text):
    # Initialize embedding function and vector store
    embedding_function = get_embedding()
    db = Chroma(
        persist_directory=CHROMA_PATH,
        embedding_function=embedding_function
    )

    # Retrieve relevant documents
    docs = db.similarity_search(query_text, k=3)
    context = "\n\n".join([doc.page_content for doc in docs])

    # Prepare prompt
    prompt = PromptTemplate(
        input_variables=["context", "question"],
        template=PROMPT_TEMPLATE
    ).format(context=context, question=query_text)

    # Initialize LLM
    llm = HuggingFaceHub(
        repo_id="google/flan-t5-base",
        model_kwargs={"temperature": 0.5, "max_length": 512},
        huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN
    )

    # Generate response
    response = llm(prompt)
    return response.strip()

if __name__ == "__main__":
    main()