Update app.py

app.py

@@ -5,9 +5,10 @@ import requests
 import inspect
 import subprocess
 import pandas as pd
-import torch, spaces
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from smolagents import CodeAgent, HfApiModel
+from smolagents.models import ChatMessage
 from huggingface_hub import InferenceClient, hf_hub_download
 
 subprocess.run(["playwright", "install"], check=True)
@@ -61,7 +62,7 @@ def check_token_access():
 # --- Basic Agent Definition ---
 # ----- THIS IS WHERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicAgent:
-    def __init__(self, model_id="meta-llama/Llama-3.1-8B-Instruct", hf_token=None):
+    def __init__(self, model_id="meta-llama/Llama-3.1-8B-Instruct", hf_token=""):
         print("BasicAgent initialized.")
         print("ENV-HF_TOKEN-LEN", len(hf_token), file=sys.stderr)
         check_token_access()
@@ -97,7 +98,7 @@ class BasicAgent:
             model=mod,
             tokenizer=tok,
             max_new_tokens=512,
-            temperature=1.0,
+            # temperature=1.0,
         )
         self.agent = CodeAgent(model=self, tools=[], add_base_tools=True)
         
@@ -136,10 +137,11 @@ class BasicAgent:
         print(f"Agent returning its generated answer: {response}")
         
         # wrap back into a chat message dict
-        return {
-            "role": 'assistant',
-            "content": [{"type": "text", "text": response}],
-        }
+        return ChatMessage(role="assistant", content=response)
+        # return {
+        #     "role": 'assistant',
+        #     "content": [{"type": "text", "text": response}],
+        # }
 
     __call__ = generate