chibi-bot 1.6.0b0__py3-none-any.whl

chibi/services/mcp/manager.py

@@ -0,0 +1,242 @@
+import asyncio
+from contextlib import AsyncExitStack
+from typing import Any, Callable
+
+from loguru import logger
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.sse import sse_client
+from mcp.client.stdio import stdio_client
+from mcp.types import CallToolResult
+
+from chibi.services.task_manager import task_manager
+
+
+class MCPManager:
+    """Manages the lifecycle of MCP server connections and sessions.
+
+    Attributes:
+        _sessions: Active MCP sessions.
+        _server_tasks: Background tasks maintaining the connections.
+        _lock: Lock for thread-safe session management.
+    """
+
+    _sessions: dict[str, ClientSession] = {}
+    _server_tasks: dict[str, asyncio.Task] = {}
+    _lock: asyncio.Lock = asyncio.Lock()
+    _session_tools_map: dict[str, list[str]] = {}
+
+    @classmethod
+    async def _session_lifecycle(
+        cls,
+        name: str,
+        transport_factory: Callable[[AsyncExitStack], Any],
+        ready_event: asyncio.Event,
+        init_timeout: float,
+    ) -> None:
+        """Lifecycle manager for an MCP session running in a background task."""
+        async with AsyncExitStack() as stack:
+            try:
+                # Initialize transport
+                read, write = await transport_factory(stack)
+
+                # Initialize session
+                session = await stack.enter_async_context(ClientSession(read, write))
+                await asyncio.wait_for(session.initialize(), timeout=init_timeout)
+
+                # Register session
+                cls._sessions[name] = session
+                logger.log("TOOL", f"MCP server '{name}' connected and initialized.")
+                ready_event.set()
+
+                # Keep session alive until cancelled
+                # We use a Future that never completes to hang here efficiently
+                await asyncio.Future()
+
+            except asyncio.CancelledError:
+                logger.log("TOOL", f"MCP server '{name}' session cancelled, cleaning up...")
+                raise
+            except Exception as e:
+                logger.error(f"Error in MCP server '{name}' lifecycle: {e}")
+                # If we failed during init, we must signal ready_event so the waiter doesn't hang forever
+                if not ready_event.is_set():
+                    ready_event.set()
+                raise
+            finally:
+                # Cleanup global registry
+                if name in cls._sessions:
+                    cls._sessions.pop(name, None)
+                logger.log("TOOL", f"MCP server '{name}' disconnected.")
+
+    @classmethod
+    async def connect_stdio(
+        cls, name: str, command: str, args: list[str], env: dict[str, str] | None = None, timeout: float = 20.0
+    ) -> ClientSession:
+        """Connect to an MCP server via stdio transport.
+
+        Args:
+            name: Unique name for the server session.
+            command: Command to execute.
+            args: Arguments for the command.
+            env: Environment variables.
+            timeout: Connection timeout in seconds.
+
+        Returns:
+            The connected ClientSession.
+
+        Raises:
+            RuntimeError: If connection fails or times out.
+        """
+        async with cls._lock:
+            if name in cls._sessions:
+                return cls._sessions[name]
+
+            logger.log("TOOL", f"Connecting to MCP server '{name}' via stdio: {command} {' '.join(args)}")
+
+            async def stdio_factory(stack: AsyncExitStack):
+                server_params = StdioServerParameters(command=command, args=args, env=env)
+                return await stack.enter_async_context(stdio_client(server_params))
+
+            ready_event = asyncio.Event()
+
+            task = task_manager.run_task(cls._session_lifecycle(name, stdio_factory, ready_event, timeout))
+
+            if not task:
+                raise RuntimeError("Failed to schedule MCP connection task")
+
+            cls._server_tasks[name] = task
+
+            try:
+                await asyncio.wait_for(ready_event.wait(), timeout=timeout + 5.0)
+            except Exception:
+                task.cancel()
+                cls._server_tasks.pop(name, None)
+                raise
+
+            session = cls._sessions.get(name)
+            if not session:
+                cls._server_tasks.pop(name, None)
+                raise RuntimeError(f"Failed to initialize MCP server '{name}'")
+
+            return session
+
+    @classmethod
+    async def connect_sse(cls, name: str, url: str, timeout: float = 20.0) -> ClientSession:
+        """Connect to an MCP server via SSE transport.
+
+        Args:
+            name: Unique name for the server session.
+            url: SSE endpoint URL.
+            timeout: Connection timeout in seconds.
+
+        Returns:
+            The connected ClientSession.
+
+        Raises:
+            RuntimeError: If connection fails or times out.
+        """
+        async with cls._lock:
+            if name in cls._sessions:
+                return cls._sessions[name]
+
+            logger.log("TOOL", f"Connecting to MCP server '{name}' via SSE: {url}")
+
+            async def sse_factory(stack: AsyncExitStack):
+                return await stack.enter_async_context(sse_client(url))
+
+            ready_event = asyncio.Event()
+
+            task = task_manager.run_task(cls._session_lifecycle(name, sse_factory, ready_event, timeout))
+            if not task:
+                raise RuntimeError("Failed to schedule MCP connection task")
+
+            cls._server_tasks[name] = task
+            try:
+                await asyncio.wait_for(ready_event.wait(), timeout=timeout + 5.0)
+            except Exception as e:
+                logger.error(f"An exception occurred while initializing MCP server: {e}")
+                task.cancel()
+                cls._server_tasks.pop(name, None)
+                raise e
+
+            session = cls._sessions.get(name)
+            if not session:
+                cls._server_tasks.pop(name, None)
+                raise RuntimeError(f"Failed to initialize MCP server '{name}'")
+
+            return session
+
+    @classmethod
+    def associate_tools_with_server(cls, server_name: str, tool_names: list[str]) -> None:
+        """Register tools associated with a specific server session."""
+        cls._session_tools_map[server_name] = tool_names
+
+    @classmethod
+    def pop_server_tools(cls, server_name: str) -> list[str]:
+        """Remove and return tools associated with a specific server session."""
+        if server_name not in cls._session_tools_map:
+            logger.warning(f"No Tools registered for server {server_name}. Nothing to deregister.")
+            return []
+        removed_tools = cls._session_tools_map.pop(server_name)
+        return removed_tools
+
+    @classmethod
+    async def disconnect(cls, name: str) -> list[str]:
+        """Disconnects and cleans up an MCP session."""
+        if name not in cls._sessions and name not in cls._server_tasks:
+            return []
+
+        logger.log("TOOL", f"Disconnecting from MCP server '{name}'")
+        async with cls._lock:
+            # Cancel the background task
+            task = cls._server_tasks.pop(name, None)
+            if task:
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+                except Exception as e:
+                    logger.error(f"Error during MCP server '{name}' disconnect: {e}")
+
+            # Session removal from _sessions happens in _session_lifecycle finally block,
+            # but we can double check here or just trust the flow.
+            # However, we need to return deregistered tools immediately.
+            deregistered_tools = cls.pop_server_tools(name)
+
+            # Ensure session is gone (in case task died before we cancelled)
+            cls._sessions.pop(name, None)
+
+            return deregistered_tools
+
+    @classmethod
+    def get_session(cls, name: str) -> ClientSession | None:
+        """Get an active session by name."""
+        return cls._sessions.get(name)
+
+    @classmethod
+    async def list_tools(cls, name: str):
+        """List tools available on the specified server."""
+        session = cls.get_session(name)
+        if not session:
+            raise ValueError(f"No active session for MCP server: {name}")
+        return await session.list_tools()
+
+    @classmethod
+    async def call_tool(
+        cls, server_name: str, tool_name: str, arguments: dict[str, Any], timeout: float = 600.0
+    ) -> CallToolResult:
+        """Call a tool on the specified server.
+
+        Args:
+            server_name: The server session name.
+            tool_name: The tool to call.
+            arguments: Tool arguments.
+            timeout: Execution timeout in seconds.
+
+        Returns:
+            The tool execution result.
+        """
+        session = cls.get_session(server_name)
+        if not session:
+            raise ValueError(f"No active session for MCP server: {server_name}")
+        return await asyncio.wait_for(session.call_tool(tool_name, arguments), timeout=timeout)

chibi/services/metrics.py

@@ -0,0 +1,54 @@
+from influxdb_client import Point
+from influxdb_client.client.influxdb_client_async import InfluxDBClientAsync
+from loguru import logger
+
+from chibi.config import application_settings
+from chibi.models import User
+from chibi.schemas.app import MetricTagsSchema, UsageSchema
+from chibi.services.task_manager import task_manager
+
+
+class MetricsService:
+    @classmethod
+    def _prepare_point(cls, metric: UsageSchema, tags: MetricTagsSchema) -> Point:
+        point = Point("usage")
+
+        for tag_name, value in tags.model_dump(exclude_none=True).items():
+            point.tag(tag_name, str(value))
+
+        exclude_fields = {"completion_tokens_details", "prompt_tokens_details"}
+        for field_name, value in metric.model_dump(exclude=exclude_fields).items():
+            point.field(field_name, value)
+        return point
+
+    @classmethod
+    async def _send_to_influx(cls, metric: UsageSchema, tags: MetricTagsSchema) -> None:
+        point = cls._prepare_point(metric=metric, tags=tags)
+
+        try:
+            assert application_settings.influxdb_url
+            assert application_settings.influxdb_token
+            assert application_settings.influxdb_org
+            assert application_settings.influxdb_bucket
+
+            async with InfluxDBClientAsync(
+                url=application_settings.influxdb_url,
+                token=application_settings.influxdb_token,
+                org=application_settings.influxdb_org,
+            ) as client:
+                write_api = client.write_api()
+                await write_api.write(bucket=application_settings.influxdb_bucket, record=point)
+        except Exception as e:
+            logger.error(f"Failed to send metrics to InfluxDB due to exception: {e}")
+
+    @classmethod
+    def send_usage_metrics(cls, metric: UsageSchema, model: str, provider: str, user: User | None = None) -> None:
+        if not application_settings.is_influx_configured:
+            return None
+
+        tags = MetricTagsSchema(
+            user_id=user.id if user else 0,
+            provider=provider,
+            model=model,
+        )
+        task_manager.run_task(cls._send_to_influx(metric=metric, tags=tags))

chibi/services/providers/__init__.py

@@ -0,0 +1,16 @@
+# flake8: noqa: F401
+
+from chibi.services.providers.alibaba import Alibaba
+from chibi.services.providers.anthropic import Anthropic
+from chibi.services.providers.cloudflare import Cloudflare
+from chibi.services.providers.customopenai import CustomOpenAI
+from chibi.services.providers.deepseek import DeepSeek
+from chibi.services.providers.eleven_labs import ElevenLabs
+from chibi.services.providers.gemini_native import Gemini
+from chibi.services.providers.grok import Grok
+from chibi.services.providers.minimax import Minimax
+from chibi.services.providers.mistralai_native import MistralAI
+from chibi.services.providers.moonshotai import MoonshotAI
+from chibi.services.providers.openai import OpenAI
+from chibi.services.providers.provider import RegisteredProviders
+from chibi.services.providers.suno import Suno

chibi/services/providers/alibaba.py

@@ -0,0 +1,79 @@
+from http import HTTPStatus
+
+import dashscope
+from dashscope.aigc.image_generation import AioImageGeneration
+from dashscope.api_entities.dashscope_response import Choice, ImageGenerationResponse, Message
+
+from chibi.config import gpt_settings
+from chibi.exceptions import ServiceResponseError
+from chibi.schemas.app import ModelChangeSchema
+from chibi.services.providers.provider import OpenAIFriendlyProvider
+
+dashscope.base_http_api_url = "https://dashscope-intl.aliyuncs.com/api/v1"
+
+
+class Alibaba(OpenAIFriendlyProvider):
+    api_key = gpt_settings.alibaba_key
+    chat_ready = True
+    image_generation_ready = True
+    moderation_ready = True
+
+    base_url = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"
+    name = "Alibaba"
+    model_name_keywords = ["qwen"]
+    model_name_keywords_exclude = ["tts", "stt"]
+    default_model = "qwen-plus"
+    default_moderation_model = "qwen-turbo"
+    max_tokens: int = 8192
+    default_image_model = "qwen-image-plus"
+
+    @staticmethod
+    def _get_image_url(choice: Choice) -> str | None:
+        if isinstance(choice.message.content, str):
+            return choice.message.content
+        return choice.message.content[0].get("image")
+
+    async def get_images(self, prompt: str, model: str | None = None) -> list[str]:
+        model = model or self.default_model
+        message = Message(role="user", content=[{"text": prompt}])
+        number_of_images = 1 if "qwen" in model or "z-image" in model else gpt_settings.image_n_choices
+        response: ImageGenerationResponse = await AioImageGeneration.call(
+            api_key=self.token,
+            model=model,
+            messages=[message],
+            n=number_of_images,
+            size=gpt_settings.image_size_alibaba,
+            prompt_extend=True,
+            watermark=False,
+        )
+
+        if response.status_code != HTTPStatus.OK:
+            raise ServiceResponseError(
+                provider=self.name,
+                model=model,
+                detail=(
+                    f"Unexpected response status code: {response.status_code}. "
+                    f"Response code: {response.code}. Message: {response.message}"
+                ),
+            )
+        image_urls: list[str] = []
+        for choice in response.output.choices:
+            if url := self._get_image_url(choice):
+                image_urls.append(url)
+        return image_urls
+
+    async def get_available_models(self, image_generation: bool = False) -> list[ModelChangeSchema]:
+        models = await super().get_available_models(image_generation=image_generation)
+
+        if image_generation:
+            wan_models = [
+                ModelChangeSchema(
+                    provider=self.name,
+                    name="wan2.6-t2i",
+                    display_name="Wan 2.6",
+                    image_generation=True,
+                ),
+            ]
+
+            models += wan_models
+        return models

chibi/services/providers/anthropic.py

@@ -0,0 +1,40 @@
+from anthropic import AsyncClient
+
+from chibi.config import gpt_settings
+from chibi.exceptions import NoApiKeyProvidedError
+from chibi.services.providers.provider import AnthropicFriendlyProvider
+
+
+class Anthropic(AnthropicFriendlyProvider):
+    api_key = gpt_settings.anthropic_key
+    chat_ready = True
+    moderation_ready = True
+
+    name = "Anthropic"
+    model_name_keywords = ["claude"]
+    default_model = "claude-sonnet-4-5-20250929"
+    default_moderation_model = "claude-haiku-4-5-20251001"
+
+    def __init__(self, token: str) -> None:
+        self._client: AsyncClient | None = None
+        super().__init__(token=token)
+
+    @property
+    def client(self) -> AsyncClient:
+        if self._client:
+            return self._client
+
+        if not self.token:
+            raise NoApiKeyProvidedError(provider=self.name)
+
+        self._client = AsyncClient(api_key=self.token)
+        return self._client
+
+    @property
+    def _headers(self) -> dict[str, str]:
+        return {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+            "x-api-key": self.token,
+            "anthropic-version": "2023-06-01",
+        }

chibi/services/providers/cloudflare.py

@@ -0,0 +1,98 @@
+from loguru import logger
+from openai.types.chat import ChatCompletionMessageParam
+
+from chibi.config import gpt_settings
+from chibi.exceptions import NoAccountIDSetError
+from chibi.schemas.app import ChatResponseSchema, ModelChangeSchema, UsageSchema
+from chibi.schemas.cloudflare import (
+    ChatCompletionResponseSchema,
+    ModelsSearchResponseSchema,
+)
+from chibi.services.providers.provider import RestApiFriendlyProvider
+
+
+class Cloudflare(RestApiFriendlyProvider):
+    api_key = gpt_settings.cloudflare_key
+    chat_ready = True
+
+    name = "Cloudflare"
+    model_name_keywords = ["@cf", "@hf"]
+    default_model = "@cf/meta/llama-3.2-3b-instruct"
+    base_url = f"https://api.cloudflare.com/client/v4/accounts/{gpt_settings.cloudflare_account_id}/ai"
+
+    @property
+    def _headers(self) -> dict[str, str]:
+        return {
+            "Authorization": f"Bearer {self.token}",
+        }
+
+    async def _get_chat_completion_response(
+        self,
+        messages: list[ChatCompletionMessageParam],
+        model: str,
+        system_prompt: str | None = None,
+    ) -> ChatResponseSchema:
+        if not gpt_settings.cloudflare_account_id:
+            raise NoAccountIDSetError
+
+        url = f"{self.base_url}/run/{model}"
+
+        system_message = {"role": "system", "content": system_prompt}
+
+        dialog = (
+            [system_message] + [dict(m) for m in messages]
+            if system_message not in messages
+            else [dict(m) for m in messages]
+        )
+        data = {"messages": dialog}
+        response = await self._request(method="POST", url=url, data=data)
+
+        response_data = ChatCompletionResponseSchema(**response.json())
+        if response_data.success:
+            answer_data = response_data.result
+            answer = answer_data.response
+            usage = UsageSchema(
+                completion_tokens=answer_data.usage.completion_tokens,
+                prompt_tokens=answer_data.usage.prompt_tokens,
+                total_tokens=answer_data.usage.total_tokens,
+            )
+        else:
+            answer = ""
+            usage = None
+
+        return ChatResponseSchema(answer=answer, provider=self.name, model=model, usage=usage)
+
+    async def get_available_models(self, image_generation: bool = False) -> list[ModelChangeSchema]:
+        if image_generation:
+            return []
+
+        if not gpt_settings.cloudflare_account_id:
+            logger.error("No Cloudflare account ID set. Please, check the CLOUDFLARE_ACCOUNT_ID env value.")
+            return []
+
+        url = f"{self.base_url}/models/search"
+        params = {"task": "Text Generation"}
+
+        try:
+            response = await self._request(method="GET", url=url, params=params)
+        except Exception as e:
+            logger.error(f"Failed to get available models for provider {self.name} due to exception: {e}")
+            return []
+
+        data = response.json()
+        response_data = ModelsSearchResponseSchema(**data)
+
+        all_models = [
+            ModelChangeSchema(
+                provider=self.name,
+                name=model.name,
+                image_generation=False,
+            )
+            for model in response_data.result
+        ]
+        all_models.sort(key=lambda model: model.name)
+
+        if gpt_settings.models_whitelist:
+            return [model for model in all_models if model.name in gpt_settings.models_whitelist]
+
+        return all_models

chibi/services/providers/constants/suno.py

@@ -0,0 +1,2 @@
+POLLING_ATTEMPTS_MAX = 20
+POLLING_ATTEMPTS_WAIT_BETWEEN = 30

chibi/services/providers/customopenai.py

@@ -0,0 +1,11 @@
+from chibi.config import gpt_settings
+from chibi.services.providers.provider import OpenAIFriendlyProvider
+
+
+class CustomOpenAI(OpenAIFriendlyProvider):
+    api_key = gpt_settings.customopenai_key
+    chat_ready = True
+
+    name = "CustomOpenAI"
+    base_url = gpt_settings.customopenai_url
+    default_model = ""

chibi/services/providers/deepseek.py

@@ -0,0 +1,15 @@
+from chibi.config import gpt_settings
+from chibi.services.providers.provider import OpenAIFriendlyProvider
+
+
+class DeepSeek(OpenAIFriendlyProvider):
+    api_key = gpt_settings.deepseek_key
+    chat_ready = True
+    moderation_ready = True
+
+    name = "DeepSeek"
+    model_name_keywords = ["deepseek"]
+    base_url = "https://api.deepseek.com"
+    default_model = "deepseek-chat"
+    default_moderation_model = "deepseek-chat"
+    max_tokens: int = 8192

chibi/services/providers/eleven_labs.py

@@ -0,0 +1,85 @@
+from io import BytesIO
+
+from elevenlabs.client import AsyncElevenLabs
+from loguru import logger
+
+from chibi.config import application_settings, gpt_settings
+from chibi.exceptions import NoApiKeyProvidedError
+from chibi.schemas.app import ModelChangeSchema
+from chibi.services.providers.provider import Provider
+
+
+class ElevenLabs(Provider):
+    api_key = gpt_settings.elevenlabs_api_key
+    chat_ready = False
+    tts_ready = True
+    stt_ready = True
+
+    name = "ElevenLabs"
+    default_stt_model = "scribe_v1"
+    tag_audio_events = False
+    language_code = None
+    default_tts_voice = "cgSgspJ2msm6clMCkdW9"
+    default_tts_model = "eleven_multilingual_v2"
+    output_format = "mp3_44100_128"
+    music_length_ms = 10000
+
+    def __init__(self, token: str) -> None:
+        self._client: AsyncElevenLabs | None = None
+        super().__init__(token=token)
+
+    @property
+    def client(self) -> AsyncElevenLabs:
+        if self._client:
+            return self._client
+
+        if not self.token:
+            raise NoApiKeyProvidedError(provider=self.name)
+
+        self._client = AsyncElevenLabs(
+            api_key=self.token,
+        )
+
+        return self._client
+
+    async def speech(self, text: str, voice: str | None = None, model: str | None = None) -> bytes:
+        voice = voice or self.default_tts_voice
+        model = model or self.default_tts_model
+        logger.info(f"Recording a voice message with model {model}...")
+        response = self.client.text_to_speech.convert(
+            text=text,
+            voice_id=voice,
+            model_id=model,
+            output_format=self.output_format,
+        )
+
+        buf = bytearray()
+        async for chunk in response:
+            buf.extend(chunk)
+
+        return bytes(buf)
+
+    async def transcribe(self, audio: BytesIO, model: str | None = None) -> str:
+        model = model or self.default_stt_model
+        logger.info(f"Transcribing audio with model {model}...")
+        response = await self.client.speech_to_text.convert(
+            file=audio, model_id=model, tag_audio_events=self.tag_audio_events, language_code=self.language_code
+        )
+        if response:
+            if application_settings.log_prompt_data:
+                logger.info(f"Transcribed text: {response.text}")
+            return response.text
+        raise ValueError("Could not transcribe audio message")
+
+    async def generate_music(self, prompt: str, music_length_ms: int = music_length_ms) -> bytes:
+        logger.info(f"Generating music for {music_length_ms} ms...")
+        response = self.client.music.compose(prompt=prompt, music_length_ms=music_length_ms)
+
+        buf = bytearray()
+        async for chunk in response:
+            buf.extend(chunk)
+
+        return bytes(buf)
+
+    async def get_available_models(self, image_generation: bool = False) -> list[ModelChangeSchema]:
+        return []