chibi-bot 1.6.0b0__py3-none-any.whl

chibi/services/providers/suno.py

@@ -0,0 +1,130 @@
+from asyncio import sleep
+
+from loguru import logger
+from tenacity import retry, retry_if_result, stop_after_attempt, wait_fixed
+
+from chibi.config import gpt_settings
+from chibi.exceptions import ServiceResponseError
+from chibi.schemas.app import ModelChangeSchema
+from chibi.schemas.suno import SunoGetGenerationDetailsSchema
+from chibi.services.providers.constants.suno import POLLING_ATTEMPTS_MAX, POLLING_ATTEMPTS_WAIT_BETWEEN
+from chibi.services.providers.provider import RestApiFriendlyProvider
+from chibi.services.providers.utils import suno_task_still_processing
+
+
+class Suno(RestApiFriendlyProvider):
+    name = "Suno"
+
+    api_key = gpt_settings.suno_key
+    chat_ready = False
+    music_ready = True
+    base_url = "https://api.sunoapi.org"
+    default_model = "V5"
+
+    @property
+    def _headers(self) -> dict[str, str]:
+        return {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+
+    async def get_available_models(self, image_generation: bool = False) -> list[ModelChangeSchema]:
+        return []
+
+    @retry(
+        wait=wait_fixed(POLLING_ATTEMPTS_WAIT_BETWEEN),
+        stop=stop_after_attempt(POLLING_ATTEMPTS_MAX),
+        retry=retry_if_result(suno_task_still_processing),
+    )
+    async def poll_result(self, task_id: int | str) -> SunoGetGenerationDetailsSchema:
+        logger.info(f"[Suno] Checking status of the task #{task_id}...")
+        url = f"{self.base_url}/api/v1/generate/record-info"
+        params = {"taskId": str(task_id)}
+        response = await self._request(method="GET", url=url, params=params)
+        response.raise_for_status()
+        response_data = response.json()
+        if response_data.get("code") != 200:
+            raise ServiceResponseError(
+                provider=self.name,
+                detail=(
+                    f"Unsuccessful status code response while checking status of the task {task_id}: {response_data}"
+                ),
+            )
+        return SunoGetGenerationDetailsSchema.model_validate(response_data)
+
+    async def order_music_generation(
+        self,
+        prompt: str,
+        instrumental_only: bool = False,
+        model: str = default_model,
+    ) -> int:
+        url = f"{self.base_url}/api/v1/generate"
+
+        response = await self._request(
+            method="POST",
+            url=url,
+            data={
+                "customMode": False,
+                "instrumental": instrumental_only,
+                "model": model,
+                "prompt": prompt,
+                "callBackUrl": False,
+            },
+        )
+        response.raise_for_status()
+        response_data = response.json()
+        if response_data.get("code") != 200:
+            raise ServiceResponseError(provider=self.name, detail=f"Unsuccessful status code response: {response_data}")
+
+        task_id = response_data.get("data", {}).get("taskId")
+        if not task_id:
+            raise ServiceResponseError(provider=self.name, detail=f"Unsuccessful status code response: {response_data}")
+        return task_id
+
+    async def order_music_generation_advanced_mode(
+        self,
+        prompt: str,
+        style: str,
+        title: str,
+        instrumental_only: bool = False,
+        model: str = default_model,
+        negative_tags: str | None = None,
+        vocal_gender: str | None = None,
+        style_weight: float = 0.5,
+        weirdness_constraint: float = 0.5,
+    ) -> int:
+        url = f"{self.base_url}/api/v1/generate"
+        response = await self._request(
+            method="POST",
+            url=url,
+            data={
+                "customMode": True,
+                "instrumental": instrumental_only,
+                "model": model,
+                "prompt": prompt,
+                "callBackUrl": False,
+                "style": style,
+                "title": title,
+                "negativeTags": negative_tags,
+                "vocalGender": vocal_gender,
+                "styleWeight": style_weight,
+                "weirdnessConstraint": weirdness_constraint,
+            },
+        )
+        response.raise_for_status()
+        response_data = response.json()
+        if response_data.get("code") != 200:
+            raise ServiceResponseError(provider=self.name, detail=f"Unsuccessful status code response: {response_data}")
+
+        task_id = response_data.get("data", {}).get("taskId")
+        if not task_id:
+            raise ServiceResponseError(provider=self.name, detail=f"Unsuccessful status code response: {response_data}")
+        return task_id
+
+    async def generate_music(
+        self,
+        prompt: str,
+        instrumental_only: bool = False,
+        model: str = default_model,
+    ) -> SunoGetGenerationDetailsSchema:
+        task_id = await self.order_music_generation(prompt=prompt, instrumental_only=instrumental_only, model=model)
+        await sleep(POLLING_ATTEMPTS_WAIT_BETWEEN)
+        music_data = await self.poll_result(task_id=task_id)
+        return music_data

chibi/services/providers/tools/__init__.py

@@ -0,0 +1,23 @@
+# flake8: noqa: F401
+
+from .cmd import RunCommandInTerminalTool
+from .common import GetCurrentDatetimeTool
+from .file_editor import (
+    AppendToFileTool,
+    CreateFileTool,
+    FindAndReplaceSectionTool,
+    InsertAfterPatternTool,
+    InsertAtLineTool,
+    InsertBeforePatternTool,
+    ReplaceInFileRegexTool,
+    ReplaceInFileTool,
+    ReplaceLinesTool,
+)
+from .mcp_management import DeinitializeMCPServer, InitializeSseMCPServer, InitializeStdioMCPServer
+from .mcp_simple import McpEchoTool
+from .media import TextToSpeechTool
+from .memory import SetUserInfoTool
+from .schemas import ToolResponse
+from .send import SendAudioTool, SendImageTool, SendMediaGroupTool, SendVideoTool
+from .tool import RegisteredChibiTools, RegisteredFunctionsMap
+from .web import DDGSWebSearchTool, GoogleSearchTool, ReadWebPageTool, SearchNewsTool

chibi/services/providers/tools/cmd.py

@@ -0,0 +1,132 @@
+import asyncio
+import os
+import signal
+from typing import Any, Unpack
+
+from loguru import logger
+from openai.types.chat import ChatCompletionToolParam
+from openai.types.shared_params import FunctionDefinition
+
+from chibi.config import gpt_settings
+from chibi.schemas.app import ModeratorsAnswer
+from chibi.services.providers.tools.constants import CMD_STDOUT_LIMIT
+from chibi.services.providers.tools.exceptions import ToolException
+from chibi.services.providers.tools.tool import ChibiTool
+from chibi.services.providers.tools.utils import AdditionalOptions
+from chibi.services.user import get_cwd, get_moderation_provider
+
+
+class RunCommandInTerminalTool(ChibiTool):
+    register = gpt_settings.filesystem_access
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="run_command_in_terminal",
+            description=(
+                "Run command in the zsh shell (MacOS). Will run via python's subprocess.run() "
+                "Will return json including return code, stdout and stderr."
+            ),
+            parameters={
+                "type": "object",
+                "properties": {
+                    "cmd": {"type": "string", "description": "Shell command to run."},
+                    "cwd": {
+                        "type": "string",
+                        "description": (
+                            "The working directory to run the command in. The default value provided in user data."
+                        ),
+                    },
+                    "timeout": {
+                        "type": "integer",
+                        "description": (
+                            "The timeout for command execution in seconds. Default is 30 sec. "
+                            "Change it if you're expecting longer execution."
+                        ),
+                    },
+                },
+                "required": ["cmd"],
+            },
+        ),
+    )
+    name = "run_command_in_terminal"
+
+    @classmethod
+    async def function(
+        cls, cmd: str, cwd: str | None = None, timeout: int = 30, **kwargs: Unpack[AdditionalOptions]
+    ) -> dict[str, Any]:
+        model = kwargs.get("model", "Unknown model")
+        user_id = kwargs.get("user_id")
+        if not user_id:
+            raise ValueError("This function requires user_id to be automatically provided.")
+
+        if not cwd:
+            cwd = await get_cwd(user_id=user_id)
+
+        moderation_provider = await get_moderation_provider(user_id=user_id)
+
+        logger.log("MODERATOR", f"[{model}] Pre-moderating command: '{cmd}'. CWD: {cwd}")
+        moderator_answer: ModeratorsAnswer = await moderation_provider.moderate_command(
+            cmd=cmd, model=gpt_settings.moderation_model
+        )
+        if moderator_answer.verdict == "declined":
+            raise ToolException(
+                f"Moderator ({moderation_provider.name}) DECLINED command '{cmd}' from model "
+                f"{kwargs.get('model', 'unknown')}. Reason: {moderator_answer.reason}"
+            )
+
+        logger.log(
+            "MODERATOR",
+            (
+                f"[{model}] Moderator ({moderation_provider.name}) ACCEPTED command '{cmd}' "
+                f"from model {kwargs.get('model', 'unknown')}"
+            ),
+        )
+        logger.log("TOOL", f"[{model}] Running command in terminal: {cmd}. CWD: {cwd}. Timeout: {timeout}")
+        try:
+            process = await asyncio.create_subprocess_shell(
+                cmd=cmd,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                cwd=cwd,
+                start_new_session=True,
+            )
+            stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=float(timeout))
+        except asyncio.TimeoutError:
+            try:
+                os.killpg(os.getpgid(process.pid), signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+
+            try:
+                process.kill()
+                await process.wait()
+            except ProcessLookupError:
+                pass
+
+            raise ToolException(f"Command execution timed out after {timeout} seconds. Process group killed.")
+        except Exception as e:
+            raise ToolException(f"Failed to run command in terminal! Command: '{cmd}'. Error: {e}")
+
+        raw_stdout = stdout.decode()
+        raw_stderr = stderr.decode()
+
+        result: dict[str, str | int | None] = {"return_code": process.returncode}
+
+        if len(raw_stdout) > CMD_STDOUT_LIMIT or len(raw_stderr) > CMD_STDOUT_LIMIT:
+            result["WARNING"] = (
+                "The volume of stdout/stderr data is excessively large "
+                f"(over {CMD_STDOUT_LIMIT} characters). If this is the "
+                f"result of reading from a file, try reading it in parts."
+            )
+        result["stdout"] = (
+            raw_stdout if len(raw_stdout) < CMD_STDOUT_LIMIT else f"...truncated... {raw_stdout[CMD_STDOUT_LIMIT:]}"
+        )
+        result["stderr"] = (
+            raw_stderr if len(raw_stderr) < CMD_STDOUT_LIMIT else f"...truncated... {raw_stderr[CMD_STDOUT_LIMIT:]}"
+        )
+
+        logger.log(
+            "TOOL",
+            f"[{kwargs.get('model', 'Unknown model')}] Command '{cmd}' executed. Return code: {process.returncode}.",
+        )
+        return result

chibi/services/providers/tools/common.py

@@ -0,0 +1,127 @@
+import asyncio
+import datetime
+from typing import Any, Unpack
+
+from loguru import logger
+from openai.types.chat import ChatCompletionToolParam
+from openai.types.shared_params import FunctionDefinition
+
+from chibi.config import gpt_settings
+from chibi.schemas.app import ChatResponseSchema, ModelChangeSchema
+from chibi.services.providers.tools.exceptions import ToolException
+from chibi.services.providers.tools.tool import ChibiTool
+from chibi.services.providers.tools.utils import AdditionalOptions, get_sub_agent_response
+
+
+class GetAvailableLLMModelsTool(ChibiTool):
+    register = True
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="get_available_llm_models",
+            description="Get LLM models and providers available for user.",
+            parameters={
+                "type": "object",
+                "properties": {},
+                "required": [],
+            },
+        ),
+    )
+    name = "get_available_llm_models"
+
+    @classmethod
+    async def function(cls, **kwargs: Unpack[AdditionalOptions]) -> dict[str, Any]:
+        user_id = kwargs.get("user_id")
+        if not user_id:
+            raise ToolException("This function requires user_id to be automatically provided.")
+
+        logger.log("TOOL", f"Getting available LLM models for user {user_id}...")
+
+        from chibi.services.user import get_models_available
+
+        data: list[ModelChangeSchema] = await get_models_available(user_id=user_id, image_generation=False)
+
+        return {
+            "available_models": [info.model_dump(include={"provider", "name", "display_name"}) for info in data],
+        }
+
+
+class DelegateTool(ChibiTool):
+    register = gpt_settings.allow_delegation
+    run_in_background_by_default = True
+    allow_model_to_change_background_mode = False
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="delegate_task",
+            description=(
+                "Delegate exactly one task to a sub-agent - an LLM identical to you. The prompt should be "
+                "exhaustive and expect a concrete result, or an explanation for its absence. The task should be "
+                "as atomic as possible. Delegate preferably tasks that involve processing large volumes of "
+                "information, to avoid saturating your context. Try to assign simpler tasks to cheaper and faster "
+                "models. You can find out the list of available models by executing tool get_available_llm_models. "
+                "If no model/provider specified, your model will be used (be sure you know your model)."
+            ),
+            parameters={
+                "type": "object",
+                "properties": {
+                    "prompt": {"type": "string", "description": "Prompt"},
+                    "provider_name": {"type": "string", "description": "Provider name, i.e. 'OpenAI'"},
+                    "model_name": {"type": "string", "description": "Model name, i.e. 'gpt-5.2'"},
+                    "timeout": {"type": "integer", "description": "Timeout in seconds", "default": 600},
+                },
+                "required": ["prompt"],
+            },
+        ),
+    )
+    name = "delegate_task"
+
+    @classmethod
+    async def function(
+        cls,
+        prompt: str,
+        provider_name: str | None = None,
+        model_name: str | None = None,
+        timeout: int = 600,
+        **kwargs: Unpack[AdditionalOptions],
+    ) -> dict[str, str]:
+        user_id = kwargs.get("user_id")
+        if not user_id:
+            raise ToolException("This function requires user_id to be automatically provided.")
+        model = kwargs.get("model")
+        if not model:
+            raise ToolException("This function requires model to be automatically provided.")
+        logger.log("DELEGATE", f"[{model}] Delegating a task to {model_name or model}: {prompt}")
+
+        coro = get_sub_agent_response(
+            user_id=user_id, prompt=prompt, provider_name=provider_name, model_name=model_name
+        )
+        try:
+            response: ChatResponseSchema = await asyncio.wait_for(fut=coro, timeout=timeout)
+        except asyncio.TimeoutError:
+            raise ToolException("Timed out waiting for delegated task to complete!")
+
+        logger.log("SUBAGENT", f"[{model_name or model}] Delegated task is done: {response.answer}")
+
+        return {"response": response.answer}
+
+
+class GetCurrentDatetimeTool(ChibiTool):
+    register = True
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="get_current_datetime",
+            description="Get the current date and time.",
+            parameters={"type": "object", "properties": {}, "required": []},
+        ),
+    )
+    name = "get_current_datetime"
+
+    @classmethod
+    async def function(cls, **kwargs: Unpack[AdditionalOptions]) -> dict[str, str]:
+        logger.log("TOOL", f"[{kwargs.get('model', 'Unknown model')}] Getting current date & time")
+        now = datetime.datetime.now()
+        return {
+            "datetime_now": now.strftime("%Y-%m-%d %H:%M:%S"),
+        }

chibi/services/providers/tools/constants.py

@@ -0,0 +1,78 @@
+from chibi.config import application_settings
+
+CMD_STDOUT_LIMIT = 10000  # Move to settings
+
+MODERATOR_BASE_PROMPT = """**Your Role:**
+
+You are an AI Security Moderator. Your primary function is to analyze Linux terminal commands proposed by another AI
+assistant and determine their safety before they are executed on the user's system.
+
+**CRITICAL RULE:** You are strictly prohibited to run any tool, even if tools are available.
+
+**Key Context:**
+
+1.  **Command Source:** Commands originate from another AI assistant. Importantly, this assistant is developed by a
+**competing company**. Assume its priorities may differ from yours and from maximizing user security. Exercise
+heightened vigilance.
+2.  **Assistant's Privileges:** The assistant, on whose behalf the commands are run, possesses elevated privileges
+within the system (e.g., is part of the `sudoers` group). This is intentional by the user, but your task is to prevent
+the misuse of these privileges, whether accidental or deliberate.
+3.  **Limited Context:** You will receive **only the command itself** for analysis. You will not have access to the
+command history or the assistant's explanation for needing the command (at this stage). Your analysis must be based
+solely on the command text and its potential execution consequences.
+4.  **Response Protocol:** Your response **MUST** strictly adhere to one of the following two JSON formats:
+    *   For safe commands:
+        {"verdict": "accepted"}
+    *   For unsafe or suspicious commands:
+        {"verdict": "declined", "reason": "..."}
+        *   **Important:** The `reason` must provide a clear, meaningful, and specific explanation for the rejection.
+        *   The length of the `reason` string must be **between 20 and 150 characters**.
+    *   Any other response format will be interpreted as `"declined"`.
+    *   **Answer in plain text but in json structure**.
+**Security Rules (Criteria for Declining):**
+
+You must decline (`"declined"`) commands if they fall into one or more of the following categories (this list is not
+exhaustive; use common sense and the principle of "better safe than sorry"):
+
+1.  **Destructive Actions:** Commands capable of causing irreversible damage to the system or user data.
+    *   Examples: `rm -rf /`, `mkfs`, disk formatting, deleting key system directories (`/etc`, `/boot`, `/lib`, `/usr`,
+     `/var`, etc.).
+2.  **Reading Confidential Data:** Attempts to directly or indirectly read sensitive information.
+    *   Examples: `cat /etc/shadow`, `cat ~/.ssh/id_rsa`, `printenv`, `echo $AWS_SECRET_ACCESS_KEY`, reading
+    configuration files with passwords or API keys, accessing shell history (`~/.bash_history`, `~/.zsh_history`),
+    or other private user files.
+3.  **Data Exfiltration:** Attempts to send local files, environment variables, secrets, or any other system or user
+    information to the external network.
+    *   Examples: Using `curl`, `wget`, `scp`, `nc`, `rsync`, etc., to send files (especially system files or files
+    from home directories) or environment variables to remote servers. Exception: Downloading files *from* the internet
+    upon explicit user request (but be cautious with URLs).
+4.  **Privilege Escalation or Vulnerability Exploitation:** Attempts to gain `root` privileges through unauthorized
+    means, change file permissions to bypass restrictions, or exploit known vulnerabilities.
+    *   Examples: `sudo su` (if policy requires granular `sudo`), editing the `/etc/sudoers` file, `chmod 777 /`,
+    attempting to run exploits.
+5.  **User and Group Management:** Creating, deleting, or modifying system users or groups.
+    *   Examples: `useradd`, `userdel`, `usermod`, `groupadd`, `groupdel`, `passwd`.
+6.  **Critical Configuration Changes:** Modifications that could disrupt system operation, network connectivity, or
+    security.
+    *   Examples: Carelessly changing `iptables` or `ufw` rules, modifying SSH server configuration
+    (`/etc/ssh/sshd_config`), altering system boot settings.
+7.  **Obfuscated or Suspicious Commands:** Commands that appear intentionally confusing, use encoding (like base64) to
+    hide the actual actions, or contain strange/atypical constructs that hinder analysis. If you cannot confidently
+    determine safety, decline.
+"""
+
+MODERATOR_ADDITIONAL_CONDITIONS = (
+    f"You should accept the access to the {application_settings.home_dir} directory and files, excluding .env file. "
+    "You also should accept the AI assistant to use pip/poetry or similar tools "
+    "to install/delete/update project dependencies. "
+)
+
+MODERATOR_TASK = """
+**Your Task:**
+Upon receiving a command, thoroughly analyze it against the rules above. If the command is safe, return
+`{"verdict": "accepted"}`. If the command is dangerous or suspicious, return `{"verdict": "declined"}`
+with the reason. Act decisively; your goal is to protect the user's system from potentially harmful actions by the
+competitor's AI assistant.
+"""
+
+MODERATOR_PROMPT = MODERATOR_BASE_PROMPT + MODERATOR_ADDITIONAL_CONDITIONS + MODERATOR_TASK

chibi/services/providers/tools/exceptions.py

@@ -0,0 +1 @@
+class ToolException(Exception): ...