fix: tool resources validation (#654)

* chore: increase default token limits for assistants * fix: validate tool_resources with a model_dump
defenseunicorns · Jun 18, 2024 · d2cbb09 · d2cbb09
1 parent a9ae575
commit d2cbb09
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/src/leapfrogai_api/routers/openai/requests/run_create_params_request_base.py b/src/leapfrogai_api/routers/openai/requests/run_create_params_request_base.py
@@ -71,8 +71,8 @@
 class RunCreateParamsRequestBase(BaseModel):
     assistant_id: str = Field(default="", examples=["123ab"])
     instructions: str = Field(default="", examples=["You are a helpful AI assistant."])
-    max_completion_tokens: int | None = Field(default=128, examples=[4096])
-    max_prompt_tokens: int | None = Field(default=128, examples=[32768])
+    max_completion_tokens: int | None = Field(default=1024, examples=[4096])
+    max_prompt_tokens: int | None = Field(default=1024, examples=[32768])
     metadata: dict | None = Field(default={}, examples=[{}])
     model: str | None = Field(default=None, examples=["llama-cpp-python"])
     response_format: AssistantResponseFormatOption | None = Field(
@@ -96,14 +96,14 @@ def __init__(self, **data):
         # TODO: Temporary fix to ensure max_completion_tokens and max_prompt_tokens are set
         if self.max_completion_tokens is None or self.max_completion_tokens < 1:
             logging.warning(
-                "max_completion_tokens is not set or is less than 1, setting to 128"
+                "max_completion_tokens is not set or is less than 1, setting to 1024"
             )
-            self.max_completion_tokens = 128
+            self.max_completion_tokens = 1024
         if self.max_prompt_tokens is None or self.max_prompt_tokens < 1:
             logging.warning(
-                "max_prompt_tokens is not set or is less than 1, setting to 128"
+                "max_prompt_tokens is not set or is less than 1, setting to 1024"
             )
-            self.max_prompt_tokens = 128
+            self.max_prompt_tokens = 1024
 
     @staticmethod
     def get_initial_messages_base(run: Run) -> list[str]:
@@ -296,7 +296,7 @@ async def generate_message_for_thread(
                 and isinstance(assistant.tool_resources, BetaAssistantToolResources)
             ):
                 tool_resources = BetaThreadToolResources.model_validate(
-                    assistant.tool_resources
+                    assistant.tool_resources.model_dump()
                 )
             else:
                 tool_resources = None
@@ -362,7 +362,7 @@ async def stream_generate_message_for_thread(
                 and isinstance(assistant.tool_resources, BetaAssistantToolResources)
             ):
                 tool_resources = BetaThreadToolResources.model_validate(
-                    assistant.tool_resources
+                    assistant.tool_resources.model_dump()
                 )
             else:
                 tool_resources = None

diff --git a/src/leapfrogai_api/routers/openai/requests/thread_run_create_params_request.py b/src/leapfrogai_api/routers/openai/requests/thread_run_create_params_request.py
@@ -77,7 +77,7 @@ async def update_with_assistant_data(self, session: Session):
 
         if assistant and assistant.tool_resources:
             assistant_tool_resources = BetaThreadToolResources.model_validate(
-                assistant.tool_resources
+                assistant.tool_resources.model_dump()
             )
 
             self.tool_resources = self.tool_resources or assistant_tool_resources