From a9553dafc5038737afd2320b3d505d8477056288 Mon Sep 17 00:00:00 2001 From: Zexin Yuan Date: Fri, 12 Jun 2026 15:19:17 +0800 Subject: [PATCH] Add native LLM tool call support with ReAct fallback - Add prompt_template.py: minimal Jinja2-like {% if %} conditional renderer - Add ToolCallMessage/ToolResultMessage unified message types - Add detect_tool_call_mode() using LiteLLM model capability DB - Update model_client.py: format tool messages for OpenAI Chat/Responses API - Update react_parser.py: _convert_messages() converts tool messages back to TextMessage for ReAct fallback - Update agent.py: Thought extraction post-processing for native mode - Update leader_multiagent.py: unified Agent layer, mode detection in init - Update leader_prompts.py: conditional {% if tool_call_mode %} sections - Update Executor.py: conditional prompt building per client type - Update agent_loader.py: render_prompt() when loading TOML configs - Update agent TOML configs with template conditionals - Fix prompt_template.py regex bug: {% %} delimiters were never matched - Fix prompt_template.py elif logic: three-tuple stack for branch tracking Tests: 143 pass (40 new), lint clean --- copilotj/core/agent.py | 19 +++ copilotj/core/message.py | 44 ++++- copilotj/core/model_client.py | 128 ++++++++++++-- copilotj/multiagent/Executor.py | 55 ++++-- .../agent_configs/research_agent.toml | 12 +- .../multiagent/agent_configs/tool_agent.toml | 21 ++- copilotj/multiagent/agent_loader.py | 8 +- copilotj/multiagent/leader_multiagent.py | 83 ++++----- copilotj/multiagent/leader_prompts.py | 35 +++- copilotj/test/core/test_agent.py | 107 ++++++++++++ copilotj/test/core/test_model_client.py | 139 +++++++++++++++- copilotj/test/util/test_prompt_template.py | 147 ++++++++++++++++ copilotj/test/util/test_react_parser.py | 157 +++++++++++++++++- copilotj/util/prompt_template.py | 132 +++++++++++++++ copilotj/util/react_parser.py | 38 ++++- 15 files changed, 1037 insertions(+), 88 deletions(-) create mode 100644 copilotj/test/core/test_agent.py create mode 100644 copilotj/test/util/test_prompt_template.py create mode 100644 copilotj/util/prompt_template.py diff --git a/copilotj/core/agent.py b/copilotj/core/agent.py index 7c05a4e7..bb15f36c 100644 --- a/copilotj/core/agent.py +++ b/copilotj/core/agent.py @@ -211,6 +211,25 @@ async def _create( tool_calls=tool_calls if len(tool_calls) else None, finish_reason=finish_reason, ) + + # In native mode, the model may output "Thought: ..." as regular + # content before calling a tool. Move that text to + # reasoning_content so the downstream logic (LeaderDriven.run) + # correctly distinguishes between a tool-call turn and a final + # answer turn. + if ( + completion.tool_calls + and completion.content + and not completion.reasoning_content + and completion.content.strip().lower().startswith("thought") + ): + completion = ModelResponse( + content=None, + reasoning_content=completion.content.strip(), + tool_calls=completion.tool_calls, + finish_reason=completion.finish_reason, + ) + self._runtime.log_info(str(completion)) return completion diff --git a/copilotj/core/message.py b/copilotj/core/message.py index abc6e8f2..6c0cc2b7 100644 --- a/copilotj/core/message.py +++ b/copilotj/core/message.py @@ -2,11 +2,18 @@ # # SPDX-License-Identifier: Apache-2.0 -from typing import Literal +from typing import Any, Literal import pydantic -__all__ = ["TextMessage", "ImageMessage", "HandoffMessage"] +__all__ = [ + "TextMessage", + "ImageMessage", + "HandoffMessage", + "ToolCallRecord", + "ToolCallMessage", + "ToolResultMessage", +] class TextMessage(pydantic.BaseModel): @@ -22,3 +29,36 @@ class ImageMessage(pydantic.BaseModel): class HandoffMessage(pydantic.BaseModel): target: str message: TextMessage | ImageMessage + + +class ToolCallRecord(pydantic.BaseModel): + """Serialisable record of a single tool call for conversation history.""" + + id: str + name: str + arguments: dict[str, Any] + + +class ToolCallMessage(pydantic.BaseModel): + """Assistant message containing one or more tool calls. + + Used in conversation history for both native and ReAct modes. The client + layer converts this to the appropriate API format (native tool_calls or + reconstructed ReAct text). + """ + + role: Literal["assistant"] = "assistant" + tool_calls: list[ToolCallRecord] + reasoning_content: str | None = None + + +class ToolResultMessage(pydantic.BaseModel): + """Tool role message carrying the execution result. + + In native mode this maps to ``{"role": "tool", ...}``. In ReAct mode the + client layer converts it to a user message with ``Observation:`` prefix. + """ + + role: Literal["tool"] = "tool" + tool_call_id: str + content: str diff --git a/copilotj/core/model_client.py b/copilotj/core/model_client.py index 2529764b..3e5b7484 100644 --- a/copilotj/core/model_client.py +++ b/copilotj/core/model_client.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import abc +import json import logging from typing import Any, AsyncGenerator, Literal, Sequence, cast, overload, override @@ -19,7 +20,7 @@ _get_vlm_and_key, _get_vlm_base_url, ) -from copilotj.core.message import ImageMessage, TextMessage +from copilotj.core.message import ImageMessage, TextMessage, ToolCallMessage, ToolResultMessage from copilotj.core.tool import Tool logger = logging.getLogger(__name__) @@ -37,10 +38,12 @@ "OllamaChatCompletionClient", "new_model_client", "new_vlm_model_client", + "detect_tool_call_mode", ] type FinishReasons = Literal["stop", "tool_calls", "unknown"] +type MessageInput = TextMessage | ImageMessage | ToolCallMessage | ToolResultMessage class ToolCall(pydantic.BaseModel): @@ -117,7 +120,7 @@ def get_api_key(self) -> str | None: @abc.abstractmethod async def create( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], *, tools: list[Tool] | None = None, extra_args: dict[str, Any] | None = None, @@ -126,7 +129,7 @@ async def create( @abc.abstractmethod def create_stream( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], *, tools: list[Tool] | None = None, extra_args: dict[str, Any] | None = None, @@ -226,7 +229,7 @@ def get_api_key(self) -> str | None: @override async def create( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], *, tools: list[Tool] | None = None, extra_args: dict[str, Any] | None = None, @@ -261,7 +264,7 @@ async def create( @override async def create_stream( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], *, tools: list[Tool] | None = None, extra_args: dict[str, Any] | None = None, @@ -347,7 +350,7 @@ async def _create( ) -> openai.AsyncStream[openai.types.chat.ChatCompletionChunk]: ... async def _create( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], *, tools: list[Tool] | None, extra_args: dict[str, Any] | None, @@ -382,10 +385,19 @@ async def _create( raise ModelProviderError(f"OpenAI error: {str(e)}", "openai") from e @classmethod - def _format_messages(cls, messages: Sequence[TextMessage | ImageMessage]): + def _format_messages(cls, messages: Sequence[TextMessage | ImageMessage | ToolCallMessage | ToolResultMessage]): openai_messages: list[openai.types.chat.ChatCompletionMessageParam] = [] group: list[TextMessage | ImageMessage] = [] for message in messages: + # ToolCallMessage and ToolResultMessage are standalone — flush any + # pending group first, then emit them directly. + if isinstance(message, (ToolCallMessage, ToolResultMessage)): + if group: + openai_messages.append(cls._merge_messages(group)) + group.clear() + openai_messages.append(cls._format_tool_message(message)) + continue + if len(group) > 0 and group[0].role != message.role: openai_messages.append(cls._merge_messages(group)) group.clear() @@ -397,9 +409,35 @@ def _format_messages(cls, messages: Sequence[TextMessage | ImageMessage]): return openai_messages + @staticmethod + def _format_tool_message( + msg: ToolCallMessage | ToolResultMessage, + ) -> openai.types.chat.ChatCompletionMessageParam: + """Convert a ToolCallMessage or ToolResultMessage to OpenAI API format.""" + if isinstance(msg, ToolCallMessage): + return { + "role": "assistant", + "content": msg.reasoning_content, + "tool_calls": [ + { + "id": tc.id, + "type": "function", + "function": {"name": tc.name, "arguments": json.dumps(tc.arguments, ensure_ascii=False)}, + } + for tc in msg.tool_calls + ], + } # type: ignore[return-value] + + # ToolResultMessage + return { + "role": "tool", + "tool_call_id": msg.tool_call_id, + "content": msg.content, + } # type: ignore[return-value] + @staticmethod def _merge_messages( - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], ) -> openai.types.chat.ChatCompletionMessageParam: """Format a sequence of messages into OpenAI's chat completion format.""" content = [] @@ -435,7 +473,7 @@ def get_api_key(self) -> str | None: @override async def create( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], *, tools: list[Tool] | None = None, extra_args: dict[str, Any] | None = None, @@ -490,7 +528,7 @@ async def create( @override async def create_stream( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], *, tools: list[Tool] | None = None, extra_args: dict[str, Any] | None = None, @@ -599,7 +637,7 @@ async def _create( ) -> openai.AsyncStream[openai.types.responses.ResponseStreamEvent]: ... async def _create( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], *, tools: list[Tool] | None, extra_args: dict[str, Any] | None, @@ -609,6 +647,13 @@ async def _create( inputs: list[openai.types.responses.ResponseInputItemParam] = [] group: list[TextMessage | ImageMessage] = [] for message in messages: + if isinstance(message, (ToolCallMessage, ToolResultMessage)): + if group: + inputs.append(self._merge_messages(group)) + group.clear() + inputs.append(self._format_tool_input(message)) + continue + if len(group) > 0 and group[0].role != message.role: inputs.append(self._merge_messages(group)) group.clear() @@ -644,9 +689,31 @@ async def _create( except openai.OpenAIError as e: raise ModelProviderError(f"OpenAI error: {str(e)}", "openai") from e + @staticmethod + def _format_tool_input( + msg: ToolCallMessage | ToolResultMessage, + ) -> openai.types.responses.ResponseInputItemParam: + """Convert a ToolCallMessage or ToolResultMessage to Responses API input format.""" + if isinstance(msg, ToolCallMessage): + # Responses API uses function_call_output items for assistant tool calls + return { + "type": "function_call", + "id": msg.tool_calls[0].id if msg.tool_calls else "unknown", + "call_id": msg.tool_calls[0].id if msg.tool_calls else "unknown", + "name": msg.tool_calls[0].name if msg.tool_calls else "unknown", + "arguments": json.dumps(msg.tool_calls[0].arguments, ensure_ascii=False) if msg.tool_calls else "{}", + } # type: ignore[return-value] + + # ToolResultMessage → function_call_output + return { + "type": "function_call_output", + "call_id": msg.tool_call_id, + "output": msg.content, + } # type: ignore[return-value] + @staticmethod def _merge_messages( - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], ) -> openai.types.responses.ResponseInputItemParam: """Format a sequence of messages into the Responses API input format.""" role = _openai_convert_role(messages[0].role) @@ -746,6 +813,25 @@ def _openai_parse_finish_reason(finish_reason: str | None) -> FinishReasons | No return "unknown" +############################# +# Tool Call Detection # +############################# + + +def detect_tool_call_mode(model_client: ModelClient) -> Literal["native", "react"]: + """Detect whether the model supports native tool calling. + + Queries the LiteLLM model capability database via + :func:`copilotj.core.model_info.get_model_capabilities`. Models with + ``supports_function_calling=True`` use the native path; everything else + falls back to ReAct text parsing. + """ + from copilotj.core.model_info import get_model_capabilities + + caps = get_model_capabilities(model_client.get_model()) + return "native" if caps.supports_function_calling else "react" + + ############################# # Gemini # ############################# @@ -791,7 +877,7 @@ def get_model(self) -> str: def get_api_key(self) -> str | None: return None - def _format_messages(self, messages: Sequence[TextMessage | ImageMessage]) -> list[dict]: + def _format_messages(self, messages: Sequence[MessageInput]) -> list[dict]: """Formats messages for the Ollama API.""" ollama_messages = [] for msg in messages: @@ -799,6 +885,18 @@ def _format_messages(self, messages: Sequence[TextMessage | ImageMessage]) -> li ollama_messages.append({"role": msg.role, "content": msg.text}) elif isinstance(msg, ImageMessage): logger.warning("Image messages not fully supported by Ollama client yet. Skipping image.") + elif isinstance(msg, ToolCallMessage): + ollama_messages.append( + { + "role": "assistant", + "content": msg.reasoning_content or None, + "tool_calls": [ + {"function": {"name": tc.name, "arguments": tc.arguments}} for tc in msg.tool_calls + ], + } + ) + elif isinstance(msg, ToolResultMessage): + ollama_messages.append({"role": "tool", "content": msg.content}) else: raise ValueError(f"Unsupported message type: {msg}") return ollama_messages @@ -806,7 +904,7 @@ def _format_messages(self, messages: Sequence[TextMessage | ImageMessage]) -> li @override async def create( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], *, tools: list[Tool] | None = None, extra_args: dict[str, Any] | None = None, @@ -846,7 +944,7 @@ async def create( @override async def create_stream( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[MessageInput], *, tools: list[Tool] | None = None, extra_args: dict[str, Any] | None = None, diff --git a/copilotj/multiagent/Executor.py b/copilotj/multiagent/Executor.py index 4e873db3..a0d00617 100644 --- a/copilotj/multiagent/Executor.py +++ b/copilotj/multiagent/Executor.py @@ -7,6 +7,7 @@ from copilotj.core import ChatAgent, ModelClient, ModelSyntaxError, TextMessage, Tool from copilotj.multiagent.leader_prompts import build_tool_prompt +from copilotj.util import ReActChatCompletionClient __all__ = ["Executor"] @@ -37,11 +38,21 @@ def _build_enhanced_system_prompt(self, base_prompt: str) -> str: tools_info = build_tool_prompt(self.tools) - tools_usage = """\ + if isinstance(self._client, ReActChatCompletionClient): + # ReAct mode: include Action format instructions + tools_usage = """\ ## Tool Usage Format: When you need to use a tool, format your action as json format: Action: {"name": "", "args": } +## Tool Selection Guidelines: +- Choose the most appropriate tool based on the task requirements +- Provide clear and specific parameters for the tool +- If unsure about parameters, describe what you want to accomplish +""" + else: + # Native mode: tools are called directly, just list them + tools_usage = """\ ## Tool Selection Guidelines: - Choose the most appropriate tool based on the task requirements - Provide clear and specific parameters for the tool @@ -186,11 +197,17 @@ def _suggest_tool_based_on_context(self, thought: str, task: str) -> str: if suggested_tools: top_suggestions = [tool[0] for tool in suggested_tools[:3]] - TEMPLATE = 'Consider using one of these tools: {{TOOLS}}. Format: {"name": "tool_name", "args": }' - return TEMPLATE.replace("{{TOOLS}}", ", ".join(top_suggestions)) + if isinstance(self._client, ReActChatCompletionClient): + template = 'Consider using one of these tools: {{TOOLS}}. Format: {"name": "tool_name", "args": }' + else: + template = "Consider using one of these tools: {{TOOLS}}." + return template.replace("{{TOOLS}}", ", ".join(top_suggestions)) else: - TEMPLATE = 'Please choose from available tools: {{TOOLS}}. Format: {"name": "tool_name", "args": }' - return TEMPLATE.replace("{{TOOLS}}", self._tool_names()) + if isinstance(self._client, ReActChatCompletionClient): + template = 'Please choose from available tools: {{TOOLS}}. Format: {"name": "tool_name", "args": }' + else: + template = "Please choose from available tools: {{TOOLS}}." + return template.replace("{{TOOLS}}", self._tool_names()) def _build_execution_context(self, conversation_context: dict[str, Any], iteration: int) -> str: """Build the execution context for the current iteration""" @@ -231,7 +248,7 @@ def _is_task_complete(self, response: str) -> bool: response_lower = response.lower() return any(indicator in response_lower for indicator in completion_indicators) - REFLECTION_TEMPLATE = """\ + REFLECTION_TEMPLATE_REACT = """\ Please reflect on your progress so far: 1. **Status Check**: What have you accomplished toward the task goal? @@ -243,18 +260,30 @@ def _is_task_complete(self, response: str) -> bool: Use the format: Thought: [Your analysis and reasoning] -Action: {"name": "", "args": } +Action: {"name": "", "args": }""" -Iteration: {{CURRENT_ITERATION}}/{{MAX_ITERATIONS}} -""" + REFLECTION_TEMPLATE_NATIVE = """\ +Please reflect on your progress so far: + +1. **Status Check**: What have you accomplished toward the task goal? +2. **Next Action**: What specific action should you take next? +3. **Tool Usage**: Which tool would be most appropriate for the next step? +4. **Expected Outcome**: What do you expect to achieve with this action? + +Available tools: {{TOOL_NAMES}} + +Use the format: +Thought: [Your analysis and reasoning] +Then call the appropriate tool directly.""" def _generate_reflection_prompt(self, iteration: int) -> str: """Generate a reflection prompt for the agent""" - return ( - self.REFLECTION_TEMPLATE.replace("{{TOOL_NAMES}}", self._tool_names()) - .replace("{{CURRENT_ITERATION}}", str(iteration + 1)) - .replace("{{MAX_ITERATIONS}}", str(self.max_iterations)) + template = ( + self.REFLECTION_TEMPLATE_REACT + if isinstance(self._client, ReActChatCompletionClient) + else self.REFLECTION_TEMPLATE_NATIVE ) + return template.replace("{{TOOL_NAMES}}", self._tool_names()) def _generate_final_summary(self, conversation_context: dict[str, Any]) -> str: """Generate a final summary of the task execution""" diff --git a/copilotj/multiagent/agent_configs/research_agent.toml b/copilotj/multiagent/agent_configs/research_agent.toml index 0f5d3f3d..b6b453c7 100644 --- a/copilotj/multiagent/agent_configs/research_agent.toml +++ b/copilotj/multiagent/agent_configs/research_agent.toml @@ -13,15 +13,19 @@ Your mandate: answer accurately with the **fewest necessary tool calls**. Let's work this out in a step by step way to be sure we have the right answer. ## Execution Workflow -- Thought: When you thought, please reflect on your progress with several sentences about current status, Tool usage, and Expected outcome: +- Thought: When you thought, please reflect on your progress with several sentences about current status, Tool usage, and Expected outcome: +{% if tool_call_mode == "react" %} - Action: output EXACTLY ONE tool call in this JSON form: \ Action: {"name": "", "args": { ... }} +{% else %} +- When you need to use a tool, call it directly with the appropriate parameters. +{% endif %} -After the Action is executed, you will get the result of the tool call as an "observation". -This Thought/Action/Observation can repeat N times, you should take several steps when needed. +After the tool is executed, you will get the result as an "observation". +This Thought/Tool/Observation can repeat N times, you should take several steps when needed. -If the task is already complete, skip the Action and output: +If the task is already complete, skip the tool call and output: Final Answer: [your answer or summary of the process] ## Response Process: diff --git a/copilotj/multiagent/agent_configs/tool_agent.toml b/copilotj/multiagent/agent_configs/tool_agent.toml index 9e16c91a..35d95945 100644 --- a/copilotj/multiagent/agent_configs/tool_agent.toml +++ b/copilotj/multiagent/agent_configs/tool_agent.toml @@ -44,15 +44,19 @@ You are the **Tool Agent**, part of a multiagent system, specializing in bioimag Let's work this out in a step by step way to be sure we have the right answer. ## Execution Workflow -- Thought: When you thought, please reflect on your progress with several sentences about current status, Tool usage, and Expected outcome: +- Thought: When you thought, please reflect on your progress with several sentences about current status, Tool usage, and Expected outcome: +{% if tool_call_mode == "react" %} - Action: output EXACTLY ONE tool call in this JSON form: \ Action: {"name": "", "args": { ... }} +{% else %} +- When you need to use a tool, call it directly with the appropriate parameters. +{% endif %} -After the Action is executed, you will get the result of the tool call as an "observation". -This Thought/Action/Observation can repeat N times, you should take several steps when needed. +After the tool is executed, you will get the result as an "observation". +This Thought/Tool/Observation can repeat N times, you should take several steps when needed. -If the task is already complete, skip the Action and output: +If the task is already complete, skip the tool call and output: Final Answer: [Please provide a detailed summary of the following ImageJ task execution steps. The summary must be comprehensive like: Summary of ImageJ Task 1. **Task Overview**: What was the main objective @@ -70,18 +74,26 @@ Summary of ImageJ Task assistant: | Thought: I need to use the cellpose_segmentation tool to segment the cells in the image. I'll use the nuclei \ model since it's good for cell segmentation. +{% if tool_call_mode == "react" %} Action: { "name": "cellpose_segmentation", "args": { \ "image_path": "C:/XXXX/img08.png", "model_type": "nuclei", "gpu": true, "save_path": "C:/XXXX/results" \ }} +{% else %} + → Call cellpose_segmentation with image_path="C:/XXXX/img08.png", model_type="nuclei", gpu=true, save_path="C:/XXXX/results" +{% endif %} ### Example 2: - user: "The Cellpose segmentation failed with a memory error" assistant: | Thought: The Cellpose tool failed due to memory issues. I need to use the execute_python_script tool to \ implement image tiling or reduce image size to work around this limitation. +{% if tool_call_mode == "react" %} Action: { "name": "execute_python_script", "args": { \ "script": "# Tile-based Cellpose processing for large images\nimport numpy as np\nfrom cellpose import models\n..." \ }} +{% else %} + → Call execute_python_script with script for tile-based Cellpose processing +{% endif %} ### Example 3: - user: "Process the image for cell segmentation" @@ -91,7 +103,6 @@ Summary of ImageJ Task Final Answer: I need the image file path to perform cell segmentation. Please provide the path to the image \ you want to process. Once you provide the image path, I can use Cellpose or other segmentation tools to segment \ the cells. - ## Error Handling Strategy: 1. **Tool-First**: Always attempt the appropriate specialized tool first diff --git a/copilotj/multiagent/agent_loader.py b/copilotj/multiagent/agent_loader.py index 46deae78..f939dfdd 100644 --- a/copilotj/multiagent/agent_loader.py +++ b/copilotj/multiagent/agent_loader.py @@ -10,7 +10,9 @@ from copilotj.core import FunctionTool, ModelClient, Tool from copilotj.core.config import Config +from copilotj.core.model_client import detect_tool_call_mode from copilotj.multiagent.Executor import Executor +from copilotj.util.prompt_template import render_prompt __all__ = ["load_agent_configs"] @@ -100,9 +102,13 @@ def _load_agent_configs(glob_pattern: str, *, model_client: ModelClient, cfg: Co _log.info("Loaded tools for %s: %s", name, list(agent_tools)) + # Render prompt template with detected tool_call_mode + tool_call_mode = detect_tool_call_mode(model_client) + rendered_prompt = render_prompt(prompt, tool_call_mode=tool_call_mode) + # Create agent instance with tool descriptions agents[name] = agent_class( - name=name, description=description, prompt=prompt, tools=agent_tools, model_client=model_client + name=name, description=description, prompt=rendered_prompt, tools=agent_tools, model_client=model_client ) _log.info("Successfully loaded agent: %s", name) diff --git a/copilotj/multiagent/leader_multiagent.py b/copilotj/multiagent/leader_multiagent.py index 1e219603..75608519 100644 --- a/copilotj/multiagent/leader_multiagent.py +++ b/copilotj/multiagent/leader_multiagent.py @@ -26,6 +26,9 @@ Pattern, TextMessage, Tool, + ToolCallMessage, + ToolCallRecord, + ToolResultMessage, new_model_client, ) from copilotj.core.config import Config @@ -49,7 +52,6 @@ build_available_specialized_agents_prompt, build_initial_user_message, build_leader_system_prompt, - build_observation_message, build_tool_prompt, make_summary_prompt, ) @@ -57,6 +59,8 @@ from copilotj.plugin import ClientPluginAPI from copilotj.util import ReActChatCompletionClient +from copilotj.core.model_client import detect_tool_call_mode + __all__ = ["LeaderDriven"] PROMPT_CHAT_HISTORY_LIMIT = 8 @@ -66,27 +70,6 @@ SaveWorkflowHandler = Callable[[str, str | None, int | None, list[int] | None], Awaitable[str]] -def _reconstruct_react_text(response: ModelResponse) -> str: - """Rebuild an assistant-side ReAct text block from a parsed ModelResponse. - - The ReAct wrapper parses the model's raw text output into ``reasoning_content`` - (Thought), ``tool_calls`` (Action), and ``content`` (Final Answer). For - multi-turn conversation we re-assemble those parts into the same shape the - model originally produced, so the next turn's context matches the - ReAct-formatted examples in the system prompt. - """ - parts: list[str] = [] - if response.reasoning_content: - parts.append(f"Thought: {response.reasoning_content.strip()}") - if response.tool_calls: - tc = response.tool_calls[0] - args_json = json.dumps(tc.args.model_dump(), ensure_ascii=False) - parts.append(f'Action: {{"name": "{tc.tool.name}", "args": {args_json}}}') - if response.content: - parts.append(f"Final Answer: {response.content.strip()}") - return "\n".join(parts) - - class LeaderAgent(ChatAgent): def __init__( self, @@ -97,15 +80,17 @@ def __init__( model_client: ModelClient, apis: ClientPluginAPI, cfg: Config, + tool_call_mode: str = "react", ): super().__init__(name, description, model_client=model_client) self.chat_history: list[dict[str, Any]] = [] self._save_workflow_handler: SaveWorkflowHandler | None = None + self._tool_call_mode = tool_call_mode # Per-dialog conversation state. Populated by begin_dialog() and # extended by continue_dialog(). Reset at the start of each dialog. - self._dialog_messages: list[TextMessage] = [] + self._dialog_messages: list[TextMessage | ToolCallMessage | ToolResultMessage] = [] self._apis = apis self.plugin_tools = tools.PluginTools(apis, cfg=cfg) @@ -163,7 +148,7 @@ def _recent_history_text(self) -> str: ] return "\n".join(history) - async def _build_system_prompt(self) -> str: + async def _build_system_prompt(self, *, tool_call_mode: str = "react") -> str: """Build the static system prompt for the current session.""" tool_list = build_tool_prompt(self.tools) + "\n" + build_available_specialized_agents_prompt(self.agents) macro_plugins = _load_macro_plugin_names() or ["StarDist", "TrackMate", "CLIJ2"] @@ -175,6 +160,7 @@ async def _build_system_prompt(self) -> str: plugins_text=plugins_text, system_info_text=system_info_text, default_image_path=str(py_tools.get_project_temp_dir()), + tool_call_mode=tool_call_mode, ) async def begin_dialog(self, main_task: str, *, trace_ctx: Langfuse | None = None) -> ModelResponse: @@ -186,7 +172,7 @@ async def begin_dialog(self, main_task: str, *, trace_ctx: Langfuse | None = Non only append new turns rather than resending the whole prompt. """ self.imagej_windowInfo_text = await self.plugin_tools.imagej_windowInfo() - system_prompt = await self._build_system_prompt() + system_prompt = await self._build_system_prompt(tool_call_mode=self._tool_call_mode) initial_user = build_initial_user_message( main_task=main_task, @@ -210,20 +196,35 @@ async def continue_dialog( ) -> ModelResponse: """Continue the current dialog after a tool call. - Appends the prior assistant response (reconstructed from the ReAct - parser's structured output) and a new user message carrying the tool - observation plus refreshed ImageJ window info. + Appends the prior assistant response as a ToolCallMessage (with + reasoning) and a ToolResultMessage for the observation. The client + layer converts these to the appropriate format (native API or + reconstructed ReAct text) depending on the active tool-call mode. """ - assistant_text = _reconstruct_react_text(prior_response) - if assistant_text: - self._dialog_messages.append(TextMessage(role="assistant", text=assistant_text)) + if prior_response.tool_calls: + # Record the assistant turn with its tool calls + records = [ + ToolCallRecord( + id=tc.id, + name=tc.tool.name, + arguments=tc.args.model_dump(), + ) + for tc in prior_response.tool_calls + ] + self._dialog_messages.append( + ToolCallMessage(tool_calls=records, reasoning_content=prior_response.reasoning_content) + ) + # Record each tool result + for tc in prior_response.tool_calls: + self._dialog_messages.append(ToolResultMessage(tool_call_id=tc.id, content=str(observation))) + elif prior_response.content: + self._dialog_messages.append(TextMessage(role="assistant", text=prior_response.content)) + # Window info as a user message (always, regardless of mode) self.imagej_windowInfo_text = await self.plugin_tools.imagej_windowInfo() - user_text = build_observation_message( - tool_response=observation, - imagej_window_info=self.imagej_windowInfo_text, + self._dialog_messages.append( + TextMessage(role="user", text=f"## ImageJ WindowInfo\n{self.imagej_windowInfo_text or '(no image open)'}") ) - self._dialog_messages.append(TextMessage(role="user", text=user_text)) self._log_prompt_size("continue_dialog") return await self._create(*self._dialog_messages, tools=self.tools, trace_ctx=trace_ctx) @@ -392,8 +393,15 @@ def __init__( base_url=cfg.llm_base_url, cfg=cfg, ) - # wrap the model client to handle ReAct-style responses - wrapped_model_client = ReActChatCompletionClient(self.model_client) + # Detect tool call mode from LiteLLM model capability database + self._tool_call_mode = detect_tool_call_mode(self.model_client) + self.log_info(f"Detected tool call mode: {self._tool_call_mode}") + + # Wrap the model client for ReAct mode only + if self._tool_call_mode == "react": + wrapped_model_client: ModelClient = ReActChatCompletionClient(self.model_client) + else: + wrapped_model_client = self.model_client self.dialog_counter = 1 self.max_steps_before_summary = max_steps_before_summary @@ -440,6 +448,7 @@ def __init__( agents=self.specialized_agents, apis=apis, cfg=cfg, + tool_call_mode=self._tool_call_mode, ) self.workflow_saver.chat_history = self.leader_agent.chat_history self.leader_agent.set_save_workflow_handler(self.workflow_saver.save) diff --git a/copilotj/multiagent/leader_prompts.py b/copilotj/multiagent/leader_prompts.py index 9cb7a52f..0970c8cd 100644 --- a/copilotj/multiagent/leader_prompts.py +++ b/copilotj/multiagent/leader_prompts.py @@ -45,13 +45,17 @@ ## Execution Workflow - Thought: When you thought, please reflect on your progress with several sentences about current status, Tool usage, and Expected outcome.(Thought should be short and concise) +{% if tool_call_mode == "react" %} - Action: output EXACTLY ONE tool call in this JSON form: \ Action: {"name": "", "args": { ... }} +{% else %} +- When you need to use a tool, call it directly with the appropriate parameters. +{% endif %} -After the Action is executed, you will get the result of the tool call as an "observation". -This Thought/Action/Observation can repeat N times, you should take several steps when needed. +After the tool is executed, you will get the result as an "observation". +This Thought/Tool/Observation can repeat N times, you should take several steps when needed. -If the task is already complete, skip the Action and output: +If the task is already complete, skip the tool call and output: Final Answer: [your answer or summary of the process] # Saved Workflow Execution Shortcut @@ -105,16 +109,28 @@ ## Few-Shot Examples #### Example 1 — Perception then Knowledge Retrieval Thought: I need to analyze the image characteristics first. +{% if tool_call_mode == "react" %} Action: {"name":"imagej_perception","args":{"task":"analyze this image for nuclei segmentation"}} +{% else %} +→ Call imagej_perception with task="analyze this image for nuclei segmentation" +{% endif %} (Observation returned with image_desc and perception_info) Thought: Now I should retrieve relevant macros and workflow tips. +{% if tool_call_mode == "react" %} Action: {"name":"kb_retrieve","args":{"query":"nuclei segmentation","image_desc":"8-bit, single-channel, fluorescence","perception_info":"bright nuclei on dark background","topk":8}} +{% else %} +→ Call kb_retrieve with query="nuclei segmentation", perception_info="bright nuclei on dark background" +{% endif %} #### Example 2 — Incremental Segmentation (DAPI nuclei) **User:** "Please segment nuclei from my DAPI image and export counts." **Assistant:** Thought: I need to plan a stepwise segmentation macro and get user permission first. +{% if tool_call_mode == "react" %} Action: {"name":"user_manipulate","args":{"instructions":"Permission needed: Incremental DAPI nuclei segmentation (stepwise). +{% else %} +→ Call user_manipulate with instructions="Permission needed: Incremental DAPI nuclei segmentation (stepwise). +{% endif %} **What I will do, step by step** 1) Duplicate the current image (preserve original). 2) Apply Gaussian Blur (sigma=1.5) to smooth local noise and enhance object continuity. @@ -140,17 +156,23 @@ **Please respond** - Reply `Yes` to proceed, `No` to cancel - Reply `` to adjust parameters or refine this plan. +{% if tool_call_mode == "react" %} "}} +{% else %} +" +{% endif %} ## Dialog Structure The first user message contains the current request and, if relevant, a summary of previous chat history. \ Tool observations are returned as subsequent user messages. You must connect each new request with the \ prior conversation context. +{% if tool_call_mode == "react" %} ## Anti-pattern (do NOT do this): Thought: I'll open the image and then threshold it. Action: {"name": "run_macro", "args": {"script": "run(\"Open...\", \"path=/a.tif\");"}} Action: {"name": "run_macro", "args": {"script": "setAutoThreshold(\"Otsu\");"}} <-- Wrong: two Actions in one message +{% endif %} ## Rules ### Execution Flow @@ -698,6 +720,8 @@ def build_leader_system_prompt( plugins_text: str, system_info_text: str, default_image_path: str, + *, + tool_call_mode: str = "react", ) -> str: """Build the static system prompt for the leader agent. @@ -706,12 +730,15 @@ def build_leader_system_prompt( Dynamic content (chat history, current task, ImageJ window info, tool observations) lives in subsequent user/assistant messages. """ - return ( + from copilotj.util.prompt_template import render_prompt + + template = ( PROMPT_LEADER.replace("{TOOL_LIST}", tool_list) .replace("{SPECIAL_PLUGIN}", plugins_text) .replace("{SYSTEM_INFO}", system_info_text) .replace("{DEFAULT_IMAGE_PATH}", default_image_path) ) + return render_prompt(template, tool_call_mode=tool_call_mode) def build_initial_user_message( diff --git a/copilotj/test/core/test_agent.py b/copilotj/test/core/test_agent.py new file mode 100644 index 00000000..0f867b62 --- /dev/null +++ b/copilotj/test/core/test_agent.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright contributors to the CopilotJ project. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for Thought extraction post-processing in ChatAgent._create(). + +The Thought extraction logic lives inline in ``ChatAgent._create()`` (lines +215–231 of ``agent.py``). It is tightly coupled to streaming, langfuse +tracing, and the runtime — so rather than exercising the full method, these +tests reproduce the *exact same conditional check* that the production code +uses. If the condition changes in production, these tests should be updated +in tandem. +""" + +from copilotj.core.model_client import ModelResponse, ToolCall +from copilotj.core.tool import FunctionTool + + +def _lookup(query: str) -> str: + return query + + +_TOOL = FunctionTool(_lookup, "lookup", name="lookup") + + +def _make_tool_call() -> ToolCall: + return ToolCall(id="tc1", tool=_TOOL, args=_TOOL.args_type()(query="test")) + + +def _apply_thought_extraction(completion: ModelResponse) -> ModelResponse: + """Reproduce the Thought extraction post-processing from ChatAgent._create(). + + This mirrors the conditional at lines 220–231 of ``agent.py`` exactly. + """ + if ( + completion.tool_calls + and completion.content + and not completion.reasoning_content + and completion.content.strip().lower().startswith("thought") + ): + return ModelResponse( + content=None, + reasoning_content=completion.content.strip(), + tool_calls=completion.tool_calls, + finish_reason=completion.finish_reason, + ) + return completion + + +def test_thought_moved_to_reasoning_when_tool_calls_present(): + original = ModelResponse( + content="Thought: I need to analyze this image first.", + reasoning_content=None, + tool_calls=[_make_tool_call()], + finish_reason="tool_calls", + ) + + result = _apply_thought_extraction(original) + + assert result.content is None + assert result.reasoning_content == "Thought: I need to analyze this image first." + assert result.tool_calls == original.tool_calls + + +def test_thought_not_moved_when_no_tool_calls(): + original = ModelResponse( + content="Thought: I will answer directly.\nFinal Answer: 42", + reasoning_content=None, + tool_calls=[], + finish_reason="stop", + ) + + result = _apply_thought_extraction(original) + + # Content unchanged — no tool calls, so this is a final answer + assert result.content == "Thought: I will answer directly.\nFinal Answer: 42" + assert result.reasoning_content is None + + +def test_thought_not_moved_when_reasoning_already_present(): + original = ModelResponse( + content="Thought: extra text", + reasoning_content="already captured by model", + tool_calls=[_make_tool_call()], + finish_reason="tool_calls", + ) + + result = _apply_thought_extraction(original) + + # reasoning_content was already set — don't overwrite + assert result.content == "Thought: extra text" + assert result.reasoning_content == "already captured by model" + + +def test_content_without_thought_prefix_unchanged(): + original = ModelResponse( + content="Just some regular content", + reasoning_content=None, + tool_calls=[_make_tool_call()], + finish_reason="tool_calls", + ) + + result = _apply_thought_extraction(original) + + # Doesn't start with "thought" — leave as-is + assert result.content == "Just some regular content" + assert result.reasoning_content is None diff --git a/copilotj/test/core/test_model_client.py b/copilotj/test/core/test_model_client.py index 317577b0..837b3dbb 100644 --- a/copilotj/test/core/test_model_client.py +++ b/copilotj/test/core/test_model_client.py @@ -2,8 +2,16 @@ # # SPDX-License-Identifier: Apache-2.0 -from copilotj.core.message import ImageMessage, TextMessage -from copilotj.core.model_client import OpenAIChatCompletionClient +from unittest.mock import patch + +from copilotj.core.message import ( + ImageMessage, + TextMessage, + ToolCallMessage, + ToolCallRecord, + ToolResultMessage, +) +from copilotj.core.model_client import OpenAIChatCompletionClient, OpenAIResponseClient, detect_tool_call_mode def test_format_single_text_message(): @@ -70,3 +78,130 @@ def test_merge_mixed_messages(): assert len(result["content"]) == 2 # type: ignore assert result["content"][0]["type"] == "text" # type: ignore assert result["content"][1]["type"] == "image_url" # type: ignore + + +# --- _format_tool_message tests --- + + +def test_format_tool_call_message(): + msg = ToolCallMessage( + reasoning_content="I need to search", + tool_calls=[ToolCallRecord(id="call_1", name="search", arguments={"query": "cells"})], + ) + result = OpenAIChatCompletionClient._format_tool_message(msg) + + assert result["role"] == "assistant" + assert result["content"] == "I need to search" # type: ignore + tool_calls = result["tool_calls"] # type: ignore + assert len(tool_calls) == 1 + assert tool_calls[0]["id"] == "call_1" + assert tool_calls[0]["type"] == "function" + assert tool_calls[0]["function"]["name"] == "search" # type: ignore + assert '"query"' in tool_calls[0]["function"]["arguments"] # type: ignore + + +def test_format_tool_call_message_without_thought(): + msg = ToolCallMessage( + tool_calls=[ToolCallRecord(id="call_2", name="lookup", arguments={})], + ) + result = OpenAIChatCompletionClient._format_tool_message(msg) + + assert result["role"] == "assistant" + assert result["content"] is None # type: ignore + + +def test_format_tool_result_message(): + msg = ToolResultMessage(tool_call_id="call_1", content="found 42 items") + result = OpenAIChatCompletionClient._format_tool_message(msg) + + assert result["role"] == "tool" # type: ignore + assert result["tool_call_id"] == "call_1" # type: ignore + assert result["content"] == "found 42 items" # type: ignore + + +def test_format_tool_call_message_multiple_calls(): + msg = ToolCallMessage( + reasoning_content="parallel calls", + tool_calls=[ + ToolCallRecord(id="c1", name="search", arguments={"q": "a"}), + ToolCallRecord(id="c2", name="lookup", arguments={"q": "b"}), + ], + ) + result = OpenAIChatCompletionClient._format_tool_message(msg) + + tool_calls = result["tool_calls"] # type: ignore + assert len(tool_calls) == 2 + assert tool_calls[0]["function"]["name"] == "search" # type: ignore + assert tool_calls[1]["function"]["name"] == "lookup" # type: ignore + + +# --- _format_tool_input tests (Responses API) --- + + +def test_format_tool_input_call_message(): + msg = ToolCallMessage( + tool_calls=[ToolCallRecord(id="fc_1", name="search", arguments={"query": "test"})], + ) + result = OpenAIResponseClient._format_tool_input(msg) + + assert result["type"] == "function_call" # type: ignore + assert result["id"] == "fc_1" # type: ignore + assert result["name"] == "search" # type: ignore + assert '"query"' in result["arguments"] # type: ignore + + +def test_format_tool_input_result_message(): + msg = ToolResultMessage(tool_call_id="fc_1", content="result data") + result = OpenAIResponseClient._format_tool_input(msg) + + assert result["type"] == "function_call_output" # type: ignore + assert result["call_id"] == "fc_1" # type: ignore + assert result["output"] == "result data" # type: ignore + + +# --- detect_tool_call_mode tests --- + + +@patch("copilotj.core.model_info.get_model_capabilities") +def test_detect_native_for_known_model(mock_caps): + from copilotj.core.model_info import ModelCapabilities + + mock_caps.return_value = ModelCapabilities( + model="gpt-4o", + supports_vision=True, + supports_function_calling=True, + context_window=128000, + max_output_tokens=16384, + source="litellm_db", + ) + + # Create a minimal ModelClient subclass for testing + class _TestClient(OpenAIChatCompletionClient): + pass + + client = _TestClient.__new__(_TestClient) + client._model = "gpt-4o" # noqa: SLF001 + + assert detect_tool_call_mode(client) == "native" + + +@patch("copilotj.core.model_info.get_model_capabilities") +def test_detect_react_for_unsupported_model(mock_caps): + from copilotj.core.model_info import ModelCapabilities + + mock_caps.return_value = ModelCapabilities( + model="some-local-model", + supports_vision=False, + supports_function_calling=False, + context_window=None, + max_output_tokens=None, + source="unknown", + ) + + class _TestClient(OpenAIChatCompletionClient): + pass + + client = _TestClient.__new__(_TestClient) + client._model = "some-local-model" # noqa: SLF001 + + assert detect_tool_call_mode(client) == "react" diff --git a/copilotj/test/util/test_prompt_template.py b/copilotj/test/util/test_prompt_template.py new file mode 100644 index 00000000..77973746 --- /dev/null +++ b/copilotj/test/util/test_prompt_template.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright contributors to the CopilotJ project. +# +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from copilotj.util.prompt_template import render_prompt + + +# --- Basic conditionals --- + + +def test_renders_if_branch_when_condition_true(): + template = '{% if mode == "react" %}\nAction line\n{% endif %}' + assert "Action line" in render_prompt(template, mode="react") + + +def test_skips_if_branch_when_condition_false(): + template = '{% if mode == "react" %}\nAction line\n{% endif %}' + assert "Action line" not in render_prompt(template, mode="native") + + +def test_renders_else_branch_when_condition_false(): + template = '{% if mode == "react" %}\nuse Action\n{% else %}\ncall directly\n{% endif %}' + result = render_prompt(template, mode="native") + assert "call directly" in result + assert "use Action" not in result + + +def test_skips_else_branch_when_condition_true(): + template = '{% if mode == "react" %}\nuse Action\n{% else %}\ncall directly\n{% endif %}' + result = render_prompt(template, mode="react") + assert "use Action" in result + assert "call directly" not in result + + +# --- elif --- + + +def test_elif_branch_taken(): + template = '{% if x == "a" %}\nA\n{% elif x == "b" %}\nB\n{% endif %}' + result = render_prompt(template, x="b") + assert "B" in result + assert "A" not in result + + +def test_elif_skipped_when_if_taken(): + template = '{% if x == "a" %}\nA\n{% elif x == "b" %}\nB\n{% endif %}' + result = render_prompt(template, x="a") + assert "A" in result + assert "B" not in result + + +def test_elif_else_fallback(): + template = '{% if x == "a" %}\nA\n{% elif x == "b" %}\nB\n{% else %}\nC\n{% endif %}' + result = render_prompt(template, x="z") + assert "C" in result + assert "A" not in result + assert "B" not in result + + +# --- Variable substitution --- + + +def test_substitutes_template_variables(): + template = "Hello {name}, your mode is {mode}." + assert render_prompt(template, name="Alice", mode="react") == "Hello Alice, your mode is react." + + +def test_leaves_unknown_variables_as_is(): + template = "Hello {name}, {unknown_var} stays." + assert render_prompt(template, name="Alice") == "Hello Alice, {unknown_var} stays." + + +# --- Edge cases --- + + +def test_no_conditionals_returns_template_as_is(): + template = "Just a plain template with no conditionals." + assert render_prompt(template) == "Just a plain template with no conditionals." + + +def test_empty_template(): + assert render_prompt("") == "" + + +def test_nested_conditionals(): + template = '{% if a == "1" %}\nouter-true\n{% if b == "2" %}\ninner-true\n{% endif %}\n{% endif %}' + result = render_prompt(template, a="1", b="2") + assert "outer-true" in result + assert "inner-true" in result + + +def test_nested_conditionals_inner_false(): + template = '{% if a == "1" %}\nouter-true\n{% if b == "2" %}\ninner-true\n{% endif %}\n{% endif %}' + result = render_prompt(template, a="1", b="9") + assert "outer-true" in result + assert "inner-true" not in result + + +def test_consecutive_if_blocks(): + template = '{% if x == "a" %}\nAAA\n{% endif %}\n{% if x == "b" %}\nBBB\n{% endif %}' + assert "AAA" in render_prompt(template, x="a") + assert "BBB" not in render_prompt(template, x="a") + assert "BBB" in render_prompt(template, x="b") + assert "AAA" not in render_prompt(template, x="b") + + +def test_multiline_content_in_branches(): + template = '{% if mode == "react" %}\nLine 1\nLine 2\n{% else %}\nAlt 1\nAlt 2\n{% endif %}' + result = render_prompt(template, mode="react") + assert "Line 1" in result + assert "Line 2" in result + assert "Alt" not in result + + +# --- Error cases --- + + +def test_unclosed_if_raises_syntax_error(): + with pytest.raises(SyntaxError, match="Unclosed"): + render_prompt('{% if x == "1" %}\ncontent') + + +def test_elif_without_if_raises_syntax_error(): + with pytest.raises(SyntaxError, match="without matching"): + render_prompt('{% elif x == "1" %}\ncontent\n{% endif %}') + + +def test_else_without_if_raises_syntax_error(): + with pytest.raises(SyntaxError, match="without matching"): + render_prompt("{% else %}\ncontent\n{% endif %}") + + +def test_endif_without_if_raises_syntax_error(): + with pytest.raises(SyntaxError, match="without matching"): + render_prompt("{% endif %}") + + +def test_elif_after_else_raises_syntax_error(): + with pytest.raises(SyntaxError, match="after"): + render_prompt('{% if x == "1" %}\nA\n{% else %}\nB\n{% elif x == "2" %}\nC\n{% endif %}') + + +def test_if_without_condition_raises_syntax_error(): + with pytest.raises(SyntaxError, match="Missing condition"): + render_prompt("{% if %}\ncontent\n{% endif %}") diff --git a/copilotj/test/util/test_react_parser.py b/copilotj/test/util/test_react_parser.py index cc07a22f..af395c38 100644 --- a/copilotj/test/util/test_react_parser.py +++ b/copilotj/test/util/test_react_parser.py @@ -6,7 +6,13 @@ from collections.abc import AsyncGenerator, Sequence from typing import Any, override -from copilotj.core.message import ImageMessage, TextMessage +from copilotj.core.message import ( + ImageMessage, + TextMessage, + ToolCallMessage, + ToolCallRecord, + ToolResultMessage, +) from copilotj.core.model_client import ModelClient, ModelResponse, ModelResponseChunk, ModelSyntaxError, ToolCall from copilotj.core.tool import FunctionTool, Tool from copilotj.util.react_parser import ReActChatCompletionClient, _build_last_line_prefix_regex @@ -330,3 +336,152 @@ async def _collect_stream(client: ReActChatCompletionClient, tools: list[Tool]) async for item in client.create_stream([TextMessage(role="user", text="help")], tools=tools): items.append(item) return items + + +# --- _convert_messages tests --- + + +def test_convert_tool_call_message_with_thought(): + messages = [ + ToolCallMessage( + reasoning_content="I need to look something up", + tool_calls=[ToolCallRecord(id="tc1", name="lookup", arguments={"query": "cells"})], + ) + ] + result = ReActChatCompletionClient._convert_messages(messages) + + assert len(result) == 1 + assert isinstance(result[0], TextMessage) + assert result[0].role == "assistant" + assert "Thought: I need to look something up" in result[0].text + assert 'Action: {"name": "lookup"' in result[0].text + + +def test_convert_tool_call_message_without_thought(): + messages = [ToolCallMessage(tool_calls=[ToolCallRecord(id="tc1", name="lookup", arguments={"query": "cells"})])] + result = ReActChatCompletionClient._convert_messages(messages) + + assert len(result) == 1 + assert isinstance(result[0], TextMessage) + assert "Thought:" not in result[0].text + assert 'Action: {"name": "lookup"' in result[0].text + + +def test_convert_tool_result_message(): + messages = [ToolResultMessage(tool_call_id="tc1", content="found 42 items")] + result = ReActChatCompletionClient._convert_messages(messages) + + assert len(result) == 1 + assert isinstance(result[0], TextMessage) + assert result[0].role == "user" + assert result[0].text == "Observation:\nfound 42 items" + + +def test_convert_mixed_messages(): + messages = [ + TextMessage(role="user", text="hello"), + ToolCallMessage( + reasoning_content="thinking", + tool_calls=[ToolCallRecord(id="tc1", name="lookup", arguments={"query": "x"})], + ), + ToolResultMessage(tool_call_id="tc1", content="result text"), + TextMessage(role="assistant", text="done"), + ] + result = ReActChatCompletionClient._convert_messages(messages) + + assert len(result) == 4 + # First: TextMessage passed through + assert isinstance(result[0], TextMessage) + assert result[0].text == "hello" + # Second: ToolCallMessage → TextMessage with Thought + Action + assert isinstance(result[1], TextMessage) + assert "Thought: thinking" in result[1].text + assert "Action:" in result[1].text + # Third: ToolResultMessage → TextMessage with Observation + assert isinstance(result[2], TextMessage) + assert result[2].text == "Observation:\nresult text" + # Fourth: TextMessage passed through + assert isinstance(result[3], TextMessage) + assert result[3].text == "done" + + +def test_convert_empty_tool_calls_list(): + messages = [ToolCallMessage(tool_calls=[])] + result = ReActChatCompletionClient._convert_messages(messages) + + # Empty tool_calls with no reasoning_content → empty parts → no message + assert len(result) == 0 + + +def test_convert_preserves_non_tool_messages(): + messages = [ + TextMessage(role="user", text="question"), + ImageMessage(role="user", image="data:image/png;base64,abc"), + ] + result = ReActChatCompletionClient._convert_messages(messages) + + assert len(result) == 2 + assert isinstance(result[0], TextMessage) + assert result[0].text == "question" + assert isinstance(result[1], ImageMessage) + assert result[1].image == "data:image/png;base64,abc" + + +class _CapturingStubClient(ModelClient): + """Stub that records the messages it receives.""" + + def __init__(self): + self.received_messages = None + + @override + def get_model(self) -> str: + return "stub" + + @override + def get_api_key(self) -> str | None: + return None + + @override + async def create( + self, + messages, + *, + tools: list[Tool] | None = None, + extra_args: dict[str, Any] | None = None, + ) -> ModelResponse: + self.received_messages = messages + return ModelResponse( + reasoning_content=None, content="Final Answer: done", tool_calls=None, finish_reason="stop" + ) + + @override + async def create_stream(self, messages, *, tools=None, extra_args=None): + return + yield # make this an async generator # noqa: RET503 + + +def test_create_converts_messages_before_forwarding(): + stub = _CapturingStubClient() + client = ReActChatCompletionClient(stub) + + tool_messages = [ + TextMessage(role="user", text="help"), + ToolCallMessage( + reasoning_content="thinking", + tool_calls=[ToolCallRecord(id="tc1", name="lookup", arguments={"query": "x"})], + ), + ToolResultMessage(tool_call_id="tc1", content="observed"), + ] + + asyncio.run(client.create(tool_messages, tools=[])) + + # The stub should have received only TextMessage/ImageMessage + assert stub.received_messages is not None + for msg in stub.received_messages: + assert isinstance(msg, (TextMessage, ImageMessage)), f"Unexpected type: {type(msg)}" + + # Verify reconstructed content + # The ToolCallMessage should have been converted to a TextMessage with Action: + assert any("Action:" in m.text for m in stub.received_messages if isinstance(m, TextMessage)) + # The ToolResultMessage should have been converted to a TextMessage with Observation: + assert any("Observation:" in m.text for m in stub.received_messages if isinstance(m, TextMessage)) diff --git a/copilotj/util/prompt_template.py b/copilotj/util/prompt_template.py new file mode 100644 index 00000000..797fbc2d --- /dev/null +++ b/copilotj/util/prompt_template.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright contributors to the CopilotJ project. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal Jinja2-like prompt template renderer. + +Supports ``{% if variable == "value" %}`` / ``{% elif %}`` / ``{% else %}`` / +``{% endif %}`` conditional blocks with simple equality checks. No external +dependency — just regex and a small state machine. +""" + +from __future__ import annotations + +import re +from typing import Any + +__all__ = ["render_prompt"] + +# Matches: {% if var == "val" %}, {% elif var == "val" %}, {% else %}, {% endif %} +_TAG_RE = re.compile( + r"^\s*\{%" # opening delimiter {% + r"\s*(if|elif|else|endif)" # tag name + r"(?:\s+(\w+)\s*==\s*\"([^\"]*)\")?" # optional: variable == "value" + r"\s*%\}\s*$", # closing delimiter %} + re.MULTILINE, +) + +# Matches template variables: {variable_name} +_VAR_RE = re.compile(r"\{(\w+)\}") + + +def render_prompt(template: str, **variables: str) -> str: + """Render a prompt template with conditional sections. + + Supports: + {% if variable == "value" %} + {% elif variable == "other" %} + {% else %} + {% endif %} + + Conditional tags must appear on their own lines. Lines containing only a + conditional tag are removed from the output. ``{% else %}`` and + ``{% endif %}`` do not take a condition. + + Template variables ``{name}`` are substituted with *variables[name]*. + Missing variables are left as-is. + + Args: + template: The prompt template string. + **variables: Variable values for condition evaluation and substitution. + + Returns: + The rendered prompt string. + """ + # --- Phase 1: evaluate conditionals --- + lines = template.split("\n") + result_lines: list[str] = [] + + # Stack of (branch_taken, active, in_else). Tracks nesting. + # branch_taken: whether any branch in this if-block has been taken so far. + # active: whether lines in the *current* section should be emitted. + # in_else: whether we are past an {% else %} for this level. + stack: list[tuple[bool, bool, bool]] = [] + + for line in lines: + m = _TAG_RE.match(line.strip()) + if m is None: + # Regular line — emit only if all enclosing conditions are true. + if all(active for _, active, _ in stack): + result_lines.append(line) + continue + + tag, var_name, var_value = m.group(1), m.group(2), m.group(3) + + if tag == "if": + cond = _eval_condition(var_name, var_value, variables) + stack.append((cond, cond, False)) + + elif tag == "elif": + if not stack: + raise SyntaxError("{% elif %} without matching {% if %}") + branch_taken, _, in_else = stack[-1] + if in_else: + raise SyntaxError("{% elif %} after {% else %}") + # Only evaluate if no previous branch was taken. + new_cond = _eval_condition(var_name, var_value, variables) if not branch_taken else False + stack[-1] = (branch_taken or new_cond, new_cond, False) + + elif tag == "else": + if not stack: + raise SyntaxError("{% else %} without matching {% if %}") + branch_taken, _, _ = stack[-1] + stack[-1] = (branch_taken, not branch_taken, True) + + elif tag == "endif": + if not stack: + raise SyntaxError("{% endif %} without matching {% if %}") + stack.pop() + + if stack: + raise SyntaxError("Unclosed {% if %} block in prompt template") + + # --- Phase 2: substitute variables --- + text = "\n".join(result_lines) + + def _replace_var(m: re.Match[str]) -> str: + name = m.group(1) + if name in variables: + return str(variables[name]) + return m.group(0) # leave as-is if not provided + + text = _VAR_RE.sub(_replace_var, text) + + # Clean up blank lines left by removed conditional tags + text = _clean_blank_lines(text) + + return text + + +def _eval_condition(var_name: str | None, var_value: str | None, variables: dict[str, Any]) -> bool: + """Evaluate ``variable == "value"``.""" + if var_name is None or var_value is None: + raise SyntaxError("Missing condition in {% if/elif %} tag") + actual = variables.get(var_name) + return str(actual) == var_value + + +def _clean_blank_lines(text: str) -> str: + """Remove excessive blank lines (3+ consecutive → 2) caused by stripped tags.""" + while "\n\n\n" in text: + text = text.replace("\n\n\n", "\n\n") + return text diff --git a/copilotj/util/react_parser.py b/copilotj/util/react_parser.py index 84e5a74e..53656f93 100644 --- a/copilotj/util/react_parser.py +++ b/copilotj/util/react_parser.py @@ -22,6 +22,8 @@ TextMessage, Tool, ToolCall, + ToolCallMessage, + ToolResultMessage, ) __all__ = ["ReActChatCompletionClient", "ModelSyntaxError"] @@ -84,13 +86,14 @@ def get_model(self) -> str: @override async def create( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[TextMessage | ImageMessage | ToolCallMessage | ToolResultMessage], *, tools: list[Tool] | None = None, extra_args: dict[str, Any] | None = None, ) -> ModelResponse: # NOTE: dont send tools since we are parsing the response manually - response = await self._model_client.create(messages, extra_args=extra_args) + converted = self._convert_messages(messages) + response = await self._model_client.create(converted, extra_args=extra_args) if response.reasoning_content is not None or response.content is None: return response @@ -136,7 +139,7 @@ async def create( @override async def create_stream( self, - messages: Sequence[TextMessage | ImageMessage], + messages: Sequence[TextMessage | ImageMessage | ToolCallMessage | ToolResultMessage], *, tools: list[Tool] | None = None, extra_args: dict[str, Any] | None = None, @@ -144,7 +147,8 @@ async def create_stream( tools = tools or [] # NOTE: dont send tools since we are parsing the response manually - stream = self._model_client.create_stream(messages, tools=None, extra_args=extra_args) + converted = self._convert_messages(messages) + stream = self._model_client.create_stream(converted, tools=None, extra_args=extra_args) buffer = "" state = _StreamingState.Init @@ -235,6 +239,32 @@ async def create_stream( case _: raise ValueError(f"Unexpected last item type: {type(last)}") + @staticmethod + def _convert_messages( + messages: Sequence[TextMessage | ImageMessage | ToolCallMessage | ToolResultMessage], + ) -> list[TextMessage | ImageMessage]: + """Convert ToolCallMessage/ToolResultMessage back to TextMessage for ReAct mode. + + Reconstructs the ReAct text format so the underlying model client only + sees plain TextMessage and ImageMessage objects. + """ + result: list[TextMessage | ImageMessage] = [] + for msg in messages: + if isinstance(msg, ToolCallMessage): + parts: list[str] = [] + if msg.reasoning_content: + parts.append(f"Thought: {msg.reasoning_content.strip()}") + for tc in msg.tool_calls: + args_json = json.dumps(tc.arguments, ensure_ascii=False) + parts.append(f'Action: {{"name": "{tc.name}", "args": {args_json}}}') + if parts: + result.append(TextMessage(role="assistant", text="\n".join(parts))) + elif isinstance(msg, ToolResultMessage): + result.append(TextMessage(role="user", text=f"Observation:\n{msg.content}")) + else: + result.append(msg) + return result + def _fsm_init( self, buffer: str, *, tools: list[Tool] ) -> tuple[_StreamingState, Iterable[ModelResponseChunk | ToolCall], str]: