From c126d2be5b924dc9a5d8565462eb8febc7892706 Mon Sep 17 00:00:00 2001
From: monoxgas <monoxgas@gmail.com>
Date: Wed, 10 Apr 2024 16:34:19 -0600
Subject: [PATCH] 0.2.1

- Tool prompt updates
- Added force to using()
- Remove stop sequence for tool calls
---
 README.md         |  2 +-
 pyproject.toml    |  2 +-
 rigging/chat.py   | 30 ++++++++++++++++++++++++++----
 rigging/prompt.py |  9 +++++----
 4 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 4988459..5b81915 100644
--- a/README.md
+++ b/README.md
@@ -150,7 +150,7 @@ chat = (
             {"role": "user", "content": "What is the weather in London?"},
         ]
     )
-    .using(WeatherTool())
+    .using(WeatherTool(), force=True)
     .run()
 )
 
diff --git a/pyproject.toml b/pyproject.toml
index 70403d4..83c1c0c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "rigging"
-version = "0.2.0"
+version = "0.2.1"
 description = "LLM Interaction Framework"
 authors = ["Nick Landers <monoxgas@gmail.com>"]
 license = "MIT"
diff --git a/rigging/chat.py b/rigging/chat.py
index 00562b1..c6d9f05 100644
--- a/rigging/chat.py
+++ b/rigging/chat.py
@@ -137,6 +137,7 @@ def __init__(self, generator: "Generator", messages: t.Sequence[Message], params
         self.until_types: list[type[Model]] = []
         self.until_tools: list[Tool] = []
         self.inject_tool_prompt: bool = True
+        self.force_tool: bool = False
 
         self.params = params
 
@@ -180,6 +181,7 @@ def using(
         self,
         tool: Tool | t.Sequence[Tool],
         *,
+        force: bool = False,
         attempt_recovery: bool = True,
         drop_dialog: bool = False,
         max_rounds: int = DEFAULT_MAX_ROUNDS,
@@ -187,8 +189,16 @@ def using(
     ) -> "PendingChat":
         self.until_tools += tool if isinstance(tool, t.Sequence) else [tool]
         self.inject_tool_prompt = inject_prompt or self.inject_tool_prompt
+        self.force_tool = force
         if next((c for c in self.until_callbacks if c[0] == self._until_tools_callback), None) is None:
-            self.until_callbacks.append((self._until_tools_callback, attempt_recovery, drop_dialog, max_rounds))
+            self.until_callbacks.append(
+                (
+                    self._until_tools_callback,
+                    attempt_recovery,
+                    drop_dialog,
+                    max_rounds,
+                )
+            )
         return self
 
     def until_parsed_as(
@@ -215,8 +225,14 @@ def _until_tools_callback(self, message: Message) -> tuple[bool, list[Message]]:
             return (True, next_messages)
 
         if tool_calls is None:
-            logger.debug("No tool calls or types, returning message")
-            return (False, next_messages)
+            if self.force_tool:
+                logger.debug("No tool calls or types, returning error")
+                next_messages.append(Message.from_model(SystemErrorModel(content="You must use a tool")))
+            else:
+                logger.debug("No tool calls or types, returning message")
+            return (self.force_tool, next_messages)
+
+        self.force_tool = False
 
         tool_results: list[ToolResult] = []
         errors: list[SystemErrorModel] = []
@@ -298,7 +314,13 @@ def _until(
 
     def _execute(self) -> list[Message]:
         if self.until_tools:
-            self.params.stop = [ToolCalls.xml_end_tag()]
+            # TODO: This can cause issues when certain APIs do not return
+            # the stop sequence as part of the response. This behavior
+            # seems like a larger issue than the model continuining after
+            # requesting a tool call, so we'll remove it for now.
+            #
+            # self.params.stop = [ToolCalls.xml_end_tag()]
+
             if self.inject_tool_prompt:
                 self.chat.inject_tool_prompt(self.until_tools)
 
diff --git a/rigging/prompt.py b/rigging/prompt.py
index b264330..e7aaa1b 100644
--- a/rigging/prompt.py
+++ b/rigging/prompt.py
@@ -1,14 +1,15 @@
 def system_tool_extension(call_format: str, tool_descriptions: str) -> str:
     return f"""\
+# Tool Use
 In this environment you have access to a set of tools you can use to improve your responses.
 
-Tool call format:
+## Tool Call Format
 {call_format}
 
-Available tools:
+## Available Tools
 {tool_descriptions}
 
-You can use tools by responding in the format above. The inputs will be parsed and the specified tool function will be executed with the parameters you provided. The results of each function call will be given before you continue the conversation. You can execute multiple steps of function calls by continuing to respond in the format above. Function calls take explicit values and are independent of each other. Function calls CANNOT use the results of other functions. DO NOT USE terms like `$result` or `TOOL_RESULT` in your parameters.
+You can use any of the available tools by responding in the call format above. The XML will be parsed and the tool(s) will be executed with the parameters you provided. The results of each tool call will be provided back to you before you continue the conversation. You can execute multiple tool calls by continuing to respond in the format above until you are finished. Function calls take explicit values and are independent of each other. Tool calls cannot share, re-use, and transfer values between eachother. The use of placeholders is forbidden.
 
-The user will not see the results of your tool calls, only the final message of your conversation. Wait to perform your full response until after you have used any required tools.
+The user will not see the results of your tool calls, only the final message of your conversation. Wait to perform your full response until after you have used any required tools. If you intend to use a tool, please do so before you continue the conversation.
 """