Update app.py

app.py

@@ -12,6 +12,9 @@ from Gradio_UI import GradioUI
 # Define the Perplexity system prompt
 Perplex_Assistant_Prompt = """You are a helpful AI assistant that searches the web for accurate information."""
 
+# Set up API key in environment variable as expected by HfApiModel
+os.environ["HUGGINGFACE_API_TOKEN"] = os.getenv("HUGGINGFACE_API_KEY", "")
+
 #@weave.op()
 def tracked_perplexity_call(prompt: str, system_messages: str, model_name: str = "sonar-pro", assistant_meta: bool = False):
     """Enhanced Perplexity API call with explicit model tracking."""
@@ -62,15 +65,12 @@ def get_current_time_in_timezone(timezone: str) -> str:
 
 final_answer = FinalAnswerTool()
 
-# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
-# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 
-
+# Remove the huggingface_api_key parameter - it's not supported
 model = HfApiModel(
     max_tokens=2096,
     temperature=0.5,
     model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',  # Using the backup endpoint
-    custom_role_conversions=None,
-    huggingface_api_key=os.getenv("HUGGINGFACE_API_KEY")  # Add this line to provide authentication
+    custom_role_conversions=None
 )
 
 # Import tool from Hub
@@ -81,7 +81,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, Sonar_Web_Search_Tool, get_current_time_in_timezone],
+    tools=[final_answer, Sonar_Web_Search_Tool, get_current_time_in_timezone, image_generation_tool],  # Added the image tool
     max_steps=6,
     verbosity_level=1,
     grammar=None,
@@ -91,6 +91,16 @@ agent = CodeAgent(
     prompt_templates=prompt_templates
 )
 
-# To fix the TypeError, modify the Gradio_UI.py file to handle potential None values
-# But for now, we can just use the agent directly
-GradioUI(agent).launch()
+# To fix the TypeError in Gradio_UI.py, you would need to modify that file
+# For now, we'll just use the agent directly
+try:
+    GradioUI(agent).launch()
+except TypeError as e:
+    if "unsupported operand type(s) for +=" in str(e):
+        print("Error: Token counting issue in Gradio UI")
+        print("To fix, edit Gradio_UI.py and change:")
+        print("total_input_tokens += agent.model.last_input_token_count")
+        print("To:")
+        print("total_input_tokens += (agent.model.last_input_token_count or 0)")
+    else:
+        raise e