AstrBot 4.5.0__py3-none-any.whl → 4.5.2__py3-none-any.whl

astrbot/core/provider/sources/sensevoice_selfhosted_source.py

@@ -1,22 +1,24 @@
-"""
-Author: diudiu62
+"""Author: diudiu62
 Date: 2025-02-24 18:04:18
 LastEditTime: 2025-02-25 14:06:30
 """
 
 import asyncio
-from datetime import datetime
 import os
 import re
+from datetime import datetime
+
 from funasr_onnx import SenseVoiceSmall
 from funasr_onnx.utils.postprocess_utils import rich_transcription_postprocess
-from ..provider import STTProvider
-from ..entities import ProviderType
-from astrbot.core.utils.io import download_file
-from ..register import register_provider_adapter
+
 from astrbot.core import logger
+from astrbot.core.utils.io import download_file
 from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav
 
+from ..entities import ProviderType
+from ..provider import STTProvider
+from ..register import register_provider_adapter
+
 
 @register_provider_adapter(
     "sensevoice_stt_selfhost",
@@ -30,7 +32,7 @@ class ProviderSenseVoiceSTTSelfHost(STTProvider):
         provider_settings: dict,
     ) -> None:
         super().__init__(provider_config, provider_settings)
-        self.set_model(provider_config.get("stt_model", None))
+        self.set_model(provider_config.get("stt_model"))
         self.model = None
         self.is_emotion = provider_config.get("is_emotion", False)
 
@@ -39,7 +41,8 @@ class ProviderSenseVoiceSTTSelfHost(STTProvider):
 
         # 将模型加载放到线程池中执行
         self.model = await asyncio.get_event_loop().run_in_executor(
-            None, lambda: SenseVoiceSmall(self.model_name, quantize=True, batch_size=16)
+            None,
+            lambda: SenseVoiceSmall(self.model_name, quantize=True, batch_size=16),
         )
 
         logger.info("SenseVoice 模型加载完成。")
@@ -55,8 +58,7 @@ class ProviderSenseVoiceSTTSelfHost(STTProvider):
 
         if silk_header in file_header:
             return True
-        else:
-            return False
+        return False
 
     async def get_text(self, audio_url: str) -> str:
         try:

astrbot/core/provider/sources/vllm_rerank_source.py

@@ -1,8 +1,10 @@
 import aiohttp
+
 from astrbot import logger
+
+from ..entities import ProviderType, RerankResult
 from ..provider import RerankProvider
 from ..register import register_provider_adapter
-from ..entities import ProviderType, RerankResult
 
 
 @register_provider_adapter(
@@ -30,7 +32,10 @@ class VLLMRerankProvider(RerankProvider):
         )
 
     async def rerank(
-        self, query: str, documents: list[str], top_n: int | None = None
+        self,
+        query: str,
+        documents: list[str],
+        top_n: int | None = None,
     ) -> list[RerankResult]:
         payload = {
             "query": query,
@@ -40,14 +45,15 @@ class VLLMRerankProvider(RerankProvider):
         if top_n is not None:
             payload["top_n"] = top_n
         async with self.client.post(
-            f"{self.base_url}/v1/rerank", json=payload
+            f"{self.base_url}/v1/rerank",
+            json=payload,
         ) as response:
             response_data = await response.json()
             results = response_data.get("results", [])
 
             if not results:
                 logger.warning(
-                    f"Rerank API 返回了空的列表数据。原始响应: {response_data}"
+                    f"Rerank API 返回了空的列表数据。原始响应: {response_data}",
                 )
 
             return [

astrbot/core/provider/sources/volcengine_tts.py

@@ -1,18 +1,23 @@
-import uuid
+import asyncio
 import base64
 import json
 import os
 import traceback
-import asyncio
+import uuid
+
 import aiohttp
-from ..provider import TTSProvider
+
+from astrbot import logger
+
 from ..entities import ProviderType
+from ..provider import TTSProvider
 from ..register import register_provider_adapter
-from astrbot import logger
 
 
 @register_provider_adapter(
-    "volcengine_tts", "火山引擎 TTS", provider_type=ProviderType.TEXT_TO_SPEECH
+    "volcengine_tts",
+    "火山引擎 TTS",
+    provider_type=ProviderType.TEXT_TO_SPEECH,
 )
 class ProviderVolcengineTTS(TTSProvider):
     def __init__(self, provider_config: dict, provider_settings: dict) -> None:
@@ -23,7 +28,8 @@ class ProviderVolcengineTTS(TTSProvider):
         self.voice_type = provider_config.get("volcengine_voice_type", "")
         self.speed_ratio = provider_config.get("volcengine_speed_ratio", 1.0)
         self.api_base = provider_config.get(
-            "api_base", "https://openspeech.bytedance.com/api/v1/tts"
+            "api_base",
+            "https://openspeech.bytedance.com/api/v1/tts",
         )
         self.timeout = provider_config.get("timeout", 20)
 
@@ -66,43 +72,44 @@ class ProviderVolcengineTTS(TTSProvider):
         logger.debug(f"请求体: {json.dumps(payload, ensure_ascii=False)[:100]}...")
 
         try:
-            async with aiohttp.ClientSession() as session:
-                async with session.post(
+            async with (
+                aiohttp.ClientSession() as session,
+                session.post(
                     self.api_base,
                     data=json.dumps(payload),
                     headers=headers,
                     timeout=self.timeout,
-                ) as response:
-                    logger.debug(f"响应状态码: {response.status}")
-
-                    response_text = await response.text()
-                    logger.debug(f"响应内容: {response_text[:200]}...")
+                ) as response,
+            ):
+                logger.debug(f"响应状态码: {response.status}")
 
-                    if response.status == 200:
-                        resp_data = json.loads(response_text)
+                response_text = await response.text()
+                logger.debug(f"响应内容: {response_text[:200]}...")
 
-                        if "data" in resp_data:
-                            audio_data = base64.b64decode(resp_data["data"])
+                if response.status == 200:
+                    resp_data = json.loads(response_text)
 
-                            os.makedirs("data/temp", exist_ok=True)
+                    if "data" in resp_data:
+                        audio_data = base64.b64decode(resp_data["data"])
 
-                            file_path = f"data/temp/volcengine_tts_{uuid.uuid4()}.mp3"
+                        os.makedirs("data/temp", exist_ok=True)
 
-                            loop = asyncio.get_running_loop()
-                            await loop.run_in_executor(
-                                None, lambda: open(file_path, "wb").write(audio_data)
-                            )
+                        file_path = f"data/temp/volcengine_tts_{uuid.uuid4()}.mp3"
 
-                            return file_path
-                        else:
-                            error_msg = resp_data.get("message", "未知错误")
-                            raise Exception(f"火山引擎 TTS API 返回错误: {error_msg}")
-                    else:
-                        raise Exception(
-                            f"火山引擎 TTS API 请求失败: {response.status}, {response_text}"
+                        loop = asyncio.get_running_loop()
+                        await loop.run_in_executor(
+                            None,
+                            lambda: open(file_path, "wb").write(audio_data),
                         )
 
+                        return file_path
+                    error_msg = resp_data.get("message", "未知错误")
+                    raise Exception(f"火山引擎 TTS API 返回错误: {error_msg}")
+                raise Exception(
+                    f"火山引擎 TTS API 请求失败: {response.status}, {response_text}",
+                )
+
         except Exception as e:
             error_details = traceback.format_exc()
             logger.debug(f"火山引擎 TTS 异常详情: {error_details}")
-            raise Exception(f"火山引擎 TTS 异常: {str(e)}")
+            raise Exception(f"火山引擎 TTS 异常: {e!s}")

astrbot/core/provider/sources/whisper_api_source.py

@@ -1,13 +1,16 @@
-import uuid
 import os
-from openai import AsyncOpenAI, NOT_GIVEN
-from ..provider import STTProvider
-from ..entities import ProviderType
-from astrbot.core.utils.io import download_file
-from ..register import register_provider_adapter
+import uuid
+
+from openai import NOT_GIVEN, AsyncOpenAI
+
 from astrbot.core import logger
-from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav
 from astrbot.core.utils.astrbot_path import get_astrbot_data_path
+from astrbot.core.utils.io import download_file
+from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav
+
+from ..entities import ProviderType
+from ..provider import STTProvider
+from ..register import register_provider_adapter
 
 
 @register_provider_adapter(
@@ -26,11 +29,11 @@ class ProviderOpenAIWhisperAPI(STTProvider):
 
         self.client = AsyncOpenAI(
             api_key=self.chosen_api_key,
-            base_url=provider_config.get("api_base", None),
+            base_url=provider_config.get("api_base"),
             timeout=provider_config.get("timeout", NOT_GIVEN),
         )
 
-        self.set_model(provider_config.get("model", None))
+        self.set_model(provider_config.get("model"))
 
     async def _is_silk_file(self, file_path):
         silk_header = b"SILK"
@@ -39,11 +42,10 @@ class ProviderOpenAIWhisperAPI(STTProvider):
 
         if silk_header in file_header:
             return True
-        else:
-            return False
+        return False
 
     async def get_text(self, audio_url: str) -> str:
-        """only supports mp3, mp4, mpeg, m4a, wav, webm"""
+        """Only supports mp3, mp4, mpeg, m4a, wav, webm"""
         is_tencent = False
 
         if audio_url.startswith("http"):

astrbot/core/provider/sources/whisper_selfhosted_source.py

@@ -1,14 +1,17 @@
-import uuid
-import os
 import asyncio
+import os
+import uuid
+
 import whisper
-from ..provider import STTProvider
-from ..entities import ProviderType
-from astrbot.core.utils.io import download_file
-from ..register import register_provider_adapter
+
 from astrbot.core import logger
-from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav
 from astrbot.core.utils.astrbot_path import get_astrbot_data_path
+from astrbot.core.utils.io import download_file
+from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav
+
+from ..entities import ProviderType
+from ..provider import STTProvider
+from ..register import register_provider_adapter
 
 
 @register_provider_adapter(
@@ -23,14 +26,16 @@ class ProviderOpenAIWhisperSelfHost(STTProvider):
         provider_settings: dict,
     ) -> None:
         super().__init__(provider_config, provider_settings)
-        self.set_model(provider_config.get("model", None))
+        self.set_model(provider_config.get("model"))
         self.model = None
 
     async def initialize(self):
         loop = asyncio.get_event_loop()
         logger.info("下载或者加载 Whisper 模型中，这可能需要一些时间 ...")
         self.model = await loop.run_in_executor(
-            None, whisper.load_model, self.model_name
+            None,
+            whisper.load_model,
+            self.model_name,
         )
         logger.info("Whisper 模型加载完成。")
 
@@ -41,8 +46,7 @@ class ProviderOpenAIWhisperSelfHost(STTProvider):
 
         if silk_header in file_header:
             return True
-        else:
-            return False
+        return False
 
     async def get_text(self, audio_url: str) -> str:
         loop = asyncio.get_event_loop()

astrbot/core/provider/sources/xinference_rerank_source.py

@@ -0,0 +1,116 @@
+from xinference_client.client.restful.async_restful_client import (
+    AsyncClient as Client,
+)
+
+from astrbot import logger
+
+from ..entities import ProviderType, RerankResult
+from ..provider import RerankProvider
+from ..register import register_provider_adapter
+
+
+@register_provider_adapter(
+    "xinference_rerank",
+    "Xinference Rerank 适配器",
+    provider_type=ProviderType.RERANK,
+)
+class XinferenceRerankProvider(RerankProvider):
+    def __init__(self, provider_config: dict, provider_settings: dict) -> None:
+        super().__init__(provider_config, provider_settings)
+        self.provider_config = provider_config
+        self.provider_settings = provider_settings
+        self.base_url = provider_config.get("rerank_api_base", "http://127.0.0.1:8000")
+        self.base_url = self.base_url.rstrip("/")
+        self.timeout = provider_config.get("timeout", 20)
+        self.model_name = provider_config.get("rerank_model", "BAAI/bge-reranker-base")
+        self.api_key = provider_config.get("rerank_api_key")
+        self.launch_model_if_not_running = provider_config.get(
+            "launch_model_if_not_running",
+            False,
+        )
+        self.client = None
+        self.model = None
+        self.model_uid = None
+
+    async def initialize(self):
+        if self.api_key:
+            logger.info("Xinference Rerank: Using API key for authentication.")
+            self.client = Client(self.base_url, api_key=self.api_key)
+        else:
+            logger.info("Xinference Rerank: No API key provided.")
+            self.client = Client(self.base_url)
+
+        try:
+            running_models = await self.client.list_models()
+            for uid, model_spec in running_models.items():
+                if model_spec.get("model_name") == self.model_name:
+                    logger.info(
+                        f"Model '{self.model_name}' is already running with UID: {uid}",
+                    )
+                    self.model_uid = uid
+                    break
+
+            if self.model_uid is None:
+                if self.launch_model_if_not_running:
+                    logger.info(f"Launching {self.model_name} model...")
+                    self.model_uid = await self.client.launch_model(
+                        model_name=self.model_name,
+                        model_type="rerank",
+                    )
+                    logger.info("Model launched.")
+                else:
+                    logger.warning(
+                        f"Model '{self.model_name}' is not running and auto-launch is disabled. Provider will not be available.",
+                    )
+                    return
+
+            if self.model_uid:
+                self.model = await self.client.get_model(self.model_uid)
+
+        except Exception as e:
+            logger.error(f"Failed to initialize Xinference model: {e}")
+            logger.debug(
+                f"Xinference initialization failed with exception: {e}",
+                exc_info=True,
+            )
+            self.model = None
+
+    async def rerank(
+        self,
+        query: str,
+        documents: list[str],
+        top_n: int | None = None,
+    ) -> list[RerankResult]:
+        if not self.model:
+            logger.error("Xinference rerank model is not initialized.")
+            return []
+        try:
+            response = await self.model.rerank(documents, query, top_n)
+            results = response.get("results", [])
+            logger.debug(f"Rerank API response: {response}")
+
+            if not results:
+                logger.warning(
+                    f"Rerank API returned an empty list. Original response: {response}",
+                )
+
+            return [
+                RerankResult(
+                    index=result["index"],
+                    relevance_score=result["relevance_score"],
+                )
+                for result in results
+            ]
+        except Exception as e:
+            logger.error(f"Xinference rerank failed: {e}")
+            logger.debug(f"Xinference rerank failed with exception: {e}", exc_info=True)
+            return []
+
+    async def terminate(self) -> None:
+        """关闭客户端会话"""
+        if self.client:
+            logger.info("Closing Xinference rerank client...")
+            try:
+                await self.client.close()
+            except Exception as e:
+                logger.error(f"Failed to close Xinference client: {e}", exc_info=True)

astrbot/core/provider/sources/xinference_stt_provider.py

@@ -0,0 +1,197 @@
+import os
+import uuid
+
+import aiohttp
+from xinference_client.client.restful.async_restful_client import (
+    AsyncClient as Client,
+)
+
+from astrbot.core import logger
+from astrbot.core.utils.astrbot_path import get_astrbot_data_path
+from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav
+
+from ..entities import ProviderType
+from ..provider import STTProvider
+from ..register import register_provider_adapter
+
+
+@register_provider_adapter(
+    "xinference_stt",
+    "Xinference STT",
+    provider_type=ProviderType.SPEECH_TO_TEXT,
+)
+class ProviderXinferenceSTT(STTProvider):
+    def __init__(self, provider_config: dict, provider_settings: dict) -> None:
+        super().__init__(provider_config, provider_settings)
+        self.provider_config = provider_config
+        self.provider_settings = provider_settings
+        self.base_url = provider_config.get("api_base", "http://127.0.0.1:9997")
+        self.base_url = self.base_url.rstrip("/")
+        self.timeout = provider_config.get("timeout", 180)
+        self.model_name = provider_config.get("model", "whisper-large-v3")
+        self.api_key = provider_config.get("api_key")
+        self.launch_model_if_not_running = provider_config.get(
+            "launch_model_if_not_running",
+            False,
+        )
+        self.client = None
+        self.model_uid = None
+
+    async def initialize(self):
+        if self.api_key:
+            logger.info("Xinference STT: Using API key for authentication.")
+            self.client = Client(self.base_url, api_key=self.api_key)
+        else:
+            logger.info("Xinference STT: No API key provided.")
+            self.client = Client(self.base_url)
+
+        try:
+            running_models = await self.client.list_models()
+            for uid, model_spec in running_models.items():
+                if model_spec.get("model_name") == self.model_name:
+                    logger.info(
+                        f"Model '{self.model_name}' is already running with UID: {uid}",
+                    )
+                    self.model_uid = uid
+                    break
+
+            if self.model_uid is None:
+                if self.launch_model_if_not_running:
+                    logger.info(f"Launching {self.model_name} model...")
+                    self.model_uid = await self.client.launch_model(
+                        model_name=self.model_name,
+                        model_type="audio",
+                    )
+                    logger.info("Model launched.")
+                else:
+                    logger.warning(
+                        f"Model '{self.model_name}' is not running and auto-launch is disabled. Provider will not be available.",
+                    )
+                    return
+
+        except Exception as e:
+            logger.error(f"Failed to initialize Xinference model: {e}")
+            logger.debug(
+                f"Xinference initialization failed with exception: {e}",
+                exc_info=True,
+            )
+
+    async def get_text(self, audio_url: str) -> str:
+        if not self.model_uid or self.client is None or self.client.session is None:
+            logger.error("Xinference STT model is not initialized.")
+            return ""
+
+        audio_bytes = None
+        temp_files = []
+        is_tencent = False
+
+        try:
+            # 1. Get audio bytes
+            if audio_url.startswith("http"):
+                if "multimedia.nt.qq.com.cn" in audio_url:
+                    is_tencent = True
+                async with aiohttp.ClientSession() as session:
+                    async with session.get(audio_url, timeout=self.timeout) as resp:
+                        if resp.status == 200:
+                            audio_bytes = await resp.read()
+                        else:
+                            logger.error(
+                                f"Failed to download audio from {audio_url}, status: {resp.status}",
+                            )
+                            return ""
+            elif os.path.exists(audio_url):
+                with open(audio_url, "rb") as f:
+                    audio_bytes = f.read()
+            else:
+                logger.error(f"File not found: {audio_url}")
+                return ""
+
+            if not audio_bytes:
+                logger.error("Audio bytes are empty.")
+                return ""
+
+            # 2. Check for conversion
+            needs_conversion = False
+            if (
+                audio_url.endswith((".amr", ".silk"))
+                or is_tencent
+                or b"SILK" in audio_bytes[:8]
+            ):
+                needs_conversion = True
+
+            # 3. Perform conversion if needed
+            if needs_conversion:
+                logger.info("Audio requires conversion, using temporary files...")
+                temp_dir = os.path.join(get_astrbot_data_path(), "temp")
+                os.makedirs(temp_dir, exist_ok=True)
+
+                input_path = os.path.join(temp_dir, str(uuid.uuid4()))
+                output_path = os.path.join(temp_dir, str(uuid.uuid4()) + ".wav")
+                temp_files.extend([input_path, output_path])
+
+                with open(input_path, "wb") as f:
+                    f.write(audio_bytes)
+
+                logger.info("Converting silk/amr file to wav ...")
+                await tencent_silk_to_wav(input_path, output_path)
+
+                with open(output_path, "rb") as f:
+                    audio_bytes = f.read()
+
+            # 4. Transcribe
+            # 官方asyncCLient的客户端似乎实现有点问题，这里直接用aiohttp实现openai标准兼容请求，提交issue等待官方修复后再改回来
+            url = f"{self.base_url}/v1/audio/transcriptions"
+            headers = {
+                "accept": "application/json",
+            }
+            if self.client and self.client._headers:
+                headers.update(self.client._headers)
+
+            data = aiohttp.FormData()
+            data.add_field("model", self.model_uid)
+            data.add_field(
+                "file",
+                audio_bytes,
+                filename="audio.wav",
+                content_type="audio/wav",
+            )
+
+            async with self.client.session.post(
+                url,
+                data=data,
+                headers=headers,
+                timeout=self.timeout,
+            ) as resp:
+                if resp.status == 200:
+                    result = await resp.json()
+                    text = result.get("text", "")
+                    logger.debug(f"Xinference STT result: {text}")
+                    return text
+                error_text = await resp.text()
+                logger.error(
+                    f"Xinference STT transcription failed with status {resp.status}: {error_text}",
+                )
+                return ""
+
+        except Exception as e:
+            logger.error(f"Xinference STT failed: {e}")
+            logger.debug(f"Xinference STT failed with exception: {e}", exc_info=True)
+            return ""
+        finally:
+            # 5. Cleanup
+            for temp_file in temp_files:
+                try:
+                    if os.path.exists(temp_file):
+                        os.remove(temp_file)
+                        logger.debug(f"Removed temporary file: {temp_file}")
+                except Exception as e:
+                    logger.error(f"Failed to remove temporary file {temp_file}: {e}")
+
+    async def terminate(self) -> None:
+        """关闭客户端会话"""
+        if self.client:
+            logger.info("Closing Xinference STT client...")
+            try:
+                await self.client.close()
+            except Exception as e:
+                logger.error(f"Failed to close Xinference client: {e}", exc_info=True)

astrbot/core/star/__init__.py

@@ -1,10 +1,11 @@
-from .star import StarMetadata, star_map, star_registry
-from .star_manager import PluginManager
-from .context import Context
-from astrbot.core.provider import Provider
-from astrbot.core.utils.command_parser import CommandParserMixin
 from astrbot.core import html_renderer
+from astrbot.core.provider import Provider
 from astrbot.core.star.star_tools import StarTools
+from astrbot.core.utils.command_parser import CommandParserMixin
+
+from .context import Context
+from .star import StarMetadata, star_map, star_registry
+from .star_manager import PluginManager
 
 
 class Star(CommandParserMixin):
@@ -36,24 +37,28 @@ class Star(CommandParserMixin):
         )
 
     async def html_render(
-        self, tmpl: str, data: dict, return_url=True, options: dict | None = None
+        self,
+        tmpl: str,
+        data: dict,
+        return_url=True,
+        options: dict | None = None,
     ) -> str:
         """渲染 HTML"""
         return await html_renderer.render_custom_template(
-            tmpl, data, return_url=return_url, options=options
+            tmpl,
+            data,
+            return_url=return_url,
+            options=options,
         )
 
     async def initialize(self):
         """当插件被激活时会调用这个方法"""
-        pass
 
     async def terminate(self):
         """当插件被禁用、重载插件时会调用这个方法"""
-        pass
 
     def __del__(self):
         """[Deprecated] 当插件被禁用、重载插件时会调用这个方法"""
-        pass
 
 
-__all__ = ["Star", "StarMetadata", "PluginManager", "Context", "Provider", "StarTools"]
+__all__ = ["Context", "PluginManager", "Provider", "Star", "StarMetadata", "StarTools"]