youtu-agent 0.1.2__py3-none-any.whl

utils/openai_utils/openai_utils.py

@@ -0,0 +1,147 @@
+import logging
+from typing import Any
+
+from openai._streaming import AsyncStream
+from openai.types.chat import (
+    ChatCompletionChunk,
+    ChatCompletionMessage,
+    ChatCompletionMessageFunctionToolCall,
+    ChatCompletionMessageToolCall,
+    ChatCompletionToolParam,
+)
+from openai.types.responses import (
+    FunctionToolParam,
+    Response,
+    ResponseStreamEvent,
+)
+from pydantic import BaseModel
+
+from ..print_utils import PrintUtils
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIUtils:
+    # --------------------------------------------------------
+    # chat completions
+    # --------------------------------------------------------
+    @staticmethod
+    def print_message(message: ChatCompletionMessage) -> None:
+        if hasattr(message, "reasoning_content") and message.reasoning_content:
+            PrintUtils.print_info(f"{message.reasoning_content}")
+        if message.content:
+            PrintUtils.print_bot(f"{message.content}", add_prefix=True)
+        if message.tool_calls:
+            for tool_call in message.tool_calls:
+                PrintUtils.print_bot(f"<{tool_call.function.name}>{tool_call.function.arguments}", add_prefix=True)
+
+    @staticmethod
+    async def print_stream(stream: AsyncStream[ChatCompletionChunk]) -> ChatCompletionMessage:
+        final_tool_calls: dict[int, ChatCompletionMessageToolCall] = {}
+        content = ""
+        async for chunk in stream:
+            delta = chunk.choices[0].delta
+            if hasattr(delta, "reasoning_content") and delta.reasoning_content:
+                PrintUtils.print_info(f"{delta.reasoning_content}", end="", color="green")
+            if delta.content:
+                content += delta.content
+                PrintUtils.print_info(f"{delta.content}", end="", color="gray")
+            if delta.tool_calls:
+                for tool_call in delta.tool_calls:
+                    index = tool_call.index
+                    if index not in final_tool_calls:
+                        final_tool_calls[index] = tool_call
+                        PrintUtils.print_info(
+                            f"<{tool_call.function.name}>{tool_call.function.arguments}", end="", color="blue"
+                        )
+                    else:
+                        if final_tool_calls[index].function.arguments:
+                            final_tool_calls[index].function.arguments += tool_call.function.arguments
+                        else:
+                            final_tool_calls[index].function.arguments = tool_call.function.arguments
+                        PrintUtils.print_info(f"{tool_call.function.arguments}", end="", color="blue")
+        PrintUtils.print_info("")  # print a newline
+        tool_calls = [
+            ChatCompletionMessageFunctionToolCall(
+                id=tool_call.id,
+                function=tool_call.function.model_dump(),
+                type=tool_call.type,  # type is always "function"
+            )
+            for tool_call in final_tool_calls.values()
+        ]
+        message = ChatCompletionMessage(role="assistant", content=content, tool_calls=tool_calls)
+        OpenAIUtils.print_message(message)
+        return message
+
+    # --------------------------------------------------------
+    # responses
+    # --------------------------------------------------------
+    @staticmethod
+    def print_response(response: Response) -> None:
+        for item in response.output:
+            # print(f"> responses item: {item}")
+            match item.type:
+                case "reasoning":
+                    content = getattr(item, "content", item.summary)
+                    PrintUtils.print_bot(f"<reasoning>{content}</reasoning>", add_prefix=True, color="gray")
+                case "message":
+                    PrintUtils.print_bot(f"{item.content}", add_prefix=True)
+                case "function_call":
+                    PrintUtils.print_info(f"<{item.name}>({item.arguments})")
+                case "file_search_call":
+                    PrintUtils.print_info(f"<{item.type}>({item.queries})")
+                case "web_search_call":
+                    PrintUtils.print_info(f"<{item.type}>({item.action})")
+                case "computer_call":
+                    PrintUtils.print_info(f"<{item.type}>({item.action})")
+                case "image_generation_call":
+                    PrintUtils.print_info(f"<{item.type}> -> {item.result[:4]}")
+                case "code_interpreter_call":
+                    PrintUtils.print_info(
+                        f"<{item.type}>(container_id={item.container_id}, code={item.code}) -> {item.outputs}"
+                    )
+                case "local_shell_call":
+                    PrintUtils.print_info(f"<{item.type}>(action={item.action})")
+                case "mcp_list_tools":
+                    PrintUtils.print_info(f"<{item.type}>(server={item.server_label}) -> {item.tools}")
+                case "mcp_call":
+                    PrintUtils.print_info(
+                        f"<{item.type}>(server={item.server_label}) {item.name}({item.arguments}) -> {item.output}"
+                    )
+                case "mcp_approval_request":
+                    PrintUtils.print_info(f"<{item.type}>(server={item.server_label}) {item.name}({item.arguments})")
+                case _:
+                    PrintUtils.print_error(f"Unknown item type: {item.type}\n{item}")
+
+    @staticmethod
+    def print_response_stream(stream: AsyncStream[ResponseStreamEvent]) -> Response:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_response_configs(response: Response, include_output: bool = False) -> dict:
+        """Get response configs from response"""
+        data = response.model_dump()
+        if not include_output:
+            del data["output"]
+        return data
+
+    @staticmethod
+    def get_response_output(response: Response) -> list[dict]:
+        """Get response output from response"""
+        return response.model_dump()["output"]
+
+    @classmethod
+    def tool_chatcompletion_to_responses(cls, tool: ChatCompletionToolParam) -> FunctionToolParam:
+        assert tool["type"] == "function"
+        return FunctionToolParam(
+            name=tool["function"]["name"],
+            description=tool["function"].get("description", ""),
+            parameters=tool["function"].get("parameters", None),
+            type="function",
+        )
+
+    @staticmethod
+    def maybe_basemodel_to_dict(obj: Any) -> dict | None:
+        if isinstance(obj, BaseModel):
+            return obj.model_dump()
+        return obj

utils/openai_utils/simplified_client.py

@@ -0,0 +1,119 @@
+import logging
+import os
+from typing import Literal
+
+from openai import AsyncOpenAI, AsyncStream
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+from openai.types.responses import Response, ResponseStreamEvent
+
+from .types import (
+    OpenAIChatCompletionParams,
+    OpenAIChatCompletionParamsKeys,
+    OpenAIResponsesParams,
+    OpenAIResponsesParamsKeys,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class SimplifiedAsyncOpenAI(AsyncOpenAI):
+    """Simplified OpenAI client for chat.completions and responses API, with default config"""
+
+    def __init__(
+        self,
+        *,
+        type: Literal["chat.completions", "responses"] = None,
+        # openai client kwargs
+        api_key: str | None = None,
+        base_url: str | None = None,
+        # default configs
+        **kwargs: dict,
+    ) -> None:
+        logger.info(f"> type: {type}, base_url: {base_url}, kwargs: {kwargs}")
+        super().__init__(
+            api_key=api_key or os.getenv("UTU_LLM_API_KEY") or "xxx", base_url=base_url or os.getenv("UTU_LLM_BASE_URL")
+        )
+        self.type = type or os.getenv("UTU_LLM_TYPE", "chat.completions")
+        self.type_create_params = (
+            OpenAIChatCompletionParamsKeys if self.type == "chat.completions" else OpenAIResponsesParamsKeys
+        )
+        self.default_config = self._process_kwargs(kwargs)
+
+    def _process_kwargs(self, kwargs: dict) -> dict:
+        # parse kwargs for ChatCompletionParams
+        default_config = {}
+        for k, v in kwargs.items():
+            if k in self.type_create_params:
+                default_config[k] = v
+        default_config["model"] = default_config.get("model", os.getenv("UTU_LLM_MODEL"))
+        return default_config
+
+    async def query_one(self, **kwargs) -> str:
+        """Simplified chat.complete / responses API
+        WARNING: Only for basic text i/o usage! You should not use the method with querying with customized configs!
+        """
+        if "stream" in kwargs:
+            assert kwargs["stream"] is False, "stream is not supported in `query_one`"
+
+        if self.type == "chat.completions":
+            chat_completion: ChatCompletion = await self.chat_completions_create(**kwargs)
+            return chat_completion.choices[0].message.content
+        elif self.type == "responses":
+            response: Response = await self.responses_create(**kwargs)
+            return response.output_text  # NOTE: will not return toolcall or reasoning
+        else:
+            raise ValueError(f"Unknown type: {self.type}")
+
+    async def chat_completions_create(self, **kwargs) -> ChatCompletion | AsyncStream[ChatCompletionChunk]:
+        assert self.type == "chat.completions", "`chat_completions_create` is not supported for responses API"
+        unknown_params = self.check_known_keys(kwargs, self.type_create_params)
+        if unknown_params:
+            logger.warning(f"Unknown parameters: {unknown_params} for {self.type} API!")
+        kwargs = self.process_chat_completion_params(kwargs, self.default_config)
+        return await self.chat.completions.create(**kwargs)
+
+    async def responses_create(self, **kwargs) -> Response | AsyncStream[ResponseStreamEvent]:
+        unknown_params = self.check_known_keys(kwargs, self.type_create_params)
+        if unknown_params - {"messages"}:  # ignore
+            logger.warning(f"Unknown parameters: {unknown_params} for {self.type} API!")
+        assert self.type == "responses", "`responses_create` is not supported for chat.completions API"
+        kwargs = self.process_responses_params(kwargs, self.default_config)
+        return await self.responses.create(**kwargs)
+
+    def process_chat_completion_params(
+        self, kwargs: OpenAIChatCompletionParams, default_config: OpenAIChatCompletionParams
+    ) -> OpenAIChatCompletionParams:
+        """Process chat completion params, convert str to list of messages, merge default config"""
+        assert "messages" in kwargs
+        if isinstance(kwargs["messages"], str):
+            kwargs["messages"] = [{"role": "user", "content": kwargs["messages"]}]
+        return self._merge_default_config(kwargs, default_config)
+
+    def process_responses_params(
+        self, kwargs: OpenAIResponsesParams, default_config: OpenAIResponsesParams
+    ) -> OpenAIResponsesParams:
+        """Process responses params, convert str to list of messages, merge default config"""
+        if "input" not in kwargs:
+            # try parse query for chat.completions
+            assert "messages" in kwargs
+            input = kwargs.pop("messages")
+            if isinstance(input, str):
+                kwargs["input"] = [{"role": "user", "content": input}]
+            else:
+                kwargs["input"] = input
+        else:
+            if isinstance(kwargs["input"], str):
+                kwargs["input"] = [{"role": "user", "content": kwargs["input"]}]
+        return self._merge_default_config(kwargs, default_config)
+
+    def _merge_default_config(self, kwargs: dict, default_config: dict) -> dict:
+        """Merge default config"""
+        for k, v in default_config.items():
+            if k not in kwargs:
+                kwargs[k] = v
+        return kwargs
+
+    def check_known_keys(self, kwargs: dict, known_keys: set[str]) -> set:
+        """Check if all keys in kwargs are in known_keys"""
+        unknown_keys = set(kwargs.keys()) - known_keys
+        return unknown_keys

utils/openai_utils/types.py

@@ -0,0 +1,66 @@
+from collections.abc import Iterable
+from typing import Literal, TypedDict
+
+import httpx
+from openai._types import NOT_GIVEN, Body, Headers, NotGiven, Query
+from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolParam
+from openai.types.chat.completion_create_params import ResponseFormat
+from openai.types.responses import ResponseInputParam, ResponseTextConfigParam, ToolParam
+from openai.types.responses.response_create_params import ToolChoice
+from openai.types.shared import ChatModel, Reasoning, ReasoningEffort, ResponsesModel
+
+
+class OpenAICreateBaseParams(TypedDict):
+    stream: bool | None = False
+    # from openai.resources.chat.completions.Completions.create
+    extra_headers: Headers | None = None
+    extra_query: Query | None = None
+    extra_body: Body | None = None
+    timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN
+
+
+# CompletionCreateParams | https://platform.openai.com/docs/api-reference/chat/create
+class OpenAIChatCompletionParams(TypedDict, OpenAICreateBaseParams):
+    # NOTE: only for typing
+    messages: Iterable[ChatCompletionMessageParam]  # required
+    model: str | ChatModel  # required
+    frequency_penalty: float | None
+    logit_bias: dict[str, int] | None
+    logprobs: bool | None
+    max_completion_tokens: int | None
+    max_tokens: int | None
+    n: int | None
+    presence_penalty: float | None
+    reasoning_effort: ReasoningEffort | None
+    response_format: ResponseFormat
+    seed: int | None
+    temperature: float | None
+    top_p: float | None
+    tools: list[ChatCompletionToolParam] | None
+    tool_choice: Literal["none", "auto", "required"] | None
+    parallel_tool_calls: bool | None
+    stop: str | list[str] | None
+    top_logprobs: int | None
+
+
+# ResponseCreateParams | https://platform.openai.com/docs/api-reference/responses/create
+class OpenAIResponsesParams(TypedDict, OpenAICreateBaseParams):
+    input: str | ResponseInputParam
+    instructions: str | None
+    max_output_tokens: int | None
+    max_tool_calls: int | None
+    model: ResponsesModel
+    parallel_tool_calls: bool | None
+    previous_response_id: str | None
+    reasoning: Reasoning | None
+    temperature: float | None
+    text: ResponseTextConfigParam
+    tool_choice: ToolChoice
+    tools: Iterable[ToolParam]
+    top_logprobs: int | None
+    top_p: float | None
+    truncation: Literal["auto", "disabled"] | None
+
+
+OpenAIChatCompletionParamsKeys = OpenAIChatCompletionParams.__annotations__.keys()
+OpenAIResponsesParamsKeys = OpenAIResponsesParams.__annotations__.keys()

utils/path.py

@@ -0,0 +1,130 @@
+import hashlib
+import json
+import os
+import pathlib
+import tempfile
+from typing import Any
+from urllib.parse import urlparse
+
+import requests
+import yaml
+from jinja2 import Environment, FileSystemLoader, Template
+
+
+def get_package_path() -> pathlib.Path:
+    return pathlib.Path(__file__).parent.parent.parent
+
+
+DIR_ROOT = get_package_path()
+CACHE_DIR = DIR_ROOT / ".cache"
+
+
+class FileUtils:
+    @staticmethod
+    def is_web_url(url: str) -> bool:
+        parsed_url = urlparse(url)
+        return all([parsed_url.scheme, parsed_url.netloc])
+
+    @staticmethod
+    def get_file_ext(file_path: str) -> str:
+        if FileUtils.is_web_url(file_path):
+            return pathlib.Path(urlparse(file_path).path).suffix
+        return pathlib.Path(file_path).suffix
+
+    @staticmethod
+    def download_file(url: str, save_path: str = None) -> str:
+        """Download file from web. Return the saved path"""
+        # if not save_path, use tempfile
+        if not save_path:
+            save_path = tempfile.NamedTemporaryFile(
+                suffix=FileUtils.get_file_ext(url),
+                delete=False,
+            ).name
+        response = requests.get(url)
+        response.raise_for_status()
+        with open(save_path, "wb") as f:
+            f.write(response.content)
+        return save_path
+
+    @staticmethod
+    def get_file_md5(file_path: str) -> str:
+        """Clac md5 for local or web file"""
+        hash_md5 = hashlib.md5()
+        if FileUtils.is_web_url(file_path):
+            with requests.get(file_path, stream=True) as r:
+                r.raise_for_status()
+                for chunk in r.iter_content(chunk_size=4096):
+                    hash_md5.update(chunk)
+        else:
+            with open(file_path, "rb") as f:
+                for chunk in iter(lambda: f.read(4096), b""):
+                    hash_md5.update(chunk)
+        return hash_md5.hexdigest()
+
+    @staticmethod
+    def load_yaml(file_path: pathlib.Path | str) -> dict[str, Any]:
+        if isinstance(file_path, str):
+            file_path = pathlib.Path(file_path)
+        if not file_path.exists():
+            raise FileNotFoundError(f"File {file_path} does not exist")
+        with file_path.open() as f:
+            return yaml.safe_load(f)
+
+    @staticmethod
+    def load_prompts(fn: str | pathlib.Path) -> dict[str, str]:
+        """Load prompts from yaml file.
+
+        - Default path: `DIR_ROOT / "utu/prompts" / fn`
+        """
+        if isinstance(fn, str):
+            if not fn.endswith(".yaml"):
+                fn += ".yaml"
+            fn = DIR_ROOT / "utu" / "prompts" / fn
+        assert fn.exists(), f"File {fn} does not exist!"
+        with fn.open(encoding="utf-8") as f:
+            return yaml.safe_load(f)
+
+    @staticmethod
+    def get_jinja_env(directory: str | pathlib.Path) -> Environment:
+        if isinstance(directory, str):
+            directory = DIR_ROOT / "utu" / "prompts" / directory
+        if not directory.exists():
+            raise FileNotFoundError(f"Directory {directory} does not exist")
+        return Environment(loader=FileSystemLoader(directory))
+
+    @staticmethod
+    def get_jinja_template(template_path: str | pathlib.Path) -> Template:
+        if isinstance(template_path, str):
+            if not template_path.endswith(".j2"):
+                template_path += ".j2"
+            template_path = DIR_ROOT / "utu" / "prompts" / template_path
+        if not template_path.exists():
+            raise FileNotFoundError(f"File {template_path} does not exist")
+        with template_path.open(encoding="utf-8") as f:
+            return Template(f.read())
+
+    @staticmethod
+    def get_jinja_template_str(template_str: str) -> Template:
+        return Template(template_str)
+
+    @staticmethod
+    def load_json(file_path: str | pathlib.Path) -> dict[str, Any]:
+        if isinstance(file_path, str):
+            file_path = pathlib.Path(file_path)
+        if not file_path.exists():
+            raise FileNotFoundError(f"File {file_path} does not exist")
+        with file_path.open() as f:
+            return json.load(f)
+
+    @staticmethod
+    def load_json_data(file_path: str | pathlib.Path) -> list[dict[str, Any]]:
+        if isinstance(file_path, str) and not os.path.exists(file_path):
+            file_path = DIR_ROOT / "utu" / "data" / file_path
+        return FileUtils.load_json(file_path)
+
+    @staticmethod
+    def save_json(file_path: str | pathlib.Path, data: dict[str, Any]) -> None:
+        if isinstance(file_path, str):
+            file_path = pathlib.Path(file_path)
+        with file_path.open("w", encoding="utf-8") as f:
+            json.dump(data, f, ensure_ascii=False, indent=4)

utils/print_utils.py

@@ -0,0 +1,112 @@
+import json
+from collections import defaultdict
+
+import prompt_toolkit
+from colorama import Fore, Style, init
+from prompt_toolkit.patch_stdout import patch_stdout
+
+init(autoreset=True)  # Reset color to default (autoreset=True handles this automatically)
+
+COLOR_DICT = defaultdict(lambda: Style.RESET_ALL)
+COLOR_DICT.update(
+    {
+        "gray": Fore.LIGHTBLACK_EX,
+        "orange": Fore.LIGHTYELLOW_EX,
+        "red": Fore.RED,
+        "green": Fore.GREEN,
+        "blue": Fore.BLUE,
+        "yellow": Fore.YELLOW,
+        "magenta": Fore.MAGENTA,
+        "cyan": Fore.CYAN,
+        "white": Fore.WHITE,
+        "bold_blue": Style.BRIGHT + Fore.BLUE,
+    }
+)
+
+
+class PrintUtils:
+    @staticmethod
+    def print_input(prompt_text: str):
+        """styled user input"""
+        # https://github.com/prompt-toolkit/python-prompt-toolkit
+        # user_input = input(COLOR_DICT[prompt_color] + prompt_text + COLOR_DICT[input_color])
+        # print(Style.RESET_ALL, end="")
+        user_input = prompt_toolkit.prompt(prompt_text)
+        return user_input
+
+    @staticmethod
+    async def async_print_input(prompt_text: str) -> str:
+        # https://python-prompt-toolkit.readthedocs.io/en/master/pages/asking_for_input.html#prompt-in-an-asyncio-application  # noqa: E501
+        session = prompt_toolkit.PromptSession()
+        while True:
+            with patch_stdout():
+                result = await session.prompt_async(prompt_text)
+                return result
+
+    @staticmethod
+    def print_info(
+        msg: str,
+        color: str = "gray",
+        add_prefix: bool = False,
+        prefix: str = "",
+        end: str = "\n",
+        flush: bool = True,
+    ):
+        if add_prefix:
+            msg = prefix + " " + msg
+        print(COLOR_DICT[color] + msg + Style.RESET_ALL, end=end, flush=flush)
+
+    @staticmethod
+    def print_bot(
+        msg: str,
+        color: str = "orange",
+        add_prefix: bool = False,
+        prefix: str = "[BOT]",
+        end: str = "\n",
+        flush: bool = True,
+    ):
+        PrintUtils.print_info(msg, color=color, add_prefix=add_prefix, prefix=prefix, end=end, flush=flush)
+
+    @staticmethod
+    def print_tool(
+        msg: str,
+        color: str = "green",
+        add_prefix: bool = False,
+        prefix: str = "[TOOL]",
+        end: str = "\n",
+        flush: bool = True,
+    ):
+        PrintUtils.print_info(msg, color=color, add_prefix=add_prefix, prefix=prefix, end=end, flush=flush)
+
+    @staticmethod
+    def print_error(
+        msg: str,
+        color: str = "red",
+        add_prefix: bool = False,
+        prefix: str = "",
+        end: str = "\n",
+        flush: bool = True,
+    ):
+        PrintUtils.print_info(msg, color=color, add_prefix=add_prefix, prefix=prefix, end=end, flush=flush)
+
+    @staticmethod
+    def format_json(obj: dict, indent: int = None) -> str:
+        return json.dumps(obj, indent=indent, ensure_ascii=False)
+
+    @staticmethod
+    def truncate_text(text: str | dict, max_length: int = 100, oneline: bool = True) -> str:
+        """Truncate text to max_length, adding ellipsis if truncated.
+
+        Args:
+            text: The text to truncate
+            max_length: Maximum length of the returned text
+            oneline: Whether to convert text to a single line by removing newlines
+        """
+        if isinstance(text, dict):
+            text = PrintUtils.format_json(text)
+        if oneline:
+            text = " ".join(text.splitlines())
+        if len(text) <= max_length:
+            return text
+        else:
+            return text[: max_length - 3] + "..."

utils/script_utils.py

@@ -0,0 +1,22 @@
+import argparse
+
+from ..config import AgentConfig, ConfigLoader
+from .print_utils import PrintUtils
+
+
+def parse_cli_args() -> AgentConfig:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_name", type=str, default="simple/base", help="Configuration name")
+    parser.add_argument("--agent_model", type=str, default=None, help="Agent model.")
+    args = parser.parse_args()
+
+    config: AgentConfig = ConfigLoader.load_agent_config(args.config_name)
+    # Override basic configs
+    if args.agent_model:
+        assert config.type == "simple", f"--agent_model only support SimpleAgent, get {config.type}"
+        config.model.model_provider.model = args.agent_model
+    if config.type == "workforce":
+        PrintUtils.print_info("Error: Workforce agent is not supported in CLI mode yet.")
+        raise NotImplementedError
+
+    return config