fix: support func type in code and class in schemas

ai21/clients/common/beta/assistant/plans.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
+import inspect
 from abc import ABC, abstractmethod
-from typing import Any, Dict
+from typing import Any, Dict, Type
+
+from pydantic import BaseModel
 
 from ai21.models.responses.plan_response import PlanResponse, ListPlanResponse
 from ai21.types import NOT_GIVEN, NotGiven
@@ -16,23 +19,33 @@ def create(
         self,
         *,
         assistant_id: str,
-        code: str,
-        schemas: list[dict] | NotGiven = NOT_GIVEN,
+        code: str | callable,
+        schemas: list[dict] | list[Type[BaseModel]] | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> PlanResponse:
         pass
 
     def _create_body(
         self,
         *,
-        code: str,
-        schemas: list[dict] | NotGiven = NOT_GIVEN,
+        code: str | callable,
+        schemas: list[dict] | list[Type[BaseModel]] | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> Dict[str, Any]:
+        if callable(code):
+            code = inspect.getsource(code).strip()
+
+        schema_dicts = []
+        for schema in schemas:
+            if inspect.isclass(schema) and issubclass(schema, BaseModel):
+                schema_dicts.append(schema.model_json_schema())
+            else:
+                schema_dicts.append(schema)
+
         return remove_not_given(
             {
                 "code": code,
-                "schemas": schemas,
+                "schemas": schema_dicts,
                 **kwargs,
             }
         )

examples/studio/assistant/user_defined_plans.py

@@ -20,6 +20,10 @@
 ask_for_approval=True,\n    )\n    return summary"""
 
 
+def test_func():
+    pass
+
+
 class ExampleSchema(BaseModel):
     name: str
     id: str
@@ -30,10 +34,7 @@ def main():
 
     assistant = ai21_client.beta.assistants.create(name="My Assistant")
 
-    plan = ai21_client.beta.assistants.plans.create(
-        assistant_id=assistant.id, code=CODE_STR, schemas=[ExampleSchema.model_json_schema()]
-    )
-
+    plan = ai21_client.beta.assistants.plans.create(assistant_id=assistant.id, code=test_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"
     )