app.py

"""
Emerald-RAG: A Retrieval-Augmented Generation system for Pokémon Emerald.
This module serves as the entry point for the Hugging Face Space, 
handling the Gradio/Streamlit UI and orchestrating the RAG pipeline 
to answer queries based on the game's mechanics and manual.
"""
import gradio as gr
from helpers import create_prompt, find_similar_documents, load_vector_store
from model import CHAT_MODEL_ID, generate_answer, load_chat_model
import pytest
import sys

# Trigger tests on start up. 
retcode = pytest.main(["tests/"])
if retcode != 0:
    print("Tests failed!")
else:
    print("Tests passed!") 
    
# Initialiations and globals
print(f"Loading {CHAT_MODEL_ID} model...")
chat_model, tokenizer = load_chat_model()
print("Model loaded!")

MANUAL_PATH = "emerald_manual.txt"
print(f"Loading vector store from {MANUAL_PATH}...")
vector_store = load_vector_store(MANUAL_PATH)
print("Vector store loaded!")

def respond(message, history):
    print("Enter app.respond...")
    context = find_similar_documents(vector_store, message)
    prompt = create_prompt(message, history, context)
    response = generate_answer(chat_model, tokenizer, prompt)
    if not response:
        return "Query failed. Please try again."
    formatted_output = f"**Answer:** {response.answer}\n\n"

    emoji = "✅" if response.confidence_score > 80 else "⚠️"
    formatted_output += f"**Confidence:** {response.confidence_score}% {emoji}\n\n"
    
    # 4. Add citations as a list of quotes
    formatted_output += "**Sources from Manual:**"
    for quote in response.citations:
        formatted_output += f"\n> *\"{quote}\"*\n"
        
    return formatted_output

app = gr.ChatInterface(
    fn=respond,
    title="RAG system for Pokémon Emerald Q&A",
    description=(
        "Ask me anything about Pokémon Emerald mechanics, items, or walkthroughs! "
        "This high-precision Retrieval-Augmented Generation (RAG) system eliminates "
        "LLM hallucinations by grounding responses directly in an  game reference "
        "manual ([link](https://gamefaqs.gamespot.com/gba/921905-pokemon-emerald-version/faqs/44694)). "
        "Built using [Mistral-Small-Instruct-2409](https://huggingface.co/mistralai/Mistral-Small-Instruct-2409), "
        "[nomic-ai/nomic-embed-text-v1.5](https://huggingface.co/nomic-ai/nomic-embed-text-v1.5), "
        "and a ChromaDB vector store, the pipeline outputs structured JSON complete with direct source citations "
        "and confidence scores."
    ),
    examples=[
        "Where can I find a super rod?", 
        "Who is the best starter pokemon in Emerald?", 
        "What does the Mach Bike do?",
        "What are the most important items?"
    ],
    cache_examples=False,
)

custom_css = """
.p, .md p {
    font-size: 16px !important;
}
.message-text {
    font-size: 16px !important;
}
.bubble-wrap {
    max-height: 400px !important;
    overflow-y: auto !important;
}
"""

if __name__ == "__main__":   
    app.launch(css=custom_css)