unique_toolkit 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

unique_toolkit/language_model/infos.py

@@ -1,16 +1,14 @@
 from datetime import date
 from enum import StrEnum
-from typing import ClassVar, Optional, Type, cast
+from typing import ClassVar, Optional
 
 from pydantic import BaseModel
+from typing_extensions import deprecated
 
 from unique_toolkit.language_model.schemas import LanguageModelTokenLimits
 
 
 class LanguageModelName(StrEnum):
-    AZURE_GPT_35_TURBO_0613 = "AZURE_GPT_35_TURBO_0613"
-    AZURE_GPT_35_TURBO = "AZURE_GPT_35_TURBO"
-    AZURE_GPT_35_TURBO_16K = "AZURE_GPT_35_TURBO_16K"
     AZURE_GPT_35_TURBO_0125 = "AZURE_GPT_35_TURBO_0125"
     AZURE_GPT_4_0613 = "AZURE_GPT_4_0613"
     AZURE_GPT_4_TURBO_1106 = "AZURE_GPT_4_TURBO_1106"
@@ -30,12 +28,7 @@ class EncoderName(StrEnum):
 def get_encoder_name(model_name: LanguageModelName) -> Optional[EncoderName]:
     LMN = LanguageModelName
     match model_name:
-        case (
-            LMN.AZURE_GPT_35_TURBO
-            | LMN.AZURE_GPT_35_TURBO_16K
-            | LMN.AZURE_GPT_35_TURBO_0613
-            | LMN.AZURE_GPT_35_TURBO_0125
-        ):
+        case LMN.AZURE_GPT_35_TURBO_0125:
             return EncoderName.CL100K_BASE
         case (
             LMN.AZURE_GPT_4_0613
@@ -76,7 +69,147 @@ class LanguageModelInfo(BaseModel):
     deprecated_at: Optional[date] = None
     retirement_text: Optional[str] = None
 
+    @classmethod
+    def from_name(cls, model_name: LanguageModelName):
+        match model_name:
+            case LanguageModelName.AZURE_GPT_35_TURBO_0125:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.AZURE,
+                    version="0125",
+                    encoder_name=get_encoder_name(model_name),
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=16385, token_limit_output=4096
+                    ),
+                    info_cutoff_at=date(2021, 9, 1),
+                    published_at=date(2023, 1, 25),
+                    retirement_at=date(5, 3, 31),
+                )
+            case LanguageModelName.AZURE_GPT_4_0613:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.AZURE,
+                    version="0613",
+                    encoder_name=get_encoder_name(model_name),
+                    token_limits=LanguageModelTokenLimits(token_limit=8192),
+                    info_cutoff_at=date(2021, 9, 1),
+                    published_at=date(2023, 6, 13),
+                    deprecated_at=date(2024, 10, 1),
+                    retirement_at=date(2025, 6, 6),
+                )
+            case LanguageModelName.AZURE_GPT_4_TURBO_1106:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.AZURE,
+                    version="1106-preview",
+                    encoder_name=get_encoder_name(model_name),
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=128000, token_limit_output=4096
+                    ),
+                    info_cutoff_at=date(2023, 4, 1),
+                    published_at=date(2023, 11, 6),
+                )
+            case LanguageModelName.AZURE_GPT_4_VISION_PREVIEW:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.AZURE,
+                    version="vision-preview",
+                    encoder_name=get_encoder_name(model_name),
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=128000, token_limit_output=4096
+                    ),
+                    info_cutoff_at=date(2023, 4, 1),
+                    published_at=date(2023, 11, 6),
+                )
+            case LanguageModelName.AZURE_GPT_4_32K_0613:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.AZURE,
+                    version="1106-preview",
+                    encoder_name=get_encoder_name(model_name),
+                    token_limits=LanguageModelTokenLimits(token_limit=32768),
+                    info_cutoff_at=date(2021, 9, 1),
+                    published_at=date(2023, 6, 13),
+                    deprecated_at=date(2024, 10, 1),
+                    retirement_at=date(2025, 6, 6),
+                )
+            case LanguageModelName.AZURE_GPT_4_TURBO_2024_0409:
+                return cls(
+                    name=model_name,
+                    encoder_name=get_encoder_name(model_name),
+                    provider=LanguageModelProvider.AZURE,
+                    version="turbo-2024-04-09",
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=128000, token_limit_output=4096
+                    ),
+                    info_cutoff_at=date(2023, 12, 1),
+                    published_at=date(2024, 4, 9),
+                )
+            case LanguageModelName.AZURE_GPT_4o_2024_0513:
+                return cls(
+                    name=model_name,
+                    encoder_name=get_encoder_name(model_name),
+                    provider=LanguageModelProvider.AZURE,
+                    version="2024-05-13",
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=128000, token_limit_output=4096
+                    ),
+                    info_cutoff_at=date(2023, 10, 1),
+                    published_at=date(2024, 5, 13),
+                )
+            case LanguageModelName.AZURE_GPT_4o_2024_0806:
+                return cls(
+                    name=model_name,
+                    encoder_name=get_encoder_name(model_name),
+                    provider=LanguageModelProvider.AZURE,
+                    version="2024-08-06",
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=128000, token_limit_output=16384
+                    ),
+                    info_cutoff_at=date(2023, 10, 1),
+                    published_at=date(2024, 8, 6),
+                )
+            case LanguageModelName.AZURE_GPT_4o_MINI_2024_0718:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.AZURE,
+                    version="2024-07-18",
+                    encoder_name=get_encoder_name(model_name),
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=128000, token_limit_output=16384
+                    ),
+                    info_cutoff_at=date(2023, 10, 1),
+                    published_at=date(2024, 7, 18),
+                )
+            case _:
+                if isinstance(model_name, LanguageModelName):
+                    raise ValueError(
+                        f"{model_name} is not supported. Please add model information in toolkit."
+                    )
+
+                return cls(
+                    name=model_name,
+                    version="custom",
+                    provider=LanguageModelProvider.CUSTOM,
+                )
+
+    @property
+    def display_name(self) -> str:
+        """
+        Returns the name of the model as a string.
+        """
+
+        if isinstance(self.name, LanguageModelName):
+            return self.name.value
+        else:
+            return self.name
+
 
+@deprecated(
+    """
+Use `LanguageModelInfo` instead of `LanguageModel`
+"""
+)
 class LanguageModel:
     _info: ClassVar[LanguageModelInfo]
 
@@ -199,17 +332,8 @@ class LanguageModel:
 
     @classmethod
     def get_model_info(cls, model_name: LanguageModelName | str) -> LanguageModelInfo:
-        if not model_name:
-            raise ValueError("Model name must be provided to get the model info.")
-
-        for subclass in cls.__subclasses__():
-            if hasattr(subclass, "info") and subclass._info.name == model_name:
-                # TODO find alternative solution for warning
-                # if subclass._info.retirement_at:
-                #     warning_text = f"WARNING: {subclass._info.name} will be retired on {subclass._info.retirement_at.isoformat()} and from then on not accessible anymore. {subclass._info.retirement_text}"
-                #     print(warning_text)
-                #     warnings.warn(warning_text, DeprecationWarning, stacklevel=2)
-                return subclass._info
+        if isinstance(model_name, LanguageModelName):
+            return LanguageModelInfo.from_name(model_name)
 
         return LanguageModelInfo(
             name=model_name,
@@ -224,190 +348,5 @@ class LanguageModel:
         """
 
         return [
-            cast(LanguageModelInfo, subclass._info)
-            for subclass in cls.__subclasses__()
-            if hasattr(subclass, "_info")
+            LanguageModelInfo.from_name(model_name=name) for name in LanguageModelName
         ]
-
-
-def create_language_model(
-    name: LanguageModelName,
-    version: str,
-    provider: LanguageModelProvider,
-    info_cutoff_at: date,
-    published_at: date,
-    encoder_name: Optional[EncoderName] = None,
-    token_limit: Optional[int] = None,
-    token_limit_input: Optional[int] = None,
-    token_limit_output: Optional[int] = None,
-    retirement_at: Optional[date] = None,
-    deprecated_at: Optional[date] = None,
-    retirement_text: Optional[str] = None,
-) -> Type[LanguageModel]:
-    info = LanguageModelInfo(
-        name=name,
-        version=version,
-        provider=provider,
-        encoder_name=encoder_name,
-        token_limits=LanguageModelTokenLimits(
-            token_limit=token_limit,
-            token_limit_input=token_limit_input,
-            token_limit_output=token_limit_output,
-        ),
-        info_cutoff_at=info_cutoff_at,
-        published_at=published_at,
-        retirement_at=retirement_at,
-        deprecated_at=deprecated_at,
-        retirement_text=retirement_text,
-    )
-
-    class Model(LanguageModel):
-        _info = info
-
-    return Model
-
-
-############################################################################################################
-# Define the models here
-############################################################################################################
-
-
-AzureGpt35Turbo0613 = create_language_model(
-    name=LanguageModelName.AZURE_GPT_35_TURBO_0613,
-    provider=LanguageModelProvider.AZURE,
-    version="0613",
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_35_TURBO_0613),
-    token_limit=8192,
-    info_cutoff_at=date(2021, 9, 1),
-    published_at=date(2023, 6, 13),
-    retirement_at=date(2025, 2, 13),
-)
-
-AzureGpt35Turbo = create_language_model(
-    name=LanguageModelName.AZURE_GPT_35_TURBO,
-    provider=LanguageModelProvider.AZURE,
-    version="0301",
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_35_TURBO),
-    token_limit=4096,
-    info_cutoff_at=date(2021, 9, 1),
-    published_at=date(2023, 3, 1),
-    retirement_at=date(2025, 2, 13),
-)
-
-
-AzureGpt35Turbo16k = create_language_model(
-    name=LanguageModelName.AZURE_GPT_35_TURBO_16K,
-    provider=LanguageModelProvider.AZURE,
-    version="0613",
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_35_TURBO_16K),
-    token_limit=16382,
-    info_cutoff_at=date(2021, 9, 1),
-    published_at=date(2023, 6, 13),
-    retirement_at=date(2025, 2, 13),
-)
-
-
-AzureGpt35Turbo0125 = create_language_model(
-    name=LanguageModelName.AZURE_GPT_35_TURBO_0125,
-    provider=LanguageModelProvider.AZURE,
-    version="0125",
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_35_TURBO_0125),
-    token_limit_input=16385,
-    token_limit_output=4096,
-    info_cutoff_at=date(2021, 9, 1),
-    published_at=date(2023, 1, 25),
-    retirement_at=date(5, 3, 31),
-)
-
-
-AzureGpt40613 = create_language_model(
-    name=LanguageModelName.AZURE_GPT_4_0613,
-    provider=LanguageModelProvider.AZURE,
-    version="0613",
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_4_0613),
-    token_limit=8192,
-    info_cutoff_at=date(2021, 9, 1),
-    published_at=date(2023, 6, 13),
-    deprecated_at=date(2024, 10, 1),
-    retirement_at=date(2025, 6, 6),
-)
-
-
-AzureGpt4Turbo1106 = create_language_model(
-    name=LanguageModelName.AZURE_GPT_4_TURBO_1106,
-    provider=LanguageModelProvider.AZURE,
-    version="1106-preview",
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_4_TURBO_1106),
-    token_limit_input=128000,
-    token_limit_output=4096,
-    info_cutoff_at=date(2023, 4, 1),
-    published_at=date(2023, 11, 6),
-)
-
-
-AzureGpt4VisionPreview = create_language_model(
-    name=LanguageModelName.AZURE_GPT_4_VISION_PREVIEW,
-    provider=LanguageModelProvider.AZURE,
-    version="vision-preview",
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_4_VISION_PREVIEW),
-    token_limit_input=128000,
-    token_limit_output=4096,
-    info_cutoff_at=date(2023, 4, 1),
-    published_at=date(2023, 11, 6),
-)
-
-AzureGpt432k0613 = create_language_model(
-    name=LanguageModelName.AZURE_GPT_4_32K_0613,
-    provider=LanguageModelProvider.AZURE,
-    version="1106-preview",
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_4_32K_0613),
-    token_limit=32768,
-    info_cutoff_at=date(2021, 9, 1),
-    published_at=date(2023, 6, 13),
-    deprecated_at=date(2024, 10, 1),
-    retirement_at=date(2025, 6, 6),
-)
-
-AzureGpt4Turbo20240409 = create_language_model(
-    name=LanguageModelName.AZURE_GPT_4_TURBO_2024_0409,
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_4_TURBO_2024_0409),
-    provider=LanguageModelProvider.AZURE,
-    version="turbo-2024-04-09",
-    token_limit_input=128000,
-    token_limit_output=4096,
-    info_cutoff_at=date(2023, 12, 1),
-    published_at=date(2024, 4, 9),
-)
-
-AzureGpt4o20240513 = create_language_model(
-    name=LanguageModelName.AZURE_GPT_4o_2024_0513,
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_4o_2024_0513),
-    provider=LanguageModelProvider.AZURE,
-    version="2024-05-13",
-    token_limit_input=128000,
-    token_limit_output=4096,
-    info_cutoff_at=date(2023, 10, 1),
-    published_at=date(2024, 5, 13),
-)
-
-AzureGpt4o20240806 = create_language_model(
-    name=LanguageModelName.AZURE_GPT_4o_2024_0806,
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_4o_2024_0806),
-    provider=LanguageModelProvider.AZURE,
-    version="2024-08-06",
-    token_limit_input=128000,
-    token_limit_output=16384,
-    info_cutoff_at=date(2023, 10, 1),
-    published_at=date(2024, 8, 6),
-)
-
-AzureGpt4oMini20240718 = create_language_model(
-    name=LanguageModelName.AZURE_GPT_4o_MINI_2024_0718,
-    provider=LanguageModelProvider.AZURE,
-    version="2024-07-18",
-    encoder_name=get_encoder_name(LanguageModelName.AZURE_GPT_4o_MINI_2024_0718),
-    token_limit_input=128000,
-    token_limit_output=16384,
-    info_cutoff_at=date(2023, 10, 1),
-    published_at=date(2024, 7, 18),
-)

unique_toolkit/language_model/schemas.py

@@ -9,11 +9,13 @@ from pydantic import (
     BaseModel,
     ConfigDict,
     Field,
+    PrivateAttr,
     RootModel,
     field_validator,
     model_serializer,
     model_validator,
 )
+from typing_extensions import deprecated
 
 from unique_toolkit.language_model.utils import format_message
 
@@ -205,30 +207,50 @@ class LanguageModelStreamResponse(BaseModel):
 
 
 class LanguageModelTokenLimits(BaseModel):
-    token_limit: Optional[int] = None
-    token_limit_input: Optional[int] = None
-    token_limit_output: Optional[int] = None
-
-    fraction_input: float = Field(default=0.4, le=1, ge=0)
-
-    @model_validator(mode="after")
-    def check_required_fields(self):
-        # Best case input and output is determined
-        if self.token_limit_input and self.token_limit_output:
-            self.token_limit = self.token_limit_input + self.token_limit_output
-            self.fraction_input = self.token_limit_input / self.token_limit
-            return self
-
-        # Deal with case where only token_limit and optional fraction_input is given
-        if self.token_limit:
-            if not self.fraction_input:
-                self.fraction_input = 0.4
-
-            self.token_limit_input = math.floor(self.fraction_input * self.token_limit)
-            self.token_limit_output = math.floor(
-                (1 - self.fraction_input) * self.token_limit
-            )
-            return self
+    token_limit_input: int
+    token_limit_output: int
+
+    _fraction_adaptable = PrivateAttr(default=False)
+
+    @property
+    @deprecated("""
+    Deprecated: Use the more specific `token_limit_input` and `token_limit_output` instead.
+    """)
+    def token_limit(self):
+        return self.token_limit_input + self.token_limit_output
+
+    @token_limit.setter
+    @deprecated("""
+    Deprecated: Token limit can only be reduced
+    """)
+    def token_limit(self, token_limit, fraction_input=0.4):
+        if self.token_limit > token_limit:
+            self.token_limit_input = math.floor(fraction_input * token_limit)
+            self.token_limit_output = math.floor((1 - fraction_input) * token_limit)
+            self._fraction_adaptable = True
+
+    def adapt_fraction(self, fraction_input: float):
+        if self._fraction_adaptable:
+            token_limit = self.token_limit_input + self.token_limit_output
+            self.token_limit_input = math.floor(fraction_input * token_limit)
+            self.token_limit_output = math.floor((1 - fraction_input) * token_limit)
+
+    @model_validator(mode="before")
+    def check_required_fields(cls, data):
+        if isinstance(data, dict):
+            if {"token_limit_input", "token_limit_output"}.issubset(data.keys()):
+                return data
+
+            if {"token_limit"}.issubset(data.keys()):
+                token_limit = data.get("token_limit")
+                fraction_input = data.get("fraction_input", 0.4)
+
+                data["token_limit_input"] = math.floor(fraction_input * token_limit)
+                data["token_limit_output"] = math.floor(
+                    (1 - fraction_input) * token_limit
+                )
+                data["_fraction_adaptpable"] = True
+            return data
 
         raise ValueError(
             'Either "token_limit_input" and "token_limit_output" must be provided together, or "token_limit" must be provided.'