feat: make assistant work with claude2 (#36)

Update prompts to make them compatible with claude2.

aidial_assistant/application/prompts.py

@@ -28,82 +28,88 @@ def build(self, **kwargs) -> Template:
         )
 
 
-_REQUEST_EXAMPLE_TEXT = """
-ALWAYS reply with a JSON containing an array of available commands. You must not use natural language:
+_REQUEST_FORMAT_TEXT = """
+You should ALWAYS reply with a JSON containing an array of commands:
 {
   "commands": [
     {
       "command": "<command name>",
       "args": [
-        <array of arguments>
+        // <array of arguments>
       ]
     }
   ]
 }
+The commands are invoked by system on user's behalf.
+""".strip()
+
+_PROTOCOL_FOOTER = """
+* reply
+The command delivers final response to the user.
+Arguments:
+ - MESSAGE is a string containing the final and complete result for the user.
 
-Example:
+Your goal is to answer user questions. Use relevant commands when they help to achieve the goal.
+
+## Example
 {"commands": [{"command": "reply", "args": ["Hello, world!"]}]}
 """.strip()
 
 _SYSTEM_TEXT = """
 Today's date is {{today_date}}.
+This message defines the following communication protocol.
 
 {%- if system_prefix %}
 {{system_prefix}}
 {%- endif %}
 
-Protocol
-The following commands are available to reply to user or find out the answer to the user's question:
+# Protocol
+{{request_format}}
+
+## Commands
 {%- if tools %}
-> run-plugin
-The command runs a specified plugin to solve a one-shot task written in natural language.
-Plugins do not see current conversation and require all details to be provided in the query to solve the task.
-The command returns the result of the plugin call.
+* run-addon
+This command executes a specified addon to address a one-time task described in natural language.
+Addons do not see current conversation and require all details to be provided in the query to solve the task.
 Arguments:
- - NAME is one of the following plugins:
+ - NAME is one of the following addons:
 {%- for name, description in tools.items() %}
     * {{name}} - {{description | decap}}
 {%- endfor %}
- - QUERY is a string formulating the query to the plugin.
+ - QUERY is the query string.
 {%- endif %}
-> reply
-The command delivers ultimate result to the user.
-Arguments:
- - MESSAGE is a string containing response for user.
-
-{{request_response}}
+{{protocol_footer}}
 """.strip()
 
-_PLUGIN_SYSTEM_TEXT = """
+_ADDON_SYSTEM_TEXT = """
 Today's date is {{today_date}}.
+This message defines the following communication protocol.
 
-Service
+# Service
 API_DESCRIPTION:
 {{api_description}}
 
-API_SCHEMA:
+# API Schema
 ```typescript
 {{api_schema}}}
 ```
 
-Protocol
-The following commands are available to reply to user or find out the answer to the user's question:
+# Protocol
+{{request_format}}
+
+## Commands
 {%- for command_name in command_names %}
-> {{command_name}}
+* {{command_name}}
 Arguments:
- - <JSON dict according to the API_SCHEMA>
+ - <JSON dict according to the API Schema>
 {%- endfor %}
-> reply
-The command delivers ultimate result to the user
-Arguments:
- - MESSAGE is a string containing response for user.
-
-{{request_response}}
+{{protocol_footer}}
 """.strip()
 
 _ENFORCE_JSON_FORMAT_TEXT = """
 {{response}}
-**Remember to reply with a JSON with commands**
+
+**Protocol reminder: reply with commands**
 """.strip()
 
 _MAIN_BEST_EFFORT_TEXT = (
@@ -134,7 +140,7 @@ def build(self, **kwargs) -> Template:
 """
 ).strip()
 
-_PLUGIN_BEST_EFFORT_TEXT = (
+_ADDON_BEST_EFFORT_TEXT = (
     """
 You were allowed to use the following API to answer the query below.
 
@@ -165,18 +171,24 @@ def build(self, **kwargs) -> Template:
 
 MAIN_SYSTEM_DIALOG_MESSAGE = PartialTemplate(
     _SYSTEM_TEXT,
-    globals={"request_response": _REQUEST_EXAMPLE_TEXT},
+    globals={
+        "request_format": _REQUEST_FORMAT_TEXT,
+        "protocol_footer": _PROTOCOL_FOOTER,
+    },
     template_class=DateAwareTemplate,
 )
 
-PLUGIN_SYSTEM_DIALOG_MESSAGE = PartialTemplate(
-    _PLUGIN_SYSTEM_TEXT,
-    globals={"request_response": _REQUEST_EXAMPLE_TEXT},
+ADDON_SYSTEM_DIALOG_MESSAGE = PartialTemplate(
+    _ADDON_SYSTEM_TEXT,
+    globals={
+        "request_format": _REQUEST_FORMAT_TEXT,
+        "protocol_footer": _PROTOCOL_FOOTER,
+    },
     template_class=DateAwareTemplate,
 )
 
 ENFORCE_JSON_FORMAT_TEMPLATE = JINJA2_ENV.from_string(_ENFORCE_JSON_FORMAT_TEXT)
 
 MAIN_BEST_EFFORT_TEMPLATE = PartialTemplate(_MAIN_BEST_EFFORT_TEXT)
 
-PLUGIN_BEST_EFFORT_TEMPLATE = PartialTemplate(_PLUGIN_BEST_EFFORT_TEXT)
+ADDON_BEST_EFFORT_TEMPLATE = PartialTemplate(_ADDON_BEST_EFFORT_TEXT)

aidial_assistant/chain/command_chain.py

@@ -150,7 +150,8 @@ async def _run_with_protocol_failure_retries(
                     retry += 1
                     last_error = e
                     retries.append(
-                        chunk_stream.buffer, json.dumps({"error": str(e)})
+                        chunk_stream.buffer,
+                        "Failed to parse JSON commands: " + str(e),
                     )
                 finally:
                     self._log_message(Role.ASSISTANT, chunk_stream.buffer)

aidial_assistant/commands/run_plugin.py

@@ -5,8 +5,8 @@
 from typing_extensions import override
 
 from aidial_assistant.application.prompts import (
-    PLUGIN_BEST_EFFORT_TEMPLATE,
-    PLUGIN_SYSTEM_DIALOG_MESSAGE,
+    ADDON_BEST_EFFORT_TEMPLATE,
+    ADDON_SYSTEM_DIALOG_MESSAGE,
 )
 from aidial_assistant.chain.command_chain import (
     CommandChain,
@@ -50,7 +50,7 @@ def __init__(
 
     @staticmethod
     def token():
-        return "run-plugin"
+        return "run-addon"
 
     @override
     async def execute(
@@ -78,7 +78,7 @@ async def _run_plugin(
     ) -> ResultObject:
         if name not in self.plugins:
             raise ValueError(
-                f"Unknown plugin: {name}. Available plugins: {[*self.plugins.keys()]}"
+                f"Unknown addon: {name}. Available addons: {[*self.plugins.keys()]}"
             )
 
         plugin = self.plugins[name]
@@ -98,12 +98,12 @@ def create_command(op: APIOperation):
         command_dict[Reply.token()] = Reply
 
         history = History(
-            assistant_system_message_template=PLUGIN_SYSTEM_DIALOG_MESSAGE.build(
+            assistant_system_message_template=ADDON_SYSTEM_DIALOG_MESSAGE.build(
                 command_names=ops.keys(),
                 api_description=info.ai_plugin.description_for_model,
                 api_schema=api_schema,
             ),
-            best_effort_template=PLUGIN_BEST_EFFORT_TEMPLATE.build(
+            best_effort_template=ADDON_BEST_EFFORT_TEMPLATE.build(
                 api_schema=api_schema
             ),
             scoped_messages=[ScopedMessage(message=Message.user(query))],

pyproject.toml

@@ -24,7 +24,7 @@ openai = "^0.28.0"
 pydantic = "1.10.13"
 pyyaml = "^6.0.1"
 typing-extensions = "^4.8.0"
-aidial-sdk = "^0.1.0"
+aidial-sdk = "^0.2.0"
 aiohttp = "^3.9.0"
 openapi-schema-pydantic = "^1.2.4"
 openapi-pydantic = "^0.3.2"

tests/unit_tests/application/test_prompts.py

@@ -1,6 +1,6 @@
 from aidial_assistant.application.prompts import (
+    ADDON_BEST_EFFORT_TEMPLATE,
     MAIN_BEST_EFFORT_TEMPLATE,
-    PLUGIN_BEST_EFFORT_TEMPLATE,
 )
 
 
@@ -65,9 +65,7 @@ def test_main_best_effort_prompt_with_empty_dialogue():
 
 
 def test_plugin_best_effort_prompt():
-    actual = PLUGIN_BEST_EFFORT_TEMPLATE.build(
-        api_schema="<api schema>"
-    ).render(
+    actual = ADDON_BEST_EFFORT_TEMPLATE.build(api_schema="<api schema>").render(
         error="<error>",
         message="<message>",
         dialogue=[{"role": "<role>", "content": "<content>"}],
@@ -99,9 +97,7 @@ def test_plugin_best_effort_prompt():
 
 
 def test_plugin_best_effort_prompt_with_empty_dialogue():
-    actual = PLUGIN_BEST_EFFORT_TEMPLATE.build(
-        api_schema="<api schema>"
-    ).render(
+    actual = ADDON_BEST_EFFORT_TEMPLATE.build(api_schema="<api schema>").render(
         error="<error>",
         message="<message>",
         dialogue=[],

tests/unit_tests/chain/test_command_chain_best_effort.py

@@ -20,7 +20,7 @@
 
 SYSTEM_MESSAGE = "<system message>"
 USER_MESSAGE = "<user message>"
-ENFORCE_JSON_FORMAT = "**Remember to reply with a JSON with commands**"
+ENFORCE_JSON_FORMAT = "\n\n**Protocol reminder: reply with commands**"
 BEST_EFFORT_ANSWER = "<best effort answer>"
 NO_TOKENS_ERROR = "No tokens left"
 FAILED_PROTOCOL_ERROR = "The next constructed API request is incorrect."
@@ -78,7 +78,7 @@ async def test_model_doesnt_support_protocol():
         call(
             [
                 Message.system(f"system_prefix={SYSTEM_MESSAGE}"),
-                Message.user(f"{USER_MESSAGE}\n{ENFORCE_JSON_FORMAT}"),
+                Message.user(f"{USER_MESSAGE}{ENFORCE_JSON_FORMAT}"),
             ],
             usage_publisher,
         ),
@@ -132,7 +132,7 @@ async def test_model_partially_supports_protocol():
         call(
             [
                 Message.system(f"system_prefix={SYSTEM_MESSAGE}"),
-                Message.user(f"{USER_MESSAGE}\n{ENFORCE_JSON_FORMAT}"),
+                Message.user(f"{USER_MESSAGE}{ENFORCE_JSON_FORMAT}"),
             ],
             usage_publisher,
         ),
@@ -141,7 +141,7 @@ async def test_model_partially_supports_protocol():
                 Message.system(f"system_prefix={SYSTEM_MESSAGE}"),
                 Message.user(USER_MESSAGE),
                 Message.assistant(TEST_COMMAND_REQUEST),
-                Message.user(f"{TEST_COMMAND_RESPONSE}\n{ENFORCE_JSON_FORMAT}"),
+                Message.user(f"{TEST_COMMAND_RESPONSE}{ENFORCE_JSON_FORMAT}"),
             ],
             usage_publisher,
         ),
@@ -191,7 +191,7 @@ async def test_no_tokens_for_tools():
         call(
             [
                 Message.system(f"system_prefix={SYSTEM_MESSAGE}"),
-                Message.user(f"{USER_MESSAGE}\n{ENFORCE_JSON_FORMAT}"),
+                Message.user(f"{USER_MESSAGE}{ENFORCE_JSON_FORMAT}"),
             ],
             usage_publisher,
         ),
@@ -200,7 +200,7 @@ async def test_no_tokens_for_tools():
                 Message.system(f"system_prefix={SYSTEM_MESSAGE}"),
                 Message.user(USER_MESSAGE),
                 Message.assistant(TEST_COMMAND_REQUEST),
-                Message.user(f"{TEST_COMMAND_RESPONSE}\n{ENFORCE_JSON_FORMAT}"),
+                Message.user(f"{TEST_COMMAND_RESPONSE}{ENFORCE_JSON_FORMAT}"),
             ],
             usage_publisher,
         ),