Update text_generator.py

text_generator.py

@@ -16,15 +16,17 @@ class TextGenerationTool(Tool):
     def __call__(self, prompt: str):
         API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-v0.1"
         headers = {"Authorization":  os.environ['hf']}
+
+        payload = {
+            "inputs": prompt  # Adjust this based on your model's input format
+        }
         
-        def query(payload):
-            response = requests.post(API_URL, headers=headers, json=payload)
-            print(response)
-            return response.json()
+        #def query(payload):
+        response = requests.post(API_URL, headers=headers, json=payload)
+        print(response)
+        return response.json()
             
-        output = query({
-            "inputs": prompt,
-        })
+
         
         # Define the payload for the request
         #payload = {