Update app.py

app.py

@@ -11,6 +11,9 @@ import pandas as pd  # If you're working with DataFrames
 import matplotlib.figure  # If you're using matplotlib figures
 import numpy as np
 
+from custom_agent import CustomHfAgent
+
+
 # For Altair charts
 import altair as alt
 
@@ -33,19 +36,7 @@ transformers_logger = logging.getLogger("transformers.file_utils")
 transformers_logger.setLevel(logging.INFO)  # Set the desired logging level
   
 
-
- 
-
-
-
-
-
-
-
-
-
 import time  
-from transformers import load_tool, Agent
 import torch
 
 class ToolLoader:
@@ -62,39 +53,6 @@ class ToolLoader:
                 log_response(f"Error loading tool '{tool_name}': {e}")
         return loaded_tools
 
-class CustomHfAgent(Agent):
-    def __init__(self, url_endpoint, token, chat_prompt_template=None, run_prompt_template=None, additional_tools=None, input_params=None):
-        super().__init__(
-            chat_prompt_template=chat_prompt_template,
-            run_prompt_template=run_prompt_template,
-            additional_tools=additional_tools,
-        )
-        self.url_endpoint = url_endpoint
-        self.token = token
-        self.input_params = input_params
-
-    def generate_one(self, prompt, stop):
-        headers = {"Authorization": self.token}
-        max_new_tokens = self.input_params.get("max_new_tokens", 192)
-        parameters = {"max_new_tokens": max_new_tokens, "return_full_text": False, "stop": stop, "padding": True, "truncation": True}
-        inputs = {
-            "inputs": prompt,
-            "parameters": parameters,
-        }
-        response = requests.post(self.url_endpoint, json=inputs, headers=headers)
-
-        if response.status_code == 429:
-            log_response("Getting rate-limited, waiting a tiny bit before trying again.")
-            time.sleep(1)
-            return self._generate_one(prompt)
-        elif response.status_code != 200:
-            raise ValueError(f"Errors {inputs} {response.status_code}: {response.json()}")
-        log_response(response)
-        result = response.json()[0]["generated_text"]
-        for stop_seq in stop:
-            if result.endswith(stop_seq):
-                return result[: -len(stop_seq)]
-        return result
 
 def handle_submission(user_message, selected_tools, url_endpoint):