rag.py

import os
import logging
import signal
import sys
import zipfile
import gdown
import time

from google import genai

from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings


# ==============================
# Logging
# ==============================
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# ==============================
# Gemini Client
# ==============================
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

if not GEMINI_API_KEY:
    raise ValueError("GEMINI_API_KEY is missing")

client = genai.Client(api_key=GEMINI_API_KEY)


MODEL_NAME = "models/gemini-2.5-flash-lite"


# ==============================
# Gemini Call (with retry)
# ==============================
def call_gemini(prompt: str, retries: int = 3):
    for attempt in range(retries):
        try:
            response = client.models.generate_content(
                model=MODEL_NAME,
                contents=prompt
            )

            if response and response.text:
                return response.text

            return "Empty response from model."

        except Exception as e:
            msg = str(e)

            # handling quota error
            if "429" in msg or "RESOURCE_EXHAUSTED" in msg:
                wait_time = 5 * (attempt + 1)
                logger.warning(f"Quota hit. Retrying in {wait_time}s...")
                time.sleep(wait_time)
                continue

            logger.error(f"Gemini error: {e}")
            break

    return "Model is temporarily unavailable. Please try again later."


# ==============================
# Download Vector DB
# ==============================
if not os.path.exists("chroma_db"):
    print("Downloading vector DB...")

    url = "https://drive.google.com/uc?id=1yelCLmiRD2Qds-_8EsnKzUNY7uPT5xMg"

    gdown.download(url, "chroma_db.zip", quiet=False)

    print("Extracting DB...")

    with zipfile.ZipFile("chroma_db.zip", "r") as zip_ref:
        zip_ref.extractall(".")

    print("Vector DB Ready ✅")


# ==============================
# Load Vector DB
# ==============================
embedding = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)

vector_db = Chroma(
    persist_directory="./chroma_db",
    embedding_function=embedding
)

print("Vector DB Loaded ✅")


# ==============================
# Conversation Memory
# ==============================
class Conversation:
    def __init__(self):
        self.history = []

    def add_user_message(self, message):
        self.history.append({"role": "user", "content": message})

    def add_bot_message(self, message):
        self.history.append({"role": "assistant", "content": message})

    def get_history(self):
        return self.history

    def clear_history(self):
        self.history = []


conversation = Conversation()


# ==============================
# Get Context
# ==============================
def get_relevant_context_from_db(query):
    try:
        results = vector_db.similarity_search(query, k=2)

        context = "\n\n".join([
            doc.page_content[:1200]
            for doc in results
        ])

        sources = list(set([
            doc.metadata.get("source", "")
            for doc in results
        ]))

        return context, sources

    except Exception as e:
        logger.error(f"Vector DB error: {e}")
        return "", []


# ==============================
# Prompt Builder
# ==============================
def build_prompt(query, context, history):

    history_text = "\n".join([
        f"{msg['role']}: {msg['content']}"
        for msg in history[-3:]
    ]) if history else ""

    prompt = f"""
You are an expert nutritionist and dietitian.

RULES:
- Answer ONLY the user question
- Keep answers clean and organized
- Use bullet points when helpful
- Do NOT repeat sentences
- Do NOT mention Ramadan unless user asks
- If user asks for diet:
  give breakfast, lunch, dinner, snacks
- If user asks for foods:
  include protein/calories if possible
- If user writes Arabic:
  answer in Arabic
- Keep answers practical and realistic

Conversation:
{history_text}
Context:
{context}
User Question:
{query}
Answer:
"""

    return prompt


# ==============================
# Main Function
# ==============================
def generate_chat_answer(query, history=[]):

    try:
        logger.info(f"Query: {query}")

        context, sources = get_relevant_context_from_db(query)

        prompt = build_prompt(
            query,
            context,
            history
        )

        answer = call_gemini(prompt).strip()

        answer = answer.replace("**", "").replace("##", "")

        if not answer:
            answer = "Sorry, no response generated."

        # memory
        conversation.add_user_message(query)
        conversation.add_bot_message(answer)

        # add sources
        if sources:
            answer += "\n\nSources:\n" + "\n".join(
                [f"- {s}" for s in sources if s]
            )

        return answer

    except Exception as e:
        logger.error(f"Error: {e}", exc_info=True)
        return "Something went wrong."


# ==============================
# Graceful shutdown
# ==============================
def signal_handler(sig, frame):
    print("\nShutting down...")
    sys.exit(0)


signal.signal(signal.SIGINT, signal_handler)