Update app.py

app.py

@@ -98,27 +98,30 @@ Do not summarize. Only copy from the context.
 
 }
 
-def generate_output(prompt_type, query, retriever, llm):
-    # Try to extract SOL code
-    sol_match = re.search(r"\bG\.[A-Z]+\.\d+\b", query)
-    matched_code = sol_match.group(0) if sol_match else None
+def generate_prompt_output(prompt_type, query, retriever, llm):
+    try:
+        import re
+        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:
-        # Retrieve and filter by metadata
-        all_docs = retriever.vectorstore._collection.get(include=['documents', 'metadatas'])
-        filtered = []
-        for doc_text, metadata in zip(all_docs['documents'], all_docs['metadatas']):
-            if metadata.get('standard') == matched_code:
-                filtered.append(doc_text)
+        if matched_code:
+            all_docs = retriever.vectorstore._collection.get(include=['documents', 'metadatas'])
+            filtered = []
+            for doc_text, metadata in zip(all_docs['documents'], all_docs['metadatas']):
+                if metadata.get('standard') == matched_code:
+                    filtered.append(doc_text)
 
-        context = "\n\n".join(filtered)
-    else:
-        # fallback to semantic retrieval
-        docs = retriever.get_relevant_documents(query)
-        context = "\n\n".join([doc.page_content for doc in docs])
+            context = "\n\n".join(filtered)
+        else:
+            docs = retriever.get_relevant_documents(query)
+            context = "\n\n".join([doc.page_content for doc in docs])
+
+        chain = LLMChain(llm=llm, prompt=templates[prompt_type])
+        return chain.run({"context": context, "query": query}).strip()
+
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
 
-    chain = LLMChain(llm=llm, prompt=templates[prompt_type])
-    return chain.run({"context": context, "query": query})
 
 
 
@@ -137,6 +140,6 @@ with gr.Blocks() as demo:
     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)
+    btn.click(fn=generate_prompt_output, inputs=[prompt_type, query], outputs=output)
 
 demo.launch()