Support additional_instructions (#302)

React agent supports additional_instructions configuration option.  
Add test cases for custom system prompt and adding additional_instructions.




Closes #301

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/AIQToolkit/blob/develop/docs/source/resources/contributing.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - https://github.com/gfreeman-nvidia

Approvers:
  - David Gardner (https://github.com/dagardner-nv)

URL: #302

Author: gfreeman-nvidia

src/aiq/agent/react_agent/agent.py

@@ -26,7 +26,8 @@
 from langchain_core.messages.base import BaseMessage
 from langchain_core.messages.human import HumanMessage
 from langchain_core.messages.tool import ToolMessage
-from langchain_core.prompts.chat import ChatPromptTemplate
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.prompts import MessagesPlaceholder
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel
@@ -42,6 +43,9 @@
 from aiq.agent.dual_node import DualNodeAgent
 from aiq.agent.react_agent.output_parser import ReActOutputParser
 from aiq.agent.react_agent.output_parser import ReActOutputParserException
+from aiq.agent.react_agent.prompt import SYSTEM_PROMPT
+from aiq.agent.react_agent.prompt import USER_PROMPT
+from aiq.agent.react_agent.register import ReActAgentWorkflowConfig
 
 logger = logging.getLogger(__name__)
 
@@ -320,3 +324,32 @@ def validate_system_prompt(system_prompt: str) -> bool:
             logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
             raise ValueError(error_text)
         return True
+
+
+def create_react_agent_prompt(config: ReActAgentWorkflowConfig) -> ChatPromptTemplate:
+    """
+    Create a ReAct Agent prompt from the config.
+
+    Args:
+        config (ReActAgentWorkflowConfig): The config to use for the prompt.
+
+    Returns:
+        ChatPromptTemplate: The ReAct Agent prompt.
+    """
+    # the ReAct Agent prompt can be customized via config option system_prompt and additional_instructions.
+
+    if config.system_prompt:
+        prompt_str = config.system_prompt
+    else:
+        prompt_str = SYSTEM_PROMPT
+
+    if config.additional_instructions:
+        prompt_str += f" {config.additional_instructions}"
+
+    valid_prompt = ReActAgentGraph.validate_system_prompt(prompt_str)
+    if not valid_prompt:
+        logger.exception("%s Invalid system_prompt", AGENT_LOG_PREFIX)
+        raise ValueError("Invalid system_prompt")
+    prompt = ChatPromptTemplate([("system", prompt_str), ("user", USER_PROMPT),
+                                 MessagesPlaceholder(variable_name='agent_scratchpad', optional=True)])
+    return prompt

src/aiq/agent/react_agent/prompt.py

@@ -14,8 +14,6 @@
 # limitations under the License.
 
 # flake8: noqa
-from langchain_core.prompts.chat import ChatPromptTemplate
-from langchain_core.prompts.chat import MessagesPlaceholder
 
 SYSTEM_PROMPT = """
 Answer the following questions as best you can. You may ask the human to use the following tools:
@@ -37,10 +35,7 @@
 Thought: I now know the final answer
 Final Answer: the final answer to the original input question
 """
+
 USER_PROMPT = """
 Question: {question}
 """
-
-# This is the prompt - (ReAct Agent prompt)
-react_agent_prompt = ChatPromptTemplate([("system", SYSTEM_PROMPT), ("user", USER_PROMPT),
-                                         MessagesPlaceholder(variable_name='agent_scratchpad', optional=True)])

src/aiq/agent/react_agent/register.py

@@ -63,29 +63,13 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
 async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builder):
     from langchain.schema import BaseMessage
     from langchain_core.messages import trim_messages
-    from langchain_core.prompts import ChatPromptTemplate
-    from langchain_core.prompts import MessagesPlaceholder
     from langgraph.graph.graph import CompiledGraph
 
-    from aiq.agent.react_agent.prompt import USER_PROMPT
-
-    from .agent import ReActAgentGraph
-    from .agent import ReActGraphState
-    from .prompt import react_agent_prompt
-
-    # the ReAct Agent prompt comes from prompt.py, and can be customized there or via config option system_prompt.
-    if config.system_prompt:
-        _prompt_str = config.system_prompt
-        if config.additional_instructions:
-            _prompt_str += f" {config.additional_instructions}"
-        valid_prompt = ReActAgentGraph.validate_system_prompt(config.system_prompt)
-        if not valid_prompt:
-            logger.exception("%s Invalid system_prompt", AGENT_LOG_PREFIX)
-            raise ValueError("Invalid system_prompt")
-        prompt = ChatPromptTemplate([("system", config.system_prompt), ("user", USER_PROMPT),
-                                     MessagesPlaceholder(variable_name='agent_scratchpad', optional=True)])
-    else:
-        prompt = react_agent_prompt
+    from aiq.agent.react_agent.agent import ReActAgentGraph
+    from aiq.agent.react_agent.agent import ReActGraphState
+    from aiq.agent.react_agent.agent import create_react_agent_prompt
+
+    prompt = create_react_agent_prompt(config)
 
     # we can choose an LLM for the ReAct agent in the config file
     llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)

tests/aiq/agent/test_react.py

@@ -25,12 +25,12 @@
 from aiq.agent.react_agent.agent import TOOL_NOT_FOUND_ERROR_MESSAGE
 from aiq.agent.react_agent.agent import ReActAgentGraph
 from aiq.agent.react_agent.agent import ReActGraphState
+from aiq.agent.react_agent.agent import create_react_agent_prompt
 from aiq.agent.react_agent.output_parser import FINAL_ANSWER_AND_PARSABLE_ACTION_ERROR_MESSAGE
 from aiq.agent.react_agent.output_parser import MISSING_ACTION_AFTER_THOUGHT_ERROR_MESSAGE
 from aiq.agent.react_agent.output_parser import MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE
 from aiq.agent.react_agent.output_parser import ReActOutputParser
 from aiq.agent.react_agent.output_parser import ReActOutputParserException
-from aiq.agent.react_agent.prompt import react_agent_prompt
 from aiq.agent.react_agent.register import ReActAgentWorkflowConfig
 
 
@@ -59,7 +59,7 @@ def mock_config():
 
 def test_react_init(mock_config_react_agent, mock_llm, mock_tool):
     tools = [mock_tool('Tool A'), mock_tool('Tool B')]
-    prompt = react_agent_prompt
+    prompt = create_react_agent_prompt(mock_config_react_agent)
     agent = ReActAgentGraph(llm=mock_llm, prompt=prompt, tools=tools, detailed_logs=mock_config_react_agent.verbose)
     assert isinstance(agent, ReActAgentGraph)
     assert agent.llm == mock_llm
@@ -72,7 +72,7 @@ def test_react_init(mock_config_react_agent, mock_llm, mock_tool):
 @pytest.fixture(name='mock_react_agent', scope="module")
 def mock_agent(mock_config_react_agent, mock_llm, mock_tool):
     tools = [mock_tool('Tool A'), mock_tool('Tool B')]
-    prompt = react_agent_prompt
+    prompt = create_react_agent_prompt(mock_config_react_agent)
     agent = ReActAgentGraph(llm=mock_llm, prompt=prompt, tools=tools, detailed_logs=mock_config_react_agent.verbose)
     return agent
 
@@ -412,3 +412,30 @@ async def test_output_parser_missing_action_input(mock_react_output_parser):
         await mock_react_output_parser.aparse(mock_input)
     assert isinstance(ex.value, ReActOutputParserException)
     assert ex.value.observation == MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE
+
+
+def test_react_additional_instructions(mock_llm, mock_tool):
+    config_react_agent = ReActAgentWorkflowConfig(tool_names=['test'],
+                                                  llm_name='test',
+                                                  verbose=True,
+                                                  retry_parsing_errors=False,
+                                                  additional_instructions="Talk like a parrot and repeat the question.")
+    tools = [mock_tool('Tool A'), mock_tool('Tool B')]
+    prompt = create_react_agent_prompt(config_react_agent)
+    agent = ReActAgentGraph(llm=mock_llm, prompt=prompt, tools=tools, detailed_logs=config_react_agent.verbose)
+    assert isinstance(agent, ReActAgentGraph)
+    assert "Talk like a parrot" in agent.agent.get_prompts()[0].messages[0].prompt.template
+
+
+def test_react_custom_system_prompt(mock_llm, mock_tool):
+    config_react_agent = ReActAgentWorkflowConfig(
+        tool_names=['test'],
+        llm_name='test',
+        verbose=True,
+        retry_parsing_errors=False,
+        system_prompt="Refuse to run any of the following tools: {tools}.  or ones named: {tool_names}")
+    tools = [mock_tool('Tool A'), mock_tool('Tool B')]
+    prompt = create_react_agent_prompt(config_react_agent)
+    agent = ReActAgentGraph(llm=mock_llm, prompt=prompt, tools=tools, detailed_logs=config_react_agent.verbose)
+    assert isinstance(agent, ReActAgentGraph)
+    assert "Refuse" in agent.agent.get_prompts()[0].messages[0].prompt.template