app.py

import os
import zipfile
import gradio as gr
from langchain_openai import ChatOpenAI
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

# Unzip vector DB if not already extracted
if not os.path.exists("geometry_chroma"):
    with zipfile.ZipFile("geometry_chroma.zip", 'r') as zip_ref:
        zip_ref.extractall(".")

# Load vector DB
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectordb = Chroma(persist_directory="geometry_chroma", embedding_function=embedding_model)
retriever = vectordb.as_retriever()

# Set OpenAI key (use Secrets or .env later)
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

llm = ChatOpenAI(model_name="gpt-4.1", temperature=0.2)

# ✅ Prompt templates
templates = {
    "flashcard": PromptTemplate(
        input_variables=["context", "query"],
        template="""
{context}

Create 5 flashcards based on the topic: "{query}"
Each flashcard should include:
- A clear question
- A short answer
Focus on high school geometry understanding.
"""
    ),
    "lesson plan": PromptTemplate(
        input_variables=["context", "query"],
        template="""
Given the following retrieved SOL text:
{context}

Generate a Geometry lesson plan based on: "{query}"
Include:
1. Simple explanation of the concept.
2. Real-world example.
3. Engaging class activity.
Be concise and curriculum-aligned for high school.
"""
    ),
    "worksheet": PromptTemplate(
        input_variables=["context", "query"],
        template="""
{context}

Create a student worksheet for: "{query}"
Include:
- Concept summary
- A worked example
- 3 practice problems
"""
    ),
    "proofs": PromptTemplate(
        input_variables=["context", "query"],
        template="""
{context}

Generate a proof-focused geometry lesson plan for: "{query}"
Include:
- Student-friendly explanation
- Real-world connection
- One short class activity
"""
    ),
"general question": PromptTemplate(
    input_variables=["context", "query"],
    template="""
You are a Virginia Geometry Standards of Learning (SOL) assistant.

From the following standards content:
{context}

Identify the SOL standard (e.g., G.RLT.1, G.TR.3) that most accurately answers the user's question: "{query}"

Respond with:
1. The SOL code (e.g., G.RLT.3)
2. The **exact standard description** for that code, as written in the context

Do not summarize. Copy the official description exactly.
"""
)



}

def generate_output(prompt_type, query):
    import re

    # Check for exact SOL code in query
    sol_match = re.search(r"\bG\.[A-Z]+\.\d+\b", query)
    matched_code = sol_match.group(0) if sol_match else None

    if matched_code:
        # Find matching document by metadata
        docs_with_code = vectordb.similarity_search_with_score(query)
        filtered_docs = [doc for doc, _ in docs_with_code if doc.metadata.get("standard") == matched_code]

        if filtered_docs:
            context = "\n\n".join([doc.page_content for doc in filtered_docs])
        else:
            # fallback to regular retrieval if not found
            docs = retriever.get_relevant_documents(query)
            context = "\n\n".join([doc.page_content for doc in docs])
    else:
        # regular query
        docs = retriever.get_relevant_documents(query)
        context = "\n\n".join([doc.page_content for doc in docs])

    # Run the prompt
    chain = LLMChain(llm=llm, prompt=templates[prompt_type])
    return chain.run({"context": context, "query": query}).strip()



# ✅ Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 📐 Geometry Teaching Assistant")

    with gr.Row():
        query = gr.Textbox(label="Enter a geometry topic")
        prompt_type = gr.Dropdown(
            ["general question", "lesson plan", "worksheet", "proofs", "flashcard"],
            value="general question",
            label="Prompt Type"
        )

    output = gr.Textbox(label="Generated Output", lines=12, interactive=True)
    btn = gr.Button("Generate")

    btn.click(fn=generate_output, inputs=[prompt_type, query], outputs=output)

demo.launch()