Implement tools and outputs for the Dottxt model

Author: RobinPicard

outlines/models/dottxt.py

@@ -1,11 +1,13 @@
 """Integration with Dottxt's API."""
 
 import json
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Optional, List
 
 from pydantic import TypeAdapter
 
 from outlines.models.base import Model, ModelTypeAdapter
+from outlines.outputs import Output
+from outlines.tools import ToolDef
 from outlines.types import CFG, JsonSchema, Regex
 from outlines.types.utils import (
     is_dataclass,
@@ -44,7 +46,7 @@ def format_input(self, model_input: str) -> str:
             "The only available type is `str`."
         )
 
-    def format_output_type(self, output_type: Optional[Any] = None) -> str:
+    def format_output_type(self, output_type: Optional[Any]) -> str:
         """Format the output type to pass to the client.
 
         TODO: `int`, `float` and other Python types could be supported via
@@ -98,6 +100,13 @@ def format_output_type(self, output_type: Optional[Any] = None) -> str:
                 "Consider using a local mode instead."
             )
 
+    def format_tools(self, tools):
+        """Not implemented for Dottxt."""
+        if tools:
+            raise NotImplementedError(
+                "Dottxt does not support tools."
+            )
+
 
 class Dottxt(Model):
     """Thin wrapper around the `dottxt.client.Dottxt` client.
@@ -132,9 +141,10 @@ def __init__(
     def generate(
         self,
         model_input: str,
-        output_type: Optional[Any] = None,
+        output_type: Optional[Any],
+        tools: Optional[List[ToolDef]],
         **inference_kwargs: Any,
-    ) -> str:
+    ) -> Output:
         """Generate text using Dottxt.
 
         Parameters
@@ -145,15 +155,18 @@ def generate(
             The desired format of the response generated by the model. The
             output type must be of a type that can be converted to a JSON
             schema.
+        tools
+            The tools to use for the generation.
         **inference_kwargs
             Additional keyword arguments to pass to the client.
 
         Returns
         -------
-        str
+        Output
             The text generated by the model.
 
         """
+        self.type_adapter.format_tools(tools)
         prompt = self.type_adapter.format_input(model_input)
         json_schema = self.type_adapter.format_output_type(output_type)
 
@@ -174,22 +187,26 @@ def generate(
             json_schema,
             **inference_kwargs,
         )
-        return completion.data
+
+        return Output(content=completion.data)
 
     def generate_batch(
         self,
         model_input,
-        output_type = None,
+        output_type,
+        tools,
         **inference_kwargs,
     ):
+        """Not available for Dottxt."""
         raise NotImplementedError(
             "Dottxt does not support batch generation."
         )
 
     def generate_stream(
         self,
         model_input,
-        output_type=None,
+        output_type,
+        tools,
         **inference_kwargs,
     ):
         """Not available for Dottxt."""

tests/models/test_dottxt.py

@@ -8,6 +8,7 @@
 import outlines
 from outlines import Generator
 from outlines.models.dottxt import Dottxt
+from outlines.outputs import Output
 
 
 MODEL_NAME = "dottxt/dottxt-v1-alpha"
@@ -99,7 +100,8 @@ def test_dottxt_wrong_inference_parameters(model_no_model_name):
 @pytest.mark.api_call
 def test_dottxt_direct_pydantic_call(model_no_model_name):
     result = model_no_model_name("Create a user", User)
-    assert "first_name" in json.loads(result)
+    assert isinstance(result, Output)
+    assert "first_name" in json.loads(result.content)
 
 
 @pytest.mark.api_call
@@ -112,14 +114,16 @@ def test_dottxt_direct_jsonschema_call(
         model_name=model_name_and_revision[0],
         model_revision=model_name_and_revision[1],
     )
-    assert "first_name" in json.loads(result)
+    assert isinstance(result, Output)
+    assert "first_name" in json.loads(result.content)
 
 
 @pytest.mark.api_call
 def test_dottxt_generator_pydantic_call(model):
     generator = Generator(model, User)
     result = generator("Create a user")
-    assert "first_name" in json.loads(result)
+    assert isinstance(result, Output)
+    assert "first_name" in json.loads(result.content)
 
 
 @pytest.mark.api_call

tests/models/test_dottxt_type_adapter.py

@@ -10,6 +10,7 @@
 
 from outlines.inputs import Image
 from outlines.models.dottxt import DottxtTypeAdapter
+from outlines.tools import ToolDef
 from outlines.types import cfg, json_schema, regex
 
 if sys.version_info >= (3, 12):
@@ -58,7 +59,7 @@ def test_dottxt_type_adapter_input_text(adapter):
 
 
 def test_dottxt_type_adapter_input_invalid(adapter, image):
-    prompt = ["prompt", image]
+    prompt = ["prompt", Image(image)]
     with pytest.raises(TypeError, match="The input type"):
         _ = adapter.format_input(prompt)
 
@@ -135,3 +136,15 @@ def test_dottxt_type_adapter_json_schema_str(adapter, schema):
 def test_dottxt_type_adapter_json_schema_dict(adapter, schema):
     result = adapter.format_output_type(json_schema(schema))
     assert result == json.dumps(schema)
+
+
+def test_dottxt_type_adapter_tools(adapter):
+    with pytest.raises(
+        NotImplementedError,
+        match="Dottxt does not support tools."
+    ):
+        adapter.format_tools(
+            [ToolDef(name="test", description="test", parameters={})]
+        )
+
+    adapter.format_tools(None)