import os
from pathlib import Path
from autogen import AssistantAgent, UserProxyAgent
from autogen.coding import LocalCommandLineCodeExecutor
config_list = [
{
# Let's choose the Mixtral 8x7B model
"model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
# Provide your Together.AI API key here or put it into the TOGETHER_API_KEY environment variable.
"api_key": os.environ.get("TOGETHER_API_KEY"),
# We specify the API Type as 'together' so it uses the Together.AI client class
"api_type": "together",
"stream": False,
}
]
# Setting up the code executor
workdir = Path("coding")
workdir.mkdir(exist_ok=True)
code_executor = LocalCommandLineCodeExecutor(work_dir=workdir)
# Setting up the agents
# The UserProxyAgent will execute the code that the AssistantAgent provides
user_proxy_agent = UserProxyAgent(
name="User",
code_execution_config={"executor": code_executor},
is_termination_msg=lambda msg: "FINISH" in msg.get("content"),
)
system_message = """You are a helpful AI assistant who writes code and the user executes it.
Solve tasks using your coding and language skills.
"""
# The AssistantAgent, using Together.AI's Code Llama model, will take the coding request and return code
assistant_agent = AssistantAgent(
name="Together Assistant",
system_message=system_message,
llm_config={"config_list": config_list},
)
# Start the chat, with the UserProxyAgent asking the AssistantAgent the message
chat_result = user_proxy_agent.initiate_chat(
assistant_agent,
message="Provide code to count the number of prime numbers from 1 to 10000.",
)