import asyncio import json import re import uuid from dataclasses import dataclass from typing import Any, Dict, List, Tuple, Union from agnext.application import SingleThreadedAgentRuntime from agnext.components import FunctionCall, TypeRoutedAgent, message_handler from agnext.components.code_executor import ( CodeBlock, CodeExecutor, LocalCommandLineCodeExecutor, ) from agnext.components.models import ( AssistantMessage, AzureOpenAIChatCompletionClient, ChatCompletionClient, FunctionExecutionResult, FunctionExecutionResultMessage, LLMMessage, ModelCapabilities, OpenAIChatCompletionClient, SystemMessage, UserMessage, ) from agnext.components.tools import CodeExecutionResult, PythonCodeExecutionTool from agnext.core import AgentId, CancellationToken # from azure.identity import DefaultAzureCredential, get_bearer_token_provider @dataclass class TaskMessage: content: str @dataclass class CodeExecutionRequestMessage: session_id: str execution_request: str @dataclass class CodeExecutionResultMessage: session_id: str output: str exit_code: int class Coder(TypeRoutedAgent): """An agent that uses tools to write, execute, and debug Python code.""" DEFAULT_DESCRIPTION = "A Python coder assistant." DEFAULT_SYSTEM_MESSAGES = [ SystemMessage( """You are a helpful AI assistant. Solve tasks using your Python coding skills. The code you output must be formatted in Markdown code blocks demarcated by triple backticks (```). As an example: ```python def main(): print("Hello world.") if __name__ == "__main__": main() ``` The user cannot provide any feedback or perform any other action beyond executing the code you suggest. In particular, the user can't modify your code, and can't copy and paste anything, and can't fill in missing values. Thus, do not suggest incomplete code which requires users to perform any of these actions. Check the execution result returned by the user. If the result indicates there is an error, fix the error and output the code again. Suggest the full code instead of partial code or code changes -- code blocks must stand alone and be ready to execute without modification. If the error can't be fixed or if the task is not solved even after the code is executed successfully, analyze the problem, revisit your assumption, and think of a different approach to try. If the code has executed successfully, and the problem is stolved, reply "TERMINATE". """ ) ] def __init__( self, model_client: ChatCompletionClient, description: str = DEFAULT_DESCRIPTION, system_messages: List[SystemMessage] = DEFAULT_SYSTEM_MESSAGES, max_turns: int | None = None, ) -> None: super().__init__(description) self._model_client = model_client self._system_messages = system_messages self._session_memory: Dict[str, List[LLMMessage]] = {} self._max_turns = max_turns @message_handler async def handle_user_message( self, message: TaskMessage, cancellation_token: CancellationToken ) -> None: """Handle a user message, execute the model and tools, and returns the response.""" # Create a new session. session_id = str(uuid.uuid4()) self._session_memory.setdefault(session_id, []).append( UserMessage(content=message.content, source="user") ) # Make an inference to the model. response = await self._model_client.create( self._system_messages + self._session_memory[session_id] ) assert isinstance(response.content, str) self._session_memory[session_id].append( AssistantMessage(content=response.content, source=self.metadata["name"]) ) await self.publish_message( CodeExecutionRequestMessage( execution_request=response.content, session_id=session_id ), cancellation_token=cancellation_token, ) @message_handler async def handle_code_execution_result( self, message: CodeExecutionResultMessage, cancellation_token: CancellationToken ) -> None: execution_result = f"The script ran, then exited with Unix exit code: {message.exit_code}\nIts output was:\n{message.output}" # Store the code execution output. self._session_memory[message.session_id].append( UserMessage(content=execution_result, source="user") ) # Count the number of rounds so far if self._max_turns is not None: n_turns = sum( 1 for message in self._session_memory[message.session_id] if isinstance(message, AssistantMessage) ) if n_turns >= self._max_turns: return # Make an inference to the model. response = await self._model_client.create( self._system_messages + self._session_memory[message.session_id] ) assert isinstance(response.content, str) self._session_memory[message.session_id].append( AssistantMessage(content=response.content, source=self.metadata["name"]) ) if "TERMINATE" in response.content: return else: await self.publish_message( CodeExecutionRequestMessage( execution_request=response.content, session_id=message.session_id ), cancellation_token=cancellation_token, ) class Executor(TypeRoutedAgent): def __init__(self, description: str, executor: CodeExecutor) -> None: super().__init__(description) self._executor = executor @message_handler async def handle_code_execution( self, message: CodeExecutionRequestMessage, cancellation_token: CancellationToken, ) -> None: # Extract code block from the message. code = self._extract_execution_request(message.execution_request) if code is not None: execution_requests = [CodeBlock(code=code, language="python")] future = asyncio.get_event_loop().run_in_executor( None, self._executor.execute_code_blocks, execution_requests ) cancellation_token.link_future(future) result = await future await self.publish_message( CodeExecutionResultMessage( output=result.output, exit_code=result.exit_code, session_id=message.session_id, ) ) else: await self.publish_message( CodeExecutionResultMessage( output="No code block detected. Please provide a markdown-encoded code block to execute.", exit_code=1, session_id=message.session_id, ) ) def _extract_execution_request(self, markdown_text: str) -> Union[str, None]: pattern = r"```(\w+)\n(.*?)\n```" # Search for the pattern in the markdown text match = re.search(pattern, markdown_text, re.DOTALL) # Extract the language and code block if a match is found if match: return match.group(2) return None async def main() -> None: # Create the runtime. runtime = SingleThreadedAgentRuntime() # Create the AzureOpenAI client, with AAD auth # token_provider = get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default") client = AzureOpenAIChatCompletionClient( api_version="2024-02-15-preview", azure_endpoint="https://aif-complex-tasks-west-us-3.openai.azure.com/", model="gpt-4o-2024-05-13", model_capabilities=ModelCapabilities( function_calling=True, json_output=True, vision=True ), # azure_ad_token_provider=token_provider ) # Register agents. coder = runtime.register_and_get( "Coder", lambda: Coder(model_client=client), ) runtime.register( "Executor", lambda: Executor( "A agent for executing code", executor=LocalCommandLineCodeExecutor() ), ) prompt = "" with open("prompt.txt", "rt") as fh: prompt = fh.read() entry_point = "__ENTRY_POINT__" task = TaskMessage( f""" The following python code imports the `run_tests` function from unit_tests.py, and runs it on the function `{entry_point}`. This will run a set of automated unit tests to verify the correct implementation of `{entry_point}`. However, `{entry_point}` is only partially implemented in the code below. Complete the implementation of `{entry_point}` and then execute a new stand-alone code block that contains everything needed to run the tests, including: importing `unit_tests`, calling `run_tests({entry_point})`, as well as {entry_point}'s complete definition, such that this code block can be run directly in Python. ```python from unit_tests import run_tests {prompt} # Run the unit tests run_tests({entry_point}) ``` """.strip() ) # Run the runtime until the task is completed. run_context = runtime.start() # Send a task to the tool user. await runtime.send_message(task, coder) await run_context.stop_when_idle() if __name__ == "__main__": import logging logging.basicConfig(level=logging.WARNING) logging.getLogger("agnext").setLevel(logging.DEBUG) asyncio.run(main())