add native image support

2 files changed, 56 insertions, 13 deletions

diff --git a/main.py b/main.py
index dea297d..c9d1cfa 100644
--- a/main.py
+++ b/main.py
@@ -140,7 +140,8 @@ async def guest_middleware(handler, event, data):
             replied_text = None
             
         prompt = ""
-
+        images = []
+        
         if message.document:
             if message.document.file_size > 4096:
                 if event.guest_message:
@@ -243,11 +244,9 @@ async def guest_middleware(handler, event, data):
                 file_bytes = await bot.download(photo)
                 image_bytes = file_bytes.read()
                 b64_image = base64.b64encode(image_bytes).decode("utf-8")
-                image_text = await image2text(image=b64_image)
+                images += [b64_image]
             except Exception as e:
-                image_text = "произошла ошибка при анализе фото"
                 print(f"{RED} -- photo - {e}{RESET}")
-            prompt = f"{prompt}<фото в ответе>{image_text}</фото в ответе>\n"
         
         if message.photo and IMAGE_MODEL != "":
             try:
@@ -258,11 +257,9 @@ async def guest_middleware(handler, event, data):
                 file_bytes = await bot.download(photo)
                 image_bytes = file_bytes.read()
                 b64_image = base64.b64encode(image_bytes).decode("utf-8")
-                image_text = await image2text(image=b64_image)
+                images += [b64_image]
             except Exception as e:
-                image_text = "произошла ошибка при анализе фото"
                 print(f"{RED} -- photo - {e}{RESET}")
-            prompt = f"{prompt}<фото>{image_text}</фото>\n"
 
         if event.message:
             if msg:
@@ -301,6 +298,7 @@ async def guest_middleware(handler, event, data):
                     chat_id=message.chat.id,
                     user_id=None,
                     message_id=message_id,
+                    images=images,
                 )
             elif event.message:
                 result = await generate(
@@ -309,6 +307,7 @@ async def guest_middleware(handler, event, data):
                     chat_id=message.chat.id,
                     user_id=user_id,
                     message_id=message_id,
+                    images=images,
                 )
             else:
                 result = await generate(
@@ -317,6 +316,7 @@ async def guest_middleware(handler, event, data):
                     chat_id=None,
                     user_id=user_id,
                     message_id=None,
+                    images=images,
                 )
         
             result = parse_tools(result)
@@ -410,7 +410,7 @@ async def chosen_inline_result_handler(chosen_result: ChosenInlineResult):
     else:
         current_time = time.time()
         alo, blocked = is_blocked(user_id)
-        if not alo:
+        if not alo and blocked:
             return
         if alo and blocked:
             await bot.edit_message_text(
diff --git a/supergenerator.py b/supergenerator.py
index ef31fb4..89b0902 100644
--- a/supergenerator.py
+++ b/supergenerator.py
@@ -4,6 +4,8 @@ import shelve
 import time
 from configparser import DEFAULTSECT
 
+from pprint import pprint
+
 from openai import AsyncOpenAI
 
 from statistics import *
@@ -90,6 +92,11 @@ def parse_response(text: str):
             "type": "incomplete",
             "message": "сервер петух и оборвал соединение. подожди или смени провайдера в настройках узбекгпт.",
         }
+    elif "invalid token" in text:
+        return {
+            "type": "incomplete",
+            "message": "указан неверный токен. смени провайдера или модель.",
+        }
     else:
         return {
             "type": "unknown",
@@ -179,7 +186,7 @@ def get_user_config(user_id):
 
 
 async def generate(
-    prompt: str, bot=None, user_id=None, chat_id=None, message_id=None
+    prompt: str, bot=None, user_id=None, chat_id=None, message_id=None, images=None
 ) -> str:
     if user_id:
         if user_id not in user_contexts:
@@ -191,6 +198,8 @@ async def generate(
         else:
             print(f"{RED} -- err model - failed to get prov&model in \"{config.get('model','')}\"{RESET}")
             return "произошла ошибка при парсинге вашей выбранной модели. либо хост бота гандон либо тебе надо сменить модель в настройках!!"
+
+        settings = MODELS.get(model[0]).get(model[1])
         
         stream_cf = config["stream"]
         system_prompt_name = config["prompt"]
@@ -209,7 +218,23 @@ async def generate(
         else:
             stream = False
 
-        user_contexts[user_id].append({"role": "user", "content": prompt})
+        content = []
+        
+        if images:
+            if settings.get("vision"):
+                for image in images:
+                    content += [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image}"}}]
+            else:
+                for image in images:
+                    try:
+                        image_text = await image2text(image)
+                        prompt = f"<фото>{image_text}</фото>\n{prompt}"
+                    except Exception as e:
+                        print(f"{RED} -- err photo - {e} : {config.get('model', '')}{RESET}")
+
+        content += [{"type": "text", "text": prompt}]
+        
+        user_contexts[user_id].append({"role": "user", "content": content})
         user_contexts[user_id] = user_contexts[user_id][-MAX_CONTEXT:]
 
         messages = [{"role": "system", "content": system_prompt}] + user_contexts[
@@ -217,10 +242,25 @@ async def generate(
         ]
     else:
         model = DEFAULT_MODEL.split("*", 2)
+        settings = MODELS.get(model[0]).get(model[1])
         stream = False
         stream_cf = "none"
+        content = []
+        if images:
+            if settings.get("vision"):
+                for image in images:
+                    content += [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image}"}}]
+            else:
+                for image in images:
+                    try:
+                        image_text = await image2text(image)
+                        prompt = f"<фото>{image_text}</фото>\n{prompt}"
+                    except Exception as e:
+                        print(f"{RED} -- err photo - {e} : {model}{RESET}")
+
+        content += [{"type": "text", "text": prompt}]
         messages = [{"role": "system", "content": SYSTEM_PROMPTS[DEFAULT_PROMPT]}] + [
-            {"role": "user", "content": prompt}
+            {"role": "user", "content": content}
         ]
 
     text = ""
@@ -239,7 +279,7 @@ async def generate(
         oai=client,
         bot=bot,
         model=model[1],
-        settings=MODELS.get(model[0]).get(model[1]),
+        settings=settings,
         messages=messages,
         stream=stream,
         stream_cf=stream_cf,
@@ -249,7 +289,7 @@ async def generate(
     )
 
     if user_id:
-        user_contexts[user_id].append({"role": "assistant", "content": text})
+        user_contexts[user_id].append({"role": "assistant", "content": [{"type": "text", "text": text}]})
         user_contexts[user_id] = user_contexts[user_id][-MAX_CONTEXT:]
 
     return text
@@ -290,6 +330,7 @@ async def openai_generate_text(
     message_id,
 ) -> str:
     try:
+        
         completion = await oai.chat.completions.create(
             model=model,
             messages=messages,
@@ -349,4 +390,6 @@ async def openai_generate_text(
         text, s = clean_tail_repeats(text)
         if s > 0:
             text += f"*... <{s} обрезано>*"
+        if text == "":
+            text = "(пустой ответ)"
         return text