forked from microsoft/autogen
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest_notebook.py
92 lines (72 loc) · 2.99 KB
/
test_notebook.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
import sys
import os
import pytest
try:
import openai
skip = False
except ImportError:
skip = True
here = os.path.abspath(os.path.dirname(__file__))
def run_notebook(input_nb, output_nb="executed_openai_notebook.ipynb", save=False):
import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
from nbconvert.preprocessors import CellExecutionError
try:
nb_loc = os.path.join(here, os.pardir, "notebook")
file_path = os.path.join(nb_loc, input_nb)
with open(file_path) as nb_file:
nb = nbformat.read(nb_file, as_version=4)
preprocessor = ExecutePreprocessor(timeout=4800, kernel_name="python3")
preprocessor.preprocess(nb, {"metadata": {"path": nb_loc}})
output_file_name = "executed_openai_notebook_output.txt"
output_file = os.path.join(here, output_file_name)
with open(output_file, "a") as nb_output_file:
for cell in nb.cells:
if cell.cell_type == "code" and "outputs" in cell:
for output in cell.outputs:
if "text" in output:
nb_output_file.write(output["text"].strip() + "\n")
elif "data" in output and "text/plain" in output["data"]:
nb_output_file.write(output["data"]["text/plain"].strip() + "\n")
except CellExecutionError:
raise
finally:
if save:
with open(os.path.join(here, output_nb), "w", encoding="utf-8") as nb_executed_file:
nbformat.write(nb, nb_executed_file)
@pytest.mark.skipif(
skip or not sys.version.startswith("3.10"),
reason="do not run if openai is not installed or py!=3.10",
)
def test_agentchat_auto_feedback_from_code(save=False):
run_notebook("agentchat_auto_feedback_from_code_execution.ipynb", save=save)
@pytest.mark.skipif(
skip or not sys.version.startswith("3.10"),
reason="do not run if openai is not installed or py!=3.10",
)
def test_oai_completion(save=False):
run_notebook("oai_completion.ipynb", save=save)
@pytest.mark.skipif(
skip or not sys.version.startswith("3.10"),
reason="do not run if openai is not installed or py!=3.10",
)
def test_agentchat_function_call(save=False):
run_notebook("agentchat_function_call.ipynb", save=save)
@pytest.mark.skipif(
skip or not sys.version.startswith("3.10"),
reason="do not run if openai is not installed or py!=3.10",
)
def test_agentchat_MathChat(save=False):
run_notebook("agentchat_MathChat.ipynb", save=save)
@pytest.mark.skipif(
skip or not sys.version.startswith("3.11"),
reason="do not run if openai is not installed or py!=3.11",
)
def test_oai_chatgpt_gpt4(save=False):
run_notebook("oai_chatgpt_gpt4.ipynb", save=save)
if __name__ == "__main__":
test_agentchat_auto_feedback_from_code(save=True)
# test_oai_chatgpt_gpt4(save=True)
# test_oai_completion(save=True)
# test_agentchat_MathChat(save=True)
# test_agentchat_function_call(save=True)