Update agent.py

agent.py

@@ -17,7 +17,7 @@ def build_graph():
     # Initialize Mistral model
     llm = HuggingFaceEndpoint(
         repo_id="mistralai/Mistral-7B-Instruct-v0.3",
-        # task="conversation",
+        task="conversational",
         huggingfacehub_api_token=os.environ["HF_TOKEN"]
     )
 
@@ -43,12 +43,20 @@ def build_graph():
             "- Never justify or explain"
         )
 
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": query},
-        ]
+        # prompt = [
+        #     {"role": "system", "content": system_prompt},
+        #     {"role": "user", "content": query},
+        # ]
+        
+        prompt = f"{system_prompt}\n\nUser: {query}\nAssistant:"
+        response = llm.invoke(prompt).strip()
 
-        response = llm.invoke(messages)
+        # response = llm.invoke(messages)
+        
+        for tag in ("Final answer:", "Answer:", "assistant:"):
+            if response.lower().startswith(tag.lower()):
+                response = response[len(tag):].strip()
+        
         return {"messages": [AIMessage(content=response.strip())]}
 
     # Tool node