update foramt of answers

agent_api.py

@@ -56,8 +56,6 @@ class Message(BaseModel):
 
 class ChatRequest(BaseModel):
     message: str = Field(..., description="User's question")
-    role: str = Field(..., description="User role: 'client' or 'lawyer'")
-    jurisdiction: str = Field(..., description="Selected jurisdiction")
     conversationHistory: Optional[List[Message]] = Field(default=[], description="Previous conversation messages")
 
 class ChatResponse(BaseModel):
@@ -104,12 +102,9 @@ class CyberLegalAPI:
             })
         
         try:
-            # Create enhanced query with context
-            enhanced_query = self._create_enhanced_query(request)
-            
-            # Process through agent
+            # Process through agent with raw message and conversation history
             result = await self.agent.process_query(
-                user_query=enhanced_query,
+                user_query=request.message,
                 conversation_history=conversation_history
             )
             
@@ -131,34 +126,6 @@ class CyberLegalAPI:
                 detail=f"Processing failed: {str(e)}"
             )
     
-    def _create_enhanced_query(self, request: ChatRequest) -> str:
-        """
-        Create enhanced query with role and jurisdiction context
-        """
-        base_query = request.message
-        
-        # Add role context
-        role_context = ""
-        if request.role == "client":
-            role_context = "Answer from the perspective of advising a client who needs practical guidance."
-        elif request.role == "lawyer":
-            role_context = "Answer from the perspective of providing legal analysis for a legal professional."
-        
-        # Add jurisdiction context
-        jurisdiction_context = f"Focus on the legal framework in {request.jurisdiction}."
-        
-        # Combine into enhanced query
-        enhanced_query = f"""{base_query}
-
-Context:
-- User Role: {request.role}
-- Jurisdiction: {request.jurisdiction}
-- Special Instructions: {role_context} {jurisdiction_context}
-
-Please provide a response tailored to this context."""
-        
-        return enhanced_query
-    
     async def health_check(self) -> HealthResponse:
         """
         Check health status of the API and dependencies
@@ -234,10 +201,6 @@ async def root():
             "chat": "POST /chat - Chat with the assistant",
             "health": "GET /health - Health check"
         },
-        "supported_jurisdictions": [
-            "EU", "France", "Germany", "Italy", "Spain", "Romania", "Netherlands", "Belgium"
-        ],
-        "user_roles": ["client", "lawyer"],
         "expertise": [
             "GDPR", "NIS2", "DORA", "Cyber Resilience Act", "eIDAS 2.0"
         ]
@@ -266,4 +229,4 @@ if __name__ == "__main__":
         port=port,
         reload=False,
         log_level="info"
-    )
+    )

langraph_agent.py

@@ -8,7 +8,7 @@ from typing import Dict, Any, List, Optional
 from datetime import datetime
 from langgraph.graph import StateGraph, END
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.messages import HumanMessage, SystemMessage, AssistantMessage
 
 from agent_state import AgentState, ConversationManager
 from prompts import SYSTEM_PROMPT, ERROR_HANDLING_PROMPT
@@ -127,7 +127,27 @@ class CyberLegalAgent:
                 state["final_response"] = "I apologize, but I couldn't find relevant information for your query."
                 return state
             
-            # Create prompt for LLM to answer based on retrieved context
+            # Build message stack with conversation history
+            messages = []
+            
+            # Add conversation history messages
+            history = state.get("conversation_history", [])
+            has_system_message = False
+            
+            for msg in history:
+                if msg["role"] == "system":
+                    messages.append(SystemMessage(content=msg["content"]))
+                    has_system_message = True
+                elif msg["role"] == "user":
+                    messages.append(HumanMessage(content=msg["content"]))
+                elif msg["role"] == "assistant":
+                    messages.append(AssistantMessage(content=msg["content"]))
+            
+            # Add system message only if not present in history
+            if not has_system_message:
+                messages.insert(0, SystemMessage(content=SYSTEM_PROMPT))
+            
+            # Add final message with user query and LightRAG context
             answer_prompt = f"""Based on the following retrieved legal information, please answer the user's question accurately and comprehensively.
 
 **User Question:** {state["user_query"]}
@@ -143,12 +163,8 @@ class CyberLegalAgent:
 5. Add a disclaimer that this is for guidance purposes only
 
 Please provide a clear, well-structured response."""
-
-            # Get answer from LLM
-            messages = [
-                SystemMessage(content=SYSTEM_PROMPT),
-                HumanMessage(content=answer_prompt)
-            ]
+            
+            messages.append(HumanMessage(content=answer_prompt))
             
             response = await self.llm.ainvoke(messages)
             answer = response.content