suna/agent.py

89 lines
2.6 KiB
Python
Raw Normal View History

2024-10-17 19:44:10 +08:00
import asyncio
2024-10-28 05:04:42 +08:00
import json
2024-10-17 19:44:10 +08:00
from agentpress.thread_manager import ThreadManager
from tools.files_tool import FilesTool
2024-10-28 05:04:42 +08:00
from state_manager import StateManager
2024-10-10 22:21:39 +08:00
2024-10-23 09:42:38 +08:00
async def run_agent(
thread_manager: ThreadManager,
thread_id: int,
2024-10-28 05:04:42 +08:00
state_manager: StateManager,
2024-10-23 09:42:38 +08:00
max_iterations: int = 10
):
2024-10-23 10:34:51 +08:00
2024-10-23 09:42:38 +08:00
async def init():
pass
async def pre_iteration():
pass
async def after_iteration():
await thread_manager.add_message(thread_id, {"role": "user", "content": "CREATE MORE RANDOM FILES WITH RANDOM CONTENTS. JSUT CREATE IT NO QUESTINS PLEASE.'"})
pass
async def finalizer():
pass
await init()
iteration = 0
while iteration < max_iterations:
iteration += 1
await pre_iteration()
2024-10-28 05:04:42 +08:00
# Get entire state store
state = await state_manager.export_store()
state_info = f"Current state store:\n{json.dumps(state, indent=2)}"
system_message = {
"role": "system",
"content": f"You are a helpful assistant that can create, read, update, and delete files.\n\n{state_info}"
}
2024-10-23 09:42:38 +08:00
model_name = "gpt-4o"
2024-10-28 05:04:42 +08:00
# Include entire state in additional message
# additional_message = {
# "role": "user",
# "content": state_info
# }
2024-10-23 09:42:38 +08:00
response = await thread_manager.run_thread(
thread_id=thread_id,
system_message=system_message,
model_name=model_name,
temperature=0.7,
max_tokens=150,
tool_choice="auto",
2024-10-28 05:04:42 +08:00
# additional_message=additional_message,
2024-10-23 09:42:38 +08:00
execute_tools_async=False,
2024-10-23 10:43:57 +08:00
use_tools=True,
2024-10-23 09:42:38 +08:00
execute_model_tool_calls=True
)
await after_iteration()
await finalizer()
2024-10-17 19:44:10 +08:00
2024-10-17 04:08:46 +08:00
if __name__ == "__main__":
2024-10-23 09:42:38 +08:00
async def main():
thread_manager = ThreadManager()
2024-10-28 05:04:42 +08:00
state_manager = StateManager("state.json")
2024-10-23 09:42:38 +08:00
thread_id = await thread_manager.create_thread()
2024-10-28 05:04:42 +08:00
# Read current file contents
files = await state_manager.get("files")
2024-10-23 09:42:38 +08:00
await thread_manager.add_message(thread_id, {"role": "user", "content": "Please create a file with a random name with the content 'Hello, world!'"})
thread_manager.add_tool(FilesTool)
await run_agent(
thread_manager=thread_manager,
thread_id=thread_id,
2024-10-28 05:04:42 +08:00
state_manager=state_manager,
2024-10-23 09:42:38 +08:00
max_iterations=5
)
asyncio.run(main())