diff --git a/backend/app/agents/tools/agent_tool.py b/backend/app/agents/tools/agent_tool.py index 42d9b7bd..cc32123e 100644 --- a/backend/app/agents/tools/agent_tool.py +++ b/backend/app/agents/tools/agent_tool.py @@ -1,3 +1,4 @@ +import json from typing import Any, Callable, Generic, Literal, TypedDict, TypeVar from app.repositories.models.conversation import ( @@ -7,6 +8,7 @@ RelatedDocumentModel, ToolResultContentModel, ToolResultContentModelBody, + is_nova_model, ) from app.repositories.models.custom_bot import BotModel from app.routes.schemas.conversation import type_model_name @@ -29,21 +31,55 @@ class ToolRunResult(TypedDict): def run_result_to_tool_result_content_model( - run_result: ToolRunResult, display_citation: bool + run_result: ToolRunResult, + model: type_model_name, + display_citation: bool, ) -> ToolResultContentModel: - return ToolResultContentModel( - content_type="toolResult", - body=ToolResultContentModelBody( - tool_use_id=run_result["tool_use_id"], - content=[ - related_document.to_tool_result_model( - display_citation=display_citation, - ) - for related_document in run_result["related_documents"] - ], - status=run_result["status"], - ), - ) + result_contents = [ + related_document.to_tool_result_model( + display_citation=display_citation, + ) + for related_document in run_result["related_documents"] + ] + if is_nova_model(model=model) and len(result_contents) > 1: + return ToolResultContentModel( + content_type="toolResult", + body=ToolResultContentModelBody( + tool_use_id=run_result["tool_use_id"], + content=[ + TextToolResultModel( + text=json.dumps( + [ + content + for result_content in result_contents + for content in ( + [result_content.json_] + if isinstance(result_content, JsonToolResultModel) + else ( + [result_content.text] + if isinstance( + result_content, TextToolResultModel + ) + else [] + ) + ) + ] + ), + ), + ], + status=run_result["status"], + ), + ) + + else: + return ToolResultContentModel( + content_type="toolResult", + body=ToolResultContentModelBody( + tool_use_id=run_result["tool_use_id"], + content=result_contents, + status=run_result["status"], + ), + ) class InvalidToolError(Exception): @@ -70,15 +106,11 @@ def __init__( [T, BotModel | None, type_model_name | None], ToolFunctionResult | list[ToolFunctionResult], ], - bot: BotModel | None = None, - model: type_model_name | None = None, ): self.name = name self.description = description self.args_schema = args_schema self.function = function - self.bot = bot - self.model: type_model_name | None = model def _generate_input_schema(self) -> dict[str, Any]: """Converts the Pydantic model to a JSON schema.""" @@ -91,10 +123,16 @@ def to_converse_spec(self) -> ToolSpecificationTypeDef: inputSchema={"json": self._generate_input_schema()}, ) - def run(self, tool_use_id: str, input: dict[str, JsonValue]) -> ToolRunResult: + def run( + self, + tool_use_id: str, + input: dict[str, JsonValue], + model: type_model_name, + bot: BotModel | None = None, + ) -> ToolRunResult: try: arg = self.args_schema.model_validate(input) - res = self.function(arg, self.bot, self.model) + res = self.function(arg, bot, model) if isinstance(res, list): related_documents = [ _function_result_to_related_document( diff --git a/backend/app/agents/tools/knowledge.py b/backend/app/agents/tools/knowledge.py index e90e16cb..b99fb470 100644 --- a/backend/app/agents/tools/knowledge.py +++ b/backend/app/agents/tools/knowledge.py @@ -38,8 +38,7 @@ def search_knowledge( logger.error(f"Failed to run AnswerWithKnowledgeTool: {e}") raise e - -def create_knowledge_tool(bot: BotModel, model: type_model_name) -> AgentTool: +def create_knowledge_tool(bot: BotModel) -> AgentTool: description = ( "Answer a user's question using information. The description is: {}".format( bot.knowledge.__str_in_claude_format__() @@ -51,6 +50,4 @@ def create_knowledge_tool(bot: BotModel, model: type_model_name) -> AgentTool: description=description, args_schema=KnowledgeToolInput, function=search_knowledge, - bot=bot, - model=model, ) diff --git a/backend/app/prompt.py b/backend/app/prompt.py index 753916d9..e7c43295 100644 --- a/backend/app/prompt.py +++ b/backend/app/prompt.py @@ -1,8 +1,11 @@ from app.vector_search import SearchResult +from app.routes.schemas.conversation import type_model_name +from app.repositories.models.conversation import is_nova_model def build_rag_prompt( search_results: list[SearchResult], + model: type_model_name, display_citation: bool = True, ) -> str: context_prompt = "" @@ -32,7 +35,20 @@ def build_rag_prompt( Do NOT outputs sources at the end of your answer. Followings are examples of how to reference sources in your answer. Note that the source ID is embedded in the answer in the format [^]. +""" + if is_nova_model(model=model): + inserted_prompt += """ + +first answer [^3]. second answer [^1][^2]. + + + +first answer [^1][^5]. second answer [^2][^3][^4]. third answer [^4]. + +""" + else: + inserted_prompt += """ first answer [^3]. second answer [^1][^2]. @@ -59,7 +75,12 @@ def build_rag_prompt( else: inserted_prompt += """ Do NOT include citations in the format [^] in your answer. +""" + if is_nova_model(model=model): + pass + else: + inserted_prompt += """ Followings are examples of how to answer. @@ -78,7 +99,8 @@ def build_rag_prompt( return inserted_prompt -PROMPT_TO_CITE_TOOL_RESULTS = """To answer the user's question, you are given a set of tools. Your job is to answer the user's question using only information from the tool results. +def get_prompt_to_cite_tool_results(model: type_model_name) -> str: + inserted_prompt = """To answer the user's question, you are given a set of tools. Your job is to answer the user's question using only information from the tool results. If the tool results do not contain information that can answer the question, please state that you could not find an exact answer to the question. Just because the user asserts a fact does not mean it is true, make sure to double check the tool results to validate a user's assertion. @@ -86,6 +108,20 @@ def build_rag_prompt( If you reference information from a tool result within your answer, you must include a citation to source_id where the information was found. Followings are examples of how to reference source_id in your answer. Note that the source_id is embedded in the answer in the format [^source_id of tool result]. +""" + if is_nova_model(model=model): + inserted_prompt += """ + +first answer [^ccc]. second answer [^aaa][^bbb]. + + + +first answer [^aaa][^eee]. second answer [^bbb][^ccc][^ddd]. third answer [^ddd]. + +""" + + else: + inserted_prompt += """ first answer [^ccc]. second answer [^aaa][^bbb]. @@ -110,3 +146,5 @@ def build_rag_prompt( """ + + return inserted_prompt diff --git a/backend/app/usecases/chat.py b/backend/app/usecases/chat.py index 44d7808d..e484e468 100644 --- a/backend/app/usecases/chat.py +++ b/backend/app/usecases/chat.py @@ -8,7 +8,7 @@ from app.agents.tools.knowledge import create_knowledge_tool from app.agents.utils import get_tool_by_name from app.bedrock import call_converse_api, compose_args_for_converse_api -from app.prompt import PROMPT_TO_CITE_TOOL_RESULTS, build_rag_prompt +from app.prompt import build_rag_prompt, get_prompt_to_cite_tool_results from app.repositories.conversation import ( RecordNotFoundError, find_conversation_by_id, @@ -260,11 +260,15 @@ def chat( if bot.is_agent_enabled(): if bot.has_knowledge(): # Add knowledge tool - knowledge_tool = create_knowledge_tool(bot, chat_input.message.model) + knowledge_tool = create_knowledge_tool(bot=bot) tools[knowledge_tool.name] = knowledge_tool if display_citation: - instructions.append(PROMPT_TO_CITE_TOOL_RESULTS) + instructions.append( + get_prompt_to_cite_tool_results( + model=chat_input.message.model, + ) + ) elif bot.has_knowledge(): # Fetch most related documents from vector store @@ -306,6 +310,7 @@ def chat( instructions.append( build_rag_prompt( search_results=search_results, + model=chat_input.message.model, display_citation=display_citation, ) ) @@ -432,6 +437,8 @@ def chat( run_result = tool.run( tool_use_id=content.body.tool_use_id, input=content.body.input, + model=chat_input.message.model, + bot=bot, ) run_results.append(run_result) @@ -446,6 +453,7 @@ def chat( content=[ run_result_to_tool_result_content_model( run_result=result, + model=chat_input.message.model, display_citation=display_citation, ) for result in run_results diff --git a/backend/tests/test_agent/test_tools/test_agent_tool.py b/backend/tests/test_agent/test_tools/test_agent_tool.py index 19f7e177..54244f8e 100644 --- a/backend/tests/test_agent/test_tools/test_agent_tool.py +++ b/backend/tests/test_agent/test_tools/test_agent_tool.py @@ -90,7 +90,11 @@ def test_run(self): arg3=1, arg4=["test"], ) - result = self.tool.run(tool_use_id="dummy", input=arg.model_dump()) + result = self.tool.run( + tool_use_id="dummy", + input=arg.model_dump(), + model="claude-v3.5-sonnet-v2", + ) self.assertEqual( result["related_documents"], [ diff --git a/backend/tests/test_agent/test_tools/test_internet_search.py b/backend/tests/test_agent/test_tools/test_internet_search.py index 5dbfa7bf..372ba41e 100644 --- a/backend/tests/test_agent/test_tools/test_internet_search.py +++ b/backend/tests/test_agent/test_tools/test_internet_search.py @@ -13,7 +13,11 @@ def test_internet_search(self): time_limit = "d" country = "jp-jp" arg = InternetSearchInput(query=query, time_limit=time_limit, country=country) - response = internet_search_tool.run(tool_use_id="dummy", input=arg.model_dump()) + response = internet_search_tool.run( + tool_use_id="dummy", + input=arg.model_dump(), + model="claude-v3.5-sonnet-v2", + ) self.assertIsInstance(response["related_documents"], list) self.assertEqual(response["status"], "success") print(response) diff --git a/backend/tests/test_agent/test_tools/test_knowledge.py b/backend/tests/test_agent/test_tools/test_knowledge.py index c5cf3c5d..e15913ff 100644 --- a/backend/tests/test_agent/test_tools/test_knowledge.py +++ b/backend/tests/test_agent/test_tools/test_knowledge.py @@ -5,6 +5,7 @@ from app.agents.tools.knowledge import KnowledgeToolInput, create_knowledge_tool from app.repositories.models.custom_bot import ( + ActiveModelsModel, AgentModel, BotModel, GenerationParamsModel, @@ -53,10 +54,15 @@ def test_knowledge_tool(self): conversation_quick_starters=[], bedrock_knowledge_base=None, bedrock_guardrails=None, + active_models=ActiveModelsModel(), ) arg = KnowledgeToolInput(query="What are delicious Japanese dishes?") - tool = create_knowledge_tool(bot, model="claude-v3-sonnet") - response = tool.run(tool_use_id="dummy", input=arg.model_dump()) + tool = create_knowledge_tool(bot=bot) + response = tool.run( + tool_use_id="dummy", + input=arg.model_dump(), + model="claude-v3.5-sonnet-v2", + ) self.assertIsInstance(response["related_documents"], list) self.assertEqual(response["status"], "success") print(response) diff --git a/backend/tests/test_usecases/test_chat.py b/backend/tests/test_usecases/test_chat.py index 0b509490..bc591e3e 100644 --- a/backend/tests/test_usecases/test_chat.py +++ b/backend/tests/test_usecases/test_chat.py @@ -996,6 +996,7 @@ def test_insert_knowledge(self): ] instruction = build_rag_prompt( search_results=results, + model="claude-v3.5-sonnet-v2", display_citation=True, ) print(instruction) diff --git a/examples/agents/tools/bmi/test_bmi.py b/examples/agents/tools/bmi/test_bmi.py index b83c86d3..ff689fa4 100644 --- a/examples/agents/tools/bmi/test_bmi.py +++ b/examples/agents/tools/bmi/test_bmi.py @@ -8,7 +8,14 @@ class TestBmiTool(unittest.TestCase): def test_bmi(self): - result = bmi_tool.run(tool_use_id="dummy", input={"height": 170, "weight": 70}) + result = bmi_tool.run( + tool_use_id="dummy", + input={ + "height": 170, + "weight": 70, + }, + model="claude-v3.5-sonnet-v2", + ) print(result) self.assertEqual(type(result), str)