add native image support

1 files changed, 48 insertions, 5 deletions

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