@@ -3,19 +3,18 @@ import time
 from datetime import timedelta
 from decimal import Decimal
 from io import BytesIO
-from typing import Any, Iterator
-
-import filetype
 from PIL import Image
 
-from messages.models import Message
-from ml_model.models import NeuronModel
 from ml_model.services.base import SimpleService
 from ml_model.tasks import openrouter_run
 from tools.chats.models import Chat
 from tools.copywrite.models import Copywrite
 from tools.public_api.models import APIStore
 
+from django.db.models.fields.files import FieldFile
+
+from messages.models import Message
+
 
 class Mistral(SimpleService):
     """
@@ -31,16 +30,19 @@ class Mistral(SimpleService):
         }
     }
 
-    def calculate_price(self, version: str, input_tokens: int, output_tokens: int) -> Decimal:
+    def calculate_price(self, version: str, input_tokens: int, output_tokens: int, image: FieldFile) -> Decimal:
         price_map = self.TOKENS_COST[version.split('/')[1]]
         price = (
-            input_tokens * price_map['input'] / 1_000_000
-            + output_tokens * price_map['output'] / 1_000_000
-            + price_map['input_imgs'] / 1_000
+                input_tokens * price_map['input'] / 1_000_000
+                + output_tokens * price_map['output'] / 1_000_000
         )
+        if image:
+            price += price_map['input_imgs'] / 1_000
         return price.quantize(Decimal('0.1'), rounding='ROUND_UP')
 
-    def save_results(self, content: Iterator[Any], t: timedelta, save: bool = True) -> list[Message]:
+    def save_results(
+            self, content: str, t: timedelta, save: bool = True
+    ) -> list[Message]:
         msgs = [
             Message(
                 content=content,
@@ -53,7 +55,6 @@ class Mistral(SimpleService):
         return msgs
 
     def make(self, input_message: Message, save: bool = True) -> list[Message]:
-        start_time = time.time()
         info = input_message.info.copy()
         version = f'mistralai/{info.pop("version")}'
         callback_data = {'provider': {'order': ['Parasail']}, **input_message.info}
@@ -72,6 +73,7 @@ class Mistral(SimpleService):
                 {'type': 'text', 'text': input_message.content},
                 {'type': 'image_url', 'image_url': {'url': image_url}},
             ]
+        start_time = time.time()
         result = openrouter_run(version, messages, callback_data, self.title)
         process_time = timedelta(seconds=(time.time() - start_time))
         self.handle_invoice(
@@ -79,17 +81,11 @@ class Mistral(SimpleService):
             version=version,
             input_tokens=result[1],
             output_tokens=result[2],
+            image=image
         )
         msgs = self.save_results(result[0], process_time)
         return msgs
 
-    def __init__(self, store):
-        super().__init__(store)
-
-    @property
-    def neuron_model(self):
-        return NeuronModel.objects.get(title='Mistral')
-
     def get_chat_history(self, message_limit: int = 10, max_character_limit: int = 1500):
         if isinstance(self.store, Chat):
             air_messages = list(
@@ -120,6 +116,5 @@ class Mistral(SimpleService):
                 memory.append({'role': 'user', 'content': content})
         character_length = sum(len(content['content']) for content in memory)
         while character_length > max_character_limit:
-            memory.pop(0)
-            character_length = sum(len(content['content']) for content in memory)
+            character_length -= len(memory.pop(0))
         return memory