rag_module/rag_pipeline.py

import os
from openai import OpenAI
from rag_module.vector_store import query_store
from utils.incident_logger import load_incidents

HF_MODEL = "meta-llama/Llama-3.1-8B-Instruct"


def _get_client():
    token = os.getenv("HF_TOKEN")
    if not token:
        raise EnvironmentError(
            "HF_TOKEN not set.\n"
            "Get your free token at: https://huggingface.co/settings/tokens\n"
            "Then in PowerShell: $env:HF_TOKEN='hf_xxxxxxxxxxxxxxxx'"
        )
    return OpenAI(
        base_url="https://router.huggingface.co/v1",
        api_key=token,
    )


def _build_prompt(context_docs: list[dict], user_query: str) -> str:
    context_lines = "\n".join(
        f"- [{d.get('timestamp', '?')}] Plate: {d.get('plate', '?')} | "
        f"Type: {d.get('vehicle_class', '?')} | Zone: {d.get('zone', '?')} | "
        f"Status: {d.get('status', '?')} | Notes: {d.get('notes', '') or 'none'}"
        for d in context_docs
    )
    return (
        f"You are a smart parking security assistant. "
        f"You MUST answer based on the incident log below. "
        f"Even if there is only one record, use it to answer. "
        f"Never say there is no data if records are shown below.\n\n"
        f"--- INCIDENT LOG ---\n{context_lines}\n--------------------\n\n"
        f"Question: {user_query}\n\n"
        f"Answer directly and factually using the records above:"
    )


def ask(user_query: str, n_context: int = 5) -> dict:
    retrieved = query_store(user_query, n_results=n_context)

    # Fallback: load directly from CSV if vector store is empty
    if not retrieved:
        df = load_incidents()
        if not df.empty:
            retrieved = df.head(20).to_dict(orient="records")
            for r in retrieved:
                if hasattr(r.get("timestamp"), "strftime"):
                    r["timestamp"] = r["timestamp"].strftime("%Y-%m-%d %H:%M:%S")

    if not retrieved:
        return {
            "answer": "No incident records found. Run detection or seed sample data first.",
            "retrieved_docs": [],
        }

    prompt = _build_prompt(retrieved, user_query)

    try:
        client = _get_client()
        response = client.chat.completions.create(
            model=HF_MODEL,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=400,
            temperature=0.2,
        )
        answer = response.choices[0].message.content.strip()
    except EnvironmentError as e:
        answer = str(e)
    except Exception as e:
        answer = f"Error: {e}"

    return {"answer": answer, "retrieved_docs": retrieved}