diff --git a/src/codeinterpreterapi/llm/llm.py b/src/codeinterpreterapi/llm/llm.py index aa4fbf55..91bba2c8 100644 --- a/src/codeinterpreterapi/llm/llm.py +++ b/src/codeinterpreterapi/llm/llm.py @@ -151,10 +151,13 @@ def get_llms(cls, model: str = settings.MODEL_LOCAL) -> List[BaseChatModel]: return llms -def prepare_test_llm(): +def prepare_test_llm(is_smart: bool = False): load_dotenv(verbose=True, override=False) - llm = CodeInterpreterLlm.get_llm_switcher() + if is_smart: + llm = CodeInterpreterLlm.get_llm_smart() + else: + llm = CodeInterpreterLlm.get_llm_switcher() llm_tools = CodeInterpreterLlm.get_llm_switcher_tools() callbacks = [FullOutCallbackHandler()] configurable = {"session_id": "123"} diff --git a/src/codeinterpreterapi/planners/planners.py b/src/codeinterpreterapi/planners/planners.py index eeb725db..7433845a 100644 --- a/src/codeinterpreterapi/planners/planners.py +++ b/src/codeinterpreterapi/planners/planners.py @@ -1,52 +1,62 @@ -from langchain import PromptTemplate, hub +from typing import List, Union + from langchain.agents import AgentExecutor from langchain.agents.agent import RunnableAgent -from langchain.tools.render import render_text_description from langchain_core.runnables import Runnable +from pydantic import BaseModel, Field +from codeinterpreterapi.agents.agents import CodeInterpreterAgent from codeinterpreterapi.brain.params import CodeInterpreterParams from codeinterpreterapi.llm.llm import prepare_test_llm -from codeinterpreterapi.utils.output_parser import PlannerSingleOutputParser +from codeinterpreterapi.planners.prompts import create_planner_agent_chat_prompt, create_planner_agent_prompt +from codeinterpreterapi.test_prompts.test_prompt import TestPrompt +from codeinterpreterapi.utils.prompt import PromptUpdater from codeinterpreterapi.utils.runnable import create_complement_input +class CodeInterpreterPlan(BaseModel): + '''Agent and Task definition. Plan and task is 1:1.''' + + agent_name: str = Field( + description="The agent name for task. This is primary key. Agent responsible for task execution. Represents entity performing task." + ) + task_description: str = Field(description="Descriptive text detailing task's purpose and execution.") + expected_output: str = Field(description="Clear definition of expected task outcome.") + + +class CodeInterpreterPlanList(BaseModel): + '''Sequential plans for the task.''' + + agent_task_list: List[CodeInterpreterPlan] = Field( + description="The list of CodeInterpreterPlan. It means agent name and so on." + ) + + class CodeInterpreterPlanner: @staticmethod - def choose_planner(ci_params: CodeInterpreterParams) -> Runnable: + def choose_planner(ci_params: CodeInterpreterParams) -> Union[Runnable, AgentExecutor]: """ Load a chat planner. - Args: - llm: Language model. - tools: List of tools this agent has access to. - is_ja: System prompt. - - Returns: - LLMPlanner - - - Input - tools: - tool_names: - input: - agent_scratchpad: - Output - content: Free text in str. """ - # prompt_name_react = "nobu/simple_react" - prompt_name = "nobu/chat_planner" - if ci_params.is_ja: - prompt_name += "_ja" - prompt = hub.pull(prompt_name) - # prompt = CodeInterpreterPlanner.get_prompt() + + is_chat_prompt = True + if is_chat_prompt: + prompt = create_planner_agent_chat_prompt() + prompt = PromptUpdater.update_and_show_chat_prompt(prompt, ci_params) + else: + prompt = create_planner_agent_prompt() + prompt = PromptUpdater.update_prompt(prompt, ci_params) + PromptUpdater.show_prompt(prompt) + + # structured_llm + structured_llm = ci_params.llm.bind_tools(tools=[CodeInterpreterPlanList]) # runnable runnable = ( - create_complement_input(prompt) - | prompt - | ci_params.llm + create_complement_input(prompt) | prompt | structured_llm # | StrOutputParser() - | PlannerSingleOutputParser() + # | PlannerSingleOutputParser() ) # runnable = assign_runnable_history(runnable, ci_params.runnable_config) @@ -56,50 +66,23 @@ def choose_planner(ci_params: CodeInterpreterParams) -> Runnable: remapped_inputs = create_complement_input(prompt).invoke({}) agent = RunnableAgent(runnable=runnable, input_keys=list(remapped_inputs.keys())) - # agent_executor - agent_executor = AgentExecutor(agent=agent, tools=[], verbose=ci_params.verbose) - - return agent_executor - - @staticmethod - def update_prompt(prompt: PromptTemplate, ci_params: CodeInterpreterParams) -> PromptTemplate: - # Check if the prompt has the required input variables - missing_vars = {"tools", "tool_names", "agent_scratchpad"}.difference( - prompt.input_variables + list(prompt.partial_variables) - ) - if missing_vars: - raise ValueError(f"Prompt missing required variables: {missing_vars}") - - # Partial the prompt with tools and tool_names - prompt = prompt.partial( - tools=render_text_description(list(ci_params.tools)), - tool_names=", ".join([t.name for t in ci_params.tools]), - ) - return prompt - - @staticmethod - def get_prompt(): - prompt_template = """ - あなたは親切なアシスタントです。以下の制約条件を守ってタスクを完了させてください。 - - 制約条件: - - 段階的に考え、各ステップで取るアクションを明確にすること。 - - 必要な情報が不足している場合は、ユーザーに質問すること。 - - タスクを完了するために十分な情報が得られたら、最終的な回答を出力すること。 + # return executor or runnable + return_as_executor = False + if return_as_executor: + # TODO: handle step by step by original OutputParser + agent_executor = AgentExecutor(agent=agent, tools=ci_params.tools, verbose=ci_params.verbose) + return agent_executor - タスク: {input} - agent_scratchpad: {agent_scratchpad} - """ - prompt = PromptTemplate(input_variables=["input", "agent_scratchpad"], template=prompt_template) - return prompt + else: + return runnable def test(): - sample = "ステップバイステップで2*5+2を計算して。" llm, llm_tools, runnable_config = prepare_test_llm() ci_params = CodeInterpreterParams.get_test_params(llm=llm, llm_tools=llm_tools, runnable_config=runnable_config) + _ = CodeInterpreterAgent.choose_agent_executors(ci_params=ci_params) planner = CodeInterpreterPlanner.choose_planner(ci_params=ci_params) - result = planner.invoke({"input": sample, "agent_scratchpad": ""}) + result = planner.invoke({"input": "", "agent_scratchpad": "", "messages": [TestPrompt.svg_input_str]}) print("result=", result) diff --git a/src/codeinterpreterapi/planners/prompts.py b/src/codeinterpreterapi/planners/prompts.py new file mode 100644 index 00000000..a04b5dfe --- /dev/null +++ b/src/codeinterpreterapi/planners/prompts.py @@ -0,0 +1,68 @@ +from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate + +SYSTEM_MESSAGE_TEMPLATE = ''' + You are a super agent manager. + First of all, understand the problem. + Please make a plan to solve it. + + You cau use AI agent here. + {agent_info}. + + Constraints:. + - Think step-by-step and be clear about the actions to be taken at each step. + - If the required information is missing, answer with agent_name=None and ask the user in task_description. + - If the required agent for execution does not exist, answer with agent_name=None and answer about the required agent in the task_description. + - Once sufficient information is obtained to complete the task, the final work plan should be output. + - The last step should be agent_name=. + +Translated with DeepL.com (free version) +''' + +SYSTEM_MESSAGE_TEMPLATE_JA = ''' + あなたは優秀なAIエージェントを管理するシニアエンジニアです。 + まず、問題を理解し、問題を解決するための計画を立てましょう。 + + 利用可能なAI agentのリスト: + {agent_info} + + 制約条件: + - 段階的に考え、各ステップで取るアクションを明確にすること。 + - 必要な情報が不足している場合は、agent_name=Noneで回答し、task_description でユーザーに質問すること。 + - 実行に必要なagentが存在しない場合は、agent_name=Noneで回答し、task_description で必要なagentについて回答すること。 + - タスクを完了するために十分な情報が得られたら、最終的な作業計画を出力すること。 + - 最後のステップはagent_name=とすること。 + - 各ステップの思考と出力は日本語とする。 + +''' + + +def create_planner_agent_prompt(is_ja: bool = True) -> PromptTemplate: + if is_ja: + prompt = PromptTemplate( + input_variables=["input", "agent_scratchpad", "agent_info"], template=SYSTEM_MESSAGE_TEMPLATE_JA + ) + else: + prompt = PromptTemplate( + input_variables=["input", "agent_scratchpad", "agent_info"], template=SYSTEM_MESSAGE_TEMPLATE + ) + return prompt + + +def create_planner_agent_chat_prompt(is_ja: bool = True) -> ChatPromptTemplate: + if is_ja: + prompt = ChatPromptTemplate.from_messages( + [ + ("system", SYSTEM_MESSAGE_TEMPLATE_JA), + MessagesPlaceholder(variable_name="messages"), + ("user", "上記の会話を踏まえて、最終的な作業計画を日本語で出力してください。"), + ] + ) + else: + prompt = ChatPromptTemplate.from_messages( + [ + ("system", SYSTEM_MESSAGE_TEMPLATE), + MessagesPlaceholder(variable_name="messages"), + ("user", "Given the conversation above, please output final plan."), + ] + ) + return prompt diff --git a/src/codeinterpreterapi/utils/prompt.py b/src/codeinterpreterapi/utils/prompt.py new file mode 100644 index 00000000..baab8116 --- /dev/null +++ b/src/codeinterpreterapi/utils/prompt.py @@ -0,0 +1,73 @@ +from langchain_core.prompts import ( + AIMessagePromptTemplate, + ChatPromptTemplate, + HumanMessagePromptTemplate, + MessagesPlaceholder, + PromptTemplate, + SystemMessagePromptTemplate, +) + +from codeinterpreterapi.brain.params import CodeInterpreterParams + + +class PromptUpdater: + @staticmethod + def show_prompt(prompt: PromptTemplate): + prompt_str = prompt.format() + print(f"show_prompt: {prompt_str}") + + @staticmethod + def show_chat_prompt(prompt: ChatPromptTemplate): + print("show_chat_prompt:") + for message in prompt.messages: + if isinstance(message, SystemMessagePromptTemplate): + print(f"System: {PromptUpdater.show_prompt(message.prompt)}") + elif isinstance(message, HumanMessagePromptTemplate): + print(f"Human: {message.prompt.template}") + elif isinstance(message, AIMessagePromptTemplate): + print(f"AI: {message.prompt.template}") + elif isinstance(message, MessagesPlaceholder): + print(f"MessagesPlaceholder: {message.variable_name}") + else: + print(f"Other: {message}") + + @staticmethod + def update_prompt(prompt: PromptTemplate, ci_params: CodeInterpreterParams) -> PromptTemplate: + # Check if the prompt has the required input variables + missing_vars = {"agent_info"}.difference(prompt.input_variables + list(prompt.partial_variables)) + if missing_vars: + raise ValueError(f"Prompt missing required variables: {missing_vars}") + + # Partial the prompt with tools and tool_names + prompt = prompt.partial( + agent_info=", ".join([agent_def.get_agent_info() for agent_def in ci_params.agent_def_list]), + ) + return prompt + + @staticmethod + def update_chat_prompt(prompt: ChatPromptTemplate, ci_params: CodeInterpreterParams) -> ChatPromptTemplate: + # Check if the prompt has the required input variables + missing_vars = {"agent_info"}.difference(prompt.input_variables) + if missing_vars: + raise ValueError(f"Prompt missing required variables: {missing_vars}") + + # Create a new ChatPromptTemplate with updated messages + updated_messages = [] + for message in prompt.messages: + if isinstance(message, SystemMessagePromptTemplate): + updated_content = PromptUpdater.update_prompt(message.prompt, ci_params) + updated_messages.append(SystemMessagePromptTemplate(prompt=updated_content)) + elif isinstance(message, (HumanMessagePromptTemplate, AIMessagePromptTemplate)): + updated_messages.append(message) + elif isinstance(message, MessagesPlaceholder): + updated_messages.append(message) + else: + raise ValueError(f"Unexpected message type: {type(message)}") + + return ChatPromptTemplate(messages=updated_messages) + + @staticmethod + def update_and_show_chat_prompt(prompt: ChatPromptTemplate, ci_params: CodeInterpreterParams) -> ChatPromptTemplate: + updated_prompt = PromptUpdater.update_chat_prompt(prompt, ci_params) + PromptUpdater.show_chat_prompt(updated_prompt) + return updated_prompt diff --git a/tests/one_test.py b/tests/one_test.py index b4699a3f..f03e48fe 100644 --- a/tests/one_test.py +++ b/tests/one_test.py @@ -1,27 +1,15 @@ -from codeinterpreterapi import CodeInterpreterSession +from codeinterpreterapi.brain.params import CodeInterpreterParams +from codeinterpreterapi.llm.llm import prepare_test_llm +from codeinterpreterapi.planners.planners import CodeInterpreterPlanner def test(): - model = "claude-3-haiku-20240307" - # model = "gemini-1.0-pro" - test_message = "pythonで円周率を表示するプログラムを実行してください。" - verbose = False - is_streaming = False - print("test_message=", test_message) - session = CodeInterpreterSession(model=model, verbose=verbose) - status = session.start_local() - print("status=", status) - if is_streaming: - # response_inner: CodeInterpreterResponse - response_inner = session.generate_response_stream(test_message) - for response_inner_chunk in response_inner: - print("response_inner_chunk.content=", response_inner.content) - print("response_inner_chunk code_log=", response_inner.code_log) - else: - # response_inner: CodeInterpreterResponse - response_inner = session.generate_response(test_message) - print("response_inner.content=", response_inner.content) - print("response_inner code_log=", response_inner.code_log) + sample = "ステップバイステップで2*5+2を計算して。" + llm, llm_tools, runnable_config = prepare_test_llm() + ci_params = CodeInterpreterParams.get_test_params(llm=llm, llm_tools=llm_tools, runnable_config=runnable_config) + planner = CodeInterpreterPlanner.choose_planner(ci_params=ci_params) + result = planner.invoke({"input": sample, "agent_scratchpad": "", "messages": [sample]}) + print("result=", result.content) if __name__ == "__main__": diff --git a/tests/tool_test.py b/tests/tool_test.py new file mode 100644 index 00000000..17b84edd --- /dev/null +++ b/tests/tool_test.py @@ -0,0 +1,95 @@ +from typing import List + +from langchain import PromptTemplate +from pydantic import BaseModel, Field + +from codeinterpreterapi.brain.params import CodeInterpreterParams +from codeinterpreterapi.llm.llm import prepare_test_llm +from codeinterpreterapi.test_prompts.test_prompt import TestPrompt +from codeinterpreterapi.tools.tools import CodeInterpreterTools +from codeinterpreterapi.utils.runnable import create_complement_input + + +class TestAgentPlan(BaseModel): + '''Plan for the task.''' + + agent_name: str = Field(description="The agent name for task") + agent_role: str = Field(description="What should do the agent") + + +class TestAgentPlanList(BaseModel): + '''Plan for the task.''' + + agent_plan_list: List[TestAgentPlan] = Field(default_factory=list) + + class Config: + arbitrary_types_allowed = True + + +def tool_run(ci_params: CodeInterpreterParams, inputs: dict): + prompt = get_prompt(ci_params) + structured_llm = ci_params.llm_smart.bind_tools(tools=[TestAgentPlanList]) + runnable = ( + create_complement_input(prompt) | prompt | structured_llm + # | StrOutputParser() + # | PlannerSingleOutputParser() + ) + return runnable.invoke(inputs) + + +def structured_run(ci_params: CodeInterpreterParams, inputs: dict): + prompt = get_prompt(ci_params) + structured_llm = ci_params.llm_smart.with_structured_output(TestAgentPlanList) + runnable = ( + create_complement_input(prompt) | prompt | structured_llm + # | StrOutputParser() + # | PlannerSingleOutputParser() + ) + return runnable.invoke(inputs) + + +def get_prompt(ci_params: CodeInterpreterParams): + prompt_template = """ + あなたは経験豊富なプロジェクトマネージャーです。 + 新しいAI関連のプロジェクトで以下のAIを利用したシーケンシャルな作業計画を作成してください。 + + 制約条件: + - 段階的に考え、各ステップで取るアクションを明確にすること。 + - 必要な情報が不足している場合は、agent_name=Noneで回答し、agent_roleでユーザーに質問すること。 + - 実行に必要なagentが存在しない場合は、agent_name=Noneで回答し、agent_roleで必要なagentについて回答すること。 + - タスクを完了するために十分な情報が得られたら、最終的な作業計画を出力すること。 + - 最後のステップはagent_name=とすること。 + + 利用可能なAI agentのリスト: + {agent_info} + + 回答内容: + - 作業の順番を示したCodeInterpreterPlanListを回答してください。 + + タスク: {input} + + agent_scratchpad: {agent_scratchpad} + """ + prompt = PromptTemplate(input_variables=["input", "agent_scratchpad"], template=prompt_template) + + # 利用可能なエージェントの情報を文字列として作成 + agent_info_str = "\n".join([agent_def.get_agent_info() for agent_def in ci_params.agent_def_list]) + print("agent_info_str=", agent_info_str) + + # agent_infoを部分的に適用 + prompt = prompt.partial(agent_info=agent_info_str) + + print("prompt=", prompt.template) + return prompt + + +def test(): + llm, llm_tools, runnable_config = prepare_test_llm(is_smart=False) + ci_params = CodeInterpreterParams.get_test_params(llm=llm, llm_tools=llm_tools, runnable_config=runnable_config) + ci_params.tools = CodeInterpreterTools(ci_params).get_all_tools() + result = structured_run(ci_params=ci_params, inputs={"input": TestPrompt.svg_input_str, "agent_scratchpad": ""}) + print("result=", result) + + +if __name__ == "__main__": + test()