Update app.py

app.py

@@ -1,135 +1,112 @@
-# app.py
-import gradio as gr
 import os
-from transformers import pipeline
-from sentence_transformers import SentenceTransformer
-import faiss
-import numpy as np
-import json
 import re
+import gradio as gr
+from langchain.prompts import PromptTemplate
+from langchain_openai import ChatOpenAI
+from langchain.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+
+# Load embedding model and vector store from persisted DB
+embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+vector_store = Chroma(
+    embedding=embedding_model,
+    persist_directory="geometry_db",  # relative folder inside your Hugging Face Space
+    collection_name="geometry_sol"
+)
 
+# Load OpenAI key (you must add this in Hugging Face Space Secrets)
+os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
 
-# --- Load necessary components for the RAG system ---
-# These paths are relative to the Space's root directory
-FAISS_INDEX_PATH = "sol_faiss_index.bin"
-DOCUMENT_IDS_PATH = "sol_document_ids.json"
-
-# Load SentenceTransformer model
-# Ensure this model is downloaded or available in the environment
-# For Spaces, you might need to add it to requirements.txt or directly download if space has internet
-# It's better to declare it globally or as a shared resource.
-try:
-    model = SentenceTransformer('all-mpnet-base-v2')
-except Exception as e:
-    print(f"Error loading SentenceTransformer model: {e}")
-    print("Attempting to load from local cache or download on first use.")
-    # If running in a Space, the model will be downloaded to cache if not present.
-    # Ensure you have internet access in your Space settings.
-
-# Load FAISS index
-try:
-    index = faiss.read_index(FAISS_INDEX_PATH)
-except Exception as e:
-    print(f"Error loading FAISS index: {e}")
-    # Handle error, maybe create a dummy index or exit
-    index = None # Placeholder if loading fails
-
-# Load document IDs
-try:
-    with open(DOCUMENT_IDS_PATH, "r") as f:
-        document_ids = json.load(f)
-except Exception as e:
-    print(f"Error loading document IDs: {e}")
-    document_ids = [] # Placeholder if loading fails
-
-# Placeholder for the actual content of "10 Geometry Mathematics Instructional Guide.pdf"
-# In a real deployed scenario, this content would be loaded from a file
-# that you upload to your Hugging Face Space or fetched at runtime.
-# For now, we'll assume it's available or that 'documents' are pre-processed and loaded.
-# You would typically load the 'documents' list created in Step 2 here.
-# For deployment, it's best to save the `documents` list (sol_data) as a JSON
-# and load it back. Let's add that.
-
-# Assuming you've saved sol_data as 'sol_documents.json'
-SOL_DOCUMENTS_PATH = "sol_documents.json"
-try:
-    with open(SOL_DOCUMENTS_PATH, "r") as f:
-        documents = json.load(f)
-except Exception as e:
-    print(f"Error loading sol documents: {e}")
-    documents = [] # Placeholder
-
-# Load LLM for generation
-# For a Hugging Face Space, you need to ensure the model is available.
-# 'google/gemma-2b-it' is a good option.
-# Ensure you set up environment variables or secrets for API keys if using paid models.
-try:
-   # llm_pipeline = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
-    llm_pipeline = pipeline("text-generation", model="google/gemma-2b-it")
-except Exception as e:
-    print(f"Error loading LLM pipeline: {e}")
-    llm_pipeline = None # Placeholder
-
-
-def retrieve_and_generate_app(query, top_k=3):
-    if not model or not index or not document_ids or not documents or not llm_pipeline:
-        return "System not fully initialized. Please check logs for missing components."
+# Load the LLM (GPT-3.5)
+llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.3)
 
-    try:
-        # 1. Query Embedding
-        query_embedding = model.encode([query])
+# Prompt templates
+templates = {
+    "general": PromptTemplate(
+        input_variables=["context", "query"],
+        template="""
+You are a strict assistant for the Virginia Geometry SOL.
 
-        # 2. Retrieval using FAISS
-        D, I = index.search(query_embedding, top_k)
+Only use exact phrases from the following SOL text:
+{context}
 
-        retrieved_docs = []
-        for i in I[0]:
-            sol_id = document_ids[i]
-            retrieved_content = next((doc["content"] for doc in documents if doc["id"] == sol_id), "Content not found.")
-            retrieved_docs.append({"id": sol_id, "content": retrieved_content})
+Answer the question: "{query}"
 
-        # 3. Context Construction
-        context = "\n\n".join([f"SOL {doc['id']}: {doc['content']}" for doc in retrieved_docs])
+If the answer is in the SOL text, quote it exactly. Do not rephrase or summarize. Do not add your own explanation.
 
-        # 4. LLM Generation
-        prompt = f"""
-Given the following information about Virginia Standards of Learning (SOLs):
-{context}
-Based on this information, answer the following question:
-{query}
-If the question is about a specific SOL number, provide a direct explanation for that SOL.
-If asked for lesson plans, worksheets, or proofs, explain what the document generally entails and whether it provides such materials.
-Be concise and to the point.
+If the answer is not in the context, reply: "The answer is not found in the provided SOL text."
 """
-        print(f"\n--- PROMPT SENT TO LLM ---\n{prompt}\n--------------------------\n")
-
-        response = llm_pipeline(prompt, max_new_tokens=500, num_return_sequences=1, do_sample=True, temperature=0.7)
-
-        generated_text = response[0]['generated_text']
+    ),
+    "lesson plan": PromptTemplate(
+        input_variables=["context", "query"],
+        template="""
+Given the following retrieved SOL text:
+{context}
 
-        print(f"\n--- RAW GENERATED TEXT ---\n{generated_text}\n--------------------------\n")
+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}
 
-        answer_start_marker = f"Based on this information, answer the following question:\n{query}"
-        if answer_start_marker in generated_text:
-            answer = generated_text.split(answer_start_marker, 1)[1].strip()
-            answer = re.sub(r'If the question is about a specific SOL number,.*?$', '', answer, flags=re.DOTALL).strip()
-        else:
-            answer = generated_text
+Create a student worksheet for: "{query}"
+Include: concept summary, a worked example, and 3 practice problems.
+"""
+    ),
+    "proofs": PromptTemplate(
+        input_variables=["context", "query"],
+        template="""
+{context}
 
-        print(f"\n--- FINAL ANSWER ---\n{answer}\n--------------------\n")
-        return answer if answer else "No valid response generated. Check logs for details."
+Generate a proof-focused geometry lesson plan for: "{query}"
+Include: student-friendly explanation, real-world link, and activity.
+"""
+    )
+}
+
+# Optional: shortcut to solve simple math problems (like area of rectangle)
+def try_math_solver(query):
+    match = re.search(r"rectangle.*l\s*=\s*(\d+).+w\s*=\s*(\d+)", query.lower())
+    if match:
+        l, w = int(match.group(1)), int(match.group(2))
+        return f"The area of the rectangle is {l} × {w} = {l * w} square units."
+    return None
+
+# RAG function
+def rag_query(query, mode="general"):
+    docs = vector_store.similarity_search(query, k=2)
+    context = "\n\n".join([doc.page_content for doc in docs])
+    prompt = templates[mode].format_prompt(context=context, query=query).to_string()
+    return llm.invoke(prompt).content
+
+# Gradio app function
+def ask_geometry_sol(query, mode):
+    math_result = try_math_solver(query)
+    if math_result:
+        return math_result
+    try:
+        return rag_query(query, mode)
     except Exception as e:
-        print(f"\n--- ERROR ---\n{str(e)}\n------------\n")
-        return f"An error occurred: {str(e)}. Please check the logs for more details."
-
-# Create Gradio interface
-demo = gr.Interface(
-    fn=retrieve_and_generate_app,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your geometry-related question here..."),
-    outputs=gr.Textbox(label="Generated Answer"),
-    title="Virginia SOL Geometry Assistant",
-    description="Ask questions about the Geometry SOL Instructional Guide"
+        return f"⚠️ Error: {type(e).__name__} - {str(e)}"
+
+# Gradio UI
+iface = gr.Interface(
+    fn=ask_geometry_sol,
+    inputs=[
+        gr.Textbox(label="Enter your Geometry SOL question or topic"),
+        gr.Radio(["general", "lesson plan", "worksheet", "proofs"], value="general", label="Response type")
+    ],
+    outputs="text",
+    title="📘 Virginia Geometry SOL Assistant",
+    description="Ask about any 2023 Geometry SOL (Standards of Learning). Get exact quotes, lesson plans, worksheets, or proof-based lessons."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch()