[0.4.12] [fix] Threads API

Author: yashbonde

docs/changelog.rst

@@ -0,0 +1,9 @@
+tuneapi Changelog
+=================
+
+0.4.12
+------
+
+- Fix bug in Threads API where incorrect structure was sent by client
+- Add images support for Anthropic API
+- Add ``Message.images`` field to store all images

docs/conf.py

@@ -7,13 +7,15 @@
 
 sys.path.insert(0, os.path.abspath("../"))
 
+from tuneapi import __version__
+
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
 project = "tuneapi"
 copyright = "2024, Frello Technologies"
 author = "Frello Technologies"
-release = "0.4.11"
+release = __version__
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

docs/index.rst

@@ -15,4 +15,5 @@ TuneAPI
    source/tuneapi.apis
    source/tuneapi.types
    source/tuneapi.utils
+   changelog
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tuneapi"
-version = "0.4.11"
+version = "0.4.12"
 description = "Tune AI APIs."
 authors = ["Frello Technology Private Limited <[email protected]>"]
 license = "MIT"
@@ -17,6 +17,7 @@ cryptography = ">=42.0.5"
 tqdm = "^4.66.1"
 snowflake_id = "1.0.2"
 nutree = "0.8.0"
+pillow = "^10.2.0"
 boto3 = { version = "1.29.6", optional = true }
 
 [tool.poetry.extras]

tuneapi/__init__.py

@@ -1,3 +1,3 @@
 # Copyright © 2023- Frello Technology Private Limited
 
-__version__ = "0.4.11"
+__version__ = "0.4.12"

tuneapi/apis/model_anthropic.py

@@ -1,5 +1,6 @@
 # Copyright © 2024- Frello Technology Private Limited
 
+import re
 import json
 import requests
 from typing import Optional, Dict, Any, Tuple, List
@@ -71,60 +72,55 @@ def _process_input(self, chats, token: Optional[str] = None):
         prev_tool_id = tu.get_random_string(5)
         for m in thread.chats[int(system != "") :]:
             if m.role == Message.HUMAN:
-                claude_messages.append(
-                    {
-                        "role": "user",
-                        "content": [
+                msg = {
+                    "role": "user",
+                    "content": [{"type": "text", "text": m.value.strip()}],
+                }
+                if m.images:
+                    for i in m.images:
+                        msg["content"].append(
                             {
-                                "type": "text",
-                                "text": m.value.strip(),
+                                "type": "image",
+                                "source": {
+                                    "type": "base64",
+                                    "media_type": "image/png",
+                                    "data": i,
+                                },
                             }
-                        ],
-                    }
-                )
+                        )
             elif m.role == Message.GPT:
-                claude_messages.append(
-                    {
-                        "role": "assistant",
-                        "content": [
-                            {
-                                "type": "text",
-                                "text": m.value.strip(),
-                            }
-                        ],
-                    }
-                )
+                msg = {
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": m.value.strip()}],
+                }
             elif m.role == Message.FUNCTION_CALL:
                 _m = tu.from_json(m.value) if isinstance(m.value, str) else m.value
-                claude_messages.append(
-                    {
-                        "role": "assistant",
-                        "content": [
-                            {
-                                "type": "tool_use",
-                                "id": prev_tool_id,
-                                "name": _m["name"],
-                                "input": _m["arguments"],
-                            }
-                        ],
-                    }
-                )
+                msg = {
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "type": "tool_use",
+                            "id": prev_tool_id,
+                            "name": _m["name"],
+                            "input": _m["arguments"],
+                        }
+                    ],
+                }
             elif m.role == Message.FUNCTION_RESP:
                 _m = tu.from_json(m.value) if isinstance(m.value, str) else m.value
-                claude_messages.append(
-                    {
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "tool_result",
-                                "tool_use_id": prev_tool_id,
-                                "content": tu.to_json(_m, tight=True),
-                            }
-                        ],
-                    }
-                )
+                msg = {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "tool_result",
+                            "tool_use_id": prev_tool_id,
+                            "content": tu.to_json(_m, tight=True),
+                        }
+                    ],
+                }
             else:
                 raise Exception(f"Unknown role: {m.role}")
+            claude_messages.append(msg)
 
         headers = {
             "x-api-key": token,
@@ -255,3 +251,15 @@ def stream_chat(
             except:
                 break
         return
+
+
+# helper methods
+
+
+def get_section(tag: str, out: str) -> Optional[str]:
+    pattern = re.compile("<" + tag + ">(.*?)</" + tag + ">", re.DOTALL)
+    match = pattern.search(out)
+    if match:
+        content = match.group(1)
+        return content
+    return None

tuneapi/apis/threads.py

@@ -18,7 +18,6 @@ def get_sub(
     sess.headers.update({"x-tune-key": tune_api_key})
     if tune_org_id:
         sess.headers.update({"X-Organization-Id": tune_org_id})
-    print("::::::", sess.headers)
     return tu.Subway(base_url, sess)
 
 
@@ -65,6 +64,7 @@ def put_thread(
             body["metadata"] = {"meta": tu.to_json(m, tight=True)}
 
         tu.logger.info("Creating new thread")
+        print(tu.to_json(body))
         out = self.sub.threads("post", json=body)
         thread.id = out["id"]
         return thread

tuneapi/types/chats.py

@@ -5,6 +5,7 @@
 import json
 import copy
 import random
+from PIL.Image import Image
 from functools import partial
 from collections.abc import Iterable
 from typing import Dict, List, Any, Tuple, Optional, Generator, Union
@@ -128,8 +129,8 @@ def __init__(
         self,
         value: str | float | List[Dict[str, Any]],
         role: str,
+        images: List[str | Image] = [],
         id: str = None,
-        fn_pairs: Optional[Tuple["Message", "Message"]] = None,
         **kwargs,
     ):
         if role not in self.KNOWN_ROLES:
@@ -141,7 +142,12 @@ def __init__(
         self.value = value
         self.id = id or "msg_" + str(tu.get_snowflake())
         self.metadata = kwargs
-        self.fn_pairs = fn_pairs
+        self.images = images
+        for i, img in enumerate(self.images):
+            if isinstance(img, Image):
+                buf = io.BytesIO()
+                img.save(buf, "png")
+                self.images[i] = tu.to_b64(buf.getvalue())
 
         # validations
         if self.role == self.FUNCTION_CALL:
@@ -164,14 +170,7 @@ def __add__(self, other: str):
         return Message(self.value + other, self.role)
 
     def __repr__(self) -> str:
-        out = ""
-        if self.fn_pairs:
-            for fc, fr in self.fn_pairs:
-                out += f"[[FC] {fc} => [FR] {fr}]"
-        if out:
-            out += " " + str(self.value)
-        else:
-            out = str(self.value)
+        out = str(self.value)
         return out
 
     def __getitem__(self, x):
@@ -190,13 +189,11 @@ def to_dict(
         meta: bool = False,
     ):
         """
-        if format == ``ft`` then export to following format: ``{"from": "system/human/gpt", "value": "..."}``
-
-        elif format == ``api`` then ``{"role": "system/user/assistant", "content": [{"type": "text", "text": {"value": "..."}]}``
-
-        elif format == ``full`` then ``{"id": 1234421123, "role": "system/user/assistant", "content": [{"type": "text", "text": {"value": "..."}]}``
-
-        else export to following format: ``{"role": "system/user/assistant", "content": "..."}``
+        Serialise the Message into a dictionary of different formats:
+            - format == ``ft`` then export to following format: ``{"from": "system/human/gpt", "value": "..."}``
+            - format == ``api`` then ``{"role": "system/user/assistant", "content": [{"type": "text", "text": {"value": "..."}]}``. This is used with TuneAPI
+            - format == ``full`` then ``{"id": 1234421123, "role": "system/user/assistant", "content": [{"type": "text", "text": {"value": "..."}]}``
+            - default: ``{"role": "system/user/assistant", "content": "..."}``
         """
         role = self.role
 
@@ -219,7 +216,7 @@ def to_dict(
         if ft:
             chat_message["value"] = self.value
         elif api:
-            chat_message["content"] = [{"type": "text", "text": {"value": self.value}}]
+            chat_message["content"] = [{"type": "text", "text": self.value}]
         else:
             chat_message["content"] = self.value
 

tuneapi/utils/fs.py

@@ -40,7 +40,6 @@ def get_files_in_folder(
 
             for f in files:
                 for e in ext:
-                    print(e)
                     if f.endswith(e):
                         _fp = joinp(root, f)
                         _ig_pat = ignore_pat.search(_fp)

tuneapi/utils/subway.py

@@ -7,6 +7,7 @@
 
 from tuneapi.utils.logger import logger
 from tuneapi.utils.misc import SimplerTimes
+from tuneapi.utils.serdeser import to_json
 
 
 class SubwayClientError(Exception):
@@ -85,7 +86,7 @@ def u(self, attr: str) -> "Subway":
     def _renew_session(self):
         """Renew the session"""
         _session = Session()
-        if "token" in self._session.headers:
+        if self._session.headers:
             _session.headers.update(**self._session.headers)
         self._session = _session
 
@@ -152,7 +153,6 @@ def __call__(
 
         if SimplerTimes.get_now_i64() - self._last_update > 120:
             self._renew_session()
-        print(f"{url=}, {items=}")
         r = fn(url, **items, **kwargs)
         if 399 < r.status_code < 500:
             raise SubwayClientError(r.content.decode(), code=r.status_code)