pygpt-net 2.1.37__py3-none-any.whl → 2.4.48__py3-none-any.whl

pygpt_net/provider/gpt/vision.py

@@ -6,13 +6,21 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2024.03.12 21:00:00                  #
+# Updated Date: 2024.12.14 22:00:00                  #
 # ================================================== #
 
 import base64
-import json
 import os
 import re
+from typing import Optional, Union, Dict, Any, List
+
+from pygpt_net.core.types import (
+    MODE_VISION,
+)
+from pygpt_net.core.bridge.context import BridgeContext
+from pygpt_net.item.attachment import AttachmentItem
+from pygpt_net.item.ctx import CtxItem
+from pygpt_net.item.model import ModelItem
 
 
 class Vision:
@@ -27,95 +35,127 @@ class Vision:
         self.urls = []
         self.input_tokens = 0
 
-    def send(self, **kwargs):
+    def send(
+            self,
+            context: BridgeContext,
+            extra: Optional[Dict[str, Any]] = None
+    ):
         """
         Call OpenAI API for chat with vision
 
-        :param kwargs: keyword arguments
+        :param context: Bridge context
+        :param extra: Extra arguments
         :return: response or stream chunks
         """
-        # get kwargs
-        prompt = kwargs.get("prompt", "")
-        stream = kwargs.get("stream", False)
-        max_tokens = kwargs.get("max_tokens", 200)
-        system_prompt = kwargs.get("system_prompt", "")
-        attachments = kwargs.get("attachments", {})
-        model = kwargs.get("model", None)
+        prompt = context.prompt
+        stream = context.stream
+        max_tokens = int(context.max_tokens or 0)
+        system_prompt = context.system_prompt
+        attachments = context.attachments
+        model = context.model
         model_id = model.id
         client = self.window.core.gpt.get_client()
 
+        # extra API kwargs
+        response_kwargs = {}
+        if max_tokens > 0:
+            response_kwargs['max_tokens'] = max_tokens
+
         # build chat messages
         messages = self.build(
             prompt=prompt,
             system_prompt=system_prompt,
-            attachments=attachments,
             model=model,
+            history=context.history,
+            attachments=attachments,
         )
         response = client.chat.completions.create(
             messages=messages,
             model=model_id,
-            max_tokens=int(max_tokens),
             stream=stream,
+            **response_kwargs
         )
+
         return response
 
-    def build(self, **kwargs) -> list:
+    def build(
+            self,
+            prompt: str,
+            system_prompt: str,
+            model: ModelItem,
+            history: Optional[List[CtxItem]] = None,
+            attachments: Optional[Dict[str, AttachmentItem]] = None,
+    ) -> List[dict]:
         """
         Build chat messages list
 
-        :param kwargs: keyword arguments
+        :param prompt: user prompt
+        :param system_prompt: system prompt
+        :param model: model item
+        :param history: history
+        :param attachments: attachments
         :return: messages list
         """
-        prompt = kwargs.get("prompt", "")
-        system_prompt = kwargs.get("system_prompt", None)
-        attachments = kwargs.get("attachments", {})
-        model = kwargs.get("model", None)
         messages = []
 
         # tokens config
-        mode = 'vision'
+        mode = MODE_VISION
         used_tokens = self.window.core.tokens.from_user(
             prompt,
             system_prompt,
         )  # threshold and extra included
-        max_tokens = self.window.core.config.get('max_total_tokens')
+        max_ctx_tokens = self.window.core.config.get('max_total_tokens')
 
         # fit to max model tokens
-        if max_tokens > model.ctx:
-            max_tokens = model.ctx
+        if max_ctx_tokens > model.ctx:
+            max_ctx_tokens = model.ctx
 
         # input tokens: reset
         self.reset_tokens()
 
         # append initial (system) message
         if system_prompt is not None and system_prompt != "":
-            messages.append({"role": "system", "content": system_prompt})
+            messages.append({
+                "role": "system",
+                "content": system_prompt,
+            })
         else:
             if system_prompt is not None and system_prompt != "":
-                messages.append({"role": "system", "content": system_prompt})
+                messages.append({
+                    "role": "system",
+                    "content": system_prompt,
+                })
 
         # append messages from context (memory)
         if self.window.core.config.get('use_context'):
-            items = self.window.core.ctx.get_prompt_items(
+            items = self.window.core.ctx.get_history(
+                history,
                 model.id,
                 mode,
                 used_tokens,
-                max_tokens,
+                max_ctx_tokens,
             )
             for item in items:
                 # input
-                if item.input is not None and item.input != "":
-                    content = self.build_content(item.input)
-                    messages.append({"role": "user", "content": content})
+                if item.final_input is not None and item.final_input != "":
+                    messages.append({
+                        "role": "user",
+                        "content": item.final_input,
+                    })
 
                 # output
-                if item.output is not None and item.output != "":
-                    content = self.build_content(item.output)
-                    messages.append({"role": "assistant", "content": content})
+                if item.final_output is not None and item.final_output != "":
+                    messages.append({
+                        "role": "assistant",
+                        "content": item.final_output,
+                    })
 
         # append current prompt
         content = self.build_content(prompt, attachments)
-        messages.append({"role": "user", "content": content})
+        messages.append({
+            "role": "user",
+            "content": content,
+        })
 
         # input tokens: update
         self.input_tokens += self.window.core.tokens.from_messages(
@@ -124,20 +164,27 @@ class Vision:
         )
         return messages
 
-    def build_content(self, prompt: str, attachments: dict = None) -> list:
+    def build_content(
+            self,
+            content: Union[str, list],
+            attachments: Optional[Dict[str, AttachmentItem]] = None,
+    ) -> List[dict]:
         """
-        Build vision contents
+        Build vision content
 
-        :param prompt: prompt (user input)
+        :param content: content (str or list)
         :param attachments: attachments (dict, optional)
         :return: List of contents
         """
-        content = [
-            {
-                "type": "text",
-                "text": str(prompt)
-            }
-        ]
+        if not isinstance(content, list):
+            content = [
+                {
+                    "type": "text",
+                    "text": str(content)
+                }
+            ]
+
+        prompt = content[0]['text']
 
         self.attachments = {}  # reset attachments, only current prompt
         self.urls = []
@@ -173,10 +220,11 @@ class Vision:
                             }
                         )
                         self.attachments[id] = attachment.path
+                        attachment.consumed = True
 
         return content
 
-    def extract_urls(self, text: str) -> list:
+    def extract_urls(self, text: str) -> List[str]:
         """
         Extract img urls from text
 
@@ -218,7 +266,7 @@ class Vision:
         """Reset input tokens counter"""
         self.input_tokens = 0
 
-    def get_attachments(self) -> dict:
+    def get_attachments(self) -> Dict[str, str]:
         """
         Get attachments
 
@@ -226,7 +274,7 @@ class Vision:
         """
         return self.attachments
 
-    def get_urls(self) -> list:
+    def get_urls(self) -> List[str]:
         """
         Get urls
 
@@ -241,3 +289,19 @@ class Vision:
         :return: input tokens
         """
         return self.input_tokens
+
+    def append_images(self, ctx: CtxItem):
+        """
+        Append images content to context item
+
+        :param ctx: context
+        """
+        images = self.get_attachments()  # dict -> key: id, value: path
+        urls = self.get_urls()  # list
+
+        # store sent images in ctx
+        if len(images) > 0:
+            ctx.images = self.window.core.filesystem.make_local_list(list(images.values()))
+        if len(urls) > 0:
+            ctx.images = urls
+            ctx.urls = urls