Fixed annotations

arena-ai · May 23, 2024 · 99e9c31 · 99e9c31
1 parent 5905d83
commit 99e9c31
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/client/pyproject.toml b/client/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "arena-client"
-version = "0.3.1"
+version = "0.3.2"
 description = "A client to use arena AI"
 authors = ["Nicolas Grislain <ng@sarus.tech>"]
 license = "Apache-2"

diff --git a/client/src/arena/models/anthropic.py b/client/src/arena/models/anthropic.py
@@ -29,7 +29,7 @@ class ChatCompletionRequest(BaseModel):
     """
     max_tokens: int = 1
     messages: Sequence[Message]
-    model: str | Literal[*MODELS]
+    model: str | Literal["claude-3-opus-20240229", "claude-3-sonnet-20240229", "claude-3-haiku-20240307", "claude-2.1", "claude-2.0", "claude-instant-1.2"]
     metadata: Metadata | None = None
     stop_sequences: Sequence[str] | None = None
     system: str | None = None

diff --git a/client/src/arena/models/mistral.py b/client/src/arena/models/mistral.py
@@ -20,7 +20,7 @@ class ChatCompletionRequest(BaseModel):
     https://docs.mistral.ai/api/#operation/createChatCompletion
     """
     messages: Sequence[Message]
-    model: str | Literal[*MODELS]
+    model: str | Literal["mistral-large-latest", "mistral-medium", "mistral-medium-latest", "mistral-small", "mistral-small-latest", "mistral-tiny", "open-mistral-7b", "open-mixtral-8x7b"]
     max_tokens: int | None = None
     response_format: ResponseFormat | None = None
     safe_prompt: bool | None = None

diff --git a/client/src/arena/models/openai.py b/client/src/arena/models/openai.py
@@ -12,7 +12,7 @@ class ChatCompletionRequest(models.ChatCompletionRequest):
     https://github.com/openai/openai-python/blob/main/src/openai/types/chat/completion_create_params.py#L24
     https://platform.openai.com/docs/api-reference/chat
     """
-    model: str | Literal[*MODELS]
+    model: str | Literal["gpt-4o", "gpt-4-turbo", "gpt-4", "gpt-3.5-turbo", "gpt-3.5-turbo-16k"]
 
     @classmethod
     def from_chat_completion_request(cls, ccc: models.ChatCompletionRequest) -> "ChatCompletionRequest":