Add speed capability

ls1intum · Feb 28, 2024 · 4d2de4b · 4d2de4b
1 parent f74d39f
commit 4d2de4b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/app/llm/capability/capability_list.py b/app/llm/capability/capability_list.py
@@ -75,6 +75,7 @@ class CapabilityList(BaseModel):
     gpt_version_equivalent: OrderedNumberCapability = Field(
         default=OrderedNumberCapability(value=2)
     )
+    speed: OrderedNumberCapability = Field(default=OrderedNumberCapability(value=0))
     context_length: OrderedNumberCapability = Field(
         default=OrderedNumberCapability(value=0)
     )
@@ -100,6 +101,7 @@ def from_dict(cls, data: dict[str, any]):
 capability_weights = {
     "cost": 1,
     "gpt_version_equivalent": 4,
+    "speed": 2,
     "context_length": 0.1,
     "vendor": 1,
     "privacy_compliance": 0,

diff --git a/app/llm/capability/requirement_list.py b/app/llm/capability/requirement_list.py
@@ -1,8 +1,9 @@
 class RequirementList:
     """A class to represent the requirements you want to match against"""
 
-    cost: int | None
+    cost: float | None
     gpt_version_equivalent: float | None
+    speed: float | None
     context_length: int | None
     vendor: str | None
     privacy_compliance: bool | None
@@ -12,8 +13,9 @@ class RequirementList:
 
     def __init__(
         self,
-        cost: int | None = None,
+        cost: float | None = None,
         gpt_version_equivalent: float | None = None,
+        speed: float | None = None,
         context_length: int | None = None,
         vendor: str | None = None,
         privacy_compliance: bool | None = None,
@@ -23,6 +25,7 @@ def __init__(
     ) -> None:
         self.cost = cost
         self.gpt_version_equivalent = gpt_version_equivalent
+        self.speed = speed
         self.context_length = context_length
         self.vendor = vendor
         self.privacy_compliance = privacy_compliance