autogen/test/openai/test_notebook.py
Chi Wang 169012f3e7
ChatGPT support (#942)
* improve max_valid_n and doc

* Update README.md

Co-authored-by: Li Jiang <lijiang1@microsoft.com>

* add support for chatgpt

* notebook

* newline at end of file

* chatgpt notebook

* ChatGPT in Azure

* doc

* math

* warning, timeout, log file name

* handle import error

* doc update; default value

* paper

* doc

* docstr

* eval_func

* prompt and messages

* remove confusing words

* notebook name

---------

Co-authored-by: Li Jiang <lijiang1@microsoft.com>
Co-authored-by: Susan Xueqing Liu <liususan091219@users.noreply.github.com>
2023-03-10 19:35:36 +00:00

51 lines
1.3 KiB
Python

import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
from nbconvert.preprocessors import CellExecutionError
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_notebook.ipynb", save=False):
try:
file_path = os.path.join(here, os.pardir, os.pardir, "notebook", input_nb)
with open(file_path) as f:
nb = nbformat.read(f, as_version=4)
ep = ExecutePreprocessor(timeout=3600, kernel_name="python3")
ep.preprocess(nb, {"metadata": {"path": here}})
except CellExecutionError:
raise
finally:
if save:
with open(os.path.join(here, output_nb), "w", encoding="utf-8") as f:
nbformat.write(nb, f)
@pytest.mark.skipif(
skip,
reason="do not run openai test if openai is not installed",
)
def test_integrate_openai(save=False):
run_notebook("integrate_openai.ipynb", save=save)
@pytest.mark.skipif(
skip,
reason="do not run openai test if openai is not installed",
)
def test_integrate_chatgpt(save=False):
run_notebook("integrate_chatgpt_math.ipynb", save=save)
if __name__ == "__main__":
test_integrate_chatgpt(save=True)