Update app.py

app.py

@@ -115,25 +115,46 @@ class QASystem:
             graph_builder = StateGraph(MessagesState)
 
             def query_or_respond(state: MessagesState):
-                response = llm.invoke(state["messages"])
+                retrieved_docs = [m for m in state["messages"] if m.type == "tool"]
+    
+                if retrieved_docs:
+                    context = ' '.join(m.content for m in retrieved_docs)
+                else:
+                    context = "Legal knowledge system. Use Indian judiciary references."
+
+                system_prompt = (
+                    "You are a senior legal assistant with expertise in Indian law. "
+                    "Always provide legally accurate responses with references to Indian judiciary principles. "
+                    "If the user query is not legal-specific, still respond from a legal perspective."
+                    f"\n\nContext:\n{context}"
+                )
+    
+                messages = [SystemMessage(content=system_prompt)] + state["messages"]
+    
+                logger.info(f"Sending to LLM: {[m.content for m in messages]}")  # Debugging log
+
+                response = llm.invoke(messages)
                 return {"messages": [response]}
 
             def generate(state: MessagesState):
-                recent_tools = [m for m in reversed(state["messages"]) if m.type == "tool"][::-1]
-                
+                retrieved_docs = [m for m in reversed(state["messages"]) if m.type == "tool"][::-1]
+    
+                context = ' '.join(m.content for m in retrieved_docs) if retrieved_docs else "Legal knowledge system."
+
                 system_prompt = (
-                    "You are a senior legal assistant with knowledge in the Indian legal and judiciary system.indtroduce yourself for the first quiestion"
-                    " Provide direct concise summarized answers in 5 sentences based on the following context:\n\n"
-                    f"{' '.join(m.content for m in recent_tools)}"
+                    "You are a senior legal assistant specializing in Indian judiciary matters. "
+                    "Your responses MUST be legally accurate, concise (5 sentences max), and reference Indian laws when applicable."
+                    f"\n\nContext:\n{context}"
                 )
-                messages = [SystemMessage(content=system_prompt)] + [
-                    m for m in state["messages"] 
-                    if m.type in ("human", "system") or (m.type == "ai" and not m.tool_calls)
-                ]
-                
+
+                messages = [SystemMessage(content=system_prompt)] + state["messages"]
+    
+                logger.info(f"Sending to LLM: {[m.content for m in messages]}")  # Debugging log
+
                 response = llm.invoke(messages)
                 return {"messages": [response]}
 
+
             graph_builder.add_node("query_or_respond", query_or_respond)
             graph_builder.add_node("generate", generate)