Convert embedding size to int when loading config (microsoft#53)

mluker · Sep 18, 2024 · b402597 · b402597
2 parents 9937f25 + 92d22bb
commit b402597
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 11 deletions.
diff --git a/examples/openai_deployment_config.json b/examples/openai_deployment_config.json
@@ -46,31 +46,31 @@
 	"text-embedding-3-small": {
         "model": "text-embedding-3-small",
 		"tokensPerMinute" : 10000,
-        "embeddingSize": "512"
+        "embeddingSize": 512
     },
     "text-embedding-3-medium": {
         "model": "text-embedding-3-medium",
 		"tokensPerMinute" : 10000,
-        "embeddingSize": "768"
+        "embeddingSize": 768
     },
     "text-embedding-3-large": {
         "model": "text-embedding-3-large",
 		"tokensPerMinute" : 10000,
-        "embeddingSize": "1024"
+        "embeddingSize": 1024
     },
     "text-embedding-3-xlarge": {
         "model": "text-embedding-3-xlarge",
 		"tokensPerMinute" : 10000,
-        "embeddingSize": "1536"
+        "embeddingSize": 1536
     },
     "text-embedding-ada-001": {
         "model": "text-embedding-ada-001",
 		"tokensPerMinute" : 10000,
-        "embeddingSize": "768"
+        "embeddingSize": 768
     },
     "text-embedding-ada-002": {
         "model": "text-embedding-ada-002",
 		"tokensPerMinute" : 10000,
-        "embeddingSize": "1536"
+        "embeddingSize": 1536
     }
 }
diff --git a/src/aoai-api-simulator/src/aoai_api_simulator/config_loader.py b/src/aoai-api-simulator/src/aoai_api_simulator/config_loader.py
@@ -2,13 +2,12 @@
 import json
 import logging
 import os
-
 import sys
 
+from aoai_api_simulator.generator.manager import get_default_generators
 from aoai_api_simulator.limiters import get_default_limiters
 from aoai_api_simulator.models import Config, OpenAIDeployment
 from aoai_api_simulator.record_replay.handler import get_default_forwarders
-from aoai_api_simulator.generator.manager import get_default_generators
 
 
 def get_config_from_env_vars(logger: logging.Logger) -> Config:
@@ -56,7 +55,7 @@ def _load_openai_deployments(logger: logging.Logger) -> dict[str, OpenAIDeployme
             name=deployment_name,
             model=deployment["model"],
             tokens_per_minute=deployment["tokensPerMinute"],
-            embedding_size=deployment.get("embeddingSize", 1536),
+            embedding_size=int(deployment.get("embeddingSize", 1536)),
         )
     return deployments
 
@@ -95,7 +94,6 @@ def _default_openai_deployments() -> dict[str, OpenAIDeployment]:
 
 
 def load_extension(config: Config):
-
     extension_path = config.extension_path
     if not extension_path:
         return

diff --git a/src/aoai-api-simulator/src/aoai_api_simulator/generator/openai.py b/src/aoai-api-simulator/src/aoai_api_simulator/generator/openai.py
@@ -150,7 +150,7 @@ async def calculate_latency(context: RequestContext, status_code: int):
             context.values[constants.TARGET_DURATION_MS] = target_duration_ms
 
 
-def create_embedding_content(index: int, embedding_size):
+def create_embedding_content(index: int, embedding_size: int):
     """Generates a random embedding"""
     return {
         "object": "embedding",