Draft

Author: RobinPicard

outlines/generator.py

@@ -21,6 +21,8 @@
     get_regex_logits_processor,
 )
 from outlines.backends.base import LogitsProcessorType
+from outlines.outputs import Output, StreamingOutput
+from outlines.tools import get_formatted_tools, ToolsInput
 from outlines.types import CFG, JsonSchema
 from outlines.types.dsl import python_types_to_terms, to_regex
 
@@ -35,20 +37,30 @@ class BlackBoxGenerator:
     """
     output_type: Optional[Any]
 
-    def __init__(self, model: BlackBoxModel, output_type: Optional[Any]):
+    def __init__(
+        self,
+        model: BlackBoxModel,
+        output_type: Optional[Any],
+        *,
+        tools: Optional[ToolsInput] = None,
+    ):
         """
         Parameters
         ----------
         model
             An instance of an Outlines model.
         output_type
             The output type that will be used to constrain the generation.
+        tools
+            A list of tools to use for the generator. Can contain an MCPServer,
+            a list of ToolDef, Callable, or BaseModel instances.
 
         """
         self.model = model
         self.output_type = output_type
+        self.tools = get_formatted_tools(tools)
 
-    def __call__(self, prompt: Any, **inference_kwargs) -> Any:
+    def __call__(self, prompt: Any, **inference_kwargs) -> Output | List[Output]:
         """Generate a response from the model.
 
         Parameters
@@ -60,15 +72,17 @@ def __call__(self, prompt: Any, **inference_kwargs) -> Any:
 
         Returns
         -------
-        Any
-            The response generated by the model.
+        Output | List[Output]
+            The output generated by the model.
 
         """
         return self.model.generate(
-            prompt, self.output_type, **inference_kwargs
+            prompt, self.output_type, tools=self.tools, **inference_kwargs
         )
 
-    def batch(self, prompts: List[Any], **inference_kwargs) -> List[Any]:
+    def batch(
+        self, prompts: List[Any], **inference_kwargs
+    ) -> List[Output] | List[List[Output]]:
         """Generate a batch of responses from the model.
 
         Parameters
@@ -80,15 +94,17 @@ def batch(self, prompts: List[Any], **inference_kwargs) -> List[Any]:
 
         Returns
         -------
-        List[Any]
-            The list of responses generated by the model.
+        List[Output] | List[List[Output]]
+            The list of outputs generated by the model.
 
         """
         return self.model.generate_batch(
-            prompts, self.output_type, **inference_kwargs
+            prompts, self.output_type, tools=self.tools, **inference_kwargs
         )
 
-    def stream(self, prompt: Any, **inference_kwargs) -> Iterator[Any]:
+    def stream(
+        self, prompt: Any, **inference_kwargs
+    ) -> Iterator[StreamingOutput]:
         """Generate a stream of responses from the model.
 
         Parameters
@@ -100,12 +116,12 @@ def stream(self, prompt: Any, **inference_kwargs) -> Iterator[Any]:
 
         Returns
         -------
-        Any
-            The response generated by the model.
+        Iterator[StreamingOutput]
+            A stream of StreamingOutput generated by the model.
 
         """
         return self.model.generate_stream(
-            prompt, self.output_type, **inference_kwargs
+            prompt, self.output_type, tools=self.tools, **inference_kwargs
         )
 
 
@@ -119,20 +135,32 @@ class AsyncBlackBoxGenerator:
     """
     output_type: Optional[Any]
 
-    def __init__(self, model: AsyncBlackBoxModel, output_type: Optional[Any]):
+    def __init__(
+        self,
+        model: AsyncBlackBoxModel,
+        output_type: Optional[Any],
+        *,
+        tools: Optional[ToolsInput] = None,
+    ):
         """
         Parameters
         ----------
         model
             An instance of an Outlines model.
         output_type
             The output type that will be used to constrain the generation.
+        tools
+            A list of tools to use for the generator. Can contain an MCPServer,
+            a list of ToolDef, Callable, or BaseModel instances.
 
         """
         self.model = model
         self.output_type = output_type
+        self.tools = get_formatted_tools(tools)
 
-    async def __call__(self, prompt: Any, **inference_kwargs) -> Any:
+    async def __call__(
+        self, prompt: Any, **inference_kwargs
+    ) -> Output | List[Output]:
         """Generate a response from the model.
 
         Parameters
@@ -144,15 +172,17 @@ async def __call__(self, prompt: Any, **inference_kwargs) -> Any:
 
         Returns
         -------
-        Any
-            The response generated by the model.
+        Output | List[Output]
+            The output generated by the model.
 
         """
         return await self.model.generate(
-            prompt, self.output_type, **inference_kwargs
+            prompt, self.output_type, tools=self.tools, **inference_kwargs
         )
 
-    async def batch(self, prompts: List[Any], **inference_kwargs) -> List[Any]:
+    async def batch(
+        self, prompts: List[Any], **inference_kwargs
+    ) -> List[Output] | List[List[Output]]:
         """Generate a batch of responses from the model.
 
         Parameters
@@ -164,15 +194,17 @@ async def batch(self, prompts: List[Any], **inference_kwargs) -> List[Any]:
 
         Returns
         -------
-        List[Any]
-            The list of responses generated by the model.
+        List[Output] | List[List[Output]]
+            The list of outputs generated by the model.
 
         """
         return await self.model.generate_batch(
-            prompts, self.output_type, **inference_kwargs
+            prompts, self.output_type, tools=self.tools, **inference_kwargs
         )
 
-    async def stream(self, prompt: Any, **inference_kwargs) -> AsyncIterator[Any]:
+    async def stream(
+        self, prompt: Any, **inference_kwargs
+    ) -> AsyncIterator[StreamingOutput]:
         """Generate a stream of responses from the model.
 
         Parameters
@@ -184,12 +216,13 @@ async def stream(self, prompt: Any, **inference_kwargs) -> AsyncIterator[Any]:
 
         Returns
         -------
-        Any
-            The response generated by the model.
+        AsyncIterator[StreamingOutput]
+            A coroutine that will produce an async iterator of StreamingOutput
+            produced by the model.
 
         """
         async for chunk in self.model.generate_stream(  # pragma: no cover
-            prompt, self.output_type, **inference_kwargs
+            prompt, self.output_type, tools=self.tools, **inference_kwargs
         ):
             yield chunk
 
@@ -218,6 +251,8 @@ def __init__(
         model: SteerableModel,
         output_type: Optional[Any],
         backend_name: Optional[str] = None,
+        *,
+        tools: Optional[ToolsInput] = None,
     ):
         """
         Parameters
@@ -228,9 +263,13 @@ def __init__(
             The output type expressed as a Python type
         backend_name
             The name of the backend to use to create the logits processor.
+        tools
+            A list of tools to use for the generator. Can contain an MCPServer,
+            a list of ToolDef, Callable, or BaseModel instances.
 
         """
         self.model = model
+        self.tools = get_formatted_tools(tools)
         if output_type is None:
             self.logits_processor = None
         else:
@@ -258,7 +297,11 @@ def __init__(
 
     @classmethod
     def from_processor(
-        cls, model: SteerableModel, processor: LogitsProcessorType
+        cls,
+        model: SteerableModel,
+        processor: LogitsProcessorType,
+        *,
+        tools: Optional[ToolsInput] = None,
     ):
         """Create a generator from a logits processor.
 
@@ -270,13 +313,12 @@ def from_processor(
             An instance of a logits processor.
 
         """
-        instance = cls.__new__(cls)
-        instance.model = model
+        instance = cls(model, None, tools=tools)
         instance.logits_processor = processor
 
         return instance
 
-    def __call__(self, prompt: Any, **inference_kwargs) -> Any:
+    def __call__(self, prompt: Any, **inference_kwargs) -> Output | List[Output]:
         """Generate a response from the model.
 
         Parameters
@@ -288,17 +330,19 @@ def __call__(self, prompt: Any, **inference_kwargs) -> Any:
 
         Returns
         -------
-        Any
-            The response generated by the model.
+        Output | List[Output]
+            The output generated by the model.
 
         """
         if self.logits_processor is not None:
             self.logits_processor.reset()
         return self.model.generate(
-            prompt, self.logits_processor, **inference_kwargs
+            prompt, self.logits_processor, tools=self.tools, **inference_kwargs
         )
 
-    def batch(self, prompts: List[Any], **inference_kwargs) -> List[Any]:
+    def batch(
+        self, prompts: List[Any], **inference_kwargs
+    ) -> List[Output] | List[List[Output]]:
         """Generate a batch of responses from the model.
 
         Parameters
@@ -310,17 +354,19 @@ def batch(self, prompts: List[Any], **inference_kwargs) -> List[Any]:
 
         Returns
         -------
-        List[Any]
-            The list of responses generated by the model.
+        List[Output] | List[List[Output]]
+            The list of outputs generated by the model.
 
         """
         if self.logits_processor is not None:
             self.logits_processor.reset()
         return self.model.generate_batch(
-            prompts, self.logits_processor, **inference_kwargs
+            prompts, self.logits_processor, tools=self.tools, **inference_kwargs
         )
 
-    def stream(self, prompt: Any, **inference_kwargs) -> Iterator[Any]:
+    def stream(
+        self, prompt: Any, **inference_kwargs
+    ) -> Iterator[StreamingOutput]:
         """Generate a stream of responses from the model.
 
         Parameters
@@ -332,14 +378,14 @@ def stream(self, prompt: Any, **inference_kwargs) -> Iterator[Any]:
 
         Returns
         -------
-        Any
-            The response generated by the model.
+        Iterator[StreamingOutput]
+            A stream of StreamingOutput generated by the model.
 
         """
         if self.logits_processor is not None:
             self.logits_processor.reset()
         return self.model.generate_stream(
-            prompt, self.logits_processor, **inference_kwargs
+            prompt, self.logits_processor, tools=self.tools, **inference_kwargs
         )
 
 
@@ -348,6 +394,7 @@ def Generator(
     output_type: Optional[Any] = None,
     backend: Optional[str] = None,
     *,
+    tools: Optional[ToolsInput] = None,
     processor: Optional[LogitsProcessorType] = None,
 ) -> Union[SteerableGenerator, BlackBoxGenerator, AsyncBlackBoxGenerator]:
     """Create a generator for the given model and output parameters.
@@ -367,6 +414,9 @@ def Generator(
         The name of the backend to use to create the logits processor. Only
         used for steerable models if there is an output type and `processor` is
         not provided.
+    tools
+        A list of tools to use for the generator. Can contain an MCPServer,
+        a list of ToolDef, Callable, or BaseModel instances.
     processor
         An instance of a logits processor.
 
@@ -387,18 +437,18 @@ def Generator(
 
     if isinstance(model, SteerableModel): # type: ignore
         if processor is not None:
-            return SteerableGenerator.from_processor(model, processor) # type: ignore
+            return SteerableGenerator.from_processor(model, processor, tools=tools) # type: ignore
         else:
-            return SteerableGenerator(model, output_type, backend) # type: ignore
+            return SteerableGenerator(model, output_type, backend, tools=tools) # type: ignore
     else:
         if processor is not None:
             raise NotImplementedError(
                 "This model does not support logits processors"
             )
         if isinstance(model, AsyncBlackBoxModel): # type: ignore
-            return AsyncBlackBoxGenerator(model, output_type) # type: ignore
+            return AsyncBlackBoxGenerator(model, output_type, tools=tools) # type: ignore
         elif isinstance(model, BlackBoxModel): # type: ignore
-            return BlackBoxGenerator(model, output_type) # type: ignore
+            return BlackBoxGenerator(model, output_type, tools=tools) # type: ignore
         else:
             raise ValueError(
                 "The model argument must be an instance of "