Update dspy_wrapper.py

dspy_wrapper.py

@@ -1,3 +1,4 @@
+from query_preprocessing import preprocess_query
 import dspy 
 from typing import List, Dict
 import os
@@ -8,7 +9,6 @@ if not OPENAI_API_KEY:
     raise RuntimeError("Missing OPENAI_API_KEY env var")
 
 gpt_4o_mini = dspy.LM('openai/gpt-4o-mini', api_key=OPENAI_API_KEY)
-# using unimib credentials, switch to PeS if needed!
 dspy.configure(lm=gpt_4o_mini)
 
 
@@ -19,7 +19,7 @@ class DSPyHybridRetriever(dspy.Module):
         self.retriever = retriever
 
     def forward(self, query: str, municipality: str = "", top_k: int = 5):
-        results = self.retriever.rerank(query, top_k=top_k, municipality_filter=municipality) # remember to change to rerank 
+        results = self.retriever.rerank(query, top_k=top_k, municipality_filter=municipality)
         return {"retrieved_chunks": results}
     
 class RetrieveChunks(dspy.Signature):
@@ -37,7 +37,7 @@ class RetrieveChunks(dspy.Signature):
 
 class AnswerWithEvidence(dspy.Signature):
     """Answer the query using reasoning and retrieved chunks as context."""
-    query = dspy.InputField(desc="User's question")
+    query = dspy.InputField(desc="Rewritten question")
     retrieved_chunks = dspy.InputField(desc="Retrieved text chunks (List[dict])")
     answer = dspy.OutputField(desc="Final answer")
     rationale = dspy.OutputField(desc="Chain-of-thought reasoning")
@@ -50,21 +50,29 @@ class RAGChain(dspy.Module):
         self.retriever = retriever
         self.answerer = answerer
 
-    def forward(self, query: str, municipality: str = ""):
-        # retrieve full dicts
-        retrieved = self.retriever(query=query, municipality=municipality)
+    def forward(self, raw_query: str, municipality: str = "", feedback: str = ""):
+        pre = preprocess_query(raw_query, feedback)
+        rewritten = pre["rewritten_query"] or raw_query
+        extracted_muni = pre["municipality"] or ""
+        intent = pre["intent"]
+        muni = municipality if municipality.strip() else extracted_muni
+
+        retrieved = self.retriever(query=rewritten, municipality=muni)
         chunks = retrieved["retrieved_chunks"]
 
-        # feed only the raw text into the CoT module
+        # Answer + CoT using the rewritten query
         answer_result = self.answerer(
-            query=query,
+            query=rewritten,
             retrieved_chunks=[c["chunk_text"] for c in chunks]
         )
 
-        # return both the metadata and the LLM answer
+        # Return everything for transparency & downstream use
         return {
-            "query": query,
-            "municipality": municipality,
+            "original_query": raw_query,
+            "intent": intent,
+            "rewritten_query": rewritten,
+            "llm_municipality": extracted_muni,
+            "municipality": muni,
             "retrieved_chunks": chunks,
             "chain_of_thought": answer_result.rationale,
             "final_answer": answer_result.answer,