diff --git a/README.md b/README.md
index fb40a61..0e60418 100644
--- a/README.md
+++ b/README.md
@@ -15,13 +15,11 @@
https://github.com/dbpunk-labs/octogen/assets/8623385/7445cc4d-567e-4d1a-bedc-b5b566329c41
-
|Supported OSs|Supported Interpreters|Supported Dev Enviroment|
|----|-----|-----|
|
|
|
|
-
## Getting Started
Requirement
@@ -100,7 +98,6 @@ Use /help for help
## Supported API Service
-
|name|type|status| installation|
|----|-----|----------------|---|
|[Openai GPT 3.5/4](https://openai.com/product#made-for-developers) |LLM| ✅ fully supported|use `og_up` then choose the `OpenAI`|
@@ -133,4 +130,3 @@ if you have any feature suggestion. please create a discuession to talk about it
* [roadmap for v0.5.0](https://github.com/dbpunk-labs/octogen/issues/64)
-
diff --git a/agent/src/og_agent/codellama_agent.py b/agent/src/og_agent/codellama_agent.py
index 2f28159..bd99a58 100644
--- a/agent/src/og_agent/codellama_agent.py
+++ b/agent/src/og_agent/codellama_agent.py
@@ -67,16 +67,18 @@ async def handle_show_sample_code(
"language": json_response.get("language", "text"),
})
await queue.put(
- TaskRespond(
- state=task_context.to_task_state_proto(),
- respond_type=TaskRespond.OnAgentActionType,
- on_agent_action=OnAgentAction(
+ TaskResponse(
+ state=task_context.to_context_state_proto(),
+ response_type=TaskResponse.OnStepActionStart,
+ on_step_agent_start=OnStepActionStart(
input=tool_input, tool="show_sample_code"
),
)
)
- async def handle_bash_code(self, json_response, queue, context, task_context):
+ async def handle_bash_code(
+ self, json_response, queue, context, task_context, task_opt
+ ):
commands = json_response["action_input"]
code = f"%%bash\n {commands}"
explanation = json_response["explanation"]
@@ -88,10 +90,10 @@ async def handle_bash_code(self, json_response, queue, context, task_context):
"language": json_response.get("language"),
})
await queue.put(
- TaskRespond(
- state=task_context.to_task_state_proto(),
- respond_type=TaskRespond.OnAgentActionType,
- on_agent_action=OnAgentAction(
+ TaskResponse(
+ state=task_context.to_context_state_proto(),
+ response_type=TaskResponse.OnStepActionStart,
+ on_step_action_start=OnStepActionStart(
input=tool_input, tool="execute_bash_code"
),
)
@@ -102,7 +104,7 @@ async def handle_bash_code(self, json_response, queue, context, task_context):
logger.debug("the client has cancelled the request")
break
function_result = result
- if respond:
+ if respond and task_opt.streaming:
await queue.put(respond)
return function_result
diff --git a/agent/src/og_agent/mock_agent.py b/agent/src/og_agent/mock_agent.py
index 5996573..96b4723 100644
--- a/agent/src/og_agent/mock_agent.py
+++ b/agent/src/og_agent/mock_agent.py
@@ -8,7 +8,7 @@
import time
import logging
from .base_agent import BaseAgent, TypingState, TaskContext
-from og_proto.agent_server_pb2 import OnStepActionStart, TaskResponse, OnStepActionEnd, FinalAnswer,TypingContent
+from og_proto.agent_server_pb2 import OnStepActionStart, TaskResponse, OnStepActionEnd, FinalAnswer, TypingContent
from .tokenizer import tokenize
logger = logging.getLogger(__name__)
@@ -33,7 +33,9 @@ async def call_ai(self, prompt, queue, iteration, task_context):
TaskResponse(
state=task_context.to_context_state_proto(),
response_type=TaskResponse.OnModelTypeText,
- typing_content=TypingContent(content = message["explanation"], language="text"),
+ typing_content=TypingContent(
+ content=message["explanation"], language="text"
+ ),
)
)
if message.get("code", None):
@@ -41,7 +43,9 @@ async def call_ai(self, prompt, queue, iteration, task_context):
TaskResponse(
state=task_context.to_context_state_proto(),
response_type=TaskResponse.OnModelTypeCode,
- typing_content=TypingContent(content = message["code"], language="python"),
+ typing_content=TypingContent(
+ content=message["code"], language="python"
+ ),
)
)
return message
diff --git a/agent/src/og_agent/openai_agent.py b/agent/src/og_agent/openai_agent.py
index cd93f72..fc506ff 100644
--- a/agent/src/og_agent/openai_agent.py
+++ b/agent/src/og_agent/openai_agent.py
@@ -139,6 +139,7 @@ async def call_openai(self, messages, queue, context, task_context, task_opt):
"""
call the openai api
"""
+ logger.debug(f"call openai with messages {messages}")
input_token_count = 0
for message in messages:
if not message["content"]:
@@ -212,10 +213,10 @@ async def call_openai(self, messages, queue, context, task_context, task_opt):
code_content = code_str
if task_opt.streaming and len(typed_chars) > 0:
typing_language = "text"
- if (
- delta["function_call"].get("name", "")
- == "execute_python_code"
- ):
+ if delta["function_call"].get("name", "") in [
+ "execute_python_code",
+ "python",
+ ]:
typing_language = "python"
elif (
delta["function_call"].get("name", "")
@@ -357,6 +358,8 @@ async def arun(self, task, queue, context, task_opt):
if "function_call" in chat_message:
if "content" not in chat_message:
chat_message["content"] = None
+ if "role" not in chat_message:
+ chat_message["role"] = "assistant"
messages.append(chat_message)
function_name = chat_message["function_call"]["name"]
if function_name not in [
diff --git a/agent/tests/codellama_agent_tests.py b/agent/tests/codellama_agent_tests.py
index f4d2a1e..ee73bc2 100644
--- a/agent/tests/codellama_agent_tests.py
+++ b/agent/tests/codellama_agent_tests.py
@@ -10,7 +10,8 @@
import json
import logging
import pytest
-from og_sdk.agent_sdk import AgentSDK
+
+from og_sdk.kernel_sdk import KernelSDK
from og_agent.codellama_agent import CodellamaAgent
from og_proto.agent_server_pb2 import ProcessOptions, TaskResponse
import asyncio
@@ -21,6 +22,14 @@
logger = logging.getLogger(__name__)
+@pytest.fixture
+def kernel_sdk():
+ endpoint = (
+ "localhost:9527" # Replace with the actual endpoint of your test gRPC server
+ )
+ return KernelSDK(endpoint, "ZCeI9cYtOCyLISoi488BgZHeBkHWuFUH")
+
+
class PayloadStream:
def __init__(self, payload):
@@ -49,24 +58,157 @@ def done(self):
class CodellamaMockClient:
- def __init__(self, payload):
- self.payload = payload
+ def __init__(self, payloads):
+ self.payloads = payloads
+ self.index = 0
async def prompt(self, question, chat_history=[]):
- async for line in PayloadStream(self.payload):
+ if self.index >= len(self.payloads):
+ raise StopAsyncIteration
+ self.index += 1
+ payload = self.payloads[self.index - 1]
+ async for line in PayloadStream(payload):
yield line
-@pytest_asyncio.fixture
-async def agent_sdk():
- sdk = AgentSDK(api_base, api_key)
- sdk.connect()
- yield sdk
- await sdk.close()
+@pytest.mark.asyncio
+async def test_codellama_agent_execute_bash_code(kernel_sdk):
+ kernel_sdk.connect()
+ sentence1 = {
+ "explanation": "print a hello world using python",
+ "action": "execute_bash_code",
+ "action_input": "echo 'hello world'",
+ "saved_filenames": [],
+ "language": "python",
+ "is_final_answer": False,
+ }
+ sentence2 = {
+ "explanation": "the output matchs the goal",
+ "action": "no_action",
+ "action_input": "",
+ "saved_filenames": [],
+ "language": "en",
+ "is_final_answer": False,
+ }
+ client = CodellamaMockClient([json.dumps(sentence1), json.dumps(sentence2)])
+ agent = CodellamaAgent(client, kernel_sdk)
+ task_opt = ProcessOptions(
+ streaming=True,
+ llm_name="codellama",
+ input_token_limit=100000,
+ output_token_limit=100000,
+ timeout=5,
+ )
+ queue = asyncio.Queue()
+ await agent.arun("write a hello world in bash", queue, MockContext(), task_opt)
+ responses = []
+ while True:
+ try:
+ response = await queue.get()
+ if not response:
+ break
+ responses.append(response)
+ except asyncio.QueueEmpty:
+ break
+ logger.info(responses)
+ console_output = list(
+ filter(
+ lambda x: x.response_type == TaskResponse.OnStepActionStreamStdout,
+ responses,
+ )
+ )
+ assert len(console_output) == 1, "bad console output count"
+ assert console_output[0].console_stdout == "hello world\n", "bad console output"
+
+
+@pytest.mark.asyncio
+async def test_codellama_agent_execute_python_code(kernel_sdk):
+ kernel_sdk.connect()
+ sentence1 = {
+ "explanation": "print a hello world using python",
+ "action": "execute_python_code",
+ "action_input": "print('hello world')",
+ "saved_filenames": [],
+ "language": "python",
+ "is_final_answer": False,
+ }
+ sentence2 = {
+ "explanation": "the output matchs the goal",
+ "action": "no_action",
+ "action_input": "",
+ "saved_filenames": [],
+ "language": "en",
+ "is_final_answer": False,
+ }
+ client = CodellamaMockClient([json.dumps(sentence1), json.dumps(sentence2)])
+ agent = CodellamaAgent(client, kernel_sdk)
+ task_opt = ProcessOptions(
+ streaming=True,
+ llm_name="codellama",
+ input_token_limit=100000,
+ output_token_limit=100000,
+ timeout=5,
+ )
+ queue = asyncio.Queue()
+ await agent.arun("write a hello world in python", queue, MockContext(), task_opt)
+ responses = []
+ while True:
+ try:
+ response = await queue.get()
+ if not response:
+ break
+ responses.append(response)
+ except asyncio.QueueEmpty:
+ break
+ logger.info(responses)
+ console_output = list(
+ filter(
+ lambda x: x.response_type == TaskResponse.OnStepActionStreamStdout,
+ responses,
+ )
+ )
+ assert len(console_output) == 1, "bad console output count"
+ assert console_output[0].console_stdout == "hello world\n", "bad console output"
+
+
+@pytest.mark.asyncio
+async def test_codellama_agent_show_demo_code(kernel_sdk):
+ sentence = {
+ "explanation": "Hello, how can I help you?",
+ "action": "show_demo_code",
+ "action_input": "echo 'hello world'",
+ "saved_filenames": [],
+ "language": "shell",
+ "is_final_answer": True,
+ }
+ client = CodellamaMockClient([json.dumps(sentence)])
+ agent = CodellamaAgent(client, kernel_sdk)
+ task_opt = ProcessOptions(
+ streaming=True,
+ llm_name="codellama",
+ input_token_limit=100000,
+ output_token_limit=100000,
+ timeout=5,
+ )
+ queue = asyncio.Queue()
+ await agent.arun("hello", queue, MockContext(), task_opt)
+ responses = []
+ while True:
+ try:
+ response = await queue.get()
+ if not response:
+ break
+ responses.append(response)
+ except asyncio.QueueEmpty:
+ break
+ logger.info(responses)
+ assert (
+ responses[-1].response_type == TaskResponse.OnFinalAnswer
+ ), "bad response type"
@pytest.mark.asyncio
-async def test_codellama_agent_smoke_test(agent_sdk):
+async def test_codellama_agent_smoke_test(kernel_sdk):
sentence = {
"explanation": "Hello, how can I help you?",
"action": "no_action",
@@ -75,8 +217,8 @@ async def test_codellama_agent_smoke_test(agent_sdk):
"language": "en",
"is_final_answer": True,
}
- client = CodellamaMockClient(json.dumps(sentence))
- agent = CodellamaAgent(client, agent_sdk)
+ client = CodellamaMockClient([json.dumps(sentence)])
+ agent = CodellamaAgent(client, kernel_sdk)
task_opt = ProcessOptions(
streaming=True,
llm_name="codellama",
diff --git a/agent/tests/openai_agent_tests.py b/agent/tests/openai_agent_tests.py
index 1433231..d7e92cc 100644
--- a/agent/tests/openai_agent_tests.py
+++ b/agent/tests/openai_agent_tests.py
@@ -7,9 +7,10 @@
""" """
+import json
import logging
import pytest
-from og_sdk.agent_sdk import AgentSDK
+from og_sdk.kernel_sdk import KernelSDK
from og_agent import openai_agent
from og_proto.agent_server_pb2 import ProcessOptions, TaskResponse
from openai.openai_object import OpenAIObject
@@ -19,7 +20,6 @@
api_base = "127.0.0.1:9528"
api_key = "ZCeI9cYtOCyLISoi488BgZHeBkHWuFUH"
-
logger = logging.getLogger(__name__)
@@ -48,28 +48,166 @@ async def __anext__(self):
raise StopAsyncIteration
+class FunctionCallPayloadStream:
+
+ def __init__(self, name, arguments):
+ self.name = name
+ self.arguments = arguments
+
+ def __aiter__(self):
+ # create an iterator of the input keys
+ self.iter_keys = iter(self.arguments)
+ return self
+
+ async def __anext__(self):
+ try:
+ k = next(self.iter_keys)
+ obj = OpenAIObject()
+ delta = OpenAIObject()
+ function_para = OpenAIObject()
+ function_para.name = self.name
+ function_para.arguments = k
+ function_call = OpenAIObject()
+ function_call.function_call = function_para
+ delta.delta = function_call
+ obj.choices = [delta]
+ return obj
+ except StopIteration:
+ # raise stopasynciteration at the end of iterator
+ raise StopAsyncIteration
+
+
class MockContext:
def done(self):
return False
-@pytest_asyncio.fixture
-async def agent_sdk():
- sdk = AgentSDK(api_base, api_key)
- sdk.connect()
- yield sdk
- await sdk.close()
+class MultiCallMock:
+
+ def __init__(self, responses):
+ self.responses = responses
+ self.index = 0
+
+ def call(self, *args, **kwargs):
+ if self.index >= len(self.responses):
+ raise Exception("no more response")
+ self.index += 1
+ logger.debug("call index %d", self.index)
+ return self.responses[self.index - 1]
+
+
+@pytest.fixture
+def kernel_sdk():
+ endpoint = (
+ "localhost:9527" # Replace with the actual endpoint of your test gRPC server
+ )
+ return KernelSDK(endpoint, "ZCeI9cYtOCyLISoi488BgZHeBkHWuFUH")
+
+
+@pytest.mark.asyncio
+async def test_openai_agent_call_execute_bash_code(mocker, kernel_sdk):
+ kernel_sdk.connect()
+ arguments = {
+ "explanation": "the hello world in bash",
+ "code": "echo 'hello world'",
+ "saved_filenames": [],
+ }
+ stream1 = FunctionCallPayloadStream("execute_bash_code", json.dumps(arguments))
+ sentence = "The output 'hello world' is the result"
+ stream2 = PayloadStream(sentence)
+ call_mock = MultiCallMock([stream1, stream2])
+ with mocker.patch(
+ "og_agent.openai_agent.openai.ChatCompletion.acreate",
+ side_effect=call_mock.call,
+ ) as mock_openai:
+ agent = openai_agent.OpenaiAgent("gpt4", "prompt", kernel_sdk, is_azure=False)
+ queue = asyncio.Queue()
+ task_opt = ProcessOptions(
+ streaming=True,
+ llm_name="gpt4",
+ input_token_limit=100000,
+ output_token_limit=100000,
+ timeout=5,
+ )
+ await agent.arun("write a hello world in bash", queue, MockContext(), task_opt)
+ responses = []
+ while True:
+ try:
+ response = await queue.get()
+ if not response:
+ break
+ responses.append(response)
+ except asyncio.QueueEmpty:
+ break
+ logger.info(responses)
+ console_output = list(
+ filter(
+ lambda x: x.response_type == TaskResponse.OnStepActionStreamStdout,
+ responses,
+ )
+ )
+ assert len(console_output) == 1, "bad console output count"
+ assert console_output[0].console_stdout == "hello world\n", "bad console output"
+
+
+@pytest.mark.asyncio
+async def test_openai_agent_call_execute_python_code(mocker, kernel_sdk):
+ kernel_sdk.connect()
+ arguments = {
+ "explanation": "the hello world in python",
+ "code": "print('hello world')",
+ "saved_filenames": [],
+ }
+ stream1 = FunctionCallPayloadStream("execute_python_code", json.dumps(arguments))
+ sentence = "The output 'hello world' is the result"
+ stream2 = PayloadStream(sentence)
+ call_mock = MultiCallMock([stream1, stream2])
+ with mocker.patch(
+ "og_agent.openai_agent.openai.ChatCompletion.acreate",
+ side_effect=call_mock.call,
+ ) as mock_openai:
+ agent = openai_agent.OpenaiAgent("gpt4", "prompt", kernel_sdk, is_azure=False)
+ queue = asyncio.Queue()
+ task_opt = ProcessOptions(
+ streaming=True,
+ llm_name="gpt4",
+ input_token_limit=100000,
+ output_token_limit=100000,
+ timeout=5,
+ )
+
+ await agent.arun(
+ "write a hello world in python", queue, MockContext(), task_opt
+ )
+ responses = []
+ while True:
+ try:
+ response = await queue.get()
+ if not response:
+ break
+ responses.append(response)
+ except asyncio.QueueEmpty:
+ break
+ logger.info(responses)
+ console_output = list(
+ filter(
+ lambda x: x.response_type == TaskResponse.OnStepActionStreamStdout,
+ responses,
+ )
+ )
+ assert len(console_output) == 1, "bad console output count"
+ assert console_output[0].console_stdout == "hello world\n", "bad console output"
@pytest.mark.asyncio
-async def test_openai_agent_smoke_test(mocker, agent_sdk):
+async def test_openai_agent_smoke_test(mocker, kernel_sdk):
sentence = "Hello, how can I help you?"
stream = PayloadStream(sentence)
with mocker.patch(
"og_agent.openai_agent.openai.ChatCompletion.acreate", return_value=stream
) as mock_openai:
- agent = openai_agent.OpenaiAgent("gpt4", "prompt", agent_sdk, is_azure=False)
+ agent = openai_agent.OpenaiAgent("gpt4", "prompt", kernel_sdk, is_azure=False)
queue = asyncio.Queue()
task_opt = ProcessOptions(
streaming=True,
diff --git a/chat/src/og_terminal/terminal_chat.py b/chat/src/og_terminal/terminal_chat.py
index c500b49..d3895e3 100644
--- a/chat/src/og_terminal/terminal_chat.py
+++ b/chat/src/og_terminal/terminal_chat.py
@@ -305,6 +305,31 @@ def render_image(images, sdk, image_dir, console):
return False
+def upload_file(prompt, console, history_prompt, sdk, values):
+ filepaths = parse_file_path(prompt)
+ if not filepaths:
+ return prompt
+ task_blocks = TaskBlocks(values)
+ task_blocks.begin()
+ real_prompt = prompt.replace("/up", "")
+ with Live(Group(*[]), console=console) as live:
+ mk = """The following files will be uploaded
+"""
+ task_blocks.add_markdown(mk)
+ refresh(live, task_blocks, title=SYSTEM_TITLE)
+ for file in filepaths:
+ filename = file.split("/")[-1]
+ sdk.upload_file(file, filename)
+ mk = "* ✅%s\n" % file
+ task_blocks.add_markdown(mk)
+ real_prompt = real_prompt.replace(file, "uploaded %s" % filename)
+ refresh(live, task_blocks, title=SYSTEM_TITLE)
+ task_blocks.get_last_block().finish()
+ refresh(live, task_blocks, title=SYSTEM_TITLE)
+ history_prompt.append_string(real_prompt)
+ return real_prompt
+
+
def run_chat(prompt, sdk, session, console, values, filedir=None):
"""
run the chat
@@ -406,26 +431,7 @@ def app(octogen_dir):
console.print(f"❌ /cc{number} was not found!")
continue
# try to upload first⌛⏳❌
- filepaths = parse_file_path(real_prompt)
- if filepaths:
- real_prompt = real_prompt.replace("/up", "")
- spinner = Spinner(octopus_config.get("spinner", "dots2"), text="Upload...")
- segments = [spinner]
- mk = """The following files will be uploaded
-"""
- with Live(Group(*segments), console=console) as live:
- live.update(spinner)
- for file in filepaths:
- filename = file.split("/")[-1]
- sdk.upload_file(file, filename)
- mk += "* ✅%s\n" % file
- live.update(Group(*[Markdown(mk), spinner]))
- real_prompt = real_prompt.replace(file, "uploaded %s" % filename)
- # clear the spinner
- live.update(Group(*[Markdown(mk)]))
- # add the prompt to history
- # TODO remove the last prompt
- history.append_string(real_prompt)
+ real_prompt = upload_file(real_prompt, console, history, sdk, values)
run_chat(
real_prompt,
sdk,
diff --git a/chat/tests/test_chat_function.py b/chat/tests/test_chat_function.py
index 827a171..3e1a29d 100644
--- a/chat/tests/test_chat_function.py
+++ b/chat/tests/test_chat_function.py
@@ -35,7 +35,9 @@ def test_handle_final_answer_smoke_test():
respond_content = agent_server_pb2.TaskResponse(
state=task_state,
response_type=agent_server_pb2.TaskResponse.OnModelTypeText,
- typing_content=agent_server_pb2.TypingContent(content="hello world!", language="text")
+ typing_content=agent_server_pb2.TypingContent(
+ content="hello world!", language="text"
+ ),
)
respond_final = agent_server_pb2.TaskResponse(
state=task_state,
@@ -82,6 +84,7 @@ def test_handle_action_end_boundary_test():
assert len(images) == 1000
assert all(image == "test.png" for image in images)
+
def test_handle_action_end_smoke_test():
images = []
values = []
@@ -152,6 +155,7 @@ def test_error_handle_action_end():
assert len(images) == 0
assert values[0] == "\nerror"
+
def test_handle_action_end_performance_test():
# Setup
images = []
@@ -173,7 +177,9 @@ def test_handle_action_end_performance_test():
response_type=agent_server_pb2.TaskResponse.OnStepActionEnd,
on_step_action_end=agent_server_pb2.OnStepActionEnd(
output="",
- output_files=[f"test{i}.png"], # Modify this line to create unique filenames
+ output_files=[
+ f"test{i}.png"
+ ], # Modify this line to create unique filenames
has_error=False,
),
)