-
Notifications
You must be signed in to change notification settings - Fork 0
/
code_exec_jupyter_llm_stocks.py
166 lines (136 loc) · 7.49 KB
/
code_exec_jupyter_llm_stocks.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
# Testing the Code Executor - Local Jupyter - Local LLM - Stock Prices
# Based on: https://github.com/microsoft/autogen/blob/tutorial/website/docs/getting-started/code-executors.ipynb
# Installed: pip install -qqq 'pyautogen[local-jupyter-exec]'
import os
import tempfile
import shutil
from autogen import ConversableAgent, UserProxyAgent, AssistantAgent
import datetime
import sys # Redirecting standard output to a file instead
# Duplicate output to file and screen
class Tee:
def __init__(self, file_name, mode='w'):
self.file = open(file_name, mode)
self.stdout = sys.stdout
def __enter__(self):
sys.stdout = self
def __exit__(self, exc_type, exc_value, traceback):
self.close()
def write(self, data):
self.file.write(data)
self.stdout.write(data)
def flush(self):
self.file.flush()
def close(self):
if self.file:
self.file.close()
sys.stdout = self.stdout
# The code writer agent can be powered by an LLM such as GPT-4 with code-writing capability.
# And the code executor agent is powered by a code executor.
# The following is an agent with a code writer role specified using system_message.
# The code writer agent's system message is to instruct the LLM on how to
# use the Jupyter code executor with IPython kernel.
code_writer_system_message = (
"You have been given coding capability to solve tasks using Python code in a stateful IPython kernel.\n"
"You are responsible for writing the code, and the user is responsible for executing the code.\n"
"\n"
"When you write Python code, put the code in a markdown code block with the language set to Python.\n"
"For example:\n"
"```python\n"
"x = 3\n"
"```\n"
"You can use the variable `x` in subsequent code blocks.\n"
"```python\n"
"print(x)\n"
"```\n"
"\n"
"Write code incrementally and leverage the statefulness of the kernel to avoid repeating code.\n"
"Import libraries in a separate code block.\n"
"Pip install any required libraries.\n" # MS ADDED
"Define a function or a class in a separate code block.\n"
"Run code that produces output in a separate code block.\n"
"Run code that involves expensive operations like download, upload, and call external APIs in a separate code block.\n"
"\n"
"When your code produces an output, the output will be returned to you.\n"
"Because you have limited conversation memory, if your code creates an image,\n"
"the output will be a path to the image instead of the image itself.\n"
)
ollama_models = [
{"model_name": "codellama:7b-python", "display_name": "CodeLlama_7b_Python"},
{"model_name": "codellama:13b-python", "display_name": "CodeLlama_13b_Python"},
{"model_name": "codellama:34b-instruct", "display_name": "CodeLlama_34b_Instruct"},
{"model_name": "codellama:34b-python", "display_name": "CodeLlama_34b_Python"},
{"model_name": "deepseek-coder:6.7b-instruct-q6_K", "display_name": "DeepSeek_Coder"},
{"model_name": "llama2:13b-chat", "display_name" : "Llama2_13b_Chat"},
{"model_name": "llama2:7b-chat-q6_K", "display_name": "Llama2_7b_Chat"},
{"model_name": "mistral:7b-instruct-v0.2-q6_K", "display_name" : "Mistral_7b_Instruct_v2"},
{"model_name": "mixtralq4", "display_name" : "Mixtral_8x7b_Q4"},
{"model_name": "mixtralq5", "display_name" : "Mixtral_8x7b_Q5"},
{"model_name": "neural-chat:7b-v3.3-q6_K", "display_name" : "Neural_Chat_7b"},
{"model_name": "nexusraven", "display_name" : "Nexus_Raven"},
{"model_name": "openhermes:7b-mistral-v2.5-q6_K", "display_name" : "OpenHermes_7b_Mistral_v25"},
{"model_name": "orca2:13b-q5_K_S", "display_name" : "Orca2_13b"},
{"model_name": "phi", "display_name" : "Phi"},
{"model_name": "phind-codellama:34b-v2", "display_name" : "Phind_CodeLlama_34b_v2"},
{"model_name": "qwen:14b-chat-q6_K", "display_name" : "Qwen_14b_Chat"},
{"model_name": "solar:10.7b-instruct-v1-q5_K_M", "display_name" : "Solar_107b_Instruct"},
# {"model_name": "starcoder2:3b", "display_name" : "StarCoder2_3b"},
# {"model_name": "starcoder2:7b", "display_name" : "StarCoder2_7b"},
# {"model_name": "starcoder2:15b", "display_name" : "StarCoder2_15b"},
{"model_name": "yi:34b-chat-q4_K_M", "display_name" : "Yi_34b_Chat_Q4"},
]
test_prefix = "Stocks"
for ollama_model in ollama_models:
model_name = ollama_model["model_name"]
display_name = ollama_model["display_name"]
# Two iterations per model.
for iteration in range(1, 3):
output_file = f"/home/autogen/autogen/ms_working/results/{test_prefix}_{display_name}_i{iteration}.txt"
if not os.path.exists(output_file):
# Clear the terminal (Unix/Linux/MacOS)
os.system('clear')
with Tee(output_file, 'w'):
# MS - LLM will be the one running on the host machine (Ollama)
code_writer_agent = ConversableAgent(
"code_writer",
system_message=code_writer_system_message,
llm_config={
"config_list": [{"model": model_name, "api_key": "NotRequired", "base_url": "http://192.168.0.115:11434/v1"}],
"cache_seed": None,
}, ## CRITICAL - ENSURE THERE'S NO CACHING FOR TESTING
code_execution_config=False, # Turn off code execution for this agent.
)
# Create a temporary directory to store the output images.
temp_dir = tempfile.mkdtemp()
# Create an agent with code executor configuration for a local Jupyter
# code executor.
code_executor_agent = ConversableAgent(
name="code_executor_agent",
llm_config=False,
code_execution_config={
"executor": "jupyter-local", # Use the local Jupyter executor.
# On Windows use "ipython-embedded" instead of "jupyter-local"
# due to a known issue with Jupyter Kernel Gateway on Windows.
"jupyter-local": {
"timeout": 60, # Timeout for each code execution in seconds.
"kernel_name": "python3", # Use the Python 3 kernel, which is the default.
"output_dir": temp_dir, # Use the temporary directory to store output images.
},
},
human_input_mode="ALWAYS",
)
# Now we can try a more complex example that involves querying the web. Let's say we want to get the the stock
# prices year-to-date for Tesla and Meta (formerly Facebook). We can also use the two agents with several iterations
# of conversation.
today = datetime.datetime.now().strftime("%Y-%m-%d")
print(f"-----\n\n{display_name}\n\n{today}\n\nIteration {iteration}\n\n-----\n")
chat_result = code_executor_agent.initiate_chat(
code_writer_agent, message=f"Today is {today}. Write Python code to plot TSLA's and META's stock prices YTD."
)
# Clean up the temporary directory and restart the Jupyter server to free up the kernel resources.
# Clean up for next round
shutil.rmtree(temp_dir)
code_executor_agent.code_executor.restart()
else:
print(f"{output_file} already exists, ignoring.")
# break