fix: change schemas field in PlanResponse to be optional

AI21Labs · Jan 13, 2025 · 2f8172c · 2f8172c
1 parent 4a7dfe0
commit 2f8172c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/ai21/models/responses/plan_response.py b/ai21/models/responses/plan_response.py
@@ -10,7 +10,7 @@ class PlanResponse(AI21BaseModel):
     updated_at: datetime
     assistant_id: str
     code: str
-    schemas: List[Dict[str, Any]]
+    schemas: Optional[List[Dict[str, Any]]] = None
 
 
 class ListPlanResponse(AI21BaseModel):

diff --git a/examples/studio/assistant/user_defined_plans.py b/examples/studio/assistant/user_defined_plans.py
@@ -4,7 +4,7 @@
 TIMEOUT = 20
 
 
-def test_func():
+def func():
     pass
 
 
@@ -18,8 +18,15 @@ 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])
-    route = ai21_client.beta.assistants.routes.create(
+    plan = ai21_client.beta.assistants.plans.create(assistant_id=assistant.id, code=func, schemas=[ExampleSchema])
+    ai21_client.beta.assistants.routes.create(
         assistant_id=assistant.id, plan_id=plan.id, name="My Route", examples=["hi"], description="My Route Description"
     )
-    print(f"Route: {route}")
+    routes = ai21_client.beta.assistants.routes.list(assistant_id=assistant.id)
+    print(f"Routes: {routes}")
+    plans = ai21_client.beta.assistants.plans.list(assistant_id=assistant.id)
+    print(f"Plans: {plans}")
+
+
+if __name__ == "__main__":
+    main()