From 1f30217aec31eb83f886a48e5426e868c59a9851 Mon Sep 17 00:00:00 2001
From: Maziyar Panahi <maziyar.panahi@iscpif.fr>
Date: Sun, 28 Apr 2024 22:04:19 +0200
Subject: [PATCH] Add the missing `tool_prompt` to Python client

---
 clients/python/text_generation/client.py | 8 ++++++++
 clients/python/text_generation/types.py  | 2 ++
 2 files changed, 10 insertions(+)

diff --git a/clients/python/text_generation/client.py b/clients/python/text_generation/client.py
index 95d23901..0e86901d 100644
--- a/clients/python/text_generation/client.py
+++ b/clients/python/text_generation/client.py
@@ -80,6 +80,7 @@ class Client:
         temperature: Optional[float] = None,
         top_p: Optional[float] = None,
         tools: Optional[List[Tool]] = None,
+        tool_prompt: Optional[str] = None,
         tool_choice: Optional[str] = None,
     ):
         """
@@ -119,6 +120,8 @@ class Client:
                 higher are kept for generation
             tools (`List[Tool]`):
                 List of tools to use
+            tool_prompt (`str`):
+                A prompt to be appended before the tools
             tool_choice (`str`):
                 The tool to use
 
@@ -139,6 +142,7 @@ class Client:
             temperature=temperature,
             top_p=top_p,
             tools=tools,
+            tool_prompt=tool_prompt,
             tool_choice=tool_choice,
         )
         if not stream:
@@ -466,6 +470,7 @@ class AsyncClient:
         temperature: Optional[float] = None,
         top_p: Optional[float] = None,
         tools: Optional[List[Tool]] = None,
+        tool_prompt: Optional[str] = None,
         tool_choice: Optional[str] = None,
     ) -> Union[ChatComplete, AsyncIterator[ChatCompletionChunk]]:
         """
@@ -505,6 +510,8 @@ class AsyncClient:
                 higher are kept for generation
             tools (`List[Tool]`):
                 List of tools to use
+            tool_prompt (`str`):
+                A prompt to be appended before the tools
             tool_choice (`str`):
                 The tool to use
 
@@ -525,6 +532,7 @@ class AsyncClient:
             temperature=temperature,
             top_p=top_p,
             tools=tools,
+            tool_prompt=tool_prompt,
             tool_choice=tool_choice,
         )
         if not stream:
diff --git a/clients/python/text_generation/types.py b/clients/python/text_generation/types.py
index cfa2a9ed..5e32bc6f 100644
--- a/clients/python/text_generation/types.py
+++ b/clients/python/text_generation/types.py
@@ -159,6 +159,8 @@ class ChatRequest(BaseModel):
     top_p: Optional[float] = None
     # List of tools to be used
     tools: Optional[List[Tool]] = None
+    # A prompt to be appended before the tools
+    tool_prompt: Optional[str] = None
     # Choice of tool to be used
     tool_choice: Optional[str] = None