AstrBot 3.5.6__py3-none-any.whl → 4.7.0__py3-none-any.whl

astrbot/core/provider/sources/openai_tts_api_source.py

@@ -1,12 +1,19 @@
+import os
 import uuid
-from openai import AsyncOpenAI, NOT_GIVEN
-from ..provider import TTSProvider
+
+from openai import NOT_GIVEN, AsyncOpenAI
+
+from astrbot.core.utils.astrbot_path import get_astrbot_data_path
+
 from ..entities import ProviderType
+from ..provider import TTSProvider
 from ..register import register_provider_adapter
 
 
 @register_provider_adapter(
-    "openai_tts_api", "OpenAI TTS API", provider_type=ProviderType.TEXT_TO_SPEECH
+    "openai_tts_api",
+    "OpenAI TTS API",
+    provider_type=ProviderType.TEXT_TO_SPEECH,
 )
 class ProviderOpenAITTSAPI(TTSProvider):
     def __init__(
@@ -24,16 +31,20 @@ class ProviderOpenAITTSAPI(TTSProvider):
 
         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=timeout,
         )
 
-        self.set_model(provider_config.get("model", None))
+        self.set_model(provider_config.get("model", ""))
 
     async def get_audio(self, text: str) -> str:
-        path = f"data/temp/openai_tts_api_{uuid.uuid4()}.wav"
+        temp_dir = os.path.join(get_astrbot_data_path(), "temp")
+        path = os.path.join(temp_dir, f"openai_tts_api_{uuid.uuid4()}.wav")
         async with self.client.audio.speech.with_streaming_response.create(
-            model=self.model_name, voice=self.voice, response_format="wav", input=text
+            model=self.model_name,
+            voice=self.voice,
+            response_format="wav",
+            input=text,
         ) as response:
             with open(path, "wb") as f:
                 async for chunk in response.iter_bytes(chunk_size=1024):

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

@@ -0,0 +1,71 @@
+import aiohttp
+
+from astrbot import logger
+
+from ..entities import ProviderType, RerankResult
+from ..provider import RerankProvider
+from ..register import register_provider_adapter
+
+
+@register_provider_adapter(
+    "vllm_rerank",
+    "VLLM Rerank 适配器",
+    provider_type=ProviderType.RERANK,
+)
+class VLLMRerankProvider(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.auth_key = provider_config.get("rerank_api_key", "")
+        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 = provider_config.get("rerank_model", "BAAI/bge-reranker-base")
+
+        h = {}
+        if self.auth_key:
+            h["Authorization"] = f"Bearer {self.auth_key}"
+        self.client = aiohttp.ClientSession(
+            headers=h,
+            timeout=aiohttp.ClientTimeout(total=self.timeout),
+        )
+
+    async def rerank(
+        self,
+        query: str,
+        documents: list[str],
+        top_n: int | None = None,
+    ) -> list[RerankResult]:
+        payload = {
+            "query": query,
+            "documents": documents,
+            "model": self.model,
+        }
+        if top_n is not None:
+            payload["top_n"] = top_n
+        async with self.client.post(
+            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}",
+                )
+
+            return [
+                RerankResult(
+                    index=result["index"],
+                    relevance_score=result["relevance_score"],
+                )
+                for result in results
+            ]
+
+    async def terminate(self) -> None:
+        """关闭客户端会话"""
+        if self.client:
+            await self.client.close()
+            self.client = None

astrbot/core/provider/sources/volcengine_tts.py

@@ -0,0 +1,115 @@
+import asyncio
+import base64
+import json
+import os
+import traceback
+import uuid
+
+import aiohttp
+
+from astrbot import logger
+
+from ..entities import ProviderType
+from ..provider import TTSProvider
+from ..register import register_provider_adapter
+
+
+@register_provider_adapter(
+    "volcengine_tts",
+    "火山引擎 TTS",
+    provider_type=ProviderType.TEXT_TO_SPEECH,
+)
+class ProviderVolcengineTTS(TTSProvider):
+    def __init__(self, provider_config: dict, provider_settings: dict) -> None:
+        super().__init__(provider_config, provider_settings)
+        self.api_key = provider_config.get("api_key", "")
+        self.appid = provider_config.get("appid", "")
+        self.cluster = provider_config.get("volcengine_cluster", "")
+        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",
+        )
+        self.timeout = provider_config.get("timeout", 20)
+
+    def _build_request_payload(self, text: str) -> dict:
+        return {
+            "app": {
+                "appid": self.appid,
+                "token": self.api_key,
+                "cluster": self.cluster,
+            },
+            "user": {"uid": str(uuid.uuid4())},
+            "audio": {
+                "voice_type": self.voice_type,
+                "encoding": "mp3",
+                "speed_ratio": self.speed_ratio,
+                "volume_ratio": 1.0,
+                "pitch_ratio": 1.0,
+            },
+            "request": {
+                "reqid": str(uuid.uuid4()),
+                "text": text,
+                "text_type": "plain",
+                "operation": "query",
+                "with_frontend": 1,
+                "frontend_type": "unitTson",
+            },
+        }
+
+    async def get_audio(self, text: str) -> str:
+        """异步方法获取语音文件路径"""
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer; {self.api_key}",
+        }
+
+        payload = self._build_request_payload(text)
+
+        logger.debug(f"请求头: {headers}")
+        logger.debug(f"请求 URL: {self.api_base}")
+        logger.debug(f"请求体: {json.dumps(payload, ensure_ascii=False)[:100]}...")
+
+        try:
+            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]}...")
+
+                if response.status == 200:
+                    resp_data = json.loads(response_text)
+
+                    if "data" in resp_data:
+                        audio_data = base64.b64decode(resp_data["data"])
+
+                        os.makedirs("data/temp", exist_ok=True)
+
+                        file_path = f"data/temp/volcengine_tts_{uuid.uuid4()}.mp3"
+
+                        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 异常: {e!s}")

astrbot/core/provider/sources/whisper_api_source.py

@@ -1,13 +1,17 @@
-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.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(
     "openai_whisper_api",
@@ -25,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"
@@ -38,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"):
@@ -50,7 +53,8 @@ class ProviderOpenAIWhisperAPI(STTProvider):
                 is_tencent = True
 
             name = str(uuid.uuid4())
-            path = os.path.join("data/temp", name)
+            temp_dir = os.path.join(get_astrbot_data_path(), "temp")
+            path = os.path.join(temp_dir, name)
             await download_file(audio_url, path)
             audio_url = path
 
@@ -61,7 +65,8 @@ class ProviderOpenAIWhisperAPI(STTProvider):
             is_silk = await self._is_silk_file(audio_url)
             if is_silk:
                 logger.info("Converting silk file to wav ...")
-                output_path = os.path.join("data/temp", str(uuid.uuid4()) + ".wav")
+                temp_dir = os.path.join(get_astrbot_data_path(), "temp")
+                output_path = os.path.join(temp_dir, str(uuid.uuid4()) + ".wav")
                 await tencent_silk_to_wav(audio_url, output_path)
                 audio_url = output_path
 

astrbot/core/provider/sources/whisper_selfhosted_source.py

@@ -1,14 +1,18 @@
-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.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(
     "openai_whisper_selfhost",
@@ -22,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 模型加载完成。")
 
@@ -40,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()
@@ -53,7 +58,8 @@ class ProviderOpenAIWhisperSelfHost(STTProvider):
                 is_tencent = True
 
             name = str(uuid.uuid4())
-            path = os.path.join("data/temp", name)
+            temp_dir = os.path.join(get_astrbot_data_path(), "temp")
+            path = os.path.join(temp_dir, name)
             await download_file(audio_url, path)
             audio_url = path
 
@@ -64,7 +70,8 @@ class ProviderOpenAIWhisperSelfHost(STTProvider):
             is_silk = await self._is_silk_file(audio_url)
             if is_silk:
                 logger.info("Converting silk file to wav ...")
-                output_path = os.path.join("data/temp", str(uuid.uuid4()) + ".wav")
+                temp_dir = os.path.join(get_astrbot_data_path(), "temp")
+                output_path = os.path.join(temp_dir, str(uuid.uuid4()) + ".wav")
                 await tencent_silk_to_wav(audio_url, output_path)
                 audio_url = output_path
 

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)