rework / clean up JSONSchema (#70)

docs/docs/extending-openassistants/create-a-custom-function.md

@@ -14,7 +14,6 @@ The following Python function:
 
 class PythonEvalFunction(BaseFunction):
     type: Literal["PythonEvalFunction"] = "PythonEvalFunction"
-    parameters: BaseJSONSchema
     python_code: str
 
     async def execute(

examples/fast-api-server/fast_api_server/find_email_by_name_function.py

@@ -3,9 +3,9 @@
 import anyio
 import pandas as pd
 from openassistants.contrib.python_callable import PythonCallableFunction
-from openassistants.data_models.function_input import BaseJSONSchema
 from openassistants.data_models.function_output import FunctionOutput, TextOutput
 from openassistants.functions.base import (
+    BaseFunctionParameters,
     Entity,
     EntityConfig,
     FunctionExecutionDependency,
@@ -64,7 +64,7 @@ async def _get_entity_configs() -> dict[str, EntityConfig]:
         "Find the email address for {employee}",
         "What is {employee}'s email address?",
     ],
-    parameters=BaseJSONSchema(
+    parameters=BaseFunctionParameters(
         json_schema={
             "type": "object",
             "properties": {

packages/openassistants/openassistants/contrib/index_function.py

@@ -23,6 +23,3 @@ async def execute(
             output += f"""**{function.get_display_name()}**
 {function.get_description()}\n\n"""
             yield [TextOutput(text=output)]
-
-    def get_parameters_json_schema(self) -> dict:
-        return {"type": "object", "properties": {}}

packages/openassistants/openassistants/contrib/langchain_ddg_tool.py

@@ -1,7 +1,6 @@
 import os
 from typing import Dict, List, Literal, Optional, Sequence
 
-from openassistants.data_models.function_input import BaseJSONSchema
 from openassistants.data_models.function_output import FunctionOutput, TextOutput
 from openassistants.functions.base import BaseFunction, FunctionExecutionDependency
 from openassistants.functions.utils import AsyncStreamVersion
@@ -38,7 +37,6 @@ def ddgs_text(query: str, max_results: Optional[int] = None) -> List[Dict[str, s
 
 class DuckDuckGoToolFunction(BaseFunction):
     type: Literal["DuckDuckGoToolFunction"] = "DuckDuckGoToolFunction"
-    parameters: BaseJSONSchema
 
     async def execute(
         self, deps: FunctionExecutionDependency
@@ -65,6 +63,3 @@ async def execute(
             raise RuntimeError(
                 f"Error while executing action function {self.id}. function raised: {e}"
             ) from e
-
-    def get_parameters_json_schema(self) -> dict:
-        return self.parameters.json_schema

packages/openassistants/openassistants/contrib/python_callable.py

@@ -1,6 +1,5 @@
 from typing import Awaitable, Callable, Mapping, Sequence
 
-from openassistants.data_models.function_input import BaseJSONSchema
 from openassistants.data_models.function_output import FunctionOutput
 from openassistants.functions.base import (
     BaseFunction,
@@ -15,8 +14,6 @@ class PythonCallableFunction(BaseFunction):
         [FunctionExecutionDependency], AsyncStreamVersion[Sequence[FunctionOutput]]
     ]
 
-    parameters: BaseJSONSchema
-
     get_entity_configs_callable: Callable[[], Awaitable[Mapping[str, IEntityConfig]]]
 
     async def execute(
@@ -25,8 +22,5 @@ async def execute(
         async for output in self.execute_callable(deps):
             yield output
 
-    def get_parameters_json_schema(self) -> dict:
-        return self.parameters.json_schema
-
     async def get_entity_configs(self) -> Mapping[str, IEntityConfig]:
         return await self.get_entity_configs_callable()

packages/openassistants/openassistants/contrib/python_eval.py

@@ -1,7 +1,6 @@
 import inspect
 from typing import Any, Callable, Dict, List, Literal, Sequence
 
-from openassistants.data_models.function_input import BaseJSONSchema
 from openassistants.data_models.function_output import FunctionOutput
 from openassistants.functions.base import (
     BaseFunction,
@@ -13,7 +12,6 @@
 
 class PythonEvalFunction(BaseFunction):
     type: Literal["PythonEvalFunction"] = "PythonEvalFunction"
-    parameters: BaseJSONSchema
     python_code: str
 
     async def execute(
@@ -47,6 +45,3 @@ async def execute(
             raise RuntimeError(
                 f"Error while executing action function {self.id}. function raised: {e}"
             ) from e
-
-    def get_parameters_json_schema(self) -> dict:
-        return self.parameters.json_schema

packages/openassistants/openassistants/contrib/sqlalchemy_query.py

@@ -1,6 +1,7 @@
 import asyncio
 from typing import Annotated, Any, List, Literal, Sequence
 
+import jsonschema
 import pandas as pd
 from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage
 from openassistants.data_models.chat_messages import (
@@ -9,7 +10,6 @@
     OpasMessage,
 )
 from openassistants.data_models.function_input import (
-    BaseJSONSchema,
     FunctionCall,
 )
 from openassistants.data_models.function_output import (
@@ -69,7 +69,6 @@ def _opas_to_summarization_lc(
 
 class QueryFunction(BaseFunction):
     type: Literal["QueryFunction"] = "QueryFunction"
-    parameters: BaseJSONSchema
     sqls: List[str]
     visualizations: List[str]
     summarization: str
@@ -153,7 +152,10 @@ async def execute(
         self,
         deps: FunctionExecutionDependency,
     ) -> AsyncStreamVersion[Sequence[FunctionOutput]]:
-        self.parameters.validate_args(deps.arguments)
+        try:
+            jsonschema.validate(deps.arguments, self.get_parameters_json_schema())
+        except jsonschema.ValidationError as e:
+            raise ValueError(f"Invalid arguments:\n{str(e)}") from e
 
         results: List[FunctionOutput] = []
 
@@ -203,6 +205,3 @@ async def execute(
         )
 
         yield results
-
-    def get_parameters_json_schema(self) -> dict:
-        return self.parameters.json_schema

packages/openassistants/openassistants/data_models/function_input.py

@@ -1,47 +1,14 @@
 from typing import Any, Dict
 
-import jsonschema
 from pydantic import BaseModel
 
-
-class BaseJSONSchema(BaseModel):
-    """
-    Validates a json_schema. top level must of the schema must be type object
-    """
-
-    json_schema: Dict[str, Any]
-
-    def schema_validator(cls, values):
-        jsonschema.validate(
-            values["json_schema"], jsonschema.Draft202012Validator.META_SCHEMA
-        )
-        jsonschema.validate(
-            values["json_schema"],
-            {
-                "type": "object",
-                "properties": {
-                    "type": {"type": "string", "enum": ["object"]},
-                    "properties": {"type": "object"},
-                    "required": {"type": "array", "items": {"type": "string"}},
-                },
-                "required": [
-                    "type",
-                ],
-            },
-        )
-        return values
-
-    def validate_args(self, args: dict):
-        try:
-            jsonschema.validate(args, self.json_schema)
-        except jsonschema.exceptions.ValidationError as e:
-            raise ValueError(f"invalid function arguments\n{e}")
+from openassistants.data_models.json_schema import JSONSchema
 
 
 class FunctionCall(BaseModel):
     name: str
     arguments: Dict[str, Any]
 
 
-class FunctionInputRequest(FunctionCall, BaseJSONSchema):
-    pass
+class FunctionInputRequest(FunctionCall):
+    json_schema: JSONSchema

packages/openassistants/openassistants/data_models/json_schema.py

@@ -0,0 +1,40 @@
+from typing import Annotated, Any, Dict
+
+import jsonschema
+from pydantic import AfterValidator, TypeAdapter
+
+
+def _json_schema_meta_validator(value: Any):
+    try:
+        jsonschema.validate(value, jsonschema.Draft202012Validator.META_SCHEMA)
+    except jsonschema.exceptions.ValidationError as e:
+        raise ValueError(f"Invalid JSONSchema:\n{str(e)}\n") from e
+    if value.get("type") != "object":
+        raise ValueError(
+            f"JSONSchema must have type='object'. Got '{value.get('type')}'"
+        )
+    return value
+
+
+JSONSchema = Annotated[Dict[str, Any], AfterValidator(_json_schema_meta_validator)]
+"""
+A JSONSchema is a dict that conforms to the JSONSchema specification.
+
+In order to validate an arbitrary dict
+
+```
+from pydantic import TypeAdapter
+json_schema = TypeAdapter(JSONSchema).validate_python(some_dict)
+```
+
+When used in a pydantic model as a field, the JSONSchema will be validated automatically.
+"""  # noqa: E501
+
+
+EMPTY_JSON_SCHEMA = TypeAdapter(JSONSchema).validate_python(
+    {
+        "type": "object",
+        "properties": {},
+        "required": [],
+    }
+)

packages/openassistants/openassistants/functions/base.py

@@ -6,6 +6,7 @@
 from langchain_core.language_models import BaseChatModel
 from openassistants.data_models.chat_messages import OpasMessage
 from openassistants.data_models.function_output import FunctionOutput
+from openassistants.data_models.json_schema import EMPTY_JSON_SCHEMA, JSONSchema
 from openassistants.functions.utils import AsyncStreamVersion
 from openassistants.utils.json_schema import PyRepr
 from pydantic import BaseModel
@@ -56,7 +57,7 @@ def get_sample_questions(self) -> Sequence[str]:
         pass
 
     @abc.abstractmethod
-    def get_parameters_json_schema(self) -> dict:
+    def get_parameters_json_schema(self) -> JSONSchema:
         """
         Get the json schema of the function's parameters
         """
@@ -120,13 +121,18 @@ def get_entities(self) -> Sequence[IEntity]:
         return self.entities
 
 
+class BaseFunctionParameters(BaseModel):
+    json_schema: JSONSchema = EMPTY_JSON_SCHEMA
+
+
 class BaseFunction(IBaseFunction, BaseModel, abc.ABC):
     id: str
     type: str
     display_name: Optional[str] = None
     description: str
     sample_questions: List[str] = []
     confirm: bool = False
+    parameters: BaseFunctionParameters = BaseFunctionParameters()
 
     def get_id(self) -> str:
         return self.id
@@ -146,5 +152,8 @@ def get_sample_questions(self) -> Sequence[str]:
     def get_confirm(self) -> bool:
         return self.confirm
 
+    def get_parameters_json_schema(self) -> JSONSchema:
+        return self.parameters.json_schema
+
     async def get_entity_configs(self) -> Mapping[str, IEntityConfig]:
         return {}