fix: PR comments

AI21Labs · Dec 16, 2024 · 54c8e0e · 54c8e0e
1 parent 2cd5541
commit 54c8e0e
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 14 deletions.
diff --git a/ai21/clients/common/beta/assistant/plans.py b/ai21/clients/common/beta/assistant/plans.py
@@ -51,15 +51,12 @@ def _create_body(
     ) -> Dict[str, Any]:
         code_str = self._parse_code(code)
 
-        if schemas is NOT_GIVEN:
-            schema_dicts = NOT_GIVEN
-        else:
-            schema_dicts = [self._parse_schema(schema) for schema in schemas]
-
         return remove_not_given(
             {
                 "code": code_str,
-                "schemas": schema_dicts,
+                "schemas": [self._parse_schema(schema) for schema in schemas]
+                if schemas is not NOT_GIVEN
+                else NOT_GIVEN,
                 **kwargs,
             }
         )

diff --git a/ai21/errors.py b/ai21/errors.py
@@ -115,6 +115,6 @@ def __init__(self, details: Optional[str] = None):
 
 class CodeParsingError(AI21Error):
     def __init__(self, details: Optional[str] = None):
-        message = f"Code can't be parsed: {details}"
+        message = f"Code can't be parsed{'' if details is None else f': {details}'}"
         super().__init__(message)
         self.message = message
diff --git a/ai21/models/responses/plan_response.py b/ai21/models/responses/plan_response.py
@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import List, Optional
+from typing import List, Optional, Dict, Any
 
 from ai21.models.ai21_base_model import AI21BaseModel
 
@@ -10,7 +10,7 @@ class PlanResponse(AI21BaseModel):
     updated_at: datetime
     assistant_id: str
     code: str
-    schemas: List[dict]
+    schemas: List[Dict[str, Any]]
 
 
 class ListPlanResponse(AI21BaseModel):

diff --git a/examples/studio/assistant/user_defined_plans.py b/examples/studio/assistant/user_defined_plans.py
@@ -19,10 +19,7 @@ def main():
     assistant = ai21_client.beta.assistants.create(name="My Assistant")
 
     plan = ai21_client.beta.assistants.plans.create(assistant_id=assistant.id, code=test_func, schemas=[ExampleSchema])
-    ai21_client.beta.assistants.routes.create(
+    route = ai21_client.beta.assistants.routes.create(
         assistant_id=assistant.id, plan_id=plan.id, name="My Route", examples=["hi"], description="My Route Description"
     )
-
-
-if __name__ == "__main__":
-    main()
+    print(f"Route: {route}")