Update app.py

app.py

@@ -96,32 +96,27 @@ Do not summarize. Copy the official description exactly.
 
 }
 
-def generate_output(prompt_type, query):
-    import re
-
-    # Check for exact SOL code in query
+def generate_prompt_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
 
     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])
+        # 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)
+
+        context = "\n\n".join(filtered)
     else:
-        # regular query
+        # fallback to semantic retrieval
         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()
+    return chain.run({"context": context, "query": query})