rag_engine.py

"""
RAG engine for the chatbot demo.
Uses sentence-transformers for embeddings and FAISS for retrieval.
"""

import os
import numpy as np
from faq_data import FAQ_LIST


def build_retriever(model_name="all-MiniLM-L6-v2"):
    """Load sentence transformer model and build FAISS index."""
    from sentence_transformers import SentenceTransformer
    import faiss

    # Prepare texts for embedding — embed question + answer together so all terms are searchable
    texts = [f"{item['question']} {item['answer']}" for item in FAQ_LIST]
    answers = [item["answer"] for item in FAQ_LIST]
    categories = [item["category"] for item in FAQ_LIST]

    # Load model (cached locally after first download)
    model = SentenceTransformer(model_name)

    # Create embeddings
    embeddings = model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
    dimension = embeddings.shape[1]

    # Build FAISS index
    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings.astype(np.float32))

    return {
        "model": model,
        "index": index,
        "texts": texts,
        "answers": answers,
        "categories": categories,
    }


def search(retriever, query, top_k=3):
    """Search for the most relevant FAQ entries."""
    query_vec = retriever["model"].encode([query], convert_to_numpy=True)
    distances, indices = retriever["index"].search(query_vec.astype(np.float32), top_k)

    results = []
    for i, idx in enumerate(indices[0]):
        results.append({
            "question": retriever["texts"][idx],
            "answer": retriever["answers"][idx],
            "category": retriever["categories"][idx],
            "score": float(distances[0][i]),
        })

    return results


def generate_response(query, retrieved, api_key, model="openrouter/owl-alpha"):
    """Generate a natural response using OpenRouter API with retrieved context."""
    import httpx

    if not api_key:
        return "I need an OpenRouter API key to generate responses. Set it in Settings above.", []

    # Build context from retrieved docs
    context_parts = []
    for i, r in enumerate(retrieved, 1):
        context_parts.append(f"[{i}] Q: {r['question']}\n    A: {r['answer']}")
    context = "\n\n".join(context_parts)

    system_prompt = (
        "You are a helpful customer support assistant for FlowBar, a project management SaaS."
        " Answer the user's question using ONLY the information provided in the context below."
        " If the context doesn't contain the answer, say 'I don't have that information in my knowledge base."
        " Would you like me to connect you with a human agent?'"
        " Be concise, friendly, and professional."
        " Reference the source number when appropriate."
    )

    user_prompt = f"Context:\n{context}\n\nUser Question: {query}"

    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        "temperature": 0.3,
        "max_tokens": 500,
    }

    try:
        resp = httpx.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
                "HTTP-Referer": "https://huggingface.co/spaces",
            },
            json=payload,
            timeout=30,
        )
        resp.raise_for_status()
        data = resp.json()
        return data["choices"][0]["message"]["content"], retrieved
    except Exception as e:
        return f"Sorry, I couldn't reach the AI. Error: {str(e)}", retrieved


def generate_response_light(query, retrieved):
    """Fallback: generate response without external API using template."""
    best = retrieved[0]
    return (
        f"Based on our knowledge base, here's what I found:\n\n"
        f"**{best['question']}**\n\n{best['answer']}\n\n"
        f"*Did that answer your question? If not, please rephrase or ask something else.*"
    ), retrieved