Update app.py

app.py

@@ -97,11 +97,32 @@ Do not summarize. Copy the official description exactly.
 }
 
 def generate_output(prompt_type, query):
-    docs = retriever.get_relevant_documents(query)
-    context = "\n\n".join([doc.page_content for doc in docs])
+    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])
-    output = chain.run({"context": context, "query": query})
-    return output.strip()
+    return chain.run({"context": context, "query": query}).strip()
+
 
 
 # ✅ Gradio UI
@@ -112,7 +133,7 @@ with gr.Blocks() as demo:
         query = gr.Textbox(label="Enter a geometry topic")
         prompt_type = gr.Dropdown(
             ["general question", "lesson plan", "worksheet", "proofs", "flashcard"],
-            value="flashcard",
+            value="general question",
             label="Prompt Type"
         )