Made the agent conversational

agent.py

@@ -1,28 +1,76 @@
 import os
-from PIL import Image
-from smolagents import CodeAgent, HfApiModel, InferenceClientModel
+from dotenv import load_dotenv
+from huggingface_hub import InferenceClient
+from smolagents import CodeAgent, Model, ChatMessage
+import tools.tools as tls  # Your tool definitions
 
-import tools.tools as tls
+load_dotenv()
 
-# --- Basic Agent Definition ---
-# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
+# ✅ Utility: Ensure role alternation (user/assistant/user...)
+def enforce_strict_role_alternation(messages):
+    """
+    Fixes message history to enforce 'user/assistant/user/assistant' alternation.
+    Keeps first 'system' if present.
+    """
+    cleaned = []
+    last_role = None
 
+    for msg in messages:
+        role = msg["role"]
+        if role not in ("user", "assistant", "system"):
+            continue  # skip invalid roles
+
+        if role == "system" and not cleaned:
+            cleaned.append(msg)
+            continue
+
+        if role == last_role:
+            continue  # skip consecutive same-role messages
+
+        cleaned.append(msg)
+        last_role = role
+
+    return cleaned
+
+
+# ✅ Custom Model Wrapper using InferenceClient
+class HuggingFaceChatModel(Model):
+    def __init__(self):
+        model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+        self.client = InferenceClient(model=model_id, token=os.getenv("HF_TOKEN"))
+
+    def generate(self, messages, stop_sequences=None):
+        if stop_sequences is None:
+            stop_sequences = ["Task"]
+
+        # 💡 Enforce correct message order
+        cleaned_messages = enforce_strict_role_alternation(messages)
+
+        # 🔧 Hugging Face call
+        response = self.client.chat_completion(
+            messages=cleaned_messages,
+            stop=stop_sequences,
+            max_tokens=1024
+        )
+        content = response.choices[0].message["content"]
+        return ChatMessage(role="assistant", content=content)
+
+
+# ✅ Basic Agent with SmolAgents
 class BasicAgent:
     def __init__(self):
-        print("BasicAgent initialized.")
-    def __call__(self, question: str) -> str:
-        #model = HfApiModel(model_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud", provider="together", use_auth_token=True)
+        print("✅ BasicAgent initialized with Hugging Face chat model.")
+        self.model = HuggingFaceChatModel()
 
-        agent = CodeAgent(
+        self.agent = CodeAgent(
             tools=[tls.search_tool, tls.calculate_cargo_travel_time],
-            model=InferenceClientModel(),
+            model=self.model,
             additional_authorized_imports=["pandas"],
             max_steps=20,
         )
 
-        fixed_answer = agent.run(question)
-
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        # fixed_answer = "This is a default answer."
-        print(f"Agent returning fixed answer: {fixed_answer}")
-        return str(fixed_answer)
+    def __call__(self, question: str) -> str:
+        print(f"📥 Received question: {question[:60]}...")
+        response = self.agent.run(question)
+        print(f"📤 Response generated: {response[:60]}...")
+        return response

app.py

@@ -6,7 +6,7 @@ import pandas as pd
 from huggingface_hub import login
 from dotenv import load_dotenv
 
-import agent
+from agent import BasicAgent
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -134,7 +134,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         results_df = pd.DataFrame(results_log)
         return status_message, results_df
 '''
-
+'''
 def test_init_agent_for_chat(text_input, history):
     # 1. Instantiate Agent ( modify this part to create your agent)
     try:
@@ -146,6 +146,18 @@ def test_init_agent_for_chat(text_input, history):
     submitted_answer = basicAgent(text_input)
 
     return submitted_answer
+'''
+def test_init_agent_for_chat(text_input, history):
+    try:
+        basicAgent = BasicAgent()
+    except Exception as e:
+        return f"[Error initializing agent]: {e}"
+
+    return basicAgent(text_input)
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🤖 Conversational Cargo Agent")
+    gr.ChatInterface(test_init_agent_for_chat, type="messages")
 
 # --- Build Gradio Interface using Blocks ---
 with gr.Blocks() as demo: