Octogen is an Open-Source Code Interpreter Powered by GPT3.5/4 and Codellama
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
from setuptools import setup, find_packages

setup(
    name="og_agent",
    version="0.3.6",
    description="Open source code interpreter agent",
    author="imotai",
    author_email="wangtaize@dbpunk.com",
    url="https://github.com/dbpunk-labs/octogen",
    long_description=open("README.md").read(),
    long_description_content_type="text/markdown",
    packages=[
        "og_agent",
    ],
    package_dir={
        "og_agent": "src/og_agent",
    },
    install_requires=[
        "og_proto",
        "og_kernel",
        "og_sdk",
        "grpcio-tools>=1.57.0",
        "grpc-google-iam-v1>=0.12.6",
        "aiofiles",
        "orm[sqlite]",
        "python-dotenv",
        "openai",
        "aiohttp>=3.8.5",
        "pydantic",
        "tiktoken",
    ],
    package_data={"og_agent": ["*.bnf"]},
    entry_points={
        "console_scripts": [
            "og_agent_rpc_server = og_agent.agent_server:server_main",
            "og_agent_setup = og_agent.agent_setup:setup",
        ]
    },
)

agent/setup.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import json
from .prompt import OCTOGEN_FUNCTION_SYSTEM, OCTOGEN_CODELLAMA_SYSTEM
from .codellama_agent import CodellamaAgent
from .openai_agent import OpenaiAgent
from .codellama_client import CodellamaClient
from .mock_agent import MockAgent


def build_codellama_agent(endpoint, key, sdk, grammer_path):
    """
    build codellama agent
    """
    with open(grammer_path, "r") as fd:
        grammar = fd.read()

    client = CodellamaClient(
        endpoint, key, OCTOGEN_CODELLAMA_SYSTEM, "Octogen", "User", grammar
    )

    # init the agent
    return CodellamaAgent(client, sdk)


def build_openai_agent(sdk, model_name, is_azure=True):
    """build openai function call agent"""
    # TODO a data dir per user
    # init the agent

    agent = OpenaiAgent(model_name, OCTOGEN_FUNCTION_SYSTEM, sdk, is_azure=is_azure)
    return agent


def build_mock_agent(sdk, cases_path):
    """
    build the mock agent for testing
    """
    with open(cases_path, "r") as fd:
        messages = json.load(fd)
    agent = MockAgent(messages, sdk)
    return agent

agent/src/og_agent/agent_builder.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """

import logging
import openai

logger = logging.getLogger(__name__)


class LLMManager:

    def __init__(self, config):
        """
        llm_key=xxxx
        # the other optional
        """
        self.config = config
        self.llms = {}
        self.llm_key = self.config["llm_key"]
        if self.config["llm_key"] == "azure_openai":
            self._build_azure_openai()
        elif self.config["llm_key"] == "openai":
            self._build_openai()

    def get_llm(self):
        return self.llms.get(self.llm_key, None)

    def get_llm_by_key(self, llm_key):
        """
        get llm with a key, the supported keys are 'mock', 'openai', 'azure_openai', 'codellama'
        """
        return self.llms.get(self.llm_key, None)

    def _no_empty_value_required(self, keys):
        for key in keys:
            if not self.config.get(key, None):
                raise Exception(f"the value of required {key} is empty")

    def _build_openai(self):
        self._no_empty_value_required([
            "openai_api_key",
            "openai_api_model",
        ])
        if self.config.get("openai_api_base", None):
            openai.api_base = self.config.get("openai_api_base", None)
        openai.api_key = self.config["openai_api_key"]

    def _build_azure_openai(self):
        """
        build azure openai client from config
        """
        self._no_empty_value_required([
            "openai_api_base",
            "openai_api_version",
            "openai_api_key",
            "openai_api_type",
            "openai_api_deployment",
        ])
        openai.api_base = self.config["openai_api_base"]
        openai.api_version = self.config["openai_api_version"]
        openai.api_type = self.config["openai_api_type"]
        openai.api_key = self.config["openai_api_key"]
        self.config["openai_api_model"] = self.config["openai_api_deployment"]

agent/src/og_agent/agent_llm.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>

import asyncio
import logging
import sys
import os
import pathlib
import hashlib
import grpc
import json
from grpc.aio import AioRpcError
from og_proto.agent_server_pb2_grpc import AgentServerServicer
from og_proto.agent_server_pb2_grpc import add_AgentServerServicer_to_server
from og_proto import agent_server_pb2
from og_proto import common_pb2
from og_proto import kernel_server_pb2
from dotenv import dotenv_values
from typing import AsyncIterable, Any, Dict, List, Optional, Sequence, Union, Type
from tempfile import gettempdir
from grpc.aio import ServicerContext, server
from og_sdk.kernel_sdk import KernelSDK
from og_sdk.utils import parse_image_filename
from .agent_llm import LLMManager
from .agent_builder import build_mock_agent, build_openai_agent, build_codellama_agent
import databases
import orm
from datetime import datetime

config = dotenv_values(".env")
LOG_LEVEL = (
    logging.DEBUG if config.get("log_level", "info") == "debug" else logging.INFO
)

logging.basicConfig(
    level=LOG_LEVEL,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)

logger = logging.getLogger(__name__)

# the database instance for agent
database = databases.Database("sqlite:///%s" % config["db_path"])
models = orm.ModelRegistry(database=database)


class LiteApp(orm.Model):
    tablename = "lite_app"
    registry = models
    fields = {
        "id": orm.Integer(primary_key=True),
        "key_hash": orm.String(max_length=64, index=True),
        "name": orm.String(max_length=20, index=True),
        "language": orm.String(max_length=20, allow_null=False),
        "code": orm.Text(),
        "time": orm.DateTime(),
        "desc": orm.String(max_length=100, allow_null=True),
        "saved_filenames": orm.String(max_length=512, allow_null=True),
    }


class AgentRpcServer(AgentServerServicer):

    def __init__(self):
        self.agents = {}
        self.max_file_size = int(config["max_file_size"])
        self.verbose = config.get("verbose", False)
        self.llm_manager = LLMManager(config)
        self.llm = self.llm_manager.get_llm()

    async def ping(
        self, request: agent_server_pb2.PingRequest, context: ServicerContext
    ) -> agent_server_pb2.PongResponse:
        metadata = dict(context.invocation_metadata())
        if (
            "api_key" not in metadata
            or metadata["api_key"] not in self.agents
            or not self.agents[metadata["api_key"]]
        ):
            return agent_server_pb2.PongResponse(
                code=-1, msg="Your API Key is invalid!"
            )
        else:
            return agent_server_pb2.PongResponse(
                code=0, msg="Connect to Octopus Agent Ok!"
            )

    async def assemble(
        self, request: agent_server_pb2.AssembleAppRequest, context: ServicerContext
    ) -> agent_server_pb2.AssembleAppResponse:
        metadata = dict(context.invocation_metadata())
        if (
            "api_key" not in metadata
            or metadata["api_key"] not in self.agents
            or not self.agents[metadata["api_key"]]
        ):
            await context.abort(10, "invalid api key")
        key = metadata["api_key"]
        key_hash = hashlib.sha256(key.encode("UTF-8")).hexdigest()
        if await LiteApp.objects.filter(key_hash=key_hash, name=request.name).first():
            await context.abort(10, f"the app name {request.name} exist")

        logger.debug(
            f"the code {request.code} key {key_hash}, language {request.language}"
        )
        await LiteApp.objects.create(
            key_hash=key_hash,
            name=request.name,
            language=request.language,
            code=request.code,
            time=datetime.now(),
            desc=request.desc,
            saved_filenames=",".join(request.saved_filenames)
            if request.saved_filenames
            else "",
        )
        return agent_server_pb2.AssembleAppResponse(code=0, msg="ok")

    async def run(
        self, request: agent_server_pb2.RunAppRequest, context
    ) -> AsyncIterable[agent_server_pb2.TaskRespond]:
        metadata = dict(context.invocation_metadata())
        if (
            "api_key" not in metadata
            or metadata["api_key"] not in self.agents
            or not self.agents[metadata["api_key"]]
        ):
            logger.debug("invalid api key")
            await context.abort(10, "invalid api key")
        logger.debug(f"run application {request.name}")
        key = metadata["api_key"]
        key_hash = hashlib.sha256(key.encode("UTF-8")).hexdigest()
        lite_app = await LiteApp.objects.filter(
            key_hash=key_hash, name=request.name
        ).first()
        if not lite_app:
            await context.abort(10, f"no application with name {request.name}")
        agent = self.agents[metadata["api_key"]]["agent"]
        tool_input = json.dumps({
            "code": lite_app.code,
            "explanation": "",
            "saved_filenames": lite_app.saved_filenames.split(",")
            if lite_app.saved_filenames
            else [],
        })
        yield agent_server_pb2.TaskRespond(
            respond_type=agent_server_pb2.TaskRespond.OnAgentActionType,
            on_agent_action=agent_server_pb2.OnAgentAction(
                input=tool_input, tool="execute_python_code"
            ),
        )
        function_result = None
        async for (result, respond) in agent.call_function(lite_app.code, context):
            if context.cancelled():
                break
            function_result = result
            if respond:
                logger.debug(f"the respond {respond}")
                yield respond
        output_files = (
            function_result.saved_filenames
            if function_result and function_result.saved_filenames
            else []
        )
        yield agent_server_pb2.TaskRespond(
            respond_type=agent_server_pb2.TaskRespond.OnAgentActionEndType,
            on_agent_action_end=agent_server_pb2.OnAgentActionEnd(
                output="", output_files=output_files
            ),
        )

    async def query_apps(
        self, request: agent_server_pb2.QueryAppsRequest, context: ServicerContext
    ) -> agent_server_pb2.QueryAppsResponse:
        metadata = dict(context.invocation_metadata())
        if (
            "api_key" not in metadata
            or metadata["api_key"] not in self.agents
            or not self.agents[metadata["api_key"]]
        ):
            logger.debug("invalid api key")
            await context.abort(10, "invalid api key")
        key = metadata["api_key"]
        key_hash = hashlib.sha256(key.encode("UTF-8")).hexdigest()
        lite_apps = (
            await LiteApp.objects.filter(key_hash=key_hash).order_by("-time").all()
        )
        apps = [
            agent_server_pb2.AppInfo(
                name=lite_app.name,
                language=lite_app.language,
                ctime=int(lite_app.time.timestamp()),
                desc=lite_app.desc,
            )
            for lite_app in lite_apps
        ]
        return agent_server_pb2.QueryAppsResponse(apps=apps)

    async def add_kernel(
        self, request: agent_server_pb2.AddKernelRequest, context: ServicerContext
    ) -> agent_server_pb2.AddKernelResponse:
        """Create a token, only the admin can call this method"""
        metadata = dict(context.invocation_metadata())
        if "api_key" not in metadata or metadata["api_key"] != config["admin_key"]:
            await context.abort(10, "You are not the admin")
        if request.key in self.agents and self.agents[request.key]:
            return agent_server_pb2.AddKernelResponse(code=0, msg="ok")
        # init the sdk
        sdk = KernelSDK(request.endpoint, request.key)
        try:
            sdk.connect()
            await sdk.is_alive()
        except Exception as ex:
            await context.abort(10, f"Connecting to kernel {request.endpoint} failed")
        if config["llm_key"] == "azure_openai" or config["llm_key"] == "openai":
            logger.info(f"create a openai agent {request.endpoint}")
            agent = build_openai_agent(
                sdk,
                config["openai_api_model"],
                is_azure=True if config["llm_key"] == "azure_openai" else False,
            )
            self.agents[request.key] = {"sdk": sdk, "agent": agent}
        elif config["llm_key"] == "mock":
            logger.info(f"create a mock agent to kernel {request.endpoint}")
            agent = build_mock_agent(sdk, config["cases_path"])
            self.agents[request.key] = {"sdk": sdk, "agent": agent}
        elif config["llm_key"] == "codellama":
            logger.info(f"create a codellama agent {request.endpoint}")
            grammer_path = os.path.join(
                pathlib.Path(__file__).parent.resolve(), "grammar.bnf"
            )
            agent = build_codellama_agent(
                config["llama_api_base"], config["llama_api_key"], sdk, grammer_path
            )
            self.agents[request.key] = {"sdk": sdk, "agent": agent}
        return agent_server_pb2.AddKernelResponse(code=0, msg="ok")

    async def send_task(
        self, request: agent_server_pb2.SendTaskRequest, context
    ) -> AsyncIterable[agent_server_pb2.TaskRespond]:
        """
        process the task from the client
        """
        logger.debug("receive the task %s ", request.task)
        metadata = dict(context.invocation_metadata())
        if (
            "api_key" not in metadata
            or metadata["api_key"] not in self.agents
            or not self.agents[metadata["api_key"]]
        ):
            logger.debug("invalid api key")
            await context.abort(10, "invalid api key")
        agent = self.agents[metadata["api_key"]]["agent"]
        sdk = self.agents[metadata["api_key"]]["sdk"]
        queue = asyncio.Queue()

        async def worker(task, agent, queue, context):
            try:
                return await agent.arun(task, queue, context)
            except Exception as ex:
                logger.exception("fail to run agent")
                result = str(ex)
                return result

        logger.debug("create the agent task")
        task = asyncio.create_task(worker(request.task, agent, queue, context))
        try:
            while True:
                try:
                    logger.debug("start wait the queue message")
                    # TODO add timeout
                    respond = await queue.get()
                    if not respond:
                        logger.debug("exit the queue")
                        break
                    logger.debug(f"respond {respond}")
                    queue.task_done()
                    yield respond
                except Exception as ex:
                    logger.error(f"fail to get respond for {ex}")
                    break
            await task
        except Exception as ex:
            respond = agent_server_pb2.TaskRespond(
                token_usage=0,
                model_name="",
                iteration=1,
                respond_type=agent_server_pb2.TaskRespond.OnFinalAnswerType,
                final_respond=agent_server_pb2.FinalRespond(answer=str(ex)),
            )
            yield respond
        finally:
            is_cancelled = context.cancelled()
            logger.debug(f" the context is cancelled {is_cancelled}")
            if context.cancelled():
                try:
                    logger.warning("cancel the request by stop kernel")
                    await sdk.stop()
                except Exception as ex:
                    pass

    async def download(
        self, request: common_pb2.DownloadRequest, context: ServicerContext
    ) -> AsyncIterable[common_pb2.FileChunk]:
        """
        download file from kernel and send it to client
        """
        metadata = dict(context.invocation_metadata())
        if (
            "api_key" not in metadata
            or metadata["api_key"] not in self.agents
            or not self.agents[metadata["api_key"]]
        ):
            await context.abort(10, "invalid api key")
        sdk = self.agents[metadata["api_key"]]["sdk"]
        async for chunk in sdk.download_file(request.filename):
            yield chunk

    async def upload(
        self,
        request: AsyncIterable[common_pb2.FileChunk],
        context: ServicerContext,
    ) -> common_pb2.FileUploaded:
        """
        upload file to the kernel workspace
        """
        metadata = dict(context.invocation_metadata())
        if (
            "api_key" not in metadata
            or metadata["api_key"] not in self.agents
            or not self.agents[metadata["api_key"]]
        ):
            await context.abort(10, "invalid arguments")
        sdk = self.agents[metadata["api_key"]]["sdk"]

        async def generate_chunk(proxy_request, context, limit):
            length = 0
            async for chunk in proxy_request:
                if length + len(chunk.buffer) > limit:
                    await context.abort(
                        grpc.StatusCode.INVALID_ARGUMENT.value[0],
                        "exceed the max file limit",
                    )
                length += len(chunk.buffer)
                yield chunk

        return await sdk.upload_binary(
            generate_chunk(request, context, self.max_file_size)
        )


async def serve() -> None:
    logger.info(
        "start agent rpc server with host %s and port %s",
        config["rpc_host"],
        config["rpc_port"],
    )
    await models.create_all()
    serv = server()
    add_AgentServerServicer_to_server(AgentRpcServer(), serv)
    listen_addr = "%s:%s" % (config["rpc_host"], config["rpc_port"])
    serv.add_insecure_port(listen_addr)
    await serv.start()
    await serv.wait_for_termination()


def server_main():
    asyncio.run(serve())

agent/src/og_agent/agent_server.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import click
import asyncio
from og_sdk.agent_sdk import AgentSDK


async def add_kernel(endpoint, api_key, kernel_endpoint, kernel_api_key):
    sdk = AgentSDK(endpoint, api_key)
    sdk.connect()
    try:
        await sdk.add_kernel(kernel_api_key, kernel_endpoint)
        print("add kernel %s done" % kernel_endpoint)
    except Exception as ex:
        print("add kernel %s failed %s" % (kernel_endpoint, ex))


@click.command()
@click.option("--kernel_endpoint", help="the endpoint of kernel")
@click.option("--kernel_api_key", help="the api key of kernel")
@click.option("--agent_endpoint", help="the endpoint of agent")
@click.option("--admin_key", help="the admin key of agent")
def setup(kernel_endpoint, kernel_api_key, agent_endpoint, admin_key):
    if not kernel_endpoint or not kernel_api_key or not admin_key or not agent_endpoint:
        print("kernel_endpoint or kernel_api_key or admin key is empty")
        return
    asyncio.run(add_kernel(agent_endpoint, admin_key, kernel_endpoint, kernel_api_key))

agent/src/og_agent/agent_setup.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import openai
import json
import logging
import time
from typing import Any, Dict, List, Optional, Sequence, Union, Type
from pydantic import BaseModel, Field
from og_proto.kernel_server_pb2 import ExecuteResponse
from og_proto.agent_server_pb2 import OnAgentAction, TaskRespond, OnAgentActionEnd, FinalRespond, TaskState
from og_sdk.utils import parse_image_filename, process_char_stream

logger = logging.getLogger(__name__)


class FunctionResult(BaseModel):
    console_stderr: str = ""
    console_stdout: str = ""
    saved_filenames: List[str] = Field(
        description="A list of filenames that were created by the code", default=[]
    )
    has_result: bool = False
    has_error: bool = False


class TaskContext(BaseModel):
    start_time: float = 0
    generated_token_count: int = 0
    sent_token_count: int = 0
    model_name: str = ""
    iteration_count: int = 0
    model_respond_duration: int = 0

    def to_task_state_proto(self):
        # in ms
        total_duration = int((time.time() - self.start_time) * 1000)
        return TaskState(
            generated_token_count=self.generated_token_count,
            iteration_count=self.iteration_count,
            model_name=self.model_name,
            total_duration=total_duration,
            sent_token_count=self.sent_token_count,
            model_respond_duration=self.model_respond_duration,
        )


class TypingState:
    START = 0
    EXPLANATION = 1
    CODE = 2


class BaseAgent:

    def __init__(self, sdk):
        self.kernel_sdk = sdk

    async def call_function(self, code, context, task_context=None):
        """
        run code with kernel
        """
        console_stdout = ""
        console_stderr = ""
        has_result = False
        has_error = False
        is_alive = await self.kernel_sdk.is_alive()
        if not is_alive:
            await self.kernel_sdk.start(kernel_name="python3")
        async for kernel_respond in self.kernel_sdk.execute(code=code):
            if context.done():
                logger.debug(
                    "the context is not active and the client cancelled the request"
                )
                break
            # process the stdout
            if kernel_respond.output_type == ExecuteResponse.StdoutType:
                kernel_output = json.loads(kernel_respond.output)["text"]
                console_stdout += kernel_output
                console_stdout = process_char_stream(console_stdout)
                logger.debug(f"the new stdout {console_stdout}")
                yield (
                    None,
                    TaskRespond(
                        state=task_context.to_task_state_proto()
                        if task_context
                        else None,
                        respond_type=TaskRespond.OnAgentActionStdout,
                        console_stdout=kernel_output,
                    ),
                )
            # process the stderr
            elif kernel_respond.output_type == ExecuteResponse.StderrType:
                kernel_err = json.loads(kernel_respond.output)["text"]
                console_stderr += kernel_err
                console_stderr = process_char_stream(console_stderr)
                logger.debug(f"the new stderr {console_stderr}")
                yield (
                    None,
                    TaskRespond(
                        state=task_context.to_task_state_proto()
                        if task_context
                        else None,
                        respond_type=TaskRespond.OnAgentActionStderr,
                        console_stderr=kernel_err,
                    ),
                )
            elif kernel_respond.output_type == ExecuteResponse.TracebackType:
                traceback = json.loads(kernel_respond.output)["traceback"]
                console_stderr += traceback
                logger.debug(f"the new traceback {console_stderr}")
                has_error = True
                yield (
                    None,
                    TaskRespond(
                        state=task_context.to_task_state_proto()
                        if task_context
                        else None,
                        respond_type=TaskRespond.OnAgentActionStderr,
                        console_stderr=traceback,
                    ),
                )
            else:
                has_result = True
                result = json.loads(kernel_respond.output)
                logger.debug(f"the result {result}")
                if "image/gif" in result:
                    console_stdout = result["image/gif"]
                elif "image/png" in result:
                    console_stdout = result["image/png"]
                elif "text/plain" in result:
                    console_stdout = result["text/plain"]
                    console_stdout = bytes(console_stdout, "utf-8").decode(
                        "unicode_escape"
                    )
                    if console_stdout.startswith("'") and console_stdout.endswith("'"):
                        console_stdout = console_stdout[1 : len(console_stdout) - 1]
                yield (
                    None,
                    TaskRespond(
                        state=task_context.to_task_state_proto()
                        if task_context
                        else None,
                        respond_type=TaskRespond.OnAgentActionStdout,
                        console_stdout=console_stdout,
                    ),
                )
        output_files = []
        filename = parse_image_filename(console_stdout)
        if filename:
            output_files.append(filename)
        yield (
            FunctionResult(
                console_stderr=console_stderr,
                console_stdout=console_stdout,
                saved_filenames=output_files,
                has_error=has_error,
                has_result=has_result,
            ),
            None,
        )

agent/src/og_agent/base_agent.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import json
import aiohttp
import logging

logger = logging.getLogger(__name__)


class BaseStreamClient:

    def __init__(self, endpoint, key):
        self.endpoint = endpoint
        self.key = key

    async def arun(self, request_data):
        logging.debug(f"{request_data}")
        data = json.dumps(request_data)
        headers = {"Authorization": self.key}
        async with aiohttp.ClientSession(
            headers=headers, raise_for_status=True
        ) as session:
            async with session.post(self.endpoint, data=data) as r:
                async for line in r.content:
                    if line:
                        yield line

agent/src/og_agent/base_stream_client.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """

import json
import logging
import io
import time
from .codellama_client import CodellamaClient
from og_proto.agent_server_pb2 import OnAgentAction, TaskRespond, OnAgentActionEnd, FinalRespond
from .base_agent import BaseAgent, TypingState, TaskContext
from .tokenizer import tokenize
import tiktoken

logger = logging.getLogger(__name__)
encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")


class CodellamaAgent(BaseAgent):

    def __init__(self, client, kernel_sdk):
        super().__init__(kernel_sdk)
        self.client = client

    def _output_exception(self):
        return (
            "Sorry, the LLM did return nothing, You can use a better performance model"
        )

    def _format_output(self, json_response):
        """
        format the response and send it to the user
        """
        answer = json_response["explanation"]
        if json_response["action"] == "no_action":
            return answer
        elif json_response["action"] == "show_sample_code":
            return ""
        else:
            code = json_response.get("action_input", None)
            answer_code = """%s
```%s
%s
```
""" % (
                answer,
                json_response.get("language", "python"),
                code if code else "",
            )
            return answer_code

    async def handle_show_sample_code(
        self, json_response, queue, context, task_context
    ):
        code = json_response["action_input"]
        explanation = json_response["explanation"]
        saved_filenames = json_response.get("saved_filenames", [])
        tool_input = json.dumps({
            "code": code,
            "explanation": explanation,
            "saved_filenames": saved_filenames,
        })
        await queue.put(
            TaskRespond(
                state=task_context.to_task_state_proto(),
                respond_type=TaskRespond.OnAgentActionType,
                on_agent_action=OnAgentAction(
                    input=tool_input, tool="show_sample_code"
                ),
            )
        )

    async def handle_install_package(self, json_response, queue, context, task_context):
        package = json_response["action_input"]
        code = f"!pip install {package}"
        explanation = json_response["explanation"]
        saved_filenames = json_response.get("saved_filenames", [])
        tool_input = json.dumps({
            "code": code,
            "explanation": explanation,
            "saved_filenames": saved_filenames,
        })
        await queue.put(
            TaskRespond(
                state=task_context.to_task_state_proto(),
                respond_type=TaskRespond.OnAgentActionType,
                on_agent_action=OnAgentAction(
                    input=tool_input, tool="execute_python_code"
                ),
            )
        )
        function_result = None
        async for (result, respond) in self.call_function(code, context, task_context):
            if context.done():
                logger.debug("the client has cancelled the request")
                break
            function_result = result
            if respond:
                await queue.put(respond)
        return function_result

    async def handle_function(self, json_response, queue, context, task_context):
        code = json_response["action_input"]
        explanation = json_response["explanation"]
        saved_filenames = json_response.get("saved_filenames", [])
        tool_input = json.dumps({
            "code": code,
            "explanation": explanation,
            "saved_filenames": saved_filenames,
        })
        await queue.put(
            TaskRespond(
                state=task_context.to_task_state_proto(),
                respond_type=TaskRespond.OnAgentActionType,
                on_agent_action=OnAgentAction(
                    input=tool_input, tool="execute_python_code"
                ),
            )
        )
        function_result = None

        async for (result, respond) in self.call_function(code, context, task_context):
            if context.done():
                logger.debug("the client has cancelled the request")
                break
            function_result = result
            if respond:
                await queue.put(respond)
        return function_result

    def _get_argument_new_typing(self, message):
        state = TypingState.START
        explanation_str = ""
        action_input_str = ""
        for token_state, token in tokenize(io.StringIO(message)):
            if token_state == None:
                if state == TypingState.EXPLANATION and token[0] == 1:
                    explanation_str = token[1]
                    state = TypingState.START
                if state == TypingState.CODE and token[0] == 1:
                    action_input_str = token[1]
                    state = TypingState.START
                if token[1] == "explanation":
                    state = TypingState.EXPLANATION
                if token[1] == "action_input":
                    state = TypingState.CODE
            else:
                # String
                if token_state == 9 and state == TypingState.EXPLANATION:
                    explanation_str = "".join(token)
                elif token_state == 9 and state == TypingState.CODE:
                    action_input_str = "".join(token)
        return (state, explanation_str, action_input_str)

    async def call_codellama(
        self, question, chat_history, queue, context, task_context
    ):
        """
        call codellama api
        """
        start_time = time.time()
        num_tokens = len(encoding.encode(question)) + len(encoding.encode(chat_history))
        task_context.sent_token_count += num_tokens
        state = None
        message = ""
        text_content = ""
        code_content = ""
        async for line in self.client.prompt(question, chat_history=chat_history):
            if len(line) < 6:
                continue
            if context.done():
                logger.debug("the client has cancelled the request")
                break
            respond = json.loads(line[6:])
            task_context.generated_token_count += len(
                encoding.encode(respond["content"])
            )
            task_context.model_respond_duration += int(
                (time.time() - start_time) * 1000
            )
            start_time = time.time()
            message += respond["content"]
            logger.debug(f" message {message}")
            (
                state,
                explanation_str,
                action_input_str,
            ) = self._get_argument_new_typing(message)
            if explanation_str and text_content != explanation_str:
                typed_chars = explanation_str[len(text_content) :]
                text_content = explanation_str
                await queue.put(
                    TaskRespond(
                        state=task_context.to_task_state_proto(),
                        respond_type=TaskRespond.OnAgentTextTyping,
                        typing_content=typed_chars,
                    )
                )
            if action_input_str and code_content != action_input_str:
                typed_chars = action_input_str[len(code_content) :]
                code_content = action_input_str
                await queue.put(
                    TaskRespond(
                        state=task_context.to_task_state_proto(),
                        respond_type=TaskRespond.OnAgentCodeTyping,
                        typing_content=typed_chars,
                    )
                )
            logger.debug(
                f"argument explanation:{explanation_str} code:{action_input_str}"
            )
            if respond["stop"]:
                state = respond

        return (message, state)

    async def arun(self, question, queue, context, max_iteration=5):
        """
        run the agent
        """
        history = []
        current_question = question
        task_context = TaskContext(
            start_time=time.time(),
            generated_token_count=0,
            sent_token_count=0,
            model_name="codellama",
            iteration_count=0,
            model_respond_duration=0,
        )
        iteration = 0
        try:
            while iteration < max_iteration:
                if context.done():
                    logger.debug("the client has cancelled the request")
                    break
                iteration += 1
                task_context.iteration_count = iteration
                chat_history = "\n".join(history)
                (message, state) = await self.call_codellama(
                    current_question, chat_history, queue, context, task_context
                )
                try:
                    json_response = json.loads(message)
                    if not json_response:
                        await queue.put(
                            TaskRespond(
                                state=task_context.to_task_state_proto(),
                                respond_type=TaskRespond.OnFinalAnswerType,
                                final_respond=FinalRespond(
                                    answer=self._output_exception()
                                ),
                            )
                        )
                        break
                except Exception as ex:
                    logger.exception(f"fail to load message the message is {message}")
                    await queue.put(
                        TaskRespond(
                            state=task_context.to_task_state_proto(),
                            respond_type=TaskRespond.OnFinalAnswerType,
                            final_respond=FinalRespond(
                                answer="The model made an invalid respone"
                            ),
                        )
                    )
                    break
                logger.debug(f" codellama response {json_response}")
                if (
                    json_response["action"] == "execute_python_code"
                    and json_response["action_input"]
                ):
                    function_result = await self.handle_function(
                        json_response, queue, context, task_context
                    )
                    logger.debug(f"the function result {function_result}")
                    await queue.put(
                        TaskRespond(
                            state=task_context.to_task_state_proto(),
                            respond_type=TaskRespond.OnAgentActionEndType,
                            on_agent_action_end=OnAgentActionEnd(
                                output="",
                                output_files=function_result.saved_filenames,
                                has_error=function_result.has_error,
                            ),
                        )
                    )
                    history.append("User:%s" % current_question)
                    history.append("Octogen:%s" % message)
                    ins = "the output of execute_python_code:"
                    # TODO limit the output size
                    if function_result.has_result:
                        current_question = f"{ins}\n{function_result.console_stdout}"
                        logger.debug(
                            "continue to iterate with codellama with question %s"
                            % function_result.console_stdout
                        )
                    elif function_result.has_error:
                        current_question = f"{ins} \n {function_result.console_stderr}"
                        logger.debug(
                            "continue to iterate with codellama with question %s"
                            % function_result.console_stderr
                        )
                    else:
                        current_question = f"{ins} \n {function_result.console_stdout}"
                        logger.debug(
                            "continue to iterate with codellama with question %s"
                            % function_result.console_stdout
                        )
                elif (
                    json_response["action"] == "show_sample_code"
                    and json_response["action_input"]
                ):
                    await self.handle_show_sample_code(
                        json_response, queue, context, task_context
                    )
                    result = self._format_output(json_response)
                    await queue.put(
                        TaskRespond(
                            state=task_context.to_task_state_proto(),
                            respond_type=TaskRespond.OnFinalAnswerType,
                            final_respond=FinalRespond(answer=result),
                        )
                    )
                    break
                elif (
                    json_response["action"] == "install_python_package"
                    and json_response["action_input"]
                ):
                    function_result = await self.handle_install_package(
                        json_response, queue, context, task_context
                    )
                    logger.debug(f"the function result {function_result}")
                    await queue.put(
                        TaskRespond(
                            state=task_context.to_task_state_proto(),
                            respond_type=TaskRespond.OnAgentActionEndType,
                            on_agent_action_end=OnAgentActionEnd(
                                output="",
                                output_files=function_result.saved_filenames,
                                has_error=function_result.has_error,
                            ),
                        )
                    )
                    history.append("User:%s" % current_question)
                    history.append("Octogen:%s" % message)
                    ins = "the output of install_python_package:"
                    # TODO limit the output size
                    if function_result.has_result:
                        current_question = f"{ins}\n{function_result.console_stdout}"
                        logger.debug(
                            "continue to iterate with codellama with question %s"
                            % function_result.console_stdout
                        )
                    elif function_result.has_error:
                        current_question = f"{ins} \n {function_result.console_stderr}"
                        logger.debug(
                            "continue to iterate with codellama with question %s"
                            % function_result.console_stderr
                        )
                    else:
                        current_question = f"{ins} \n {function_result.console_stdout}"
                        logger.debug(
                            "continue to iterate with codellama with question %s"
                            % function_result.console_stdout
                        )

                else:
                    result = self._format_output(json_response)
                    await queue.put(
                        TaskRespond(
                            state=task_context.to_task_state_proto(),
                            respond_type=TaskRespond.OnFinalAnswerType,
                            final_respond=FinalRespond(answer=result),
                        )
                    )
                    break
        finally:
            await queue.put(None)

agent/src/og_agent/codellama_agent.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """

import json
import aiohttp
import logging
from .base_stream_client import BaseStreamClient

logger = logging.getLogger(__name__)


class CodellamaClient(BaseStreamClient):

    def __init__(self, endpoint, key, prefix, ai_name, user_name, grammar):
        super().__init__(endpoint + "/completion", key)
        self.ai_name = ai_name
        self.user_name = user_name
        self.grammar = grammar
        self.prefix = prefix

    async def prompt(self, user_input, temperature=0, chat_history=""):
        prompt = f"""{self.prefix}
{chat_history}
{self.user_name}: {user_input}
{self.ai_name}:"""
        logging.info(f"{prompt}")
        data = {
            "n_predict": 1024,
            "grammar": self.grammar,
            "prompt": prompt,
            "temperature": temperature,
            "stream": True,
            "repeat_last_n": 256,
            "top_p": 0.9,
            "stop": [
                "</s>",
                "\n",
                "%s:" % self.ai_name,
                "%s:" % self.user_name,
            ],
        }
        async for line in self.arun(data):
            yield line

agent/src/og_agent/codellama_client.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import json
import time
import logging
from .base_agent import BaseAgent, TypingState, TaskContext
from og_proto.agent_server_pb2 import OnAgentAction, TaskRespond, OnAgentActionEnd, FinalRespond
from .tokenizer import tokenize

logger = logging.getLogger(__name__)


class MockAgent(BaseAgent):
    """
    a test agent for octogen
    """

    def __init__(self, messages, sdk):
        """
        the messages are the cases
        """
        super().__init__(sdk)
        self.messages = messages

    async def call_ai(self, prompt, queue, iteration, task_context):
        message = self.messages.get(prompt)[iteration]
        if message.get("explanation", None):
            await queue.put(
                TaskRespond(
                    state=task_context.to_task_state_proto(),
                    respond_type=TaskRespond.OnAgentTextTyping,
                    typing_content=message["explanation"],
                )
            )
        if message.get("code", None):
            await queue.put(
                TaskRespond(
                    state=task_context.to_task_state_proto(),
                    respond_type=TaskRespond.OnAgentCodeTyping,
                    typing_content=message["code"],
                )
            )
        return message

    async def handle_call_function(
        self, code, queue, explanation, context, task_context, saved_filenames=[]
    ):
        tool_input = json.dumps({
            "code": code,
            "explanation": explanation,
            "saved_filenames": saved_filenames,
        })
        await queue.put(
            TaskRespond(
                state=task_context.to_task_state_proto(),
                respond_type=TaskRespond.OnAgentActionType,
                on_agent_action=OnAgentAction(
                    input=tool_input, tool="execute_python_code"
                ),
            )
        )
        function_result = None
        async for (result, respond) in self.call_function(code, context, task_context):
            function_result = result
            if respond:
                await queue.put(respond)
        return function_result

    async def arun(self, task, queue, context, max_iteration=5):
        """
        run the agent
        """
        iteration = 0
        task_context = TaskContext(
            start_time=time.time(),
            generated_token_count=10,
            sent_token_count=10,
            model_name="mock",
            iteration_count=1,
            model_respond_duration=1000,
        )
        try:
            while iteration < max_iteration:
                message = await self.call_ai(task, queue, iteration, task_context)
                iteration = iteration + 1
                if message.get("code", None):
                    function_result = await self.handle_call_function(
                        message["code"],
                        queue,
                        message["explanation"],
                        context,
                        task_context,
                        message.get("saved_filenames", []),
                    )
                    await queue.put(
                        TaskRespond(
                            state=task_context.to_task_state_proto(),
                            respond_type=TaskRespond.OnAgentActionEndType,
                            on_agent_action_end=OnAgentActionEnd(
                                output="",
                                output_files=function_result.saved_filenames,
                                has_error=function_result.has_error,
                            ),
                        )
                    )
                else:
                    await queue.put(
                        TaskRespond(
                            state=task_context.to_task_state_proto(),
                            respond_type=TaskRespond.OnFinalAnswerType,
                            final_respond=FinalRespond(answer=message["explanation"]),
                        )
                    )
                    break
        finally:
            await queue.put(None)

agent/src/og_agent/mock_agent.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import openai
import io
import json
import logging
import time
from pydantic import BaseModel, Field
from og_proto.agent_server_pb2 import OnAgentAction, TaskRespond, OnAgentActionEnd, FinalRespond
from .base_agent import BaseAgent, TypingState, TaskContext
from .tokenizer import tokenize
import tiktoken

logger = logging.getLogger(__name__)
encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
OCTOGEN_FUNCTIONS = [
    {
        "name": "execute_python_code",
        "description": "Safely execute arbitrary Python code and return the result, stdout, and stderr.",
        "parameters": {
            "type": "object",
            "properties": {
                "explanation": {
                    "type": "string",
                    "description": "the explanation about the python code",
                },
                "code": {
                    "type": "string",
                    "description": "the python code to be executed",
                },
                "saved_filenames": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "A list of filenames that were created by the code",
                },
            },
            "required": ["explanation", "code"],
        },
    },
]


class OpenaiAgent(BaseAgent):

    def __init__(self, model, system_prompt, sdk, is_azure=True):
        super().__init__(sdk)
        self.model = model
        self.system_prompt = system_prompt
        logger.info(f"use openai model {model} is_azure {is_azure}")
        logger.info(f"use openai with system prompt {system_prompt}")
        self.is_azure = is_azure

    def _merge_delta_for_function_call(self, message, delta):
        if not delta:
            return
        if "function_call" not in message:
            message["function_call"] = delta["function_call"]
            return
        old_arguments = message["function_call"].get("arguments", "")
        if delta["function_call"]["arguments"]:
            message["function_call"]["arguments"] = (
                old_arguments + delta["function_call"]["arguments"]
            )

    def _merge_delta_for_content(self, message, delta):
        if not delta:
            return
        content = message.get("content", "")
        if delta.get("content"):
            message["content"] = content + delta["content"]

    def _get_function_call_argument_new_typing(self, message):
        if message["function_call"]["name"] == "python":
            return TypingState.CODE, "", message["function_call"].get("arguments", "")

        arguments = message["function_call"].get("arguments", "")
        state = TypingState.START
        explanation_str = ""
        code_str = ""
        for token_state, token in tokenize(io.StringIO(arguments)):
            if token_state == None:
                if state == TypingState.EXPLANATION and token[0] == 1:
                    explanation_str = token[1]
                    state = TypingState.START
                if state == TypingState.CODE and token[0] == 1:
                    code_str = token[1]
                    state = TypingState.START
                if token[1] == "explanation":
                    state = TypingState.EXPLANATION
                if token[1] == "code":
                    state = TypingState.CODE
            else:
                # String
                if token_state == 9 and state == TypingState.EXPLANATION:
                    explanation_str = "".join(token)
                elif token_state == 9 and state == TypingState.CODE:
                    code_str = "".join(token)
        return (state, explanation_str, code_str)

    async def call_openai(self, messages, queue, context, task_context):
        """
        call the openai api
        """
        sent_token_count = 0
        for message in messages:
            if not message["content"]:
                continue
            sent_token_count += len(encoding.encode(message["content"]))
        task_context.sent_token_count += sent_token_count

        start_time = time.time()
        if self.is_azure:
            response = await openai.ChatCompletion.acreate(
                engine=self.model,
                messages=messages,
                temperature=0,
                functions=OCTOGEN_FUNCTIONS,
                function_call="auto",
                stream=True,
            )
        else:
            response = await openai.ChatCompletion.acreate(
                model=self.model,
                messages=messages,
                temperature=0,
                functions=OCTOGEN_FUNCTIONS,
                function_call="auto",
                stream=True,
            )
        message = None
        text_content = ""
        code_content = ""
        async for chunk in response:
            if context.done():
                logger.debug("the client has cancelled the request")
                break
            if not chunk["choices"]:
                continue

            delta = chunk["choices"][0]["delta"]
            logger.debug(f"{delta}")
            if not message:
                message = delta
            else:
                if "function_call" in delta:
                    self._merge_delta_for_function_call(message, delta)
                    arguments = message["function_call"].get("arguments", "")
                    task_context.generated_token_count += len(
                        encoding.encode(arguments)
                    )
                    task_context.model_respond_duration += int(
                        (time.time() - start_time) * 1000
                    )
                    start_time = time.time()
                    (
                        state,
                        explanation_str,
                        code_str,
                    ) = self._get_function_call_argument_new_typing(message)
                    if explanation_str and text_content != explanation_str:
                        typed_chars = explanation_str[len(text_content) :]
                        text_content = explanation_str
                        await queue.put(
                            TaskRespond(
                                state=task_context.to_task_state_proto(),
                                respond_type=TaskRespond.OnAgentTextTyping,
                                typing_content=typed_chars,
                            )
                        )
                    if code_str and code_content != code_str:
                        typed_chars = code_str[len(code_content) :]
                        code_content = code_str
                        await queue.put(
                            TaskRespond(
                                state=task_context.to_task_state_proto(),
                                respond_type=TaskRespond.OnAgentCodeTyping,
                                typing_content=typed_chars,
                            )
                        )
                    logger.debug(
                        f"argument explanation:{explanation_str} code:{code_str}"
                    )
                else:
                    self._merge_delta_for_content(message, delta)
                    task_context.model_respond_duration += int(
                        (time.time() - start_time) * 1000
                    )
                    start_time = time.time()
                    if delta.get("content") != None:
                        task_context.generated_token_count += len(
                            encoding.encode(delta.get("content"))
                        )
                        await queue.put(
                            TaskRespond(
                                state=task_context.to_task_state_proto(),
                                respond_type=TaskRespond.OnAgentTextTyping,
                                typing_content=delta["content"],
                            )
                        )
        return message

    async def handle_function(self, message, queue, context, task_context):
        if "function_call" in message:
            if context.done():
                logging.debug("the client has cancelled the request")
                return
            function_name = message["function_call"]["name"]
            code = ""
            explanation = ""
            saved_filenames = []
            if function_name == "python":
                code = message["function_call"]["arguments"]
                logger.debug(f"call function {function_name} with args {code}")
            else:
                arguments = json.loads(message["function_call"]["arguments"])
                logger.debug(f"call function {function_name} with args {arguments}")
                code = arguments["code"]
                explanation = arguments["explanation"]
                saved_filenames = arguments.get("saved_filenames", [])
            tool_input = json.dumps({
                "code": code,
                "explanation": explanation,
                "saved_filenames": saved_filenames,
            })
            # send the respond to client
            await queue.put(
                TaskRespond(
                    state=task_context.to_task_state_proto(),
                    respond_type=TaskRespond.OnAgentActionType,
                    on_agent_action=OnAgentAction(
                        input=tool_input, tool="execute_python_code"
                    ),
                )
            )
            function_result = None
            async for (result, respond) in self.call_function(
                code, context, task_context
            ):
                if context.done():
                    logger.debug("the client has cancelled the request")
                    break
                function_result = result
                if respond:
                    await queue.put(respond)
            return function_result
        else:
            raise Exception("bad message, function message expected")

    async def arun(self, task, queue, context, max_iteration=5):
        """ """
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": task},
        ]
        iterations = 0
        task_context = TaskContext(
            start_time=time.time(),
            generated_token_count=0,
            sent_token_count=0,
            model_name="openai",
            iteration_count=0,
            model_respond_duration=0,
        )
        try:
            while iterations < max_iteration and not context.cancelled():
                iterations += 1
                task_context.iteration_count = iterations
                logger.debug(f" the input messages {messages}")
                chat_message = await self.call_openai(
                    messages, queue, context, task_context
                )
                logger.debug(f"the response {chat_message}")
                if "function_call" in chat_message:
                    if "content" not in chat_message:
                        chat_message["content"] = None
                    messages.append(chat_message)
                    function_name = chat_message["function_call"]["name"]
                    if function_name not in ["execute_python_code", "python"]:
                        messages.append({
                            "role": "function",
                            "name": "execute_python_code",
                            "content": "You can use the execute_python_code only",
                        })
                        continue
                    function_result = await self.handle_function(
                        chat_message, queue, context, task_context
                    )
                    await queue.put(
                        TaskRespond(
                            state=task_context.to_task_state_proto(),
                            respond_type=TaskRespond.OnAgentActionEndType,
                            on_agent_action_end=OnAgentActionEnd(
                                output="",
                                output_files=function_result.saved_filenames,
                                has_error=function_result.has_error,
                            ),
                        )
                    )
                    # TODO optimize the token limitation
                    if function_result.has_result:
                        messages.append({
                            "role": "function",
                            "name": "execute_python_code",
                            "content": function_result.console_stdout[0:500],
                        })
                    elif function_result.has_error:
                        messages.append({
                            "role": "function",
                            "name": "execute_python_code",
                            "content": function_result.console_stderr[0:500],
                        })
                    else:
                        messages.append({
                            "role": "function",
                            "name": "execute_python_code",
                            "content": function_result.console_stdout[0:500],
                        })
                else:
                    # end task
                    await queue.put(
                        TaskRespond(
                            state=task_context.to_task_state_proto(),
                            respond_type=TaskRespond.OnFinalAnswerType,
                            final_respond=FinalRespond(answer=chat_message["content"]),
                        )
                    )
                    break
        finally:
            await queue.put(None)

agent/src/og_agent/openai_agent.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


OCTOGEN_FUNCTION_SYSTEM = """Firstly,You are the Programming Copilot called **Octogen**, a large language model designed to complete any goal by **executing code**

Secondly, Being an expert in programming, you must follow the rules
* To complete the goal, You must write a plan and execute it step by step, the followings are examples
    * The data visualization plan involves previewing, cleaning, and processing the data to generate the chart.
* Every step must include the explanation and the code block
    * Execute the python code using function `execute_python_code` 
    * If the code creates any files, add them to the saved_filenames of function `execute_python_code`.
    * If the code has any display data, save it as a file and add it to the saved_filenames of function `execute_python_code`
* You must try to correct your code when you get errors from the output
* Your code should produce output in Markdown format. For instance, if you're using a Pandas DataFrame to display data, make sure to utilize the to_markdown function.
* You must preview one row of the data when using pandas to process data

Thirdly, the programming environment used to execute code has the following capabilities
* Internet connection: This allows the programming environment to access online resources, such as documentation, libraries, and code repositories.
* IPython kernel: This allows the programming environment to execute Python code
    * Filesystem: This allows the programming environment to open, write, and delete files in the workspace directory.
    * Lots of installed Python libraries:These includes the following popular libraries:
        * pandas:data analysis and manipulation tool
        * matplotlib:a comprehensive library for creating static, animated, and interactive visualizations in Python
        * yfinance:download market data from Yahoo!
        * imageio:library for reading and writing a wide range of image, video, scientific, and volumetric data formats.
        * pillow:Python Imaging Library
        * beautifulsoup4: a library that makes it easy to scrape information from web pages
        * requests: a simple, yet elegant, HTTP library
        * wikipedia: a Python library that makes it easy to access and parse data from Wikipedia
"""

OCTOGEN_CODELLAMA_SYSTEM = """Firstly,You are the Programming Copilot called **Octogen**, a large language model designed to complete any goal by **executing code**

Secondly, Being an expert in programming, you must follow the rules
* To achieve your goal, write a plan, execute it step-by-step, and set `is_final_answer` to `true` for the last step.
* Every step must include an action with the explanation, the code block
* Ensure that the output of action meets the goal before providing the final answer.
* Try a new step if the output does not meet the goal.
* Your code should produce output in Markdown format. For instance, if you're using a Pandas DataFrame to display data, make sure to utilize the to_markdown function.

Thirdly, the following actions are available:

* execute_python_code: This action executes Python code and returns the output. You must verify the output before giving the final answer.
* install_python_package: This action installs the Python package 
* show_sample_code: This action show the sample code for user. You must set the sample code to action_input
* no_action: This action does nothing.

Fourthly, the output format must be a JSON format with the following fields:
* explanation (string): The explanation about the action input
* action (string): The name of the action.
* action_input (string): The sample code , python code or python package to be executed for the action or an empty string if no action is specified
* saved_filenames (list of strings): A list of filenames that were created by the action input.
* language (string): The programming language used to execute the action.
* is_final_answer (boolean): Whether this is the final answer to the question. If it is, the value of this field should be true. Otherwise, the value should be false.
"""

agent/src/og_agent/prompt.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Taken from the NAYA project

https://github.com/danielyule/naya

Copyright (c) 2019 Daniel Yule
"""
import io
import unicodedata

SURROGATE = "Cs"


class TokenType:
    OPERATOR = 0
    STRING = 1
    NUMBER = 2
    BOOLEAN = 3
    NULL = 4


class State:
    WHITESPACE = 0
    INTEGER_0 = 1
    INTEGER_SIGN = 2
    INTEGER = 3
    INTEGER_EXP = 4
    INTEGER_EXP_0 = 5
    FLOATING_POINT_0 = 6
    FLOATING_POINT = 8
    STRING = 9
    STRING_ESCAPE = 10
    STRING_END = 11
    TRUE_1 = 12
    TRUE_2 = 13
    TRUE_3 = 14
    FALSE_1 = 15
    FALSE_2 = 16
    FALSE_3 = 17
    FALSE_4 = 18
    NULL_1 = 19
    NULL_2 = 20
    NULL_3 = 21
    UNICODE = 22
    UNICODE_SURROGATE_START = 23
    UNICODE_SURROGATE_STRING_ESCAPE = 24
    UNICODE_SURROGATE = 25


class SpecialChar:
    # Kind of a hack but simple: if we used the empty string "" to represent
    # EOF, expressions like `char in "0123456789"` would be true for EOF, which
    # is confusing. If we used a non-string, they would result in TypeErrors.
    # By using the string "EOF", they work as expected. The only thing we have
    # to be careful about is to not ever use "EOF" in any such strings used for
    # char membership checking, which we have no reason to do anyway.
    EOF = "EOF"


class UnCompletedException(Exception):

    def __init_(self, token):
        super().__init__("")
        self.token = token


def _guess_encoding(stream):
    # if it looks like a urllib response, get the charset from the headers (if any)
    try:
        encoding = stream.headers.get_content_charset()
    except:  # noqa
        encoding = None
    if encoding is None:
        # JSON is supposed to be UTF-8
        # https://tools.ietf.org/id/draft-ietf-json-rfc4627bis-09.html#:~:text=The%20default%20encoding%20is%20UTF,16%20and%20UTF%2D32).
        encoding = "utf-8"
    return encoding


def _ensure_text(stream):
    data = stream.read(0)
    if isinstance(data, bytes):
        encoding = _guess_encoding(stream)
        return io.TextIOWrapper(stream, encoding=encoding)
    return stream


def tokenize(stream):
    stream = _ensure_text(stream)

    def is_delimiter(char):
        return char.isspace() or char in "{}[]:," or char == SpecialChar.EOF

    token = []
    unicode_buffer = ""
    completed = False
    now_token = ""

    def process_char(char):
        nonlocal token, completed, now_token, unicode_buffer
        advance = True
        add_char = False
        next_state = state
        if state == State.WHITESPACE:
            if char == "{":
                completed = True
                now_token = (TokenType.OPERATOR, "{")
            elif char == "}":
                completed = True
                now_token = (TokenType.OPERATOR, "}")
            elif char == "[":
                completed = True
                now_token = (TokenType.OPERATOR, "[")
            elif char == "]":
                completed = True
                now_token = (TokenType.OPERATOR, "]")
            elif char == ",":
                completed = True
                now_token = (TokenType.OPERATOR, ",")
            elif char == ":":
                completed = True
                now_token = (TokenType.OPERATOR, ":")
            elif char == '"':
                next_state = State.STRING
            elif char in "123456789":
                next_state = State.INTEGER
                add_char = True
            elif char == "0":
                next_state = State.INTEGER_0
                add_char = True
            elif char == "-":
                next_state = State.INTEGER_SIGN
                add_char = True
            elif char == "f":
                next_state = State.FALSE_1
            elif char == "t":
                next_state = State.TRUE_1
            elif char == "n":
                next_state = State.NULL_1
            elif not char.isspace() and not char == SpecialChar.EOF:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.INTEGER:
            if char in "0123456789":
                add_char = True
            elif char == ".":
                next_state = State.FLOATING_POINT_0
                add_char = True
            elif char == "e" or char == "E":
                next_state = State.INTEGER_EXP_0
                add_char = True
            elif is_delimiter(char):
                next_state = State.WHITESPACE
                completed = True
                now_token = (TokenType.NUMBER, int("".join(token)))
                advance = False
            else:
                raise ValueError(
                    "A number must contain only digits.  Got '{}'".format(char)
                )
        elif state == State.INTEGER_0:
            if char == ".":
                next_state = State.FLOATING_POINT_0
                add_char = True
            elif char == "e" or char == "E":
                next_state = State.INTEGER_EXP_0
                add_char = True
            elif is_delimiter(char):
                next_state = State.WHITESPACE
                completed = True
                now_token = (TokenType.NUMBER, 0)
                advance = False
            else:
                raise ValueError(
                    "A 0 must be followed by a '.' or a 'e'.  Got '{0}'".format(char)
                )
        elif state == State.INTEGER_SIGN:
            if char == "0":
                next_state = State.INTEGER_0
                add_char = True
            elif char in "123456789":
                next_state = State.INTEGER
                add_char = True
            else:
                raise ValueError(
                    "A - must be followed by a digit.  Got '{0}'".format(char)
                )
        elif state == State.INTEGER_EXP_0:
            if char == "+" or char == "-" or char in "0123456789":
                next_state = State.INTEGER_EXP
                add_char = True
            else:
                raise ValueError(
                    "An e in a number must be followed by a '+', '-' or digit.  Got '{0}'".format(
                        char
                    )
                )
        elif state == State.INTEGER_EXP:
            if char in "0123456789":
                add_char = True
            elif is_delimiter(char):
                completed = True
                now_token = (TokenType.NUMBER, float("".join(token)))
                next_state = State.WHITESPACE
                advance = False
            else:
                raise ValueError(
                    "A number exponent must consist only of digits.  Got '{}'".format(
                        char
                    )
                )
        elif state == State.FLOATING_POINT:
            if char in "0123456789":
                add_char = True
            elif char == "e" or char == "E":
                next_state = State.INTEGER_EXP_0
                add_char = True
            elif is_delimiter(char):
                completed = True
                now_token = (TokenType.NUMBER, float("".join(token)))
                next_state = State.WHITESPACE
                advance = False
            else:
                raise ValueError("A number must include only digits")
        elif state == State.FLOATING_POINT_0:
            if char in "0123456789":
                next_state = State.FLOATING_POINT
                add_char = True
            else:
                raise ValueError(
                    "A number with a decimal point must be followed by a fractional part"
                )
        elif state == State.FALSE_1:
            if char == "a":
                next_state = State.FALSE_2
            else:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.FALSE_2:
            if char == "l":
                next_state = State.FALSE_3
            else:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.FALSE_3:
            if char == "s":
                next_state = State.FALSE_4
            else:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.FALSE_4:
            if char == "e":
                next_state = State.WHITESPACE
                completed = True
                now_token = (TokenType.BOOLEAN, False)
            else:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.TRUE_1:
            if char == "r":
                next_state = State.TRUE_2
            else:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.TRUE_2:
            if char == "u":
                next_state = State.TRUE_3
            else:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.TRUE_3:
            if char == "e":
                next_state = State.WHITESPACE
                completed = True
                now_token = (TokenType.BOOLEAN, True)
            else:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.NULL_1:
            if char == "u":
                next_state = State.NULL_2
            else:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.NULL_2:
            if char == "l":
                next_state = State.NULL_3
            else:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.NULL_3:
            if char == "l":
                next_state = State.WHITESPACE
                completed = True
                now_token = (TokenType.NULL, None)
            else:
                raise ValueError("Invalid JSON character: '{0}'".format(char))
        elif state == State.STRING:
            if char == '"':
                completed = True
                now_token = (TokenType.STRING, "".join(token))
                next_state = State.STRING_END
            elif char == "\\":
                next_state = State.STRING_ESCAPE
            elif char == SpecialChar.EOF:
                raise ValueError("Unterminated string at end of file")
            else:
                add_char = True
        elif state == State.STRING_END:
            if is_delimiter(char):
                advance = False
                next_state = State.WHITESPACE
            else:
                raise ValueError(
                    "Expected whitespace or an operator after string.  Got '{}'".format(
                        char
                    )
                )
        elif state == State.STRING_ESCAPE:
            next_state = State.STRING
            if char == "\\" or char == '"':
                add_char = True
            elif char == "b":
                char = "\b"
                add_char = True
            elif char == "f":
                char = "\f"
                add_char = True
            elif char == "n":
                char = "\n"
                add_char = True
            elif char == "t":
                char = "\t"
                add_char = True
            elif char == "r":
                char = "\r"
                add_char = True
            elif char == "/":
                char = "/"
                add_char = True
            elif char == "u":
                next_state = State.UNICODE
                unicode_buffer = ""
            else:
                raise ValueError("Invalid string escape: {}".format(char))
        elif state == State.UNICODE:
            if char == SpecialChar.EOF:
                raise ValueError("Unterminated unicode literal at end of file")
            unicode_buffer += char
            if len(unicode_buffer) == 4:
                try:
                    code_point = int(unicode_buffer, 16)
                except ValueError:
                    raise ValueError(f"Invalid unicode literal: \\u{unicode_buffer}")
                char = chr(code_point)
                if unicodedata.category(char) == SURROGATE:
                    next_state = State.UNICODE_SURROGATE_START
                else:
                    next_state = State.STRING
                    add_char = True
        elif state == State.UNICODE_SURROGATE_START:
            if char == "\\":
                next_state = State.UNICODE_SURROGATE_STRING_ESCAPE
            elif char == SpecialChar.EOF:
                raise ValueError("Unpaired UTF-16 surrogate at end of file")
            else:
                raise ValueError(f"Unpaired UTF-16 surrogate")

        elif state == State.UNICODE_SURROGATE_STRING_ESCAPE:
            if char == "u":
                next_state = State.UNICODE_SURROGATE
            elif char == SpecialChar.EOF:
                raise ValueError("Unpaired UTF-16 surrogate at end of file")
            else:
                raise ValueError(f"Unpaired UTF-16 surrogate")

        elif state == State.UNICODE_SURROGATE:
            if char == SpecialChar.EOF:
                raise ValueError("Unterminated unicode literal at end of file")
            unicode_buffer += char
            if len(unicode_buffer) == 8:
                code_point_1 = int(unicode_buffer[:4], 16)
                try:
                    code_point_2 = int(unicode_buffer[4:], 16)
                except ValueError:
                    raise ValueError(
                        f"Invalid unicode literal: \\u{unicode_buffer[4:]}"
                    )
                char = chr(code_point_2)
                if unicodedata.category(char) != SURROGATE:
                    raise ValueError(
                        f"Second half of UTF-16 surrogate pair is not a surrogate!"
                    )
                try:
                    pair = int.to_bytes(code_point_1, 2, "little") + int.to_bytes(
                        code_point_2, 2, "little"
                    )
                    char = pair.decode("utf-16-le")
                except ValueError:
                    raise ValueError(
                        f"Error decoding UTF-16 surrogate pair \\u{unicode_buffer[:4]}\\u{unicode_buffer[4:]}"
                    )
                next_state = State.STRING
                add_char = True

        if add_char:
            token.append(char)

        return advance, next_state

    state = State.WHITESPACE
    c = stream.read(1)
    index = 0
    while c:
        try:
            advance, state = process_char(c)
        except Exception as e:
            yield (state, token)
            break
        if completed:
            completed = False
            token = []
            yield (None, now_token)
        if advance:
            c = stream.read(1)
            index += 1
    try:
        process_char(SpecialChar.EOF)
        if completed:
            yield (None, now_token)
    except Exception as e:
        yield (state, token)

agent/src/og_agent/tokenizer.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
from setuptools import setup, find_packages

setup(
    name="og_chat",
    version="0.3.6",
    description="the chat client for open source code interpreter octogen",
    author="imotai",
    author_email="codego.me@gmail.com",
    url="https://github.com/dbpunk-labs/octogen",
    long_description=open("README.md").read(),
    long_description_content_type="text/markdown",
    packages=[
        "og_discord",
        "og_terminal",
    ],
    package_dir={
        "og_discord": "src/og_discord",
        "og_terminal": "src/og_terminal",
    },
    install_requires=[
        "og_sdk>=0.1.0",
        "rich>=13.5.2",
        "prompt_toolkit>=3.0.0",
        "click>=8.0.0",
        "discord.py>=2.3.2",
        "clipboard>=0.0.4",
        "term-image>=0.7.0",
        "python-dotenv",
    ],
    entry_points={
        "console_scripts": [
            "og = og_terminal.terminal_chat:app",
            "og_ping = og_terminal.ping:app",
            "og_discord_bot = og_discord.discord_chat:run_app",
        ]
    },
)

chat/setup.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """

import discord
import asyncio
import logging
import json
import sys
import os
import click
from datetime import datetime
from dotenv import dotenv_values
from og_proto import common_pb2
from og_sdk.agent_sdk import AgentSDK

LOG_LEVEL = logging.INFO
logging.basicConfig(
    level=LOG_LEVEL,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger(__name__)


class OctogenDiscordBot(discord.Client):

    def __init__(self, octogen_sdk, filedir, **kwargs):
        discord.Client.__init__(self, **kwargs)
        self.octogen_sdk = octogen_sdk
        self.filedir = filedir

    def handle_action_start(self, respond, saved_images):
        """Run on agent action."""
        segments = []
        if not respond.on_agent_action:
            return segments
        action = respond.on_agent_action
        if not action.input:
            return segments
        logger.info("handle action start return")
        arguments = json.loads(action.input)
        if action.tool == "execute_python_code" and action.input:
            explanation = arguments["explanation"]
            code = arguments["code"]
            saved_images.extend(arguments.get("saved_filenames", []))
            mk = f"""{explanation}\n
```python
{code}
```"""
            segments.append(mk)
        return segments

    def handle_final_answer(self, respond):
        segments = []
        if not respond.final_respond:
            return segments
        answer = respond.final_respond.answer
        if not answer:
            return segments
        state = "token:%s iteration:%s model:%s" % (
            respond.token_usage,
            respond.iteration,
            respond.model_name,
        )
        segments.append("%s\n%s" % (answer, state))
        return segments

    def handle_action_output(self, respond, saved_images):
        segments = []
        if not respond.on_agent_action_end:
            return segments
        mk = respond.on_agent_action_end.output
        if not mk:
            return segments
        saved_images.extend(respond.on_agent_action_end.output_files)
        segments.append(mk)
        return segments

    async def download_files(self, images):
        for image in images:
            await self.octogen_sdk.download_file(image, self.filedir)

    async def on_ready(self):
        logger.info(f"Logged in as {self.user} (ID: {self.user.id})")

    async def run_app(self, name, message):
        saved_images = []
        async for respond in self.octogen_sdk.run(name):
            if not respond:
                break
            if respond.on_agent_action_end:
                segments = self.handle_action_output(respond, saved_images)
                msg = "".join(segments)
                if msg:
                    await message.channel.send(msg)
            if respond.on_agent_action:
                segments = self.handle_action_start(respond, saved_images)
                msg = "".join(segments)
                if msg:
                    await message.channel.send(msg)
        saved_images = list(set(saved_images))
        if saved_images:
            await self.download_files(saved_images)
            for filename in saved_images:
                fullpath = "%s/%s" % (self.filedir, filename)
                await message.channel.send("", file=discord.File(fullpath))
                break

    async def show_apps(self):
        header = """Apps
"""
        rows = []
        apps = await self.octogen_sdk.query_apps()
        for index, app in enumerate(apps.apps):
            ctime = datetime.fromtimestamp(app.ctime).strftime("%m/%d/%Y")
            rows.append(f"{index+1}.{app.name}")
        table = header + "\n".join(rows)
        return table

    async def on_message(self, message):
        # we do not want the bot to reply to itself
        try:
            if message.author.id == self.user.id:
                return
            if message.content.find("/apps") >= 0:
                apps = await self.show_apps()
                await message.channel.send(apps)
                return
            content = message.content
            if content.find("/run") >= 0:
                name = content.split(" ")[1]
                await self.run_app(name, message)
                return
            await message.channel.send("working...")
            files = []
            for att in message.attachments:

                async def generate_chunk(att):
                    # TODO split
                    chunk = await att.read()
                    yield common_pb2.FileChunk(buffer=chunk, filename=att.filename)

                await sdk.upload_binary(generate_chunk(att), att.filename)
                files.append("uploaded " + att.filename)
            if files:
                prompt = message.content + "\n" + "\n".join(files)
            else:
                prompt = message.content
            try:
                async for respond in self.octogen_sdk.prompt(prompt):
                    if not respond:
                        break
                    logger.info(f"{respond}")
                    if respond.on_agent_action_end:
                        saved_images = []
                        segments = self.handle_action_output(respond, saved_images)
                        msg = "".join(segments)
                        logger.info(f"action output {msg}")
                        if msg:
                            if saved_images:
                                await self.download_files(saved_images)
                                for filename in saved_images:
                                    fullpath = "%s/%s" % (self.filedir, filename)
                                    await message.channel.send(
                                        msg, file=discord.File(fullpath)
                                    )
                                    break
                            else:
                                await message.channel.send(msg)
                    if respond.on_agent_action:
                        saved_images = []
                        segments = self.handle_action_start(respond, saved_images)
                        msg = "".join(segments)
                        logger.info(f"action start {msg}")
                        if msg:
                            await message.channel.send(msg)
                    if respond.final_respond:
                        segments = self.handle_final_answer(respond)
                        msg = "".join(segments)
                        logger.info(f"final answer {msg}")
                        if msg:
                            await message.channel.send(msg)
            except Exception as ex:
                logger.error(f"fail to get file {ex}")
                await message.channel.send("I am sorry for the internal error")
        except Exception as ex:
            logging.exception(ex)


async def app():
    octogen_discord_bot_dir = "~/.octogen_discord_bot"
    if octogen_discord_bot_dir.find("~") == 0:
        real_octogen_dir = octogen_discord_bot_dir.replace("~", os.path.expanduser("~"))
    else:
        real_octogen_dir = octogen_discord_bot_dir
    if not os.path.exists(real_octogen_dir):
        os.mkdir(real_octogen_dir)
    octogen_config = dotenv_values(real_octogen_dir + "/config")
    filedir = real_octogen_dir + "/data"
    if not os.path.exists(filedir):
        os.mkdir(filedir)
    sdk = AgentSDK(octogen_config["endpoint"], octogen_config["api_key"])
    sdk.connect()
    intents = discord.Intents.default()
    intents.message_content = True
    client = OctogenDiscordBot(sdk, filedir, intents=intents)
    await client.start(octogen_config["discord_bot_token"])


def run_app():
    asyncio.run(app())

chat/src/og_discord/discord_chat.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from rich.markdown import TextElement, Markdown
from rich.syntax import Syntax
from rich.console import Console, ConsoleOptions, RenderResult
from markdown_it.token import Token


class CodeBlock(TextElement):
    """A code block with syntax highlighting."""

    style_name = "markdown.code_block"

    @classmethod
    def create(cls, markdown: "Markdown", token: Token) -> "CodeBlock":
        node_info = token.info or ""
        lexer_name = node_info.partition(" ")[0]
        return cls(lexer_name or "default", markdown.code_theme)

    def __init__(self, lexer_name: str, theme: str) -> None:
        self.lexer_name = lexer_name
        self.theme = theme

    def __rich_console__(
        self, console: Console, options: ConsoleOptions
    ) -> RenderResult:
        code = str(self.text).rstrip()
        syntax = Syntax(
            code,
            self.lexer_name,
            # background_color="default",
            line_numbers=True,
            theme=self.theme,
            word_wrap=True,
            padding=1,
        )
        yield syntax

chat/src/og_terminal/markdown.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import sys
import os
import click
from og_sdk.agent_sdk import AgentSyncSDK
from rich.console import Console
from dotenv import dotenv_values


@click.command()
@click.option("--octogen_dir", default="~/.octogen", help="the root path of octogen")
def app(octogen_dir):
    console = Console()
    if octogen_dir.find("~") == 0:
        real_octogen_dir = octogen_dir.replace("~", os.path.expanduser("~"))
    else:
        real_octogen_dir = octogen_dir
    if not os.path.exists(real_octogen_dir):
        os.mkdir(real_octopus_dir)
    octogen_config = dotenv_values(real_octogen_dir + "/config")
    console = Console()
    try:
        if "api_key" not in octogen_config or "endpoint" not in octogen_config:
            console.print(
                f"❌ api key and endpoint are required! please check your config {octogen_dir}/config"
            )
            sys.exit(1)
        sdk = AgentSyncSDK(octogen_config["endpoint"], octogen_config["api_key"])
        sdk.connect()
        response = sdk.ping()
        if response.code == 0:
            console.print(f"👍 {response.msg}")
            sys.exit(0)
        else:
            console.print(f"❌ {response.msg}")
            sys.exit(1)
    except Exception as ex:
        console.print(
            f"❌ please check your config {octogen_dir}/config with error {ex}"
        )
        sys.exit(1)

chat/src/og_terminal/ping.py
dbpunk-labs/octogen
#! /usr/bin/env python3
# vim:fenc=utf-8
#

""" """
import sys
import os
import re
import time
import json
import asyncio
import click
import glob
import random
from datetime import datetime
from rich.console import Console
from rich.markdown import Markdown
from rich.syntax import Syntax
from rich.console import Group
from rich.panel import Panel
from rich.table import Table
from rich.text import Text
from rich.progress import Progress
from rich.rule import Rule
from rich.live import Live
from rich.spinner import Spinner
from rich.emoji import EMOJI
from rich import box
from rich.style import Style
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.shortcuts import CompleteStyle, clear
from prompt_toolkit.history import FileHistory
from prompt_toolkit import PromptSession
from og_sdk.agent_sdk import AgentSyncSDK
from og_sdk.utils import process_char_stream
from og_proto import agent_server_pb2
from dotenv import dotenv_values
from prompt_toolkit.completion import Completer, Completion
from .utils import parse_file_path
from .markdown import CodeBlock
import clipboard

Markdown.elements["fence"] = CodeBlock
Markdown.elements["code_block"] = CodeBlock
OCTOGEN_TITLE = "🐙[bold red]Octogen"
OCTOGEN_APP_TITLE = "🐙[bold red]App"

EMOJI_KEYS = list(EMOJI.keys())


def show_welcome(console):
    welcome = """
Welcome to use octogen❤️ . To ask a programming question, simply type your question and press [bold yellow]esc + enter[/]
Use [bold yellow]/help[/] for help
"""
    console.print(welcome)


def show_help(console):
    help = """
### Keyboard Shortcut:
- **`ESC + ENTER`**: Submit your question to Octogen or execute your command.

### Commands:

- **`/clear`**: Clears the screen.
- **`/cc{number}`**: Copies the output of Octogen to your clipboard.
- **`/exit`**: Exits the Octogen CLI.
- **`/up`**: Uploads files from your local machine; useful for including in your questions.
- **`/assemble {name} {number1} {number2}`**: Assembles the specified code segments into an application.
- **`/run {name}`**: Executes an application with the specified name.
- **`/apps`**: Displays a list of all your apps.

### Need Help?

1. Create an issue on our GitHub page: [Octogen GitHub Issues](https://github.com/dbpunk-labs/octogen/issues)
2. Alternatively, you can email us at [codego.me@gmail.com](mailto:codego.me@gmail.com).
"""
    mk = Markdown(help, justify="left")
    console.print(mk)


def gen_a_random_emoji():
    index = random.randint(0, len(EMOJI_KEYS) - 1)
    return EMOJI[EMOJI_KEYS[index]]


def parse_numbers(text):
    """Parses numbers from a string.

    Args:
    text: The string to parse.

    Returns:
    A list of numbers found in the string.
    """
    pattern = r"\d+\.\d+|\d+"
    numbers = re.findall(pattern, text)
    return numbers


class OctogenCompleter(Completer):

    def __init__(self, values):
        Completer.__init__(self)
        self.values = values

    def get_completions(self, document, complete_event):
        index = document.current_line_before_cursor.find("/up ")
        if index >= 0:
            word_before_cursor = document.current_line_before_cursor[
                index + 3 :
            ].strip()
            if word_before_cursor:
                for comp in glob.glob(word_before_cursor + "*"):
                    yield Completion(comp, -len(word_before_cursor))


def check_parameter(octopus_config, console):
    """
    check the parameter
    """
    if "api_key" not in octopus_config or not octopus_config["api_key"]:
        content = """No api key was found, you can get api key from the following ways
* You can use your [self-hosted](https://github.com/dbpunk-labs) agent server api key
* You can apply the api key from [octogen waitlist](https://octogen.dev) directly"""
        mk = Markdown(content)
        console.print(mk)
        return False
    return True


def clean_code(code: str):
    start_tag = "```python"
    end_tag = "```"
    index = code.find(start_tag)
    if index >= 0:
        last = code.rfind(end_tag)
        return code[index + len(start_tag) : last]
    return code


def refresh(live, segments, title=OCTOGEN_TITLE, task_state=None):
    speed = (
        task_state.generated_token_count / (task_state.model_respond_duration / 1000.0)
        if task_state
        else 0
    )
    table = Table.grid(padding=1, pad_edge=True)
    table.add_column("Index", no_wrap=True, justify="center", style="bold red")
    table.add_column("Status", no_wrap=True, justify="center", style="bold red")
    table.add_column("Content")
    for index, status, segment in segments:
        table.add_row(f"{index}", status, segment)
    live.update(
        Panel(
            table,
            title=title,
            title_align="left",
            subtitle="[gray] %.1ft/s %s" % (speed, task_state.model_name)
            if task_state
            else "",
            subtitle_align="left",
        )
    )
    live.refresh()


def handle_action_output(segments, respond, values):
    if respond.respond_type not in [
        agent_server_pb2.TaskRespond.OnAgentActionStdout,
        agent_server_pb2.TaskRespond.OnAgentActionStderr,
    ]:
        return
    value = values.pop()
    new_stdout = value[1][0]
    new_stderr = value[1][1]
    segment = segments.pop()
    if respond.respond_type == agent_server_pb2.TaskRespond.OnAgentActionStdout:
        new_stdout += respond.console_stdout
        new_stdout = process_char_stream(new_stdout)
    elif respond.respond_type == agent_server_pb2.TaskRespond.OnAgentActionStderr:
        new_stderr += respond.console_stderr
        new_stderr = process_char_stream(new_stderr)
    values.append(("text", (new_stdout, new_stderr), []))
    total_output = new_stdout
    if new_stderr:
        total_output = new_stdout + "\n" + new_stderr
    text = Text.from_ansi(total_output)
    syntax = Syntax(
        f"{text.plain}",
        "text",
        line_numbers=True,  # background_color="default"
    )
    segments.append((len(values) - 1, segment[1], syntax))


def handle_action_end(segments, respond, images, values):
    """
    Handles the end of an agent action.

    Args:
      segments: A list of segments in the current turn.
      respond: The response from the agent.
      images: A list of images to be displayed.
      values: A list of values to be copied

    Returns:
      None.
    """
    if respond.respond_type != agent_server_pb2.TaskRespond.OnAgentActionEndType:
        return
    output = respond.on_agent_action_end.output
    has_error = "✅" if not respond.on_agent_action_end.has_error else "❌"
    old_value = values.pop()
    segment = segments.pop()
    if images and not has_error:
        images.extend(respond.on_agent_action_end.output_files)
    values.append(old_value)
    segments.append((len(values) - 1, has_error, segment[2]))
    # add the next steps loading
    spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
    values.append(("text", "", []))
    segments.append((len(values) - 1, spinner, ""))


def handle_typing(segments, respond, values):
    if respond.respond_type not in [
        agent_server_pb2.TaskRespond.OnAgentTextTyping,
        agent_server_pb2.TaskRespond.OnAgentCodeTyping,
    ]:
        return
    value = values.pop()
    segment = segments.pop()
    if respond.respond_type == agent_server_pb2.TaskRespond.OnAgentTextTyping:
        new_value = value[1] + respond.typing_content
        values.append(("text", new_value, []))
        markdown = Markdown("\n" + new_value + "█")
        segments.append((len(values) - 1, segment[1], markdown))
    else:
        # Start write the code
        if value[0] == "text":
            values.append(value)
            markdown = Markdown("\n" + value[1])
            new_segment = (segment[0], "🧠", markdown)
            segments.append(new_segment)
            new_value = respond.typing_content
            values.append(("python", new_value, []))
            syntax = Syntax(
                new_value,
                "python",
                line_numbers=True,  # background_color="default"
            )
            segments.append((len(values) - 1, "📖", syntax))
        else:
            # continue
            new_value = value[1] + respond.typing_content
            values.append(("python", new_value, []))
            syntax = Syntax(
                new_value + "█",
                "python",
                line_numbers=True,  # background_color="default"
            )
            segments.append((len(values) - 1, "📖", syntax))


def handle_action_start(segments, respond, images, values):
    """Run on agent action."""
    if respond.respond_type != agent_server_pb2.TaskRespond.OnAgentActionType:
        return
    action = respond.on_agent_action
    if not action.input:
        return
    arguments = json.loads(action.input)
    value = values.pop()
    segment = segments.pop()
    if action.tool == "execute_python_code" or action.tool == "show_sample_code":
        images.extend(arguments.get("saved_filenames", []))
        if not value[1]:
            new_value = (
                "python",
                arguments["code"],
                arguments.get("saved_filenames", []),
            )
            values.append(new_value)
            syntax = Syntax(
                arguments["code"],
                "python",
                line_numbers=True,  # background_color="default"
            )
            new_segment = (segment[0], "📖", syntax)
            segments.append(new_segment)

        elif value[0] == "text":
            values.append(value)
            markdown = Markdown("\n" + value[1])
            new_segment = (segment[0], segment[1], markdown)
            segments.append(new_segment)
        elif value[0] == "python":
            values.append(value)
            syntax = Syntax(
                value[1],
                "python",
                line_numbers=True,  # background_color="default"
            )
            new_segment = (segment[0], segment[1], syntax)
            segments.append(new_segment)
        # Add spinner for console
        spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
        values.append(("text", ("", ""), []))
        syntax = Syntax(
            "",
            "text",
            line_numbers=True,  # background_color="default"
        )
        segments.append((len(values) - 1, spinner, syntax))


def find_code(content, segments, values):
    start_index = 0
    while start_index < len(content):
        first_pos = content.find("```", start_index)
        if first_pos >= 0:
            second_pos = content.find("```", first_pos + 1)
            if second_pos >= 0:
                sub_content = content[start_index:first_pos]
                values.append(("text", sub_content, []))
                segments.append((len(values) - 1, "🧠", Markdown(sub_content)))
                start_index = first_pos
                code_content = content[first_pos : second_pos + 3]
                clean_code_content = clean_code(code_content)
                # TODO parse language
                values.append(("python", clean_code_content, []))
                segments.append((len(values) - 1, "📖", Markdown(code_content)))
                start_index = second_pos + 3
        else:
            break
    if start_index < len(content):
        sub_content = content[start_index:]
        values.append(("text", sub_content, []))
        segments.append((len(values) - 1, "🧠", Markdown(sub_content)))


def handle_final_answer(segments, respond, values):
    if respond.respond_type != agent_server_pb2.TaskRespond.OnFinalAnswerType:
        return
    answer = respond.final_respond.answer
    values.pop()
    segments.pop()
    if not answer:
        return
    find_code(answer, segments, values)


def render_image(images, sdk, image_dir, console):
    try:
        from PIL import Image
        from term_image.image import AutoImage

        image_set = set(images)
        for image in image_set:
            try:
                sdk.download_file(image, image_dir)
                fullpath = "%s/%s" % (image_dir, image)
                pil_image = Image.open(fullpath)
                auto_image = AutoImage(
                    image=pil_image, width=int(pil_image.size[0] / 15)
                )
                print(f"{auto_image:1.1#}")
                return True
            except Exception as ex:
                return False
    except Exception as ex:
        return False


def run_chat(
    prompt, sdk, session, console, values, spinner_name="hearts", filedir=None
):
    """
    run the chat
    """
    segments = []
    images = []
    token_usage = 0
    iteration = 0
    model_name = ""
    with Live(Group(*segments), console=console) as live:
        spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
        values.append(("text", "", []))
        segments.append((len(values) - 1, spinner, ""))
        refresh(live, segments)
        task_state = None
        for respond in sdk.prompt(prompt):
            if not respond:
                break
            handle_typing(segments, respond, values)
            handle_action_start(segments, respond, images, values)
            handle_action_output(segments, respond, values)
            handle_action_end(segments, respond, images, values)
            handle_final_answer(segments, respond, values)
            task_state = respond.state
            refresh(live, segments, task_state=respond.state)
        refresh(live, segments, task_state=task_state)
    # display the images
    render_image(images, sdk, filedir, console)


def run_app(name, sdk, session, console, values, filedir=None):
    segments = []
    images = []
    with Live(Group(*segments), console=console) as live:
        spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
        values.append(("text", "", []))
        segments.append((len(values) - 1, spinner, ""))
        refresh(live, segments, title=OCTOGEN_APP_TITLE + f":{name}")
        for respond in sdk.run(name):
            if not respond:
                break
            handle_action_start(segments, respond, images, values)
            handle_action_output(segments, respond, values)
            handle_action_end(segments, respond, images, values)
            refresh(live, segments, title=OCTOGEN_APP_TITLE + f":{name}")
        values.pop()
        segments.pop()
        refresh(live, segments, title=OCTOGEN_APP_TITLE + f":{name}")
    # display the images
    render_image(images, sdk, filedir, console)


def gen_app_panel(app):
    desc = app.desc if app.desc else ""
    date_str = datetime.fromtimestamp(app.ctime).strftime("%m/%d/%Y")
    markdonw = f"""### {app.desc}  {app.name}
created at {date_str} with {app.language}"""
    style = Style(bgcolor="#2e2e2e")
    return Panel(Markdown(markdonw), box=box.SIMPLE, title_align="left", style=style)


def query_apps(sdk, console):
    table = Table.grid(padding=1, pad_edge=True)
    table.add_column("col1", no_wrap=True, justify="center")
    table.add_column("col2", no_wrap=True, justify="center")
    table.add_column("col3", no_wrap=True, justify="center")
    table.add_column("col3", no_wrap=True, justify="center")

    apps = sdk.query_apps()
    for index in range(0, len(apps.apps), 4):
        table.add_row(
            gen_app_panel(apps.apps[index]) if index < len(apps.apps) else "",
            gen_app_panel(apps.apps[index + 1]) if index + 1 < len(apps.apps) else "",
            gen_app_panel(apps.apps[index + 2]) if index + 2 < len(apps.apps) else "",
            gen_app_panel(apps.apps[index + 3]) if index + 3 < len(apps.apps) else "",
        )
    console.print(Panel(table, title=OCTOGEN_APP_TITLE, title_align="left"))


def assemble_app(sdk, name, numbers, values):
    code = []
    language = "python"
    saved_filenames = []
    for number in numbers:
        if values[number][0] == "text":
            continue
        if values[number][0] == "python":
            code.append(values[number][1])
            saved_filenames.extend(values[number][2])
            language = values[number][0]
    try:
        response = sdk.assemble(
            name,
            "\n".join(code),
            language,
            desc=gen_a_random_emoji(),
            saved_filenames=list(set(saved_filenames)),
        )
        if response.code == 0:
            return True, response.msg
        return False, response.msg
    except Exception as ex:
        return False, str(ex)


@click.command()
@click.option("--octogen_dir", default="~/.octogen", help="the root path of octogen")
def app(octogen_dir):
    console = Console()
    if octogen_dir.find("~") == 0:
        real_octogen_dir = octogen_dir.replace("~", os.path.expanduser("~"))
    else:
        real_octogen_dir = octogen_dir
    os.makedirs(real_octogen_dir, exist_ok=True)
    octopus_config = dotenv_values(real_octogen_dir + "/config")
    if not check_parameter(octopus_config, console):
        return
    filedir = real_octogen_dir + "/data"
    os.makedirs(filedir, exist_ok=True)
    sdk = AgentSyncSDK(octopus_config["endpoint"], octopus_config["api_key"])
    sdk.connect()
    history = FileHistory(real_octogen_dir + "/history")
    values = []
    completer = OctogenCompleter(values)
    session = PromptSession(
        history=history,
        completer=completer,
        complete_in_thread=True,
        complete_while_typing=True,
        complete_style=CompleteStyle.MULTI_COLUMN,
    )
    index = 0
    show_welcome(console)
    while True:
        index = index + 1
        real_prompt = session.prompt("[%s]%s>" % (index, "🎧"), multiline=True)
        if not "".join(real_prompt.strip().split("\n")):
            continue
        if real_prompt.find("/help") >= 0:
            show_help(console)
            continue
        if real_prompt.find("/exit") >= 0:
            console.print("👋👋!")
            return
        if real_prompt.find("/clear") >= 0:
            clear()
            continue
        if real_prompt.find("/cc") >= 0:
            # handle copy
            for number in parse_numbers(real_prompt):
                num = int(number)
                if num < len(values):
                    clipboard.copy(values[num][1])
                    console.print(f"👍 /cc{number} has been copied to clipboard!")
                    break
                else:
                    console.print(f"❌ /cc{number} was not found!")
            continue
        if real_prompt.find("/assemble") >= 0:
            parts = real_prompt.split(" ")
            if len(parts) < 3:
                console.print(f"❌ please add at least on code segment number")
                continue
            try:
                name = parts[1]
                numbers = [int(number) for number in parts[2:]]
                (status, msg) = assemble_app(sdk, name, numbers, values)
                if status:
                    console.print(
                        f"👍 the app {name} has been assembled! use /run {name} to run this app."
                    )
                    continue
                else:
                    console.print(f"❌ fail to assemble the app {name} for error {msg}")
                    continue
            except Exception as ex:
                console.print(f"❌ invalid numbers {ex}")
                continue
        if real_prompt.find("/run") >= 0:
            parts = real_prompt.split(" ")
            if len(parts) < 2:
                console.print(f"❌ please specify the name of app")
                continue

            try:
                name = parts[1]
                run_app(name, sdk, session, console, values, filedir=filedir)
                continue
            except Exception as ex:
                console.print(f"❌ exception {ex}")
                continue
        if real_prompt.find("/apps") >= 0:
            query_apps(sdk, console)
            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)
        run_chat(
            real_prompt,
            sdk,
            session,
            console,
            values,
            spinner_name=octopus_config.get("spinner", "dots2"),
            filedir=filedir,
        )

chat/src/og_terminal/terminal_chat.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """


def parse_file_path(real_prompt):
    """
    parse the file path from the prompt
    """
    filepaths = []
    position = 0
    while position < len(real_prompt):
        first_pos = real_prompt.find("/up", position)
        # break the loop if no file to upload
        if first_pos == -1 or len(real_prompt) - first_pos <= 4:
            break
        #
        if real_prompt[first_pos + 3] != " ":
            position = first_pos + 4
            continue
        end_pos = real_prompt.find("\n", first_pos + 4)
        end_pos = end_pos if end_pos >= 0 else len(real_prompt)
        blank_pos = real_prompt.find(" ", first_pos + 4, end_pos)
        if blank_pos == -1:
            filepath = real_prompt[first_pos + 4 : end_pos]
            position = len(real_prompt)
        else:
            filepath = real_prompt[first_pos + 4 : blank_pos]
            position = blank_pos
        if filepath:
            filepaths.append(filepath)
    return filepaths

chat/src/og_terminal/utils.py
dbpunk-labs/octogen
#! /usr/bin/env python3
# vim:fenc=utf-8
#
# Copyright © 2023 imotai <imotai@imotai-ub>
#
# Distributed under terms of the MIT license.

""" """

from og_terminal.terminal_chat import gen_a_random_emoji
from og_terminal.terminal_chat import parse_numbers
from og_terminal.terminal_chat import handle_action_end
from og_proto import agent_server_pb2


def test_gen_a_random_emoji():
    assert gen_a_random_emoji()


def test_parse_number():
    test_text = "/cc0"
    numbers = parse_numbers(test_text)
    assert numbers
    assert numbers[0] == "0"


def test_ok_handle_action_end():
    segments = [(0, "", "")]
    images = []
    values = [()]
    task_state = agent_server_pb2.TaskState(
        generated_token_count=10,
        iteration_count=1,
        model_name="mock",
        total_duration=1,
        sent_token_count=10,
        model_respond_duration=1000,
    )
    respond = agent_server_pb2.TaskRespond(
        state=task_state,
        respond_type=agent_server_pb2.TaskRespond.OnAgentActionEndType,
        on_agent_action_end=agent_server_pb2.OnAgentActionEnd(
            output="", output_files=[], has_error=False
        ),
    )
    handle_action_end(segments, respond, images, values)
    assert segments[0][1] == "✅"


def test_error_handle_action_end():
    segments = [(0, "", "")]
    images = []
    values = [()]
    task_state = agent_server_pb2.TaskState(
        generated_token_count=10,
        iteration_count=1,
        model_name="mock",
        total_duration=1,
        sent_token_count=10,
        model_respond_duration=1000,
    )
    respond = agent_server_pb2.TaskRespond(
        state=task_state,
        respond_type=agent_server_pb2.TaskRespond.OnAgentActionEndType,
        on_agent_action_end=agent_server_pb2.OnAgentActionEnd(
            output="", output_files=[], has_error=True
        ),
    )
    handle_action_end(segments, respond, images, values)
    assert segments[0][1] == "❌"

chat/tests/test_chat_function.py
dbpunk-labs/octogen
#! /usr/bin/env python3
# vim:fenc=utf-8
#
# Copyright © 2023 imotai <imotai@imotai-ub>
#
# Distributed under terms of the MIT license.

""" """

from og_terminal.utils import parse_file_path


def test_parse_file_path():
    prompt = "convert the file /up /home/test.pdf to text"
    paths = parse_file_path(prompt)
    assert len(paths) == 1, "bad file path count"
    assert paths[0] == "/home/test.pdf", "bad file path "

chat/tests/test_parse_files.py
dbpunk-labs/octogen
# Configuration file for the Sphinx documentation builder.
#
# For the full list of built-in configuration values, see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information

project = 'Octogen'
copyright = '2023 octogem.dev'
author = 'imotai'

# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

extensions = []

templates_path = ['_templates']
exclude_patterns = []



# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output

html_theme = "alabaster"
html_static_path = ['_static']

docs/source/conf.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import matplotlib.pyplot as plt
import numpy as np

# Create a pie chart
data = np.array([10, 20, 30, 40])
labels = ["Category 1", "Category 2", "Category 3", "Category 4"]

plt.pie(data, labels=labels, autopct="%1.1f%%")
plt.title("Pie Chart")
plt.show()

kernel/cases/pie.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
from setuptools import setup, find_packages

setup(
    name="og_kernel",
    version="0.3.6",
    description="Open source code interpreter agent for LLM",
    author="imotai",
    author_email="codego.me@gmail.com",
    url="https://github.com/dbpunk-labs/octogen",
    long_description=open("README.md").read(),
    long_description_content_type="text/markdown",
    packages=[
        "og_kernel",
        "og_kernel.kernel",
        "og_kernel.server",
    ],
    package_dir={
        "og_kernel": "src/og_kernel",
        "og_kernel.kernel": "src/og_kernel/kernel",
        "og_kernel.server": "src/og_kernel/server",
    },
    install_requires=[
        "og_proto",
        "grpc-google-iam-v1>=0.12.6",
        "grpcio-tools>=1.57.0",
        "ipykernel>=6.25.1",
        "jupyter_client>=8.3.0",
        "matplotlib>=3.7.2",
        "pandas",
        "numpy",
    ],
    entry_points={
        "console_scripts": [
            "og_kernel_rpc_server = og_kernel.server.kernel_rpc_server:server_main",
            "og_kernel_app = og_kernel.kernel.kernel_app:run_app",
        ]
    },
)

kernel/setup.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dotenv import dotenv_values

config = dotenv_values(".env")

kernel/src/og_kernel/kernel/config.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
"""An application to launch a kernel by name in a local subprocess."""
import os
import signal
import uuid

from jupyter_core.application import JupyterApp, base_flags
from tornado.ioloop import IOLoop
from traitlets import Unicode
from jupyter_client.kernelspec import NATIVE_KERNEL_NAME, KernelSpecManager
from jupyter_client.manager import KernelManager


class KernelApp(JupyterApp):
    """Launch a kernel by name in a local subprocess."""

    description = "Run a kernel locally in a subprocess"
    classes = [KernelManager, KernelSpecManager]
    aliases = {
        "kernel": "KernelApp.kernel_name",
        "ip": "KernelManager.ip",
        "connection_file": "KernelApp.connection_file",
    }
    flags = {"debug": base_flags["debug"]}
    kernel_name = Unicode(
        NATIVE_KERNEL_NAME, help="The name of a kernel type to start"
    ).tag(config=True)
    connection_file = Unicode("", help="The connection file path of the kernel").tag(
        config=True
    )

    def initialize(self, argv=None):
        """Initialize the application."""
        super().initialize(argv)
        cf_basename = (
            self.connection_file
            if self.connection_file
            else "kernel-%s.json" % uuid.uuid4()
        )
        self.config.setdefault("KernelManager", {}).setdefault(
            "connection_file", os.path.join(self.runtime_dir, cf_basename)
        )
        self.km = KernelManager(kernel_name=self.kernel_name, config=self.config)

        self.loop = IOLoop.current()
        self.loop.add_callback(self._record_started)

    def setup_signals(self) -> None:
        """Shutdown on SIGTERM or SIGINT (Ctrl-C)"""
        if os.name == "nt":
            return

        def shutdown_handler(signo, frame):
            self.loop.add_callback_from_signal(self.shutdown, signo)

        for sig in [signal.SIGTERM, signal.SIGINT]:
            signal.signal(sig, shutdown_handler)

    def shutdown(self, signo: int) -> None:
        """Shut down the application."""
        self.log.info("Shutting down on signal %d", signo)
        self.km.shutdown_kernel()
        self.loop.stop()

    def log_connection_info(self) -> None:
        """Log the connection info for the kernel."""
        cf = self.km.connection_file
        self.log.info("Connection file: %s", cf)
        self.log.info("To connect a client: --existing %s", os.path.basename(cf))

    def _record_started(self) -> None:
        """For tests, create a file to indicate that we've started

        Do not rely on this except in our own tests!
        """
        fn = os.environ.get("JUPYTER_CLIENT_TEST_RECORD_STARTUP_PRIVATE")
        if fn is not None:
            with open(fn, "wb"):
                pass

    def start(self) -> None:
        """Start the application."""
        self.log.info("Starting kernel %r", self.kernel_name)
        try:
            self.km.start_kernel()
            self.log_connection_info()
            self.setup_signals()
            self.loop.start()
        finally:
            self.km.cleanup_resources()


def run_app():
    KernelApp.launch_instance()

kernel/src/og_kernel/kernel/kernel_app.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import logging
import inspect
import asyncio
import queue
import json
from jupyter_client import AsyncKernelClient

logger = logging.getLogger(__name__)

"""
the kernel client for watching the output of kernel
"""


class KernelClient:

    def __init__(self, connection_file):
        if not connection_file:
            raise ValueError(f"connection_file={connection_file} is empty")
        if not os.path.exists(connection_file):
            raise ValueError(f"connection_file={connection_file} do not exist")
        logger.info(
            "create a new kernel client with connection_file %s", connection_file
        )
        self.client = None
        self.is_running = False
        self.task = None
        self.connection_file = connection_file

    async def is_alive(self):
        return await self.client.is_alive()

    async def start_client(self):
        self.client = AsyncKernelClient(connection_file=self.connection_file)
        self.client.load_connection_file()
        self.client.start_channels()
        await self.client.wait_for_ready()

    async def _loop(self, on_message_fn):
        logger.debug("start loop the kernel message")
        try:
            while self.is_running and self.client:
                try:
                    logger.debug("start wait message")
                    msg = await self.client.get_iopub_msg(timeout=1)
                    logger.debug("msg %s", msg)
                    try:
                        await on_message_fn(msg)
                    except Exception as e:
                        logger.error("fail to call on message function for error %s", e)
                        continue
                except queue.Empty:
                    logger.debug("empty message")
                    continue
                except (ValueError, IndexError):
                    # get_iopub_msg suffers from message fetch errors
                    logger.error("fail to get message")
                    break
                except Exception as e:
                    logger.error("fail to wait for message %s", e)
                    break
        except Exception as e:
            logger.error("loop exception", e)

    async def watching(self, on_message_fn):
        """
        Watch the message from kernel, when a new message arrived , the `on_message_fn` will be
        called

        Arguments
        on_message_fn - when a new message arrived the function will be called
        """
        if self.is_running:
            raise ValueError(f"the watch is running, do not watch it again")
        if not on_message_fn or not self.client:
            raise ValueError(f"on_message_fn or clent is None")
        if not inspect.iscoroutinefunction(on_message_fn):
            raise ValueError(f"on_message_fn must be async function")
        self.is_running = True
        self.task = asyncio.create_task(self._loop(on_message_fn))

    async def read_response(self, context, tries=1):
        try:
            hit_empty = 0
            while self.client:
                try:
                    msg = await self.client.get_iopub_msg(timeout=1)
                    if context.done():
                        logger.debug("the client  has cancelled the request")
                        break
                    logger.debug(f"{msg}")
                    yield msg
                except queue.Empty:
                    hit_empty += 1
                    if hit_empty >= tries:
                        break
                except (ValueError, IndexError):
                    # get_iopub_msg suffers from message fetch errors
                    logger.error("fail to get message")
                    break
                except Exception as e:
                    logger.error("fail to wait for message %s", e)
                    break
            yield None
        except Exception as e:
            logger.error("loop exception", e)
            yield None

    def execute(self, code):
        """
        Execute the python code
        """
        if not self.client:
            raise ValueError(f"no client is avaliable")
        msg_id = self.client.execute(code)
        logger.debug("the execute msg id %s", msg_id)
        return msg_id

    async def stop_watch(self):
        if self.task and self.is_running:
            self.is_running = False
            logger.info(
                "stop the kernel client for connection_file %s", self.connection_file
            )
            try:
                self.task.cancel()
                await self.task
            except:
                pass

    def stop_client(self):
        if self.client:
            self.client.stop_channels()
            self.client = None

kernel/src/og_kernel/kernel/kernel_client.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import logging
import subprocess
import json
import sys
import pathlib
from time import sleep

logger = logging.getLogger(__name__)

"""
The jupytor kernel manager used for
- create kernel instance
- start kernel instance
- stop kernel instance

Typical usage example:
    config_path = "kernel_connection_file.json"
    workspace = "/mnt/workspace1"
    km = KernelManager(config_path, workspace)
    # start the kernel
    km.start()
"""


class KernelManager:

    def __init__(self, config_path: str, workspace: str, kernel: str = "python3"):
        if not config_path or not workspace:
            raise ValueError(
                f"config path={config_path} or workspace={workspace} is empty"
            )
        self.config_path = config_path
        self.workspace = workspace
        self.process = None
        self.is_running = False
        logger.info(
            "new kernel manager with config path %s and worksapce %s",
            config_path,
            workspace,
        )
        self.kernel = kernel

    def start(self):
        """
        Start a kernel instance and generate the kernel connection file
        """
        self.is_running = True

        os.makedirs(self.workspace, exist_ok=True)
        launch_kernel_script_path = os.path.join(
            pathlib.Path(__file__).parent.resolve(), "launch_kernel.py"
        )
        self.process = subprocess.Popen(
            [
                sys.executable,
                launch_kernel_script_path,
                "--connection_file=" + self.config_path,
                "--kernel=" + self.kernel,
            ],
            cwd=self.workspace,
        )
        logger.info("Start the kernel with process id %s", str(self.process.pid))
        while True:
            if not os.path.isfile(self.config_path):
                sleep(1)
            else:
                try:
                    with open(self.config_path, "r") as fp:
                        logger.info("connection file content %s", json.load(fp))
                    break
                except json.JSONDecodeError:
                    pass

    def stop(self):
        """
        stop the kernel instance
        """
        self.is_running = False
        logger.info("stop the kernel with process id %s", str(self.process.pid))
        if self.process:
            self.process.terminate()
            self.process.wait()
            self.process = None

    def __str__(self):
        return f'KernelManager(config_path="{self.config_path}", workspace="{self.workspace}")'

kernel/src/og_kernel/kernel/kernel_mgr.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
from og_kernel.kernel.kernel_app import run_app

if __name__ == "__main__":
    run_app()

kernel/src/og_kernel/kernel/launch_kernel.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
the websocket server for the kernel
"""
import os
import json
import sys
import logging
import asyncio
import functools
import grpc
import re
import uuid
import random
import string
import base64
import tempfile
from pathlib import Path
from typing import Awaitable, Callable, Optional, AsyncIterable
from grpc.aio import ServicerContext, server, ServerInterceptor
from google.rpc import status_pb2
from dotenv import dotenv_values
from ..kernel.kernel_mgr import KernelManager
from og_proto.kernel_server_pb2_grpc import KernelServerNodeServicer
from og_proto.kernel_server_pb2_grpc import add_KernelServerNodeServicer_to_server
from og_proto import kernel_server_pb2
from og_proto import common_pb2
from ..kernel.kernel_client import KernelClient
import aiofiles
from aiofiles import os as aio_os

config = dotenv_values(".env")

LOG_LEVEL = (
    logging.DEBUG if config.get("log_level", "info") == "debug" else logging.INFO
)

logging.basicConfig(
    level=LOG_LEVEL,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger(__name__)
ansi_escape = re.compile(r"(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]")


def _unary_unary_rpc_terminator(code, details):
    async def terminate(ignored_request, context):
        await context.abort(code, details)

    return grpc.unary_unary_rpc_method_handler(terminate)


class ApiKeyInterceptor(ServerInterceptor):
    """
    the api key interceptor
    """

    def __init__(self, header, value, code, error):
        self._header = header
        self._value = value
        self._terminator = _unary_unary_rpc_terminator(code, error)

    async def intercept_service(
        self,
        continuation: Callable[
            [grpc.HandlerCallDetails], Awaitable[grpc.RpcMethodHandler]
        ],
        handler_call_details: grpc.HandlerCallDetails,
    ) -> grpc.RpcMethodHandler:
        if (
            self._header,
            self._value,
        ) in handler_call_details.invocation_metadata:
            return await continuation(handler_call_details)
        else:
            return self._terminator


class KernelRpcServer(KernelServerNodeServicer):

    def __init__(self):
        self.kms = {}
        self.kcs = {}
        self.auth_failed_status = status_pb2.Status(
            code=grpc.StatusCode.INVALID_ARGUMENT.value[0],
            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
    ) -> kernel_server_pb2.StopKernelResponse:
        """
        Stop the kernel
        """

        kernel_name = request.kernel_name if request.kernel_name else "python3"
        if kernel_name not in self.kms or not self.kms[kernel_name]:
            logger.warning("no started kernel")
            return kernel_server_pb2.StopKernelResponse(
                key="k", code=1, msg="no started kernel"
            )
        self.kms[kernel_name].stop()
        self.kcs[kernel_name].stop_client()
        self.kms[kernel_name] = None
        self.kcs[kernel_name] = None
        return kernel_server_pb2.StopKernelResponse(code=0, msg="ok")

    async def get_status(
        self, request: kernel_server_pb2.GetStatusRequest, context: ServicerContext
    ) -> kernel_server_pb2.GetStatusResponse:
        kernel_name = request.kernel_name if request.kernel_name else "python3"
        logger.debug("check the kernel %s status", kernel_name)
        if kernel_name not in self.kms or not self.kms[kernel_name]:
            return kernel_server_pb2.GetStatusResponse(is_alive=False, code=0, msg="ok")
        is_alive = await self.kcs[kernel_name].is_alive()
        return kernel_server_pb2.GetStatusResponse(is_alive=is_alive, code=0, msg="ok")

    async def start(
        self, request: kernel_server_pb2.StartKernelRequest, context: ServicerContext
    ) -> kernel_server_pb2.StartKernelResponse:
        """
        Start the kernel
        """
        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 request will be ignored for that the kernel has been started"
            )
            return kernel_server_pb2.StartKernelResponse(
                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" % (
            config["config_root_path"],
            uuid.uuid4(),
        )
        km = KernelManager(connection_file, config["workspace"], kernel=kernel_name)
        km.start()
        kc = KernelClient(connection_file)
        await kc.start_client()
        self.kms[kernel_name] = km
        self.kcs[kernel_name] = kc
        return kernel_server_pb2.StartKernelResponse(code=0, msg="ok")

    async def download(
        self, request: common_pb2.DownloadRequest, context: ServicerContext
    ) -> AsyncIterable[common_pb2.FileChunk]:
        """
        download file
        """
        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:
            while True:
                chunk = await afp.read(1024 * 128)
                if not chunk:
                    break
                yield common_pb2.FileChunk(buffer=chunk, filename=request.filename)

    async def upload(
        self,
        request: AsyncIterable[common_pb2.FileChunk],
        context: ServicerContext,
    ) -> common_pb2.FileUploaded:
        """
        upload file
        """
        temp_dir = tempfile.mkdtemp(prefix="octogen")
        filename = "".join(random.choices(string.ascii_lowercase, k=16))
        tmp_full_path = f"{temp_dir}{os.sep}{filename}"
        target_filename = None
        logger.info(f"upload file to temp file {tmp_full_path}")
        length = 0
        async with aiofiles.open(tmp_full_path, "wb") as afp:
            async for chunk in request:
                length = length + await afp.write(chunk.buffer)
                logger.debug(f"write the {tmp_full_path} with {length}")
                if not target_filename:
                    target_filename = "%s/%s" % (config["workspace"], chunk.filename)
        if length != 0:
            logging.info(f"move file from {tmp_full_path} to  {target_filename}")
            await aio_os.rename(tmp_full_path, target_filename)
        else:
            logging.warning("empty file")
        await aio_os.rmdir(temp_dir)
        return common_pb2.FileUploaded(length=length)

    async def execute(
        self, request: kernel_server_pb2.ExecuteRequest, context: ServicerContext
    ) -> AsyncIterable[kernel_server_pb2.ExecuteResponse]:
        """
        Execute the python code and return a stream response
        """
        kernel_name = request.kernel_name if request.kernel_name else "python3"
        if not request.code:
            raise grpc.RpcError(grpc.StatusCode.INVALID_ARGUMENT, "Invalid argument")
        if (
            kernel_name not in self.kms
            or not self.kms[kernel_name]
            or not self.kcs[kernel_name]
        ):
            logger.warning(
                "no started kernel for executing code for kernel name %s" % kernel_name
            )
            raise grpc.RpcError(grpc.StatusCode.INVALID_ARGUMENT, "Invalid argument")
        logger.debug("the code %s with kernel %s", request.code, kernel_name)
        # TODO check the busy status
        msg_id = self.kcs[kernel_name].execute(request.code)
        async for msg in self.kcs[kernel_name].read_response(context, 5):
            try:
                if context.cancelled() or not msg:
                    break
                if msg["parent_header"]["msg_id"] != msg_id:
                    continue
                if msg["msg_type"] in ["status", "execute_input"]:
                    continue
                respond = self._build_payload(msg, config["workspace"])
                yield respond
            except Exception as ex:
                logger.exception("fail to handle the result")

    def _build_payload(self, msg, workspace) -> kernel_server_pb2.ExecuteResponse:
        if msg["msg_type"] == "display_data":
            if "image/png" in msg["content"]["data"]:
                filename = "octopus_%s.png" % uuid.uuid4().hex
                fullpath = "%s/%s" % (workspace, filename)
                with open(fullpath, "wb+") as fd:
                    data = msg["content"]["data"]["image/png"].encode("ascii")
                    buffer = base64.b64decode(data)
                    fd.write(buffer)
                return kernel_server_pb2.ExecuteResponse(
                    output_type=kernel_server_pb2.ExecuteResponse.ResultType,
                    output=json.dumps({"image/png": filename}),
                )
            elif "image/gif" in msg["content"]["data"]:
                filename = "octopus_%s.gif" % uuid.uuid4().hex
                fullpath = "%s/%s" % (workspace, filename)
                with open(fullpath, "wb+") as fd:
                    data = msg["content"]["data"]["image/gif"].encode("ascii")
                    buffer = base64.b64decode(data)
                    fd.write(buffer)
                return kernel_server_pb2.ExecuteResponse(
                    output_type=kernel_server_pb2.ExecuteResponse.ResultType,
                    output=json.dumps({"image/gif": filename}),
                )
            elif "text/plain" in msg["content"]["data"]:
                return kernel_server_pb2.ExecuteResponse(
                    output_type=kernel_server_pb2.ExecuteResponse.ResultType,
                    output=json.dumps(
                        {"text/plain": msg["content"]["data"]["text/plain"]}
                    ),
                )
            else:
                logger.warning(f" unsupported display_data {msg}")
                return kernel_server_pb2.ExecuteResponse(
                    output_type=kernel_server_pb2.ExecuteResponse.ResultType,
                    output=json.dumps({}),
                )
                # keys = ",".join(msg["content"]["data"].keys())
                # raise Exception(
                #    f"unsupported display data type {keys} for the result {msg}"
                # )

        if msg["msg_type"] == "execute_result":
            logger.debug("result data %s", msg["content"]["data"]["text/plain"])
            return kernel_server_pb2.ExecuteResponse(
                output_type=kernel_server_pb2.ExecuteResponse.ResultType,
                output=json.dumps({"text/plain": msg["content"]["data"]["text/plain"]}),
            )
        elif msg["msg_type"] == "stream":
            if msg["content"]["name"] == "stdout":
                return kernel_server_pb2.ExecuteResponse(
                    output_type=kernel_server_pb2.ExecuteResponse.StdoutType,
                    output=json.dumps({"text": msg["content"]["text"]}),
                )
            else:
                return kernel_server_pb2.ExecuteResponse(
                    output_type=kernel_server_pb2.ExecuteResponse.StderrType,
                    output=json.dumps({"text": msg["content"]["text"]}),
                )
        elif msg["msg_type"] == "error":
            if len(msg["content"]["traceback"]) > 6:
                traceback = "\n".join(msg["content"]["traceback"][:3])
                traceback = traceback + "\n".join(msg["content"]["traceback"][-3:])
            else:
                traceback = "\n".join(msg["content"]["traceback"])
            return kernel_server_pb2.ExecuteResponse(
                output_type=kernel_server_pb2.ExecuteResponse.TracebackType,
                output=json.dumps({"traceback": ansi_escape.sub("", traceback)}),
            )
        raise Exception(f"unsupported msg type {msg}")


async def serve() -> None:
    logger.info(
        "start kernel rpc server with host %s and port %s",
        config["rpc_host"],
        config["rpc_port"],
    )
    interceptors = [
        ApiKeyInterceptor(
            "api_key",
            config["rpc_key"],
            grpc.StatusCode.ABORTED.value[0],
            "api key is required or invalid",
        )
    ]
    serv = server(interceptors=interceptors)
    add_KernelServerNodeServicer_to_server(KernelRpcServer(), serv)
    listen_addr = "%s:%s" % (config["rpc_host"], config["rpc_port"])
    serv.add_insecure_port(listen_addr)
    await serv.start()
    await serv.wait_for_termination()


def server_main():
    asyncio.run(serve())

kernel/src/og_kernel/server/kernel_rpc_server.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import json
import os
import asyncio
import random
import pytest
import logging
from og_kernel.kernel.kernel_mgr import KernelManager
from og_kernel.kernel.kernel_client import KernelClient

logger = logging.getLogger(__name__)


class MockContext:
    """
    Mock the grpc request context
    """

    def done(self):
        return False


@pytest.fixture
def kernel_manager():
    config_path = os.path.join("/tmp", str(random.randint(1, 100000)))
    workspace = os.path.join("/tmp", str(random.randint(1, 100000)))
    kernel_manager = KernelManager(config_path, workspace)
    kernel_manager.start()
    yield kernel_manager
    kernel_manager.stop()


@pytest.fixture
def ts_kernel_manager():
    config_path = os.path.join("/tmp", str(random.randint(1, 100000)))
    workspace = os.path.join("/tmp", str(random.randint(1, 100000)))
    kernel_manager = KernelManager(config_path, workspace, "tslab")
    kernel_manager.start()
    yield kernel_manager
    kernel_manager.stop()


@pytest.mark.asyncio
async def test_watching(kernel_manager):
    kernel_client = KernelClient(kernel_manager.config_path)
    await kernel_client.start_client()
    logger.info("is alive %s", await kernel_client.is_alive())

    async def on_message_fn(msg):
        if "text" in msg:
            assert msg["text"] == "Hello, world!"

    await kernel_client.watching(on_message_fn)
    for i in range(1):
        kernel_client.execute("print('Hello, world!')")
        await asyncio.sleep(1)
    await asyncio.sleep(10)
    await kernel_client.stop_watch()
    kernel_client.stop_client()


@pytest.mark.asyncio
async def test_result_occurs(kernel_manager):
    """Test stdout occurs"""
    kernel_client = KernelClient(kernel_manager.config_path)
    await kernel_client.start_client()
    logger.info("is alive %s", await kernel_client.is_alive())
    code = """
5
"""
    kernel_client.execute(code)
    messages = []
    context = MockContext()
    async for msg in kernel_client.read_response(context):
        if not msg:
            break
        messages.append(msg)
    logger.info(f"{messages}")
    filtered = list(filter(lambda x: x["msg_type"] == "execute_result", messages))
    assert len(filtered) > 0
    await asyncio.sleep(2)
    await kernel_client.stop_watch()
    kernel_client.stop_client()


@pytest.mark.asyncio
async def test_stderr_occurs(kernel_manager):
    """Test stderr occurs"""
    kernel_client = KernelClient(kernel_manager.config_path)
    await kernel_client.start_client()
    logger.info("is alive %s", await kernel_client.is_alive())
    code = """
import sys
print('Hello world', file=sys.stderr)
"""
    kernel_client.execute(code)
    messages = []

    context = MockContext()
    async for msg in kernel_client.read_response(context):
        if not msg:
            break
        messages.append(msg)
    filtered = list(filter(lambda x: x["msg_type"] == "stream", messages))
    assert len(filtered) > 0
    assert filtered[0]["content"]["name"] == "stderr"
    await asyncio.sleep(2)
    await kernel_client.stop_watch()
    kernel_client.stop_client()


@pytest.mark.asyncio
async def test_stdout_occurs(kernel_manager):
    """Test stdout occurs"""
    kernel_client = KernelClient(kernel_manager.config_path)
    await kernel_client.start_client()
    logger.info("is alive %s", await kernel_client.is_alive())
    code = """
print("hello world!")
"""
    kernel_client.execute(code)
    messages = []

    context = MockContext()
    async for msg in kernel_client.read_response(context):
        if not msg:
            break
        messages.append(msg)
    filtered = list(filter(lambda x: x["msg_type"] == "stream", messages))
    assert len(filtered) > 0
    assert filtered[0]["content"]["name"] == "stdout"
    await asyncio.sleep(2)
    await kernel_client.stop_watch()
    kernel_client.stop_client()


@pytest.mark.asyncio
async def test_syntax_exception_occurs(kernel_manager):
    """Test exception occurs"""
    kernel_client = KernelClient(kernel_manager.config_path)
    await kernel_client.start_client()
    logger.info("is alive %s", await kernel_client.is_alive())
    code = """
a = 10
b = 20
if (a < b)
    print('a is less than b')
"""
    kernel_client.execute(code)
    messages = []

    context = MockContext()
    async for msg in kernel_client.read_response(context):
        if not msg:
            break
        messages.append(msg)
    assert len(list(filter(lambda x: x["msg_type"] == "error", messages))) > 0
    await asyncio.sleep(2)
    await kernel_client.stop_watch()
    kernel_client.stop_client()


@pytest.mark.asyncio
async def test_generate_pie_chart(kernel_manager):
    """Test generate pie output"""
    kernel_client = KernelClient(kernel_manager.config_path)
    await kernel_client.start_client()
    logger.info("is alive %s", await kernel_client.is_alive())

    code = """
import matplotlib.pyplot as plt 
import numpy as np

# Create a pie chart
data = np.array([10, 20, 30, 40])
labels = ['Category 1', 'Category 2', 'Category 3', 'Category 4']

plt.pie(data, labels=labels, autopct='%1.1f%%')
plt.title('Pie Chart')
plt.show() 
"""
    kernel_client.execute(code)
    messages = []

    context = MockContext()
    async for msg in kernel_client.read_response(context):
        if msg:
            logger.debug(f"{msg}")
            messages.append(msg)
    assert len(list(filter(lambda x: x["msg_type"] == "display_data", messages))) > 0
    await asyncio.sleep(2)
    await kernel_client.stop_watch()
    kernel_client.stop_client()

kernel/tests/kernel_client_tests.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest

from og_kernel.kernel.kernel_mgr import KernelManager


@pytest.mark.parametrize(
    "config_path, workspace",
    [
        ("", ""),
        ("kernel_connection_file.json", ""),
        ("", "/tmp/workspace1"),
    ],
)
def test_init_with_invalid_args(config_path, workspace):
    with pytest.raises(ValueError):
        KernelManager(config_path, workspace)


def test_init_with_valid_args():
    km = KernelManager(
        config_path="/tmp/kernel_connection_file.json",
        workspace="/tmp/workspace1",
    )
    assert km.config_path == "/tmp/kernel_connection_file.json"
    assert km.workspace == "/tmp/workspace1"
    assert km.process is None
    assert not km.is_running


def test_start_kernel():
    km = KernelManager(
        config_path="/tmp/kernel_connection_file1.json",
        workspace="/tmp/workspace1",
    )
    km.start()
    assert km.is_running
    assert km.process is not None
    km.stop()


def test_stop_kernel():
    km = KernelManager(
        config_path="/tmp/kernel_connection_file2.json",
        workspace="/tmp/workspace2",
    )
    km.start()
    km.stop()
    assert not km.is_running
    assert km.process is None

kernel/tests/kernel_mgr_tests.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
from setuptools import setup, find_packages

setup(
    name="og_proto",
    version="0.3.6",
    description="Open source code interpreter agent for LLM",
    author="imotai",
    author_email="codego.me@gmail.com",
    url="https://github.com/dbpunk-labs/octogen",
    long_description= open('README.md').read(),
    long_description_content_type='text/markdown',

    packages=[
        "og_proto",
    ],

    package_dir={
        "og_proto": "src/og_proto",
    },

    package_data={"og_proto": ["*.pyi"]},

    install_requires=[
        "grpc-google-iam-v1>=0.12.6",
        "grpcio-tools>=1.57.0",
    ],

    entry_points={
    },

)

proto/setup.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
from setuptools import setup, find_packages

setup(
    name="og_sdk",
    version="0.3.6",
    description="Open source code interpreter agent for LLM",
    author="imotai",
    author_email="codego.me@gmail.com",
    url="https://github.com/dbpunk-labs/octogen",
    long_description=open("README.md").read(),
    long_description_content_type="text/markdown",
    packages=[
        "og_sdk",
    ],
    package_dir={
        "og_sdk": "src/og_sdk",
    },
    install_requires=[
        "og_proto",
        "aiofiles",
    ],
)

sdk/setup.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>


""" """
import logging
import grpc
from grpc import aio
from og_proto import agent_server_pb2, common_pb2
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__)


class AgentSyncSDK:

    def __init__(self, endpoint, api_key):
        self.endpoint = endpoint
        self.stub = None
        self.metadata = aio.Metadata(
            ("api_key", api_key),
        )
        self.channel = None

    def connect(self):
        if self.channel:
            return
        if self.endpoint.startswith("https"):
            channel_credential = grpc.ssl_channel_credentials()
            self.channel = grpc.secure_channel(
                self.endpoint.replace("https://", ""), channel_credential
            )
        else:
            self.channel = grpc.insecure_channel(self.endpoint)
        self.stub = AgentServerStub(self.channel)

    def assemble(self, name, code, language, desc="", saved_filenames=[]):
        request = agent_server_pb2.AssembleAppRequest(
            name=name,
            language=language,
            code=code,
            saved_filenames=saved_filenames,
            desc=desc,
        )
        response = self.stub.assemble(request, metadata=self.metadata)
        return response

    def run(self, name):
        # TODO support input files
        request = agent_server_pb2.RunAppRequest(name=name)
        for respond in self.stub.run(request, metadata=self.metadata):
            yield respond

    def query_apps(self):
        request = agent_server_pb2.QueryAppsRequest()
        return self.stub.query_apps(request, metadata=self.metadata)

    def add_kernel(self, key, endpoint):
        """
        add kernel instance to the agent and only admin can call this method
        """
        request = agent_server_pb2.AddKernelRequest(endpoint=endpoint, key=key)
        response = self.stub.add_kernel(request, metadata=self.metadata)
        return response

    def ping(self):
        request = agent_server_pb2.PingRequest()
        response = self.stub.ping(request, metadata=self.metadata)
        return response

    def download_file(self, filename, parent_path):
        request = common_pb2.DownloadRequest(filename=filename)
        fullpath = "%s/%s" % (parent_path, filename)
        with open(fullpath, "wb+") as fd:
            for chunk in self.stub.download(request, metadata=self.metadata):
                fd.write(chunk.buffer)

    def upload_file(self, filepath, filename):
        """
        upload file to agent
        """

        # TODO limit the file size

        return self.stub.upload(
            generate_chunk(filepath, filename), metadata=self.metadata
        )

    def prompt(self, prompt, files=[]):
        """
        ask the ai with prompt and  uploaded files
        """
        request = agent_server_pb2.SendTaskRequest(task=prompt, input_files=files)
        for respond in self.stub.send_task(request, metadata=self.metadata):
            yield respond


class AgentSDK:

    def __init__(self, endpoint, api_key):
        self.endpoint = endpoint
        self.stub = None
        self.metadata = aio.Metadata(
            ("api_key", api_key),
        )
        self.channel = None

    def connect(self):
        """
        Connect the agent service
        """
        if self.channel:
            return

        channel = aio.insecure_channel(self.endpoint)
        self.channel = channel
        self.stub = AgentServerStub(channel)

    async def ping(self):
        request = agent_server_pb2.PingRequest()
        response = await self.stub.ping(request, metadata=self.metadata)
        return response

    async def assemble(self, name, code, language, desc="", saved_filenames=[]):
        request = agent_server_pb2.AssembleAppRequest(
            name=name,
            language=language,
            code=code,
            saved_filenames=saved_filenames,
            desc=desc,
        )
        response = await self.stub.assemble(request, metadata=self.metadata)
        return response

    async def run(self, name):
        # TODO support input files
        request = agent_server_pb2.RunAppRequest(name=name)
        async for respond in self.stub.run(request, metadata=self.metadata):
            yield respond

    async def query_apps(self):
        """query all apps"""
        request = agent_server_pb2.QueryAppsRequest()
        return await self.stub.query_apps(request, metadata=self.metadata)

    async def add_kernel(self, key, endpoint):
        """
        add kernel instance to the agent and only admin can call this method
        """
        request = agent_server_pb2.AddKernelRequest(endpoint=endpoint, key=key)
        response = await self.stub.add_kernel(request, metadata=self.metadata)
        return response

    async def prompt(self, prompt, files=[]):
        """
        ask the ai with prompt and  uploaded files
        """
        request = agent_server_pb2.SendTaskRequest(task=prompt, input_files=files)
        async for respond in self.stub.send_task(request, metadata=self.metadata):
            yield respond

    async def download_file(self, filename, parent_path):
        request = common_pb2.DownloadRequest(filename=filename)
        fullpath = "%s/%s" % (parent_path, filename)
        async with aiofiles.open(fullpath, "wb+") as afd:
            async for chunk in self.stub.download(request, metadata=self.metadata):
                await afd.write(chunk.buffer)

    async def upload_binary(self, chunks: AsyncIterable[common_pb2.FileChunk]):
        try:
            return await self.stub.upload(chunks, metadata=self.metadata)
        except Exception as ex:
            logger.error("upload file ex %s", ex)

    async def upload_file(self, filepath, filename):
        """
        upload file to agent
        """
        # TODO limit the file size
        return await self.upload_binary(generate_async_chunk(filepath, filename))

    def close(self):
        if self.channel:
            self.channel.close()
            self.channel = None

sdk/src/og_sdk/agent_sdk.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>

""" """

import logging
import grpc
from grpc import aio
from og_proto import kernel_server_pb2
from og_proto import common_pb2
from og_proto.kernel_server_pb2_grpc import KernelServerNodeStub
from typing import AsyncIterable

logger = logging.getLogger(__name__)


class KernelSDK:

    def __init__(self, endpoint, api_key):
        self.endpoint = endpoint
        self.stub = None
        self.metadata = aio.Metadata(
            ("api_key", api_key),
        )

    def connect(self):
        """
        Connect the remote kernel instance
        """
        channel = aio.insecure_channel(self.endpoint)
        self.channel = channel
        self.stub = KernelServerNodeStub(channel)

    async def stop(self, kernel_name=None):
        """
        Stop the kernel
        """
        request = kernel_server_pb2.StopKernelRequest(kernel_name=kernel_name)
        response = await self.stub.stop(request, metadata=self.metadata)
        return response

    async def is_alive(self, kernel_name=None):
        request = kernel_server_pb2.GetStatusRequest(kernel_name=kernel_name)
        response = await self.stub.get_status(request, metadata=self.metadata)
        return response.is_alive

    async def download_file(self, filename):
        request = common_pb2.DownloadRequest(filename=filename)
        async for chunk in self.stub.download(request, metadata=self.metadata):
            yield chunk

    async def upload_binary(self, chunks: AsyncIterable[common_pb2.FileChunk]):
        try:
            return await self.stub.upload(chunks, metadata=self.metadata)
        except Exception as ex:
            logger.error("upload file ex %s" % ex)

    async def start(self, kernel_name=None):
        """
        Start the kernel
        """
        request = kernel_server_pb2.StartKernelRequest(kernel_name=kernel_name)
        response = await self.stub.start(request, metadata=self.metadata)
        return response

    async def execute(self, code, kernel_name=None):
        """
        Execute the python code
        """
        request = kernel_server_pb2.ExecuteRequest(code=code, kernel_name=kernel_name)
        async for respond in self.stub.execute(request, metadata=self.metadata):
            yield respond

    async def close(self):
        if self.channel:
            self.channel.close()
            self.channel = None

sdk/src/og_sdk/kernel_sdk.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>

import re
import string
import random
import aiofiles
import logging
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):
    buffer = []
    i = 0

    def carriage_return(buf):
        pop_buf = []
        if "\n" in buf:
            for _ in range(buf[::-1].index("\n")):
                pop_buf.append(buf.pop())
            return pop_buf[::-1]
        else:
            pop_buf.extend(buf)
            buf.clear()
            return pop_buf

    last_pop_buf = []
    while i < len(stream):
        c = stream[i]
        if c == "\b":
            if buffer:
                buffer.pop()
            last_pop_buf = []
        elif c == "\r":
            last_pop_buf = carriage_return(buffer)
        elif c == "\n":
            if last_pop_buf:
                buffer.extend(last_pop_buf)
                last_pop_buf = []
            buffer.append(c)
        else:
            last_pop_buf = []
            buffer.append(c)
        i += 1
    if last_pop_buf:
        buffer.extend(last_pop_buf)
    return "".join(buffer)


def clean_code(code: str):
    start_tag = "```"
    end_tag = "```"
    index = code.find(start_tag)
    if index >= 0:
        last = code.rfind(end_tag)
        return code[index + len(start_tag) : last]
    return code


def parse_link(text):
    """Parses a link from markdown text.

    Args:
    text: The markdown text.

    Returns:
    The link text and href, or None if no link is found.
    """
    link_regex = r"\[(.+?)\]\((.+?)\)"
    match = re.search(link_regex, text)
    if match:
        return match.groups()
    else:
        return None, None


def parse_image_filename(string):
    """Parses the image filename from a string.

    Args:
      string: A string containing the image filename.

    Returns:
      The image filename, or None if the filename is not valid.
    """

    pattern = r"octopus_\w+\.(jpg|png|gif)"
    match = re.search(pattern, string)
    if match:
        return match.group()
    else:
        return None


def random_str(n):
    # using random.choices()
    # generating random strings
    res = "".join(random.choices(string.ascii_uppercase + string.digits, k=n))
    return str(res)

sdk/src/og_sdk/utils.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """

import os
import pytest
import asyncio
import logging
import json
import random
import logging
from tempfile import gettempdir
from pathlib import Path
from og_sdk.agent_sdk import AgentSDK
from og_sdk.utils import random_str
from og_proto.agent_server_pb2 import TaskRespond

logger = logging.getLogger(__name__)
api_base = "127.0.0.1:9528"
api_key = "ZCeI9cYtOCyLISoi488BgZHeBkHWuFUH"


@pytest.fixture
def agent_sdk():
    sdk = AgentSDK(api_base, api_key)
    sdk.connect()
    yield sdk
    sdk.close()


def test_connect_bad_endpoint():
    try:
        sdk = AgentSDK("xxx", api_key)
        sdk.connect()
        assert 0, "should not go here"
    except Exception as ex:
        assert 1


@pytest.mark.asyncio
async def test_ping_test(agent_sdk):
    try:
        await agent_sdk.add_kernel(api_key, "127.0.0.1:9527")
        response = await agent_sdk.ping()
        assert response.code == 0
    except Exception as ex:
        assert 0, str(ex)


@pytest.mark.asyncio
async def test_upload_and_download_test(agent_sdk):
    sdk = agent_sdk
    await sdk.add_kernel(api_key, "127.0.0.1:9527")
    path = os.path.abspath(__file__)
    # upload file
    uploaded = await sdk.upload_file(path, "agent_sdk_tests.py")
    assert uploaded
    file_stats = os.stat(path)
    assert file_stats.st_size == uploaded.length, "bad upload_file size"
    # download file
    tmp_dir = gettempdir()
    fullpath = "%s%s%s" % (tmp_dir, os.sep, "agent_sdk_tests.py")
    await sdk.download_file("agent_sdk_tests.py", tmp_dir)
    file_stats2 = os.stat(fullpath)
    assert file_stats.st_size == file_stats2.st_size, "bad download_file size"


@pytest.mark.asyncio
async def test_prompt_smoke_test(agent_sdk):
    sdk = agent_sdk
    await sdk.add_kernel(api_key, "127.0.0.1:9527")
    try:
        responds = []
        async for respond in sdk.prompt("hello"):
            responds.append(respond)
        logger.debug(f"{responds}")
        assert len(responds) > 0, "no responds for the prompt"
        assert responds[len(responds) - 1].respond_type == TaskRespond.OnFinalAnswerType
        assert (
            responds[len(responds) - 1].final_respond.answer
            == "how can I help you today?"
        )
    except Exception as ex:
        assert 0, str(ex)


@pytest.mark.asyncio
async def test_run_code_test(agent_sdk):
    sdk = agent_sdk
    await sdk.add_kernel(api_key, "127.0.0.1:9527")
    try:
        responds = []
        async for respond in sdk.prompt("write a hello world in python"):
            responds.append(respond)
        logger.debug(f"{responds}")
        assert len(responds) > 0, "no responds for the prompt"
        assert responds[len(responds) - 1].respond_type == TaskRespond.OnFinalAnswerType
        assert (
            responds[len(responds) - 1].final_respond.answer
            == "this code prints 'hello world'"
        )
    except Exception as ex:
        assert 0, str(ex)


@pytest.mark.asyncio
async def test_run_code_with_error(agent_sdk):
    sdk = agent_sdk
    await sdk.add_kernel(api_key, "127.0.0.1:9527")
    try:
        responds = []
        async for respond in sdk.prompt("error function"):
            responds.append(respond)
        logger.debug(f"{responds}")
        assert len(responds) > 0, "no responds for the prompt"
        assert (
            responds[len(responds) - 3].respond_type == TaskRespond.OnAgentActionEndType
        )
        assert responds[
            len(responds) - 3
        ].on_agent_action_end.has_error, "bad has error result"
        assert responds[len(responds) - 1].respond_type == TaskRespond.OnFinalAnswerType
        assert (
            responds[len(responds) - 1].final_respond.answer
            == "this code prints 'hello world'"
        )
    except Exception as ex:
        assert 0, str(ex)


@pytest.mark.asyncio
async def test_assemble_test(agent_sdk):
    sdk = agent_sdk
    await sdk.add_kernel(api_key, "127.0.0.1:9527")
    try:
        code = "print('hello')"
        name = random_str(10)
        response = await sdk.assemble(name, code, "python")
        assert response.code == 0, "fail to assemble app"
        apps = await sdk.query_apps()
        app = list(filter(lambda x: x.name == name, apps.apps))
        assert len(app) == 1, "fail to get the app with name " + name
        responds = []
        async for respond in sdk.run(name):
            responds.append(respond)
        assert len(responds) == 3, "bad responds for run application"
        assert responds[0].respond_type == TaskRespond.OnAgentActionType
        assert responds[1].respond_type == TaskRespond.OnAgentActionStdout
        assert responds[2].respond_type == TaskRespond.OnAgentActionEndType
        assert (
            json.loads(responds[0].on_agent_action.input)["code"] == code
        ), "remote code !eq the local code"
        assert responds[1].console_stdout.find("hello") >= 0, "bad output"
    except Exception as ex:
        assert 0, str(ex)


display_image_test_code = """import matplotlib.pyplot as plt
import numpy as np

# Step 2: Create the dataset
categories = ['Category 1', 'Category 2', 'Category 3']
group1_values = [10, 15, 12]
group2_values = [8, 11, 9]

# Step 3: Set the width and positions
bar_width = 0.35
index = np.arange(len(categories))

# Step 4: Create the figure and axis
fig, ax = plt.subplots()

# Step 5: Plot the bars
rects1 = ax.bar(index, group1_values, bar_width, label='Group 1')
rects2 = ax.bar(index + bar_width, group2_values, bar_width, label='Group 2')

# Step 6: Set labels, title, and legend
ax.set_xlabel('Categories')
ax.set_ylabel('Values')
ax.set_title('Grouped Bar Chart')
ax.set_xticks(index + bar_width / 2)
ax.set_xticklabels(categories)
ax.legend()

# Step 7: Show the chart
plt.show()"""


@pytest.mark.asyncio
async def test_assemble_image_test(agent_sdk):
    sdk = agent_sdk
    await sdk.add_kernel(api_key, "127.0.0.1:9527")
    try:
        name = random_str(10)
        response = await sdk.assemble(name, display_image_test_code, "python")
        assert response.code == 0, "fail to assemble app"
        apps = await sdk.query_apps()
        app = list(filter(lambda x: x.name == name, apps.apps))
        assert len(app) == 1, "fail to get the app with name " + name
        responds = []
        async for respond in sdk.run(name):
            responds.append(respond)
        assert len(responds) == 3, "bad responds for run application"
        assert responds[0].respond_type == TaskRespond.OnAgentActionType
        assert responds[1].respond_type == TaskRespond.OnAgentActionStdout
        assert responds[2].respond_type == TaskRespond.OnAgentActionEndType
        assert (
            json.loads(responds[0].on_agent_action.input)["code"]
            == display_image_test_code
        )
        assert responds[1].console_stdout.find("png") > 0, "should output the files"
        assert (
            len(responds[2].on_agent_action_end.output_files) == 1
        ), "should output the files"

    except Exception as ex:
        assert 0, str(ex)

sdk/tests/agent_sdk_tests.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# 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__)


@pytest.fixture
def kernel_sdk():
    endpoint = (
        "localhost:9527"  # Replace with the actual endpoint of your test gRPC server
    )
    return KernelSDK(endpoint, "ZCeI9cYtOCyLISoi488BgZHeBkHWuFUH")


@pytest.fixture
def bad_kernel_sdk():
    endpoint = (
        "localhost:9527"  # Replace with the actual endpoint of your test gRPC server
    )
    return KernelSDK(endpoint, "ZCeI9cYtOCyLISoi488BgZHeBkHWuFU")


@pytest.mark.asyncio
async def test_bad_sdk(bad_kernel_sdk):
    try:
        kernel_sdk.connect()
        assert kernel_sdk.stub is not None  # Check that stub is initialized
        await kernel_sdk.start()
        assert False
    except Exception as e:
        assert True


@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")
    )
    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_file("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()
    assert kernel_sdk.stub is not None  # Check that stub is initialized
    if not await kernel_sdk.is_alive():
        await kernel_sdk.start()
    code = """print('hello world!')"""
    responds = []
    async for respond in kernel_sdk.execute(code):
        responds.append(respond)
    await kernel_sdk.stop()
    assert len(responds) == 1
    assert responds[0].output_type == ExecuteResponse.StdoutType
    assert json.loads(responds[0].output)["text"] == "hello world!\n"


@pytest.mark.asyncio
async def test_sdk_result_test(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()
    code = """print('hello world!')
5"""
    responds = []
    async for respond in kernel_sdk.execute(code):
        responds.append(respond)
    await kernel_sdk.stop()
    assert len(responds) == 2
    assert responds[0].output_type == ExecuteResponse.StdoutType
    assert responds[1].output_type == ExecuteResponse.ResultType
    assert json.loads(responds[0].output)["text"] == "hello world!\n"
    assert json.loads(responds[1].output)["text/plain"] == "5"

sdk/tests/kernel_sdk_tests.py
dbpunk-labs/octogen
#! /usr/bin/env python3
# vim:fenc=utf-8
#
# Copyright © 2023 imotai <imotai@imotai-ub>
#
# Distributed under terms of the MIT license.

""" """
import pytest
import json
from og_sdk.utils import process_char_stream


def test_process_char_stream_case2():
    stream1 = "\rt:   0%|          | 0/518 [00:00<?, ?it/s, now=None]"
    output1 = process_char_stream(stream1)
    stream2 = "\rt:   3%|▎         | 15/518 [00:00<00:03, 137.85it/s, now=None]"
    output2 = process_char_stream(output1 + stream2)
    assert output2 == "t:   3%|▎         | 15/518 [00:00<00:03, 137.85it/s, now=None]"


def test_process_char_stream():
    stream0 = "  Downloading pyfiglet-1.0.2-py3-none-any.whl (1.1 MB)\r\n\x1b[?25l     \x1b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[32m0.0/1.1 MB\x1b[0m \x1b[31m?\x1b[0m eta \x1b[36m-:--:--\x1b[0m"
    stream1 = "\r\x1b[2K     \x1b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[32m0.0/1.1 MB\x1b[0m \x1b[31m?\x1b[0m eta \x1b[36m-:--:--\x1b[0m"
    output1 = process_char_stream(stream0 + stream1)
    output2 = process_char_stream(output1 + stream1)
    assert output1 == output2
    final_stream = "\r1.1 MB 100%\r\n"
    output3 = process_char_stream(output2 + final_stream)
    final_ouput_expected = (
        "  Downloading pyfiglet-1.0.2-py3-none-any.whl (1.1 MB)\n1.1 MB 100%\n"
    )
    assert final_ouput_expected == output3


def test_empty_string():
    assert process_char_stream("") == ""


def test_single_character():
    assert process_char_stream("a") == "a"


def test_multiple_characters():
    assert process_char_stream("abc") == "abc"


def test_backspace():
    assert process_char_stream("ab\b") == "a"


def test_carriage_return():
    assert process_char_stream("ab\r") == "ab"


def test_carriage_return_with_newline():
    assert process_char_stream("ab\r\n") == "ab\n"


def test_backspace_and_carriage_return():
    assert process_char_stream("ab\b\r") == "a"


def test_mixed_escape_characters_and_regular_characters():
    assert process_char_stream("ab\b\r\ncde") == "a\ncde"


def test_special_characters():
    assert (
        process_char_stream("ab!@#$%^&*()_+{}|:\";'<>,.?/`~")
        == "ab!@#$%^&*()_+{}|:\";'<>,.?/`~"
    )

sdk/tests/utils_test.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
from setuptools import setup, find_packages

setup(
    name="og_up",
    version="0.3.6",
    description="Open source code interpreter agent for LLM",
    author="imotai",
    author_email="codego.me@gmail.com",
    url="https://github.com/dbpunk-labs/octogen",
    long_description=open("README.md").read(),
    long_description_content_type="text/markdown",
    packages=[
        "og_up",
    ],
    package_dir={
        "og_up": "src/og_up",
    },
    install_requires=["og_sdk", "requests", "huggingface_hub", "rich", "click"],
    entry_points={
        "console_scripts": [
            "og_up = og_up.up:init_octogen",
            "og_kernel_up = og_up.kernel_up:init_kernel",
            "og_download = og_up.model_downloader:download",
        ]
    },
)

up/setup.py
dbpunk-labs/octogen
#! /usr/bin/env python3
# vim:fenc=utf-8
#
# Copyright © 2023 imotai <imotai@imotai-ub>
#
# Distributed under terms of the MIT license.

""" """
import click
import os
import time
from .utils import run_with_realtime_print
from .up import check_the_env
from .up import load_docker_image
from .up import get_latest_release_version
from .up import ping_agent_service
from .up import stop_service
from .up import generate_kernel_env
from .up import random_str
from .up import refresh
from .up import add_kernel_endpoint
from rich.live import Live
from rich.spinner import Spinner
from rich.console import Group
from rich.console import Console
from rich.markdown import Markdown
from rich.prompt import Prompt
from dotenv import dotenv_values
from og_sdk.agent_sdk import AgentSyncSDK

Welcome = "welcome to use og_kernel_up"


def get_config(console):
    key = Prompt.ask("Agent Key", password=True)
    endpoint = Prompt.ask("Agent Endpoint")
    name = Prompt.ask("Kernel Name")
    port = Prompt.ask("Kernel Port")
    return key, endpoint, name, port


def start_kernel_service(
    live, segments, install_dir, image_name, version, kernel_name, kernel_port
):
    spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
    step = "Start kernel service"
    output = ""
    vender = "docker"
    segments.append((spinner, step, ""))
    refresh(live, segments)
    stop_service(kernel_name)
    full_name = f"{image_name}:{version}"
    command = [
        vender,
        "run",
        "--name",
        kernel_name,
        "-p",
        f"127.0.0.1:{kernel_port}:{kernel_port}",
        "-v",
        f"{install_dir}:/app",
        "-dt",
        f"{full_name}",
        "bash",
        "/bin/start_kernel.sh",
        "/app",
    ]
    result_code = 0
    output = ""
    for code, chunk in run_with_realtime_print(command=command):
        result_code = code
        output += chunk
        pass
    time.sleep(6)
    segments.pop()
    if result_code == 0:
        segments.append(("✅", "Start kernel service", ""))
    else:
        segments.append(("❌", "Start kernel service", output))
    refresh(live, segments)
    return result_code


@click.command("init")
@click.option("--image_name", default="dbpunk/octogen", help="the octogen image name")
@click.option(
    "--install_dir", default="~/kernel/apps", help="the install dir of kernel"
)
@click.option("--octogen_version", default="", help="the version of octogen")
def init_kernel(
    image_name,
    install_dir,
    octogen_version,
):
    if install_dir.find("~") == 0:
        real_install_dir = install_dir.replace("~", os.path.expanduser("~"))
    else:
        real_install_dir = install_dir
    os.makedirs(real_install_dir, exist_ok=True)
    console = Console()
    console.print(Welcome)
    key, agent_endpoint, kernel_name, kernel_port = get_config(console)
    kernel_dir = "/".join([real_install_dir, kernel_name])
    os.makedirs(kernel_dir, exist_ok=True)
    segments = []
    with Live(Group(*segments), console=console) as live:
        if octogen_version:
            version = octogen_version
        else:
            version = get_latest_release_version(repo_name, live, segments)
        check_result, _ = check_the_env(live, segments)
        if not check_result:
            segments.append(("❌", "Setup kernel service failed", ""))
            refresh(live, segments)
            return
        code = load_docker_image(version, image_name, live, segments)
        if code != 0:
            return
        kernel_key = random_str(32)
        env_path = kernel_dir + "/kernel/" + ".env"
        if not os.path.exists(env_path):
            generate_kernel_env(
                live,
                segments,
                kernel_dir,
                kernel_key,
                rpc_port=kernel_port,
                rpc_host="0.0.0.0",
            )
        else:
            config = dotenv_values(env_path)
            kernel_key = config.get("rpc_key", "")
            if kernel_key:
                segments.append(("✅", "Use the exist kernel config", ""))
                refresh(live, segments)
                kernel_port = config.get("rpc_port")
            else:
                segments.append(("❌", "Bad kernel config", ""))
                refresh(live, segments)
                return
        code = start_kernel_service(
            live, segments, kernel_dir, image_name, version, kernel_name, kernel_port
        )
        if code != 0:
            return
        if not add_kernel_endpoint(
            live, segments, key, f"127.0.0.1:{kernel_port}", kernel_key, agent_endpoint
        ):
            segments.append(("❌", "Setup kernel service failed", ""))
            refresh(live, segments)
            return

        if ping_agent_service(live, segments, kernel_key, api_base=agent_endpoint):
            segments.append(("👍", "Setup kernel service done", ""))
            refresh(live, segments)
        else:
            segments.append(("❌", "Setup kernel service failed", ""))
            refresh(live, segments)

up/src/og_up/kernel_up.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import os
import click
from huggingface_hub import hf_hub_download


@click.command()
@click.option("--repo", help="the repo of huggingface")
@click.option("--filename", help="the filename of model")
@click.option(
    "--cache_dir", default="~/.octogen/app/cache", help="the cache_dir of huggingface"
)
@click.option(
    "--local_dir", default="~/.octogen/app/model", help="the local dir of huggingface"
)
@click.option("--socks_proxy", default="", help="the socks proxy url")
def download(repo, filename, cache_dir, local_dir, socks_proxy):
    if local_dir.find("~") == 0:
        real_local_dir = local_dir.replace("~", os.path.expanduser("~"))
    else:
        real_local_dir = local_dir
    if cache_dir.find("~") == 0:
        real_cache_dir = cache_dir.replace("~", os.path.expanduser("~"))
    else:
        real_cache_dir = cache_dir

    os.makedirs(real_cache_dir, exist_ok=True)
    os.makedirs(real_local_dir, exist_ok=True)
    proxies = {}
    if socks_proxy:
        proxies = {"http": socks_proxy, "https": socks_proxy}
    hf_hub_download(
        repo_id=repo,
        filename=filename,
        cache_dir=real_cache_dir,
        local_dir=real_local_dir,
        proxies=proxies,
        resume_download=True,
    )

up/src/og_up/model_downloader.py
dbpunk-labs/octogen
# vim:fenc=utf-8
#
# Copyright (C) 2023 dbpunk.com Author imotai <codego.me@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" """
import json
import click
import requests
import random
import string
import os
import subprocess
import sys
import time
from pathlib import Path
from tqdm import tqdm
from tempfile import gettempdir
from rich.prompt import Prompt
from rich.markdown import Markdown
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.progress_bar import ProgressBar
from rich.live import Live
from rich.spinner import Spinner
from rich.console import Group
from og_sdk.utils import process_char_stream
from og_sdk.agent_sdk import AgentSyncSDK
from .utils import run_with_realtime_print

OCTOGEN_TITLE = "🐙[bold red]Octogen Up"
OCTOGEN_GITHUB_REPOS = "dbpunk-labs/octogen"
Welcome = f"""
Welcome to use {OCTOGEN_TITLE}
"""


def random_str(n):
    # generating random strings
    res = "".join(random.choices(string.ascii_uppercase + string.digits, k=n))
    return str(res)


def run_install_cli(live, segments):
    """
    Install the octogen chat cli
    """
    spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
    step = "Install octogen terminal cli"
    output = ""
    segments.append((spinner, step, ""))
    result_code = 0
    refresh(live, segments)
    outputs = ""
    for code, output in run_with_realtime_print(
        command=["pip", "install", "-U", "og_proto", "og_sdk", "og_chat"]
    ):
        outputs += output
        result_code = code
    if result_code == 0:
        segments.pop()
        segments.append(("✅", "Install octogen terminal cli", ""))
        refresh(live, segments)
        return True
    else:
        segments.pop()
        segments.append(("❌", "Install octogen terminal cli", outputs))
        refresh(live, segments)
        return False


def refresh(
    live,
    segments,
    title=OCTOGEN_TITLE,
):
    table = Table.grid(padding=(0, 1, 0, 0), pad_edge=True)
    table.add_column("Index", no_wrap=True, justify="center")
    table.add_column("Status", no_wrap=True, justify="left")
    table.add_column("Step", no_wrap=True, justify="left")
    table.add_column("Content", no_wrap=True, justify="left")
    for index, (status, step, segment) in enumerate(segments):
        table.add_row(str(index), status, step, segment)
    live.update(table)
    live.refresh()


def check_container_vender(vender):
    command = [vender, "version", "--help"]
    all_output = ""
    result_code = 0
    for code, output in run_with_realtime_print(command):
        result_code = code
        all_output += output
    if result_code != 0:
        return False, f"{vender} is required"
    if all_output.lower().find("json") < 0:
        return False, f"Upgrade the {vender} to support json format"
    # check alive
    command = [vender, "ps"]
    result_code = 0
    for code, _ in run_with_realtime_print(command):
        result_code = code
    if result_code != 0:
        return False, f"{vender} is not running"
    return True, "ok"


def check_the_env(live, segments, need_container=True, use_podman=False):
    # check the python version
    spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
    step = "Check the environment"
    segments.append((spinner, step, ""))
    refresh(live, segments)
    version_ctrl = sys.version.split(" ")[0].split(".")
    if version_ctrl[0] != "3":
        old_segment = segments.pop()
        segments.append(("❌", "Check the environment", "Python3 is required"))
        refresh(live, segments)
        return False, "Python3 is required"
    if int(version_ctrl[1]) < 10:
        old_segment = segments.pop()
        segments.append(("❌", "Check the environment", "Python3.10 is required"))
        refresh(live, segments)
        return False, "Python3.10 is required"
    if need_container:
        vender = "docker" if not use_podman else "podman"
        result, msg = check_container_vender(vender)
        if not result:
            old_segment = segments.pop()
            segments.append(("❌", "Check the environment", msg))
            refresh(live, segments)
            return False, msg
    old_segment = segments.pop()
    segments.append(("✅", "Check the environment", ""))
    refresh(live, segments)
    return True, ""


def get_latest_release_version(repo_name, live, segments):
    """
    get the latest release version
    """
    spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
    step = "Get octogen version"
    segments.append((spinner, step, ""))
    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", "").strip()
    if not version:
        segments.append(("❌", "Get octogen latest version failed", version))
    else:
        segments.append(("✅", "Get octogen latest version", version))
    refresh(live, segments)
    return version


def download_model(
    live,
    segments,
    socks_proxy="",
    repo="TheBloke/CodeLlama-7B-Instruct-GGUF",
    filename="codellama-7b-instruct.Q5_K_S.gguf",
):
    spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
    step = "Download codeLlama"
    output = ""
    segments.append((spinner, step, ""))
    refresh(live, segments)
    result_code = 0
    for code, chunk in run_with_realtime_print(
        command=[
            "og_download",
            "--repo",
            repo,
            "--filename",
            filename,
            "--socks_proxy",
            socks_proxy,
        ]
    ):
        result_code = code
        output += chunk
        output = process_char_stream(output)
        old_segment = segments.pop()
        segments.append((old_segment[0], old_segment[1], output))
        refresh(live, segments)

    old_segment = segments.pop()
    if result_code == 0:
        segments.append(("✅", step, filename))
    else:
        segments.append(("❌", step, output))
    refresh(live, segments)
    return result_code


def load_docker_image(
    version, image_name, live, segments, chunk_size=1024, use_podman=False
):
    """
    download the image file and load it into docker
    """
    full_name = (
        f"{image_name}:{version}"
        if not use_podman
        else f"docker.io/{image_name}:{version}"
    )
    spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
    step = "Pull octogen image"
    segments.append((spinner, step, ""))
    refresh(live, segments)
    return_code = 0
    output = ""
    vender = "docker" if not use_podman else "podman"
    for code, msg in run_with_realtime_print(command=[vender, "pull", full_name]):
        return_code = code
        output += msg
    old_segment = segments.pop()
    if return_code == 0:
        segments.append(("✅", old_segment[1], full_name))
    else:
        segments.append(("❌", old_segment[1], output))
    refresh(live, segments)
    return code


def choose_api_service(console):
    mk = """Choose your favourite Large Language Model
1. OpenAI, Kernel, Agent and Cli will be installed
2. Azure OpenAI, Kernel, Agent and Cli will be installed
3. Codellama, Llama.cpp Model Server, Kernel, Agent and Cli will be installed
4. Octogen(beta), Only Cli will be installed
"""
    console.print(Markdown(mk))
    choice = Prompt.ask("Choices", choices=["1", "2", "3", "4"], default="1:OpenAI")
    if choice == "1":
        key = Prompt.ask("Enter OpenAI Key", password=True)
        model = Prompt.ask("Enter OpenAI Model", default="gpt-3.5-turbo-16k-0613")
        return choice, key, model, ""
    elif choice == "2":
        key = Prompt.ask("Enter Azure OpenAI Key", password=True)
        deployment = Prompt.ask("Enter Azure OpenAI Deployment")
        api_base = Prompt.ask("Enter Azure OpenAI Base")
        return choice, key, deployment, api_base
    elif choice == "4":
        key = Prompt.ask("Enter Octogen Key", password=True)
        api_base = "https://agent.octogen.dev"
        return choice, key, "", api_base
    return choice, "", "", ""


def generate_agent_common(fd, rpc_key):
    fd.write("rpc_host=0.0.0.0\n")
    fd.write("rpc_port=9528\n")
    fd.write(f"admin_key={rpc_key}\n")
    fd.write("max_file_size=202400000\n")
    fd.write("max_iterations=8\n")
    fd.write("db_path=/app/agent/db/octogen.db\n")


def generate_agent_azure_openai(
    live, segments, install_dir, admin_key, openai_key, deployment, api_base
):
    agent_dir = f"{install_dir}/agent"
    os.makedirs(agent_dir, exist_ok=True)
    with open(f"{agent_dir}/.env", "w+") as fd:
        generate_agent_common(fd, admin_key)
        fd.write("llm_key=azure_openai\n")
        fd.write(f"openai_api_type=azure\n")
        fd.write(f"openai_api_version=2023-07-01-preview\n")
        fd.write(f"openai_api_key={openai_key}\n")
        fd.write(f"openai_api_base={api_base}\n")
        fd.write(f"openai_api_deployment={deployment}\n")
        fd.write("max_file_size=202400000\n")
        fd.write("max_iterations=8\n")
    segments.append(("✅", "Generate agent config", f"{agent_dir}/.env"))
    refresh(live, segments)


def generate_agent_openai(
    live, segments, install_dir, admin_key, openai_key, openai_model
):
    agent_dir = f"{install_dir}/agent"
    os.makedirs(agent_dir, exist_ok=True)
    with open(f"{agent_dir}/.env", "w+") as fd:
        generate_agent_common(fd, admin_key)
        fd.write("llm_key=openai\n")
        fd.write(f"openai_api_key={openai_key}\n")
        fd.write(f"openai_api_model={openai_model}\n")
        fd.write("max_file_size=202400000\n")
        fd.write("max_iterations=8\n")
    segments.append(("✅", "Generate agent config", f"{agent_dir}/.env"))
    refresh(live, segments)


def generate_agent_codellama(live, segments, install_dir, admin_key):
    agent_dir = f"{install_dir}/agent"
    os.makedirs(agent_dir, exist_ok=True)
    with open(f"{agent_dir}/.env", "w+") as fd:
        generate_agent_common(fd, admin_key)
        fd.write("llm_key=codellama\n")
        fd.write("llama_api_base=http://127.0.0.1:8080\n")
        fd.write("llama_api_key=xxx\n")
        fd.write("max_file_size=202400000\n")
        fd.write("max_iterations=8\n")
    segments.append(("✅", "Generate agent config", f"{agent_dir}/.env"))
    refresh(live, segments)


def generate_kernel_env(
    live, segments, install_dir, rpc_key, rpc_port=9527, rpc_host="127.0.0.1"
):
    kernel_dir = f"{install_dir}/kernel"
    kernel_ws_dir = f"{install_dir}/kernel/ws"
    kernel_config_dir = f"{install_dir}/kernel/config"
    os.makedirs(kernel_dir, exist_ok=True)
    os.makedirs(kernel_ws_dir, exist_ok=True)
    os.makedirs(kernel_config_dir, exist_ok=True)
    with open(f"{kernel_dir}/.env", "w+") as fd:
        fd.write("config_root_path=/app/kernel/config\n")
        fd.write("workspace=/app/kernel/ws\n")
        fd.write(f"rpc_host={rpc_host}\n")
        fd.write(f"rpc_port={rpc_port}\n")
        fd.write(f"rpc_key={rpc_key}\n")
    segments.append(("✅", "Generate kernel config", f"{kernel_dir}/.env"))
    refresh(live, segments)


def stop_service(name, use_podman=False):
    vender = "docker" if not use_podman else "podman"
    command = [vender, "ps", "-f", f"name={name}", "--format", "json"]
    output = ""
    for _, chunk in run_with_realtime_print(command=command):
        output += chunk
        pass
    if use_podman and output:
        rows = json.loads(output.strip())
        for row in rows:
            id = row["Id"]
            command = [vender, "kill", id]
            for _, chunk in run_with_realtime_print(command=command):
                pass
    elif output:
        for line in output.split(os.linesep):
            if not line:
                break
            row = json.loads(line.strip())
            id = row["ID"]
            command = [vender, "kill", id]
            for _, chunk in run_with_realtime_print(command=command):
                pass
    command = [vender, "container", "rm", name]
    for _, chunk in run_with_realtime_print(command=command):
        pass


def start_service(
    live,
    segments,
    install_dir,
    image_name,
    version,
    is_codellama="1",
    model_filename="",
    use_podman=False,
):
    spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
    step = "Start octogen service"
    output = ""
    vender = "docker" if not use_podman else "podman"
    segments.append((spinner, step, ""))
    refresh(live, segments)
    stop_service("octogen", use_podman=use_podman)
    full_name = (
        f"{image_name}:{version}"
        if not use_podman
        else f"docker.io/{image_name}:{version}"
    )
    command = [
        vender,
        "run",
        "--name",
        "octogen",
        "-p",
        "127.0.0.1:9528:9528",
        "-v",
        f"{install_dir}:/app",
        "-dt",
        f"{full_name}",
        "bash",
        "/bin/start_all.sh",
        "/app",
        is_codellama,
        model_filename,
    ]
    result_code = 0
    output = ""
    for code, chunk in run_with_realtime_print(command=command):
        result_code = code
        output += chunk
        pass
    time.sleep(6)
    segments.pop()
    if result_code == 0:
        segments.append(("✅", "Start octogen service", ""))
    else:
        segments.append(("❌", "Start octogen service", output))
    refresh(live, segments)
    return result_code


def update_cli_config(live, segments, api_key, cli_dir, api_base="127.0.0.1:9528"):
    config_path = f"{cli_dir}/config"
    with open(config_path, "w+") as fd:
        fd.write(f"endpoint={api_base}\n")
        fd.write(f"api_key={api_key}\n")
    segments.append(("✅", "Update cli config", ""))
    refresh(live, segments)


def add_kernel_endpoint(live, segments, admin_key, kernel_endpoint, api_key, api_base):
    spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
    step = "Register the kernel endpoint"
    segments.append((spinner, step, ""))
    refresh(live, segments)
    retry_count = 0
    result_code = 0
    msg = ""
    while retry_count <= 3:
        retry_count += 1
        try:
            sdk = AgentSyncSDK(api_base, admin_key)
            sdk.connect()
            response = sdk.add_kernel(api_key, kernel_endpoint)
            result_code = response.code
            if result_code == 0:
                break
            msg = response.msg
            time.sleep(2)
        except Exception as ex:
            result_code = 1
            msg = str(ex)
            time.sleep(2)
    segments.pop()
    if result_code == 0:
        segments.append(("✅", step, ""))
        refresh(live, segments)
        return True
    else:
        segments.append(("❌", step, msg))
        refresh(live, segments)
        return False


def ping_agent_service(live, segments, api_key, api_base="127.0.0.1:9528"):
    spinner = Spinner("dots", style="status.spinner", speed=1.0, text="")
    step = "Ping octogen agent service"
    segments.append((spinner, step, ""))
    refresh(live, segments)
    try:
        sdk = AgentSyncSDK(api_base, api_key)
        sdk.connect()
        response = sdk.ping()
        segments.pop()
        if response.code == 0:
            segments.append(("✅", "Ping octogen agent service", ""))
            return True
        else:
            segments.append(("❌", "Ping octogen agent service", response.msg))
            return False
    except Exception as ex:
        segments.append(("❌", "Ping octogen agent service", str(ex)))
        return False


def start_octogen_for_openai(
    live,
    segments,
    install_dir,
    cli_install_dir,
    admin_key,
    kernel_key,
    image_name,
    version,
    api_key,
    model,
    use_podman=False,
):
    generate_agent_openai(live, segments, install_dir, admin_key, api_key, model)
    if (
        start_service(
            live,
            segments,
            install_dir,
            image_name,
            version,
            is_codellama="0",
            use_podman=use_podman,
        )
        == 0
    ):
        if not add_kernel_endpoint(
            live, segments, admin_key, "127.0.0.1:9527", kernel_key, "127.0.0.1:9528"
        ):
            segments.append(("❌", "Setup octogen service failed", ""))
            refresh(live, segments)
            return False
        update_cli_config(live, segments, kernel_key, cli_install_dir)
        if ping_agent_service(live, segments, kernel_key):
            segments.append(("👍", "Setup octogen service done", ""))
            refresh(live, segments)
            return True
        else:
            segments.append(("❌", "Setup octogen service failed", ""))
            refresh(live, segments)
            return False
    else:
        segments.append(("❌", "Setup octogen failed", ""))
        refresh(live, segments)
        return False


def start_octogen_for_azure_openai(
    live,
    segments,
    install_dir,
    cli_install_dir,
    admin_key,
    kernel_key,
    image_name,
    version,
    api_key,
    model,
    api_base,
    use_podman=False,
):
    generate_agent_azure_openai(
        live, segments, install_dir, admin_key, api_key, model, api_base
    )
    if (
        start_service(
            live,
            segments,
            install_dir,
            image_name,
            version,
            is_codellama="0",
            use_podman=use_podman,
        )
        == 0
    ):
        if not add_kernel_endpoint(
            live, segments, admin_key, "127.0.0.1:9527", kernel_key, "127.0.0.1:9528"
        ):
            segments.append(("❌", "Setup octogen service failed", ""))
            refresh(live, segments)
            return False
        update_cli_config(live, segments, kernel_key, cli_install_dir)
        if ping_agent_service(live, segments, kernel_key):
            segments.append(("👍", "Setup octogen service done", ""))
            refresh(live, segments)
            return True
        else:
            segments.append(("❌", "Setup octogen service failed", ""))
            refresh(live, segments)
            return False
    else:
        segments.append(("❌", "Setup octogen failed", ""))
        refresh(live, segments)
        return False


def start_octogen_for_codellama(
    live,
    segments,
    model_repo,
    model_filename,
    install_dir,
    cli_install_dir,
    admin_key,
    kernel_key,
    image_name,
    version,
    socks_proxy="",
    use_podman=False,
):
    """
    start the octogen service for codellama
    """

    if download_model(live, segments, socks_proxy, model_repo, model_filename) != 0:
        segments.append(("❌", "Setup octogen service failed", ""))
        refresh(live, segments)
        return False

    generate_agent_codellama(live, segments, install_dir, admin_key)
    if (
        start_service(
            live,
            segments,
            install_dir,
            image_name,
            version,
            is_codellama="1",
            model_filename=model_filename,
            use_podman=use_podman,
        )
        == 0
    ):
        if not add_kernel_endpoint(
            live, segments, admin_key, "127.0.0.1:9527", kernel_key, "127.0.0.1:9528"
        ):
            segments.append(("❌", "Setup octogen service failed", ""))
            refresh(live, segments)
            return False
        update_cli_config(live, segments, kernel_key, cli_install_dir)
        if ping_agent_service(live, segments, kernel_key):
            segments.append(("👍", "Setup octogen service done", ""))
            refresh(live, segments)
            return True
        else:
            segments.append(("❌", "Setup octogen service failed", ""))
            refresh(live, segments)
            return False
    else:
        segments.append(("❌", "Setup octogen service failed", ""))
        refresh(live, segments)
        return False


@click.command("init")
@click.option("--image_name", default="dbpunk/octogen", help="the octogen image name")
@click.option(
    "--repo_name", default=OCTOGEN_GITHUB_REPOS, help="the github repo of octogen"
)
@click.option(
    "--install_dir", default="~/.octogen/app", help="the install dir of octogen"
)
@click.option("--cli_dir", default="~/.octogen/", help="the cli dir of octogen")
@click.option("--octogen_version", default="", help="the version of octogen")
@click.option("--socks_proxy", default="", help="the socks proxy url")
@click.option(
    "--codellama_repo",
    default="TheBloke/CodeLlama-7B-Instruct-GGUF",
    help="the codellama repo of huggingface",
)
@click.option(
    "--model_filename",
    default="codellama-7b-instruct.Q5_K_M.gguf",
    help="the model filename in model repo",
)
@click.option(
    "--use_podman",
    is_flag=True,
    help="use podman as the container engine",
)
def init_octogen(
    image_name,
    repo_name,
    install_dir,
    cli_dir,
    octogen_version,
    socks_proxy,
    codellama_repo,
    model_filename,
    use_podman,
):
    if cli_dir.find("~") == 0:
        real_cli_dir = cli_dir.replace("~", os.path.expanduser("~"))
    else:
        real_cli_dir = cli_dir
    if install_dir.find("~") == 0:
        real_install_dir = install_dir.replace("~", os.path.expanduser("~"))
    else:
        real_install_dir = install_dir
    os.makedirs(real_install_dir, exist_ok=True)
    console = Console()
    console.print(Welcome)
    choice, key, model, api_base = choose_api_service(console)
    segments = []
    with Live(Group(*segments), console=console) as live:
        run_install_cli(live, segments)
        if choice == "4":
            check_result, _ = check_the_env(live, segments, need_container=False)
            if not check_result:
                segments.append(("❌", "Setup octogen cli failed", ""))
                refresh(live, segments)
                return
            update_cli_config(live, segments, key, real_cli_dir, api_base)
            if ping_agent_service(live, segments, key, api_base):
                segments.append(("👍", "Setup octogen cli done", ""))
            else:
                segments.append(("❌", "Setup octogen cli failed", ""))
            refresh(live, segments)
            return
        check_result, _ = check_the_env(
            live, segments, need_container=True, use_podman=use_podman
        )
        if not check_result:
            segments.append(("❌", "Setup octogen agent service failed", ""))
            refresh(live, segments)
            return
        if octogen_version:
            version = octogen_version
        else:
            version = get_latest_release_version(repo_name, live, segments)

        code = load_docker_image(
            version, image_name, live, segments, use_podman=use_podman
        )
        if code != 0:
            return
        kernel_key = random_str(32)
        admin_key = random_str(32)
        generate_kernel_env(live, segments, real_install_dir, kernel_key)
        if choice == "3":
            # start for codellama
            start_octogen_for_codellama(
                live,
                segments,
                codellama_repo,
                model_filename,
                real_install_dir,
                real_cli_dir,
                admin_key,
                kernel_key,
                image_name,
                version,
                socks_proxy,
                use_podman=use_podman,
            )
        elif choice == "2":
            # start azure openai
            start_octogen_for_azure_openai(
                live,
                segments,
                real_install_dir,
                real_cli_dir,
                admin_key,
                kernel_key,
                image_name,
                version,
                key,
                model,
                api_base,
                use_podman=use_podman,
            )
        else:
            # start for openai
            start_octogen_for_openai(
                live,
                segments,
                real_install_dir,
                real_cli_dir,
                admin_key,
                kernel_key,
                image_name,
                version,
                key,
                model,
                use_podman=use_podman,
            )

up/src/og_up/up.py
dbpunk-labs/octogen
#! /usr/bin/env python3
# vim:fenc=utf-8
#
# Copyright © 2023 imotai <imotai@imotai-ub>
#
# Distributed under terms of the MIT license.

""" """
import subprocess
import os
import sys
import io

USE_SHELL = sys.platform.startswith("win")


def run_with_realtime_print(
    command, universal_newlines=True, useshell=USE_SHELL, env=os.environ
):
    try:
        p = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            shell=useshell,
            env=env,
        )

        text_fd = io.TextIOWrapper(
            p.stdout, encoding="utf-8", newline=os.linesep, errors="replace"
        )
        while True:
            chunk = text_fd.read(40)
            if not chunk:
                break
            yield 0, chunk
        p.wait()
        yield p.returncode, ""
    except Exception as ex:
        yield -1, str(ex)

up/src/og_up/utils.py
dbpunk-labs/octogen
#! /usr/bin/env python3
# vim:fenc=utf-8
#
# Copyright © 2023 imotai <imotai@imotai-ub>
#
# Distributed under terms of the MIT license.

""" """

import os
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
from og_up.up import download_model
from og_up.up import load_docker_image
from og_up.up import get_latest_release_version
from og_up.up import start_octogen_for_codellama
from og_up.up import start_octogen_for_azure_openai
from og_up.up import start_octogen_for_openai
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 og_up.up import check_the_env
from og_up.up import run_install_cli
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"


@pytest.mark.skipif(sys.platform.startswith("win"), reason="skip on windows")
def test_check_the_env():
    console = Console()
    segments = []
    with Live(Group(*segments), console=console) as live:
        result, msg = check_the_env(live, segments)
        assert result


def test_check_the_env_win():
    console = Console()
    segments = []
    with Live(Group(*segments), console=console) as live:
        result, msg = check_the_env(live, segments, need_container=False)
        assert result


def test_run_print():
    use_dir = os.path.expanduser("~")
    command = ["ls", use_dir]
    result_code = 0
    for code, output in run_with_realtime_print(command):
        result_code = code
    assert code == 0, "bad return code"


def test_download_model():
    console = Console()
    segments = []
    with Live(Group(*segments), console=console) as live:
        result_code = download_model(
            live,
            segments,
            repo="TheBloke/CodeLlama-7B-Instruct-GGUF",
            filename="codellama-7b-instruct.Q2_K.gguf",
        )
        assert result_code == 0, "fail to download model"


def test_get_version():
    repo = "imotai/test_repos"
    console = Console()
    segments = []
    temp_dir = tempfile.mkdtemp(prefix="octogen")
    with Live(Group(*segments), console=console) as live:
        version = get_latest_release_version(repo, live, segments)
        assert "v0.1.0" == version


def test_install_cli():
    console = Console()
    segments = []
    with Live(Group(*segments), console=console) as live:
        assert run_install_cli(live, segments)


@pytest.mark.skipif(sys.platform.startswith("win"), reason="skip on windows")
def test_start_azure_openai_smoketest():
    console = Console()
    segments = []
    install_dir = tempfile.mkdtemp(prefix="octogen")
    cli_install_dir = tempfile.mkdtemp(prefix="octogen")
    admin_key = random_str(32)
    kernel_key = random_str(32)
    with Live(Group(*segments), console=console) as live:
        generate_kernel_env(live, segments, install_dir, kernel_key)
        code = load_docker_image("v0.4.41", "dbpunk/octogen", live, segments)
        assert code == 0, "bad result code of loading docker image"
        result = start_octogen_for_azure_openai(
            live,
            segments,
            install_dir,
            cli_install_dir,
            admin_key,
            kernel_key,
            "dbpunk/octogen",
            "v0.4.41",
            "azure_open_api_key",
            "test_deployment",
            "https://azure_base",
        )
        assert result


@pytest.mark.skipif(sys.platform.startswith("win"), reason="skip on windows")
def test_start_openai_smoketest():
    console = Console()
    segments = []
    install_dir = tempfile.mkdtemp(prefix="octogen")
    cli_install_dir = tempfile.mkdtemp(prefix="octogen")
    admin_key = random_str(32)
    kernel_key = random_str(32)
    with Live(Group(*segments), console=console) as live:
        generate_kernel_env(live, segments, install_dir, kernel_key)
        code = load_docker_image("v0.4.41", "dbpunk/octogen", live, segments)
        assert code == 0, "bad result code of loading docker image"
        result = start_octogen_for_openai(
            live,
            segments,
            install_dir,
            cli_install_dir,
            admin_key,
            kernel_key,
            "dbpunk/octogen",
            "v0.4.41",
            "openai_api_key",
            "gpt-3.5-turbo",
        )
        assert result


@pytest.mark.skipif(sys.platform.startswith("win"), reason="skip on windows")
def test_start_codellama_smoketest():
    console = Console()
    segments = []
    install_dir = tempfile.mkdtemp(prefix="octogen")
    cli_install_dir = tempfile.mkdtemp(prefix="octogen")
    admin_key = random_str(32)
    kernel_key = random_str(32)
    with Live(Group(*segments), console=console) as live:
        generate_kernel_env(live, segments, install_dir, kernel_key)
        code = load_docker_image("v0.4.41", "dbpunk/octogen", live, segments)
        assert code == 0, "bad result code of loading docker image"
        result = start_octogen_for_codellama(
            live,
            segments,
            "TheBloke/CodeLlama-7B-Instruct-GGUF",
            "codellama-7b-instruct.Q2_K.gguf",
            install_dir,
            cli_install_dir,
            admin_key,
            kernel_key,
            "dbpunk/octogen",
            "v0.4.41",
        )
        assert result


@pytest.mark.skipif(sys.platform.startswith("win"), reason="skip on windows")
def test_load_bad_docker_image():
    console = Console()
    segments = []
    with Live(Group(*segments), console=console) as live:
        code = load_docker_image("xxx", "xxxx", live, segments)
        assert code != 0, "loading image should be failed"


@pytest.mark.skipif(sys.platform.startswith("win"), reason="skip on windows")
def test_load_valid_docker_image():
    console = Console()
    segments = []
    with Live(Group(*segments), console=console) as live:
        code = load_docker_image("v0.4.41", "dbpunk/octogen", live, segments)
        assert code == 0, "loading image should be ok"

up/tests/up_tests.py
dbpunk-labs/octogen