From 4d2de4b2066fb8b083901ab73d64e2aad617d979 Mon Sep 17 00:00:00 2001
From: Timor Morrien <timor.morrien@tum.de>
Date: Wed, 28 Feb 2024 04:02:49 +0100
Subject: [PATCH] Add speed capability

---
 app/llm/capability/capability_list.py  | 2 ++
 app/llm/capability/requirement_list.py | 7 +++++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/app/llm/capability/capability_list.py b/app/llm/capability/capability_list.py
index 3e7405a8..62aca8c7 100644
--- 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
index 52d66e58..e26eb959 100644
--- 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