From 08964f26580ac27e92d8a4f10042ff1131e3c76e Mon Sep 17 00:00:00 2001 From: imotai Date: Tue, 3 Oct 2023 12:32:01 +0800 Subject: [PATCH 1/5] test: add test for up and kernel --- .../src/og_kernel/server/kernel_rpc_server.py | 17 ++++- requirements.txt | 2 + sdk/src/og_sdk/agent_sdk.py | 29 +------- sdk/src/og_sdk/utils.py | 31 ++++++++ sdk/tests/kernel_sdk_tests.py | 31 ++++++++ up/tests/up_tests.py | 73 +++++++++++++++++++ 6 files changed, 153 insertions(+), 30 deletions(-) diff --git a/kernel/src/og_kernel/server/kernel_rpc_server.py b/kernel/src/og_kernel/server/kernel_rpc_server.py index d955ee7..e3aa5c0 100644 --- a/kernel/src/og_kernel/server/kernel_rpc_server.py +++ b/kernel/src/og_kernel/server/kernel_rpc_server.py @@ -102,8 +102,13 @@ def __init__(self): message="api key is required", details=[], ) - os.makedirs(config["config_root_path"], exist_ok=True) + os.makedirs(config["workspace"], exist_ok=True) + config_root_path = config["config_root_path"] + workspace = config["workspace"] + logger.info( + f"start kernel rpc with config root path {config_root_path} and workspace {workspace}" + ) async def stop( self, request: kernel_server_pb2.StopKernelRequest, context: ServicerContext @@ -142,9 +147,11 @@ async def start( """ kernel_name = request.kernel_name if request.kernel_name else "python3" if kernel_name in self.kms and self.kms[kernel_name]: - logger.warning("the kernel has been started") + logger.warning( + "the request will be ignored for that the kernel has been started" + ) return kernel_server_pb2.StartKernelResponse( - code=1, msg="the kernel has been started" + code=0, msg="the kernel has been started" ) logging.info("create a new kernel with kernel_name %s" % kernel_name) connection_file = "%s/kernel-%s.json" % ( @@ -165,7 +172,9 @@ async def download( """ download file """ - target_filename = "%s/%s" % (config["workspace"], request.filename) + filename = request.filename + workspace = config["workspace"] + target_filename = f"{workspace}{os.sep}{filename}" if not await aio_os.path.exists(target_filename): await context.abort(10, "%s filename do not exist" % request.filename) async with aiofiles.open(target_filename, "rb") as afp: diff --git a/requirements.txt b/requirements.txt index d7c6721..837aa22 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,3 +18,5 @@ click==8.1.7 discord.py==2.3.2 openai==0.28.0 build==1.0.3 +python-dotenv==1.0.0 + diff --git a/sdk/src/og_sdk/agent_sdk.py b/sdk/src/og_sdk/agent_sdk.py index 3f06a90..bd18b26 100644 --- a/sdk/src/og_sdk/agent_sdk.py +++ b/sdk/src/og_sdk/agent_sdk.py @@ -11,6 +11,7 @@ from og_proto.agent_server_pb2_grpc import AgentServerStub import aiofiles from typing import AsyncIterable +from .utils import generate_chunk, generate_async_chunk logger = logging.getLogger(__name__) @@ -84,19 +85,9 @@ def upload_file(self, filepath, filename): """ # TODO limit the file size - def generate_trunk(filepath, filename) -> common_pb2.FileChunk: - try: - with open(filepath, "rb") as fp: - while True: - chunk = fp.read(1024 * 128) - if not chunk: - break - yield common_pb2.FileChunk(buffer=chunk, filename=filename) - except Exception as ex: - logger.error("fail to read file %s" % ex) return self.stub.upload( - generate_trunk(filepath, filename), metadata=self.metadata + generate_chunk(filepath, filename), metadata=self.metadata ) def prompt(self, prompt, files=[]): @@ -189,22 +180,8 @@ async def upload_file(self, filepath, filename): """ upload file to agent """ - # TODO limit the file size - async def generate_trunk( - filepath, filename - ) -> AsyncIterable[common_pb2.FileChunk]: - try: - async with aiofiles.open(filepath, "rb") as afp: - while True: - chunk = await afp.read(1024 * 128) - if not chunk: - break - yield common_pb2.FileChunk(buffer=chunk, filename=filename) - except Exception as ex: - logger.error("fail to read file %s", ex) - - return await self.upload_binary(generate_trunk(filepath, filename)) + return await self.upload_binary(generate_async_chunk(filepath, filename)) def close(self): if self.channel: diff --git a/sdk/src/og_sdk/utils.py b/sdk/src/og_sdk/utils.py index cd295d4..a6757b0 100644 --- a/sdk/src/og_sdk/utils.py +++ b/sdk/src/og_sdk/utils.py @@ -5,6 +5,37 @@ import re import string import random +import aiofiles +from og_proto import agent_server_pb2, common_pb2 +from typing import AsyncIterable + +logger = logging.getLogger(__name__) + + +def generate_chunk(filepath, filename) -> common_pb2.FileChunk: + try: + with open(filepath, "rb") as fp: + while True: + chunk = fp.read(1024 * 128) + if not chunk: + break + yield common_pb2.FileChunk(buffer=chunk, filename=filename) + except Exception as ex: + logger.error("fail to read file %s" % ex) + + +async def generate_async_chunk( + filepath, filename +) -> AsyncIterable[common_pb2.FileChunk]: + try: + async with aiofiles.open(filepath, "rb") as afp: + while True: + chunk = await afp.read(1024 * 128) + if not chunk: + break + yield common_pb2.FileChunk(buffer=chunk, filename=filename) + except Exception as ex: + logger.error("fail to read file %s", ex) def process_char_stream(stream): diff --git a/sdk/tests/kernel_sdk_tests.py b/sdk/tests/kernel_sdk_tests.py index 77b5d0b..5688259 100644 --- a/sdk/tests/kernel_sdk_tests.py +++ b/sdk/tests/kernel_sdk_tests.py @@ -15,12 +15,16 @@ # limitations under the License. """ """ +import os import asyncio import pytest import logging import json from og_sdk.kernel_sdk import KernelSDK +from og_sdk.utils import generate_async_chunk from og_proto.kernel_server_pb2 import ExecuteResponse +import aiofiles +from typing import AsyncIterable logger = logging.getLogger(__name__) @@ -52,6 +56,33 @@ async def test_bad_sdk(bad_kernel_sdk): assert True +@pytest.mark.asyncio +async def test_upload_and_download_smoke_test(kernel_sdk): + path = os.path.abspath(__file__) + response = await kernel_sdk.upload_binary( + generate_async_chunk(path, "kernel_sdk_tests.py") + ) + assert response + file_stats = os.stat(path) + assert response.length == file_stats.st_size, "bad upload file size" + length = 0 + async for chunk in kernel_sdk.download("kernel_sdk_tests.py"): + length += len(chunk.buffer) + assert length == file_stats.st_size, "bad upload file size" + + +@pytest.mark.asyncio +async def test_stop_kernel(kernel_sdk): + kernel_sdk.connect() + assert kernel_sdk.stub is not None # Check that stub is initialized + if not await kernel_sdk.is_alive(): + await kernel_sdk.start() + assert await kernel_sdk.is_alive() + response = await kernel_sdk.stop() + assert response.code == 0 + assert not await kernel_sdk.is_alive() + + @pytest.mark.asyncio async def test_sdk_smoke_test(kernel_sdk): kernel_sdk.connect() diff --git a/up/tests/up_tests.py b/up/tests/up_tests.py index 5dd8971..7f31be6 100644 --- a/up/tests/up_tests.py +++ b/up/tests/up_tests.py @@ -11,6 +11,7 @@ import sys import pytest import tempfile +import logging from rich.live import Live from rich.console import Console from og_up.up import run_with_realtime_print @@ -19,7 +20,79 @@ from og_up.up import get_latest_release_version from og_up.up import start_octogen_for_codellama from og_up.up import random_str +from og_up.up import generate_agent_common, generate_agent_azure_openai, generate_agent_openai, generate_agent_codellama +from og_up.up import generate_kernel_env from rich.console import Group +from dotenv import dotenv_values + +logger = logging.getLogger(__name__) + + +def test_generate_kernel_env(): + console = Console() + segments = [] + with Live(Group(*segments), console=console) as live: + temp_dir = tempfile.mkdtemp(prefix="octogen") + rpc_key = "rpc_key" + generate_kernel_env(live, segments, temp_dir, rpc_key) + fullpath = f"{temp_dir}/kernel/.env" + config = dotenv_values(fullpath) + assert config["rpc_key"] == rpc_key, "bad rpc key" + assert config["rpc_port"] == "9527", "bad rpc port" + + +def test_generate_agent_codellama(): + console = Console() + segments = [] + with Live(Group(*segments), console=console) as live: + temp_dir = tempfile.mkdtemp(prefix="octogen") + admin_key = "admin_key" + generate_agent_codellama(live, segments, temp_dir, admin_key) + fullpath = f"{temp_dir}/agent/.env" + config = dotenv_values(fullpath) + assert config["llm_key"] == "codellama", "bad llm key" + assert ( + config["llama_api_base"] == "http://127.0.0.1:8080" + ), "bad codellama server endpoint" + assert config["admin_key"] == admin_key, "bad admin key" + + +def test_generate_agent_env_openai(): + console = Console() + segments = [] + with Live(Group(*segments), console=console) as live: + temp_dir = tempfile.mkdtemp(prefix="octogen") + admin_key = "admin_key" + openai_key = "openai_key" + model = "gpt-4-0613" + generate_agent_openai(live, segments, temp_dir, admin_key, openai_key, model) + fullpath = f"{temp_dir}/agent/.env" + config = dotenv_values(fullpath) + assert config["llm_key"] == "openai", "bad llm key" + assert config["openai_api_key"] == openai_key, "bad api key" + assert config["openai_api_model"] == model, "bad model" + assert config["admin_key"] == admin_key, "bad admin key" + + +def test_generate_agent_env_azure_openai(): + console = Console() + segments = [] + with Live(Group(*segments), console=console) as live: + temp_dir = tempfile.mkdtemp(prefix="octogen") + admin_key = "admin_key" + openai_key = "openai_key" + deployment = "octogen" + api_base = "azure" + generate_agent_azure_openai( + live, segments, temp_dir, admin_key, openai_key, deployment, api_base + ) + fullpath = f"{temp_dir}/agent/.env" + config = dotenv_values(fullpath) + assert config["llm_key"] == "azure_openai", "bad llm key" + assert config["openai_api_base"] == api_base, "bad api base" + assert config["openai_api_key"] == openai_key, "bad api key" + assert config["openai_api_deployment"] == deployment, "bad deployment" + assert config["admin_key"] == admin_key, "bad admin key" def test_run_print(): From 142132eb3bc7fd0f77ecfbc055de9d782f6327f8 Mon Sep 17 00:00:00 2001 From: imotai Date: Tue, 3 Oct 2023 12:41:32 +0800 Subject: [PATCH 2/5] fix: import logging for utils --- sdk/src/og_sdk/utils.py | 1 + sdk/tests/agent_sdk_tests.py | 1 + 2 files changed, 2 insertions(+) diff --git a/sdk/src/og_sdk/utils.py b/sdk/src/og_sdk/utils.py index a6757b0..6489274 100644 --- a/sdk/src/og_sdk/utils.py +++ b/sdk/src/og_sdk/utils.py @@ -6,6 +6,7 @@ import string import random import aiofiles +import logging from og_proto import agent_server_pb2, common_pb2 from typing import AsyncIterable diff --git a/sdk/tests/agent_sdk_tests.py b/sdk/tests/agent_sdk_tests.py index a898f1c..8e98629 100644 --- a/sdk/tests/agent_sdk_tests.py +++ b/sdk/tests/agent_sdk_tests.py @@ -22,6 +22,7 @@ import logging import json import random +import logging from tempfile import gettempdir from pathlib import Path from og_sdk.agent_sdk import AgentSDK From 84acde4b5cbc952c89b5da453bbacbc5ea711246 Mon Sep 17 00:00:00 2001 From: imotai Date: Tue, 3 Oct 2023 12:48:35 +0800 Subject: [PATCH 3/5] fix: fix the outofindex bug --- chat/src/og_terminal/terminal_chat.py | 2 +- sdk/tests/kernel_sdk_tests.py | 1 + start_sandbox.sh | 1 + 3 files changed, 3 insertions(+), 1 deletion(-) diff --git a/chat/src/og_terminal/terminal_chat.py b/chat/src/og_terminal/terminal_chat.py index 2fe2a93..7b790bd 100644 --- a/chat/src/og_terminal/terminal_chat.py +++ b/chat/src/og_terminal/terminal_chat.py @@ -82,7 +82,7 @@ def show_help(console): def gen_a_random_emoji(): - index = random.randint(0, len(EMOJI_KEYS)) + index = random.randint(0, len(EMOJI_KEYS) - 1) return EMOJI[EMOJI_KEYS[index]] diff --git a/sdk/tests/kernel_sdk_tests.py b/sdk/tests/kernel_sdk_tests.py index 5688259..732cd92 100644 --- a/sdk/tests/kernel_sdk_tests.py +++ b/sdk/tests/kernel_sdk_tests.py @@ -58,6 +58,7 @@ async def test_bad_sdk(bad_kernel_sdk): @pytest.mark.asyncio async def test_upload_and_download_smoke_test(kernel_sdk): + kernel_sdk.connect() path = os.path.abspath(__file__) response = await kernel_sdk.upload_binary( generate_async_chunk(path, "kernel_sdk_tests.py") diff --git a/start_sandbox.sh b/start_sandbox.sh index d95a8dc..ff9d65a 100644 --- a/start_sandbox.sh +++ b/start_sandbox.sh @@ -50,3 +50,4 @@ cat < ~/.octogen/config endpoint=127.0.0.1:9528 api_key=${KERNEL_RPC_KEY} EOF +og_ping From 178b90dd64c84529539adc6675b1c2be0c8da263 Mon Sep 17 00:00:00 2001 From: imotai Date: Tue, 3 Oct 2023 13:17:42 +0800 Subject: [PATCH 4/5] fix: fix the method name error --- sdk/tests/kernel_sdk_tests.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/tests/kernel_sdk_tests.py b/sdk/tests/kernel_sdk_tests.py index 732cd92..d8f54d0 100644 --- a/sdk/tests/kernel_sdk_tests.py +++ b/sdk/tests/kernel_sdk_tests.py @@ -67,7 +67,7 @@ async def test_upload_and_download_smoke_test(kernel_sdk): file_stats = os.stat(path) assert response.length == file_stats.st_size, "bad upload file size" length = 0 - async for chunk in kernel_sdk.download("kernel_sdk_tests.py"): + async for chunk in kernel_sdk.download_file("kernel_sdk_tests.py"): length += len(chunk.buffer) assert length == file_stats.st_size, "bad upload file size" From 8faff7539b58f0f9ecbdfbd59db659808ea9d128 Mon Sep 17 00:00:00 2001 From: imotai Date: Tue, 3 Oct 2023 13:27:36 +0800 Subject: [PATCH 5/5] fix: add strip --- chat/tests/test_chat_function.py | 2 ++ up/src/og_up/up.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/chat/tests/test_chat_function.py b/chat/tests/test_chat_function.py index 75dd557..07aaa69 100644 --- a/chat/tests/test_chat_function.py +++ b/chat/tests/test_chat_function.py @@ -20,3 +20,5 @@ def test_parse_number(): numbers = parse_numbers(test_text) assert numbers assert numbers[0] == "0" + + diff --git a/up/src/og_up/up.py b/up/src/og_up/up.py index 01703a7..c7618e7 100644 --- a/up/src/og_up/up.py +++ b/up/src/og_up/up.py @@ -127,7 +127,7 @@ def get_latest_release_version(repo_name, live, segments): refresh(live, segments) r = requests.get(f"https://api.github.com/repos/{repo_name}/releases/latest") old_segment = segments.pop() - version = r.json().get("name", "") + version = r.json().get("name", "").strip() if not version: segments.append(("❌", "Get octogen latest version failed", version)) else: