Implement tools and outputs for the Ollama model

Author: RobinPicard

outlines/models/ollama.py

@@ -2,12 +2,22 @@
 
 import json
 from functools import singledispatchmethod
-from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncIterator,
+    Iterator,
+    List,
+    Optional,
+    Union,
+)
 
 from pydantic import TypeAdapter
 
 from outlines.inputs import Chat, Image
 from outlines.models.base import AsyncModel, Model, ModelTypeAdapter
+from outlines.outputs import Output, StreamingOutput
+from outlines.tools import ToolDef
 from outlines.types import CFG, JsonSchema, Regex
 from outlines.types.utils import (
     is_dataclass,
@@ -74,7 +84,7 @@ def format_chat_model_input(self, model_input: Chat) -> list:
 
         """
         return [
-            self._create_message(message["role"], message["content"])
+            self._create_message(message["role"], message["content"])  # type: ignore
             for message in model_input.messages
         ]
 
@@ -107,9 +117,7 @@ def _create_message(self, role: str, content: str | list) -> dict:
                 "and a list of images."
             )
 
-    def format_output_type(
-        self, output_type: Optional[Any] = None
-    ) -> Optional[str]:
+    def format_output_type(self, output_type: Optional[Any]) -> Optional[str]:
         """Format the output type to pass to the client.
 
         TODO: `int`, `float` and other Python types could be supported via
@@ -159,6 +167,13 @@ def format_output_type(
                 "Consider using a local model instead."
             )
 
+    def format_tools(self, tools):
+        """Not available for Ollama."""
+        if tools:
+            raise NotImplementedError(
+                "Tools are not available for Ollama."
+            )
+
 
 class Ollama(Model):
     """Thin wrapper around the `ollama.Client` client.
@@ -184,9 +199,10 @@ def __init__(self, client: "Client", model_name: Optional[str] = None):
 
     def generate(self,
         model_input: Chat | str | list,
-        output_type: Optional[Any] = None,
+        output_type: Optional[Any],
+        tools: Optional[List[ToolDef]],
         **kwargs: Any,
-    ) -> str:
+    ) -> Output:
         """Generate text using Ollama.
 
         Parameters
@@ -197,15 +213,19 @@ def generate(self,
             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.
         **kwargs
             Additional keyword arguments to pass to the client.
 
         Returns
         -------
-        str
+        Output
             The text generated by the model.
 
         """
+        self.type_adapter.format_tools(tools)
+
         if "model" not in kwargs and self.model_name is not None:
             kwargs["model"] = self.model_name
 
@@ -214,12 +234,14 @@ def generate(self,
             format=self.type_adapter.format_output_type(output_type),
             **kwargs,
         )
-        return response.message.content
+
+        return Output(content=response.message.content)
 
     def generate_batch(
         self,
         model_input,
-        output_type = None,
+        output_type,
+        tools,
         **kwargs,
     ):
         raise NotImplementedError(
@@ -229,9 +251,10 @@ def generate_batch(
     def generate_stream(
         self,
         model_input: Chat | str | list,
-        output_type: Optional[Any] = None,
+        output_type: Optional[Any],
+        tools: Optional[List[ToolDef]],
         **kwargs: Any,
-    ) -> Iterator[str]:
+    ) -> Iterator[StreamingOutput]:
         """Stream text using Ollama.
 
         Parameters
@@ -242,15 +265,19 @@ def generate_stream(
             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.
         **kwargs
             Additional keyword arguments to pass to the client.
 
         Returns
         -------
-        Iterator[str]
+        Iterator[StreamingOutput]
             An iterator that yields the text generated by the model.
 
         """
+        self.type_adapter.format_tools(tools)
+
         if "model" not in kwargs and self.model_name is not None:
             kwargs["model"] = self.model_name
 
@@ -260,8 +287,9 @@ def generate_stream(
             stream=True,
             **kwargs,
         )
+
         for chunk in response:
-            yield chunk.message.content
+            yield StreamingOutput(content=chunk.message.content)
 
 
 class AsyncOllama(AsyncModel):
@@ -290,9 +318,10 @@ def __init__(
 
     async def generate(self,
         model_input: Chat | str | list,
-        output_type: Optional[Any] = None,
+        output_type: Optional[Any],
+        tools: Optional[List[ToolDef]],
         **kwargs: Any,
-    ) -> str:
+    ) -> Output:
         """Generate text using Ollama.
 
         Parameters
@@ -303,15 +332,19 @@ async def generate(self,
             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.
         **kwargs
             Additional keyword arguments to pass to the client.
 
         Returns
         -------
-        str
+        Output
             The text generated by the model.
 
         """
+        self.type_adapter.format_tools(tools)
+
         if "model" not in kwargs and self.model_name is not None:
             kwargs["model"] = self.model_name
 
@@ -320,12 +353,14 @@ async def generate(self,
             format=self.type_adapter.format_output_type(output_type),
             **kwargs,
         )
-        return response.message.content
+
+        return Output(content=response.message.content)
 
     async def generate_batch(
         self,
         model_input,
-        output_type = None,
+        output_type,
+        tools,
         **kwargs,
     ):
         raise NotImplementedError(
@@ -335,9 +370,10 @@ async def generate_batch(
     async def generate_stream( # type: ignore
         self,
         model_input: Chat | str | list,
-        output_type: Optional[Any] = None,
+        output_type: Optional[Any],
+        tools: Optional[List[ToolDef]],
         **kwargs: Any,
-    ) -> AsyncIterator[str]:
+    ) -> AsyncIterator[StreamingOutput]:
         """Stream text using Ollama.
 
         Parameters
@@ -348,15 +384,19 @@ async def generate_stream( # type: ignore
             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.
         **kwargs
             Additional keyword arguments to pass to the client.
 
         Returns
         -------
-        Iterator[str]
+        Iterator[StreamingOutput]
             An iterator that yields the text generated by the model.
 
         """
+        self.type_adapter.format_tools(tools)
+
         if "model" not in kwargs and self.model_name is not None:
             kwargs["model"] = self.model_name
 
@@ -366,8 +406,9 @@ async def generate_stream( # type: ignore
             stream=True,
             **kwargs,
         )
+
         async for chunk in stream:
-            yield chunk.message.content
+            yield StreamingOutput(content=chunk.message.content)
 
 
 def from_ollama(