[0.5.4] standardise model interface

Author: yashbonde

cookbooks/baby_server.py

@@ -77,7 +77,7 @@ async def chat_completions(request: Request, data: ChatCompletionRequest):
             max_tokens=data.max_tokens,
         )
         api_resp = tu.generator_to_api_events(
-            model=model.tune_model_id,
+            model=model.model_id,
             generator=stream_resp,
         )
         return StreamingResponse(api_resp, media_type="text/event-stream")

docs/changelog.rst

@@ -7,6 +7,11 @@ minor versions.
 
 All relevant steps to be taken will be mentioned here.
 
+0.5.4
+-----
+
+- Standardise ``tuneapi.types.chats.ModelInterface`` to have ``model_id``, ``api_token`` added to the base class.
+
 0.5.3
 -----
 

docs/conf.py

@@ -13,7 +13,7 @@
 project = "tuneapi"
 copyright = "2024, Frello Technologies"
 author = "Frello Technologies"
-release = "0.5.3"
+release = "0.5.4"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tuneapi"
-version = "0.5.3"
+version = "0.5.4"
 description = "Tune AI APIs."
 authors = ["Frello Technology Private Limited <[email protected]>"]
 license = "MIT"

tuneapi/__main__.py

@@ -7,10 +7,3 @@
 from tuneapi.apis.model_groq import Groq
 from tuneapi.apis.model_mistral import Mistral
 from tuneapi.apis.model_gemini import Gemini
-
-
-# other imports
-import os
-import random
-from time import time
-from typing import List, Optional

tuneapi/apis/__init__.py

@@ -10,75 +10,44 @@
 from typing import Optional, Dict, Any, Tuple, List
 
 import tuneapi.utils as tu
-from tuneapi.types import Thread, human, Message
+import tuneapi.types as tt
 
 
-class Anthropic:
+class Anthropic(tt.ModelInterface):
     def __init__(
         self,
         id: Optional[str] = "claude-3-haiku-20240307",
         base_url: str = "https://api.anthropic.com/v1/messages",
     ):
-        self.anthropic_model = id
+        self.model_id = id
         self.base_url = base_url
-        self.anthropic_api_token = tu.ENV.ANTHROPIC_TOKEN("")
+        self.api_token = tu.ENV.ANTHROPIC_TOKEN("")
 
     def set_api_token(self, token: str) -> None:
-        self.anthropic_api_token = token
-
-    def tool_to_claude_xml(self, tool):
-        """
-        Deprecated: was written when function calling did not exist in Anthropic API.
-        """
-        tool_signature = ""
-        if len(tool["parameters"]) > 0:
-            for name, p in tool["parameters"]["properties"].items():
-                param = f"""<parameter>
-                <name> {name} </name>
-                <type> {p['type']} </type>
-                <description> {p['description']} </description>
-                """
-                if name in tool["parameters"]["required"]:
-                    param += "<required> true </required>\n"
-                param += "</parameter>"
-                tool_signature += param + "\n"
-        tool_signature = tool_signature.strip()
-
-        constructed_prompt = (
-            "<tool_description>\n"
-            f"<tool_name> {tool['name']} </tool_name>\n"
-            "<description>\n"
-            f"{tool['description']}\n"
-            "</description>\n"
-            "<parameters>\n"
-            f"{tool_signature}\n"
-            "</parameters>\n"
-            "</tool_description>"
-        )
-        return constructed_prompt
+        self.api_token = token
 
     def _process_input(self, chats, token: Optional[str] = None):
-        if not token and not self.anthropic_api_token:  # type: ignore
+        if not token and not self.api_token:  # type: ignore
             raise Exception(
                 "Please set ANTHROPIC_TOKEN environment variable or pass through function"
             )
-        token = token or self.anthropic_api_token
-        if isinstance(chats, Thread):
+        token = token or self.api_token
+        if isinstance(chats, tt.Thread):
             thread = chats
         elif isinstance(chats, str):
-            thread = Thread(human(chats))
+            thread = tt.Thread(tt.human(chats))
         else:
             raise Exception("Invalid input")
 
         # create the anthropic style data
         system = ""
-        if thread.chats[0].role == Message.SYSTEM:
+        if thread.chats[0].role == tt.Message.SYSTEM:
             system = thread.chats[0].value
 
         claude_messages = []
         prev_tool_id = tu.get_random_string(5)
         for m in thread.chats[int(system != "") :]:
-            if m.role == Message.HUMAN:
+            if m.role == tt.Message.HUMAN:
                 msg = {
                     "role": "user",
                     "content": [{"type": "text", "text": m.value.strip()}],
@@ -95,12 +64,12 @@ def _process_input(self, chats, token: Optional[str] = None):
                                 },
                             }
                         )
-            elif m.role == Message.GPT:
+            elif m.role == tt.Message.GPT:
                 msg = {
                     "role": "assistant",
                     "content": [{"type": "text", "text": m.value.strip()}],
                 }
-            elif m.role == Message.FUNCTION_CALL:
+            elif m.role == tt.Message.FUNCTION_CALL:
                 _m = tu.from_json(m.value) if isinstance(m.value, str) else m.value
                 msg = {
                     "role": "assistant",
@@ -113,7 +82,7 @@ def _process_input(self, chats, token: Optional[str] = None):
                         }
                     ],
                 }
-            elif m.role == Message.FUNCTION_RESP:
+            elif m.role == tt.Message.FUNCTION_RESP:
                 # _m = tu.from_json(m.value) if isinstance(m.value, str) else m.value
                 msg = {
                     "role": "user",
@@ -139,7 +108,7 @@ def _process_input(self, chats, token: Optional[str] = None):
 
     def chat(
         self,
-        chats: Thread | str,
+        chats: tt.Thread | str,
         model: Optional[str] = None,
         max_tokens: int = 1024,
         temperature: Optional[float] = None,
@@ -170,7 +139,7 @@ def chat(
 
     def stream_chat(
         self,
-        chats: Thread | str,
+        chats: tt.Thread | str,
         model: Optional[str] = None,
         max_tokens: int = 1024,
         temperature: Optional[float] = None,
@@ -182,14 +151,14 @@ def stream_chat(
     ) -> Any:
 
         tools = []
-        if isinstance(chats, Thread):
+        if isinstance(chats, tt.Thread):
             tools = [x.to_dict() for x in chats.tools]
             for t in tools:
                 t["input_schema"] = t.pop("parameters")
         headers, system, claude_messages = self._process_input(chats=chats, token=token)
 
         data = {
-            "model": model or self.anthropic_model,
+            "model": model or self.model_id,
             "max_tokens": max_tokens,
             "messages": claude_messages,
             "system": system,

tuneapi/apis/model_anthropic.py

@@ -13,21 +13,21 @@
 import tuneapi.types as tt
 
 
-class Gemini:
+class Gemini(tt.ModelInterface):
     def __init__(
         self,
         id: Optional[str] = "gemini-1.5-pro-latest",
         base_url: str = "https://generativelanguage.googleapis.com/v1beta/models/{id}:{rpc}",
     ):
-        self._gemeni_model_id = id
+        self.model_id = id
         self.base_url = base_url
-        self.gemini_token = tu.ENV.GEMINI_TOKEN("")
+        self.api_token = tu.ENV.GEMINI_TOKEN("")
 
     def set_api_token(self, token: str) -> None:
-        self.gemini_token = token
+        self.api_token = token
 
     def _process_input(self, chats, token: Optional[str] = None):
-        if not token and not self.gemini_token:  # type: ignore
+        if not token and not self.api_token:  # type: ignore
             raise Exception(
                 "Gemini API key not found. Please set GEMINI_TOKEN environment variable or pass through function"
             )
@@ -98,7 +98,7 @@ def _process_input(self, chats, token: Optional[str] = None):
 
         # create headers
         headers = self._process_header()
-        params = {"key": self.gemini_token}
+        params = {"key": self.api_token}
         return headers, system.strip(), messages, params
 
     def _process_header(self):
@@ -197,7 +197,7 @@ def stream_chat(
 
         response = requests.post(
             self.base_url.format(
-                id=model or self._gemeni_model_id,
+                id=model or self.model_id,
                 rpc="streamGenerateContent",
             ),
             headers=headers,
@@ -225,12 +225,14 @@ def stream_chat(
                 continue
             block_lines += line
 
-            # is the block done?
-            if line == "{":
-                done = False
-            elif line == "}":
+            done = False
+            try:
+                tu.from_json(block_lines)
                 done = True
+            except Exception as e:
+                pass
 
+            # print(f"{block_lines=}")
             if done:
                 part_data = json.loads(block_lines)["candidates"][0]["content"][
                     "parts"

tuneapi/apis/model_gemini.py

@@ -12,25 +12,25 @@
 import tuneapi.types as tt
 
 
-class Groq:
+class Groq(tt.ModelInterface):
     def __init__(
         self,
         id: Optional[str] = "llama3-70b-8192",
         base_url: str = "https://api.groq.com/openai/v1/chat/completions",
     ):
-        self.groq_model_id = id
+        self.model_id = id
         self.base_url = base_url
-        self.groq_api_token = tu.ENV.GROQ_TOKEN("")
+        self.api_token = tu.ENV.GROQ_TOKEN("")
 
     def set_api_token(self, token: str) -> None:
-        self.groq_api_token = token
+        self.api_token = token
 
     def _process_input(self, chats, token: Optional[str] = None):
-        if not token and not self.groq_api_token:  # type: ignore
+        if not token and not self.api_token:  # type: ignore
             raise Exception(
                 "Please set GROQ_TOKEN environment variable or pass through function"
             )
-        token = token or self.groq_api_token
+        token = token or self.api_token
         if isinstance(chats, tt.Thread):
             thread = chats
         elif isinstance(chats, str):
@@ -132,7 +132,7 @@ def stream_chat(
         data = {
             "temperature": temperature,
             "messages": messages,
-            "model": model or self.groq_model_id,
+            "model": model or self.model_id,
             "stream": True,
             "max_tokens": max_tokens,
             "tools": tools,

tuneapi/apis/model_groq.py

@@ -14,25 +14,25 @@
 from tuneapi.types import Thread, human, Message
 
 
-class Mistral:
+class Mistral(tt.ModelInterface):
     def __init__(
         self,
         id: Optional[str] = "mistral-small-latest",
         base_url: str = "https://api.mistral.ai/v1/chat/completions",
     ):
-        self.mistral_model_id = id
+        self.model_id = id
         self.base_url = base_url
-        self.mistral_api_token = ENV.MISTRAL_TOKEN("")
+        self.api_token = ENV.MISTRAL_TOKEN("")
 
     def set_api_token(self, token: str) -> None:
-        self.mistral_api_token = token
+        self.api_token = token
 
     def _process_input(self, chats, token: Optional[str] = None):
-        if not token and not self.mistral_api_token:  # type: ignore
+        if not token and not self.api_token:  # type: ignore
             raise Exception(
                 "Please set MISTRAL_TOKEN environment variable or pass through function"
             )
-        token = token or self.mistral_api_token
+        token = token or self.api_token
 
         if isinstance(chats, tt.Thread):
             thread = chats
@@ -134,7 +134,7 @@ def stream_chat(
         headers, messages = self._process_input(chats, token)
         data = {
             "messages": messages,
-            "model": model or self.mistral_model_id,
+            "model": model or self.model_id,
             "stream": True,
             "max_tokens": max_tokens,
             "tools": tools,