import inspect
from pydantic import BaseModel
from rich.json import JSON
[docs]
class BaseIOSchema(BaseModel):
"""Base schema for input/output in the Atomic Agents framework."""
def __str__(self):
return self.model_dump_json()
def __rich__(self):
json_str = self.model_dump_json()
return JSON(json_str)
@classmethod
def __pydantic_init_subclass__(cls, **kwargs):
super().__pydantic_init_subclass__(**kwargs)
cls._validate_description()
@classmethod
def _validate_description(cls):
description = cls.__doc__
if not description or not description.strip():
if cls.__module__ != "instructor.function_calls" and not hasattr(cls, "from_streaming_response"):
raise ValueError(f"{cls.__name__} must have a non-empty docstring to serve as its description")
[docs]
@classmethod
def model_json_schema(cls, *args, **kwargs):
schema = super().model_json_schema(*args, **kwargs)
if "description" not in schema and cls.__doc__:
schema["description"] = inspect.cleandoc(cls.__doc__)
if "title" not in schema:
schema["title"] = cls.__name__
return schema