Update agent.py

agent.py

@@ -18,7 +18,7 @@ qa_dict = {
 def build_graph():
     llm = HuggingFaceEndpoint(
         repo_id="mistralai/Mistral-7B-Instruct-v0.3",
-        task="text-generation",
+        # task="text-generation",
         huggingfacehub_api_token=os.environ["HF_TOKEN"]
     )
 
@@ -44,7 +44,15 @@ def build_graph():
         )
 
         full_prompt = f"{prompt}\n\nQuestion: {query}\nAnswer:"
-        response = llm.invoke(full_prompt).strip()
+        
+
+        chat_input = [
+            {"role": "system", "content": prompt},
+            {"role": "user", "content": query}
+        ]
+        response = llm.invoke(chat_input).strip()
+        
+        # response = llm.invoke(full_prompt).strip()
         print(f"🧠 LLM said: {response}")
 
         if response.startswith("use_tool:"):