From c1d64a087c2c3f08e575801dbf3ae8679ba797b0 Mon Sep 17 00:00:00 2001
From: DavdGao <gaodawei.gdw@alibaba-inc.com>
Date: Tue, 30 Apr 2024 08:58:31 +0800
Subject: [PATCH] Init a prompt tuning module

---
 .../conversation_prompt_tuning/main.ipynb     | 35 ++++++++
 src/agentscope/prompt_tuning.py               | 80 +++++++++++++++++++
 2 files changed, 115 insertions(+)
 create mode 100644 examples/conversation_prompt_tuning/main.ipynb
 create mode 100644 src/agentscope/prompt_tuning.py

diff --git a/examples/conversation_prompt_tuning/main.ipynb b/examples/conversation_prompt_tuning/main.ipynb
new file mode 100644
index 000000000..29cd7437e
--- /dev/null
+++ b/examples/conversation_prompt_tuning/main.ipynb
@@ -0,0 +1,35 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "initial_id",
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/src/agentscope/prompt_tuning.py b/src/agentscope/prompt_tuning.py
new file mode 100644
index 000000000..19b4a6ea2
--- /dev/null
+++ b/src/agentscope/prompt_tuning.py
@@ -0,0 +1,80 @@
+from typing import List, Dict, Union
+
+from agentscope.agents import DialogAgent, UserAgent, AgentBase
+from agentscope.message import Msg
+
+PROMPT_SYSTEM_PROMPT_GENERATION = (
+    "Given a task description, try to generate a system prompt for large "
+    "language models (LLMs). \n\n"
+    "## Note:\n"
+    "1. A system prompt should clearly describe the role that the LLM is "
+    "playing, its target, and notes."
+    "2. The system prompt should be concise and clear, and should not contain "
+    "any irrelevant information.\n\n"
+    "## Task Description:\n"
+    "{task_description}\n\n"
+    "## Generated System Prompt:\n"
+)
+
+PROMPT_FEW_SHOT_EXAMPLES = (
+    "## Example {index}:\n"
+    "Task description: {task_description}"
+    "Generated system prompt: {system_prompt}"
+)
+
+
+class PromptFactory(object):
+
+    def __init__(self, model_config_name: str) -> None:
+        """Initialize a PromptFactory object."""
+        self.model_config_name = model_config_name
+
+    def generate_system_prompt(
+        self,
+        task_description: str,
+        few_shot_examples: List[Dict] = None,
+        interactive_mode: bool = False
+    ) -> str:
+        """Generate a system prompt for the given task description.
+
+        Args:
+            task_description (`str`):
+                The task description.
+            few_shot_examples (`List[Dict]`, defaults to `None`):
+                Few-shot examples for the system prompt generation.
+            interactive_mode (`bool`, defaults to `False`):
+                Whether to run in interactive mode.
+
+        Returns:
+            `str`: The generated system prompt.
+        """
+
+        system_prompt = PROMPT_SYSTEM_PROMPT_GENERATION.format(task_description)
+        agent = DialogAgent(
+            "assistant",
+            sys_prompt=system_prompt,
+            model_config_name=self.model_config_name,
+        )
+        user = UserAgent("user")
+
+        msg = Msg(
+            "system",
+            PROMPT_SYSTEM_PROMPT_GENERATION.format(task_description),
+            "system"
+        )
+        while True:
+            msg = agent(msg)
+            # record the generated system prompt
+            generated_system_prompt = msg.content
+            if interactive_mode:
+                break
+
+            msg = user(msg)
+            if msg.content == "exit":
+                break
+
+        return generated_system_prompt
+
+    def tune_system_prompt(self, agent: Union[AgentBase, List[Msg]]):
+        """Tune the system prompt according to the user's feedback."""
+        pass