Use the JsonSchema conversion methods in models

Author: RobinPicard

outlines/models/dottxt.py

@@ -1,18 +1,9 @@
 """Integration with Dottxt's API."""
 
-import json
-from typing import TYPE_CHECKING, Any, Optional
-
-from pydantic import TypeAdapter
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 from outlines.models.base import Model, ModelTypeAdapter
 from outlines.types import CFG, JsonSchema, Regex
-from outlines.types.utils import (
-    is_dataclass,
-    is_genson_schema_builder,
-    is_pydantic_model,
-    is_typed_dict,
-)
 
 if TYPE_CHECKING:
     from dottxt import Dottxt as DottxtClient
@@ -77,20 +68,8 @@ def format_output_type(self, output_type: Optional[Any] = None) -> str:
                 "CFG-based structured outputs will soon be available with "
                 "Dottxt. Use an open source model in the meantime."
             )
-
-        elif isinstance(output_type, JsonSchema):
-            return output_type.schema
-        elif is_dataclass(output_type):
-            schema = TypeAdapter(output_type).json_schema()
-            return json.dumps(schema)
-        elif is_typed_dict(output_type):
-            schema = TypeAdapter(output_type).json_schema()
-            return json.dumps(schema)
-        elif is_pydantic_model(output_type):
-            schema = output_type.model_json_schema()
-            return json.dumps(schema)
-        elif is_genson_schema_builder(output_type):
-            return output_type.to_json()
+        elif JsonSchema.is_json_schema(output_type):
+            return cast(str, JsonSchema.convert_to(output_type, ["str"]))
         else:
             type_name = getattr(output_type, "__name__", output_type)
             raise TypeError(

outlines/models/gemini.py

@@ -14,14 +14,11 @@
 from outlines.models.base import Model, ModelTypeAdapter
 from outlines.types import CFG, Choice, JsonSchema, Regex
 from outlines.types.utils import (
-    is_dataclass,
     is_enum,
     get_enum_from_choice,
     get_enum_from_literal,
     is_genson_schema_builder,
     is_literal,
-    is_pydantic_model,
-    is_typed_dict,
     is_typing_list,
 )
 
@@ -171,28 +168,18 @@ def format_output_type(self, output_type: Optional[Any] = None) -> dict:
                 "CFG-based structured outputs are not available with Gemini. "
                 "Use an open source model or dottxt instead."
             )
-        elif is_genson_schema_builder(output_type):
-            raise TypeError(
-                "The Gemini SDK does not accept Genson schema builders as an "
-                "input. Pass a Pydantic model, typed dict or dataclass "
-                "instead."
-            )
-        elif isinstance(output_type, JsonSchema):
-            raise TypeError(
-                "The Gemini SDK does not accept Json Schemas as an input. "
-                "Pass a Pydantic model, typed dict or dataclass instead."
-            )
 
         if output_type is None:
             return {}
 
-        # Structured types
-        elif is_dataclass(output_type):
-            return self.format_json_output_type(output_type)
-        elif is_typed_dict(output_type):
-            return self.format_json_output_type(output_type)
-        elif is_pydantic_model(output_type):
-            return self.format_json_output_type(output_type)
+        # JSON schema types
+        elif JsonSchema.is_json_schema(output_type):
+            return self.format_json_output_type(
+                JsonSchema.convert_to(
+                    output_type,
+                    ["dataclass", "typeddict", "pydantic"]
+                )
+            )
 
         # List of structured types
         elif is_typing_list(output_type):
@@ -233,21 +220,20 @@ def format_list_output_type(self, output_type: Optional[Any]) -> dict:
         if len(args) == 1:
             item_type = args[0]
 
-            # Check if list item type is supported
-            if (
-                is_pydantic_model(item_type)
-                or is_typed_dict(item_type)
-                or is_dataclass(item_type)
-            ):
+            if JsonSchema.is_json_schema(item_type):
                 return {
                     "response_mime_type": "application/json",
-                    "response_schema": output_type,
+                    "response_schema": list[  # type: ignore
+                        JsonSchema.convert_to(
+                            item_type,
+                            ["dataclass", "typeddict", "pydantic"]
+                        )
+                    ],
                 }
-
             else:
                 raise TypeError(
-                    "The only supported types for list items are Pydantic "
-                    + "models, typed dicts and dataclasses."
+                    "The list items output type must contain a JSON schema "
+                    "type."
                 )
 
         raise TypeError(

outlines/models/ollama.py

@@ -1,20 +1,19 @@
 """Integration with the `ollama` library."""
 
-import json
 from functools import singledispatchmethod
-from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, Optional, Union
-
-from pydantic import TypeAdapter
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncIterator,
+    Iterator,
+    Optional,
+    Union,
+    cast,
+)
 
 from outlines.inputs import Chat, Image
 from outlines.models.base import AsyncModel, Model, ModelTypeAdapter
 from outlines.types import CFG, JsonSchema, Regex
-from outlines.types.utils import (
-    is_dataclass,
-    is_genson_schema_builder,
-    is_pydantic_model,
-    is_typed_dict,
-)
 
 if TYPE_CHECKING:
     from ollama import Client
@@ -109,7 +108,7 @@ def _create_message(self, role: str, content: str | list) -> dict:
 
     def format_output_type(
         self, output_type: Optional[Any] = None
-    ) -> Optional[str]:
+    ) -> Optional[dict]:
         """Format the output type to pass to the client.
 
         TODO: `int`, `float` and other Python types could be supported via
@@ -126,7 +125,9 @@ def format_output_type(
             The formatted output type to be passed to the model.
 
         """
-        if isinstance(output_type, Regex):
+        if output_type is None:
+            return None
+        elif isinstance(output_type, Regex):
             raise TypeError(
                 "Regex-based structured outputs are not supported by Ollama. "
                 "Use an open source model in the meantime."
@@ -136,22 +137,8 @@ def format_output_type(
                 "CFG-based structured outputs are not supported by Ollama. "
                 "Use an open source model in the meantime."
             )
-
-        if output_type is None:
-            return None
-        elif isinstance(output_type, JsonSchema):
-            return json.loads(output_type.schema)
-        elif is_dataclass(output_type):
-            schema = TypeAdapter(output_type).json_schema()
-            return schema
-        elif is_typed_dict(output_type):
-            schema = TypeAdapter(output_type).json_schema()
-            return schema
-        elif is_pydantic_model(output_type):
-            schema = output_type.model_json_schema()
-            return schema
-        elif is_genson_schema_builder(output_type):
-            return output_type.to_json()
+        elif JsonSchema.is_json_schema(output_type):
+            return cast(dict, JsonSchema.convert_to(output_type, ["dict"]))
         else:
             type_name = getattr(output_type, "__name__", output_type)
             raise TypeError(

outlines/models/openai.py

@@ -1,29 +1,23 @@
 """Integration with OpenAI's API."""
 
-import json
 from typing import (
     TYPE_CHECKING,
     Any,
     AsyncIterator,
     Iterator,
     Optional,
     Union,
+    cast,
 )
 from functools import singledispatchmethod
 
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel
 
 from outlines.inputs import Chat, Image
 from outlines.models.base import AsyncModel, Model, ModelTypeAdapter
 from outlines.models.utils import set_additional_properties_false_json_schema
 from outlines.types import JsonSchema, Regex, CFG
-from outlines.types.utils import (
-    is_dataclass,
-    is_typed_dict,
-    is_pydantic_model,
-    is_genson_schema_builder,
-    is_native_dict
-)
+from outlines.types.utils import is_native_dict
 
 if TYPE_CHECKING:
     from openai import (
@@ -176,20 +170,10 @@ def format_output_type(self, output_type: Optional[Any] = None) -> dict:
             return {}
         elif is_native_dict(output_type):
             return self.format_json_mode_type()
-        elif is_dataclass(output_type):
-            output_type = TypeAdapter(output_type).json_schema()
-            return self.format_json_output_type(output_type)
-        elif is_typed_dict(output_type):
-            output_type = TypeAdapter(output_type).json_schema()
-            return self.format_json_output_type(output_type)
-        elif is_pydantic_model(output_type):
-            output_type = output_type.model_json_schema()
-            return self.format_json_output_type(output_type)
-        elif is_genson_schema_builder(output_type):
-            schema = json.loads(output_type.to_json())
-            return self.format_json_output_type(schema)
-        elif isinstance(output_type, JsonSchema):
-            return self.format_json_output_type(json.loads(output_type.schema))
+        elif JsonSchema.is_json_schema(output_type):
+            return self.format_json_output_type(
+                cast(dict, JsonSchema.convert_to(output_type, ["dict"]))
+            )
         else:
             type_name = getattr(output_type, "__name__", output_type)
             raise TypeError(

tests/models/test_gemini_type_adapter.py

@@ -3,16 +3,16 @@
 import sys
 from dataclasses import dataclass
 from enum import Enum, EnumMeta
-from typing import Literal
+from typing import Literal, get_args
 
 from PIL import Image as PILImage
 from genson import SchemaBuilder
-from google.genai import types
 from pydantic import BaseModel
 
 from outlines import cfg, json_schema, regex
 from outlines.inputs import Chat, Image
 from outlines.models.gemini import GeminiTypeAdapter
+from outlines.types.utils import is_dataclass
 
 if sys.version_info >= (3, 12):
     from typing import TypedDict
@@ -135,19 +135,29 @@ def test_gemini_type_adapter_output_invalid(adapter):
     with pytest.raises(TypeError, match="CFG-based structured outputs"):
         adapter.format_output_type(cfg(""))
 
-    with pytest.raises(TypeError, match="The Gemini SDK does not accept"):
-        adapter.format_output_type(SchemaBuilder())
-
-    with pytest.raises(TypeError, match="The Gemini SDK does not"):
-        adapter.format_output_type(json_schema(""))
-
 
 def test_gemini_type_adapter_output_none(adapter):
     result = adapter.format_output_type(None)
     assert result == {}
 
 
-def test_gemini_type_adapter_output_dataclass(adapter, schema):
+def test_gemini_type_adapter_output_json_schema(adapter, schema):
+    result = adapter.format_output_type(json_schema(schema))
+    assert isinstance(result, dict)
+    assert result["response_mime_type"] == "application/json"
+    assert is_dataclass(result["response_schema"])
+
+
+def test_gemini_type_adapter_output_list_json_schema(adapter, schema):
+    result = adapter.format_output_type(list[json_schema(schema)])
+    assert isinstance(result, dict)
+    assert result["response_mime_type"] == "application/json"
+    args = get_args(result["response_schema"])
+    assert len(args) == 1
+    assert is_dataclass(args[0])
+
+
+def test_gemini_type_adapter_output_dataclass(adapter):
     @dataclass
     class User:
         user_id: int
@@ -160,7 +170,7 @@ class User:
     }
 
 
-def test_gemini_type_adapter_output_list_dataclass(adapter, schema):
+def test_gemini_type_adapter_output_list_dataclass(adapter):
     class User(BaseModel):
         user_id: int
         name: str
@@ -172,7 +182,7 @@ class User(BaseModel):
     }
 
 
-def test_gemini_type_adapter_output_typed_dict(adapter, schema):
+def test_gemini_type_adapter_output_typed_dict(adapter):
     class User(TypedDict):
         user_id: int
         name: str
@@ -184,7 +194,7 @@ class User(TypedDict):
     }
 
 
-def test_gemini_type_adapter_output_list_typed_dict(adapter, schema):
+def test_gemini_type_adapter_output_list_typed_dict(adapter):
     class User(BaseModel):
         user_id: int
         name: str
@@ -196,7 +206,7 @@ class User(BaseModel):
     }
 
 
-def test_gemini_type_adapter_output_pydantic(adapter, schema):
+def test_gemini_type_adapter_output_pydantic(adapter):
     class User(BaseModel):
         user_id: int
         name: str
@@ -208,7 +218,7 @@ class User(BaseModel):
     }
 
 
-def test_gemini_type_adapter_output_list_pydantic(adapter, schema):
+def test_gemini_type_adapter_output_list_pydantic(adapter):
     class User(BaseModel):
         user_id: int
         name: str
@@ -220,6 +230,26 @@ class User(BaseModel):
     }
 
 
+def test_gemini_type_adapter_output_genson_schema_builder(adapter):
+    builder = SchemaBuilder()
+    builder.add_schema({"type": "object", "properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}}, "required": ["foo"]})
+    result = adapter.format_output_type(builder)
+    assert isinstance(result, dict)
+    assert result["response_mime_type"] == "application/json"
+    assert is_dataclass(result["response_schema"])
+
+
+def test_gemini_type_adapter_output_list_genson_schema_builder(adapter):
+    builder = SchemaBuilder()
+    builder.add_schema({"type": "object", "properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}}, "required": ["foo"]})
+    result = adapter.format_output_type(list[builder])
+    assert isinstance(result, dict)
+    assert result["response_mime_type"] == "application/json"
+    args = get_args(result["response_schema"])
+    assert len(args) == 1
+    assert is_dataclass(args[0])
+
+
 def test_gemini_type_adapter_output_enum(adapter):
     class Foo(Enum):
         Bar = "bar"