feat: 支持 Gemini 嵌入模型列表发现

Author: Sisyphbaous-DT-Project

astrbot/core/provider/sources/gemini_embedding_source.py

@@ -1,4 +1,4 @@
-from typing import cast
+from typing import Any, cast
 
 from google import genai
 from google.genai import types
@@ -78,10 +78,58 @@ async def get_embeddings(self, text: list[str]) -> list[list[float]]:
         except APIError as e:
             raise Exception(f"Gemini Embedding API批量请求失败: {e.message}")
 
+    async def get_models(self) -> list[str]:
+        try:
+            models = await self.client.models.list()
+            all_model_ids: list[str] = []
+            embedding_model_ids: list[str] = []
+
+            for model in getattr(models, "page", []):
+                model_id = self._extract_model_id(model)
+                if not model_id:
+                    continue
+                all_model_ids.append(model_id)
+                if self._supports_embedding(model, model_id):
+                    embedding_model_ids.append(model_id)
+
+            all_model_ids = sorted(dict.fromkeys(all_model_ids))
+            embedding_model_ids = sorted(dict.fromkeys(embedding_model_ids))
+
+            return embedding_model_ids or all_model_ids
+        except Exception as e:
+            raise Exception(f"获取 Gemini 嵌入模型列表失败: {e!s}") from e
+
     def get_dim(self) -> int:
         """获取向量的维度"""
         return int(self.provider_config.get("embedding_dimensions", 768))
 
     async def terminate(self):
         if self.client:
             await self.client.aclose()
+
+    @staticmethod
+    def _extract_model_id(model: Any) -> str:
+        model_name = getattr(model, "name", "") or getattr(model, "model", "")
+        if not model_name:
+            return ""
+        return str(model_name).removeprefix("models/")
+
+    @classmethod
+    def _supports_embedding(cls, model: Any, model_id: str) -> bool:
+        supported_actions = getattr(model, "supported_actions", None) or getattr(
+            model, "supported_generation_methods", []
+        )
+        if isinstance(supported_actions, list):
+            normalized_actions = {
+                str(action).lower().replace("_", "").replace("-", "")
+                for action in supported_actions
+            }
+            if "embedcontent" in normalized_actions:
+                return True
+
+        return cls._looks_like_embedding_model(model_id)
+
+    @staticmethod
+    def _looks_like_embedding_model(model_id: str) -> bool:
+        normalized_model_id = model_id.lower()
+        return "embedding" in normalized_model_id or "embed" in normalized_model_id

tests/test_gemini_embedding_source.py

@@ -0,0 +1,103 @@
+from types import SimpleNamespace
+
+import pytest
+
+from astrbot.core.provider.sources.gemini_embedding_source import (
+    GeminiEmbeddingProvider,
+)
+
+
+class _FakeModelsPager:
+    def __init__(self, models) -> None:
+        self.page = models
+
+
+class _FakeModelsAPI:
+    def __init__(self, models) -> None:
+        self._models = models
+
+    async def list(self):
+        return _FakeModelsPager(self._models)
+
+
+class _FakeClient:
+    def __init__(self, models) -> None:
+        self.models = _FakeModelsAPI(models)
+        self.closed = False
+
+    async def aclose(self):
+        self.closed = True
+
+
+def _make_provider() -> GeminiEmbeddingProvider:
+    provider_config = {
+        "id": "test-gemini-embedding",
+        "type": "gemini_embedding",
+        "embedding_api_key": "test-key",
+        "embedding_api_base": "https://generativelanguage.googleapis.com",
+        "embedding_model": "gemini-embedding-exp-03-07",
+    }
+    return GeminiEmbeddingProvider(
+        provider_config=provider_config,
+        provider_settings={},
+    )
+
+
+@pytest.mark.asyncio
+async def test_gemini_embedding_get_models_prefers_embedcontent_capability():
+    provider = _make_provider()
+    try:
+        provider.client = _FakeClient(
+            [
+                SimpleNamespace(
+                    name="models/gemini-2.5-flash",
+                    supported_actions=["generateContent"],
+                ),
+                SimpleNamespace(
+                    name="models/gemini-embedding-001",
+                    supported_actions=["embedContent"],
+                ),
+                SimpleNamespace(
+                    name="models/text-embedding-preview",
+                    supported_generation_methods=["embedContent"],
+                ),
+            ]
+        )
+        models = await provider.get_models()
+        assert models == ["gemini-embedding-001", "text-embedding-preview"]
+    finally:
+        await provider.terminate()
+
+
+@pytest.mark.asyncio
+async def test_gemini_embedding_get_models_falls_back_to_name_matching():
+    provider = _make_provider()
+    try:
+        provider.client = _FakeClient(
+            [
+                SimpleNamespace(name="models/chat-pro"),
+                SimpleNamespace(name="models/embed-lite"),
+                SimpleNamespace(name="models/text-embedding-preview"),
+            ]
+        )
+        models = await provider.get_models()
+        assert models == ["embed-lite", "text-embedding-preview"]
+    finally:
+        await provider.terminate()
+
+
+@pytest.mark.asyncio
+async def test_gemini_embedding_get_models_falls_back_to_all_when_no_match():
+    provider = _make_provider()
+    try:
+        provider.client = _FakeClient(
+            [
+                SimpleNamespace(name="models/gemini-2.5-pro"),
+                SimpleNamespace(name="models/gemini-2.5-flash"),
+                SimpleNamespace(name="models/gemini-2.5-pro"),
+            ]
+        )
+        models = await provider.get_models()
+        assert models == ["gemini-2.5-flash", "gemini-2.5-pro"]
+    finally:
+        await provider.terminate()