src/pyseekdb/utils/embedding_functions/litellm_base_embedding_function.py:60

• The docstring on line 60 says "Initialize." but should be more descriptive, like "Initialize LiteLLMBaseEmbeddingFunction." to match the naming convention used in other embedding functions.

src/pyseekdb/utils/embedding_functions/amazon_bedrock_embedding_function.py (2)

110-112: kwargs validation may be too restrictive for boto3 Config objects.

This validation rejects non-primitive types, but boto3 users may want to pass botocore.config.Config objects via kwargs. This was previously flagged.

---

255-257: Error message references wrong parameter name.

The error message says "kwargs must be a dictionary" but the config key is "client_kwargs".

📝 Proposed fix

         kwargs = config.get("client_kwargs", {})
         if not isinstance(kwargs, dict):
-            raise TypeError(f"kwargs must be a dictionary, but got {kwargs}")
+            raise TypeError(f"client_kwargs must be a dictionary, but got {kwargs}")

src/pyseekdb/utils/embedding_functions/jina_embedding_function.py (2)

84-85: Docstring default value is inconsistent with actual default.

The docstring states the default is "jina-embeddings-v2-base-en" but the actual default parameter on line 77 is "jina-embeddings-v3".

📝 Proposed fix

         Args:
             model_name (str, optional): Name of the Jina AI embedding model.
-                Defaults to "jina-embeddings-v2-base-en".
+                Defaults to "jina-embeddings-v3".

---

212-214: Error message references wrong parameter name.

The error message says "kwargs must be a dictionary" but the config key being validated is "client_kwargs".

📝 Proposed fix

         kwargs = config.get("client_kwargs", {})
         if not isinstance(kwargs, dict):
-            raise TypeError(f"kwargs must be a dictionary, but got {kwargs}")
+            raise TypeError(f"client_kwargs must be a dictionary, but got {kwargs}")

src/pyseekdb/utils/embedding_functions/tengxun_hunyuan_embedding_function.py (1)

81-89: Redundant assignment of dimensions variable.

The dimensions = None assignment on line 89 is never used since line 104 always passes dimensions=None explicitly. While harmless, it can be removed for clarity.

Proposed fix

         # Warn if dimensions parameter is provided since it's not supported
         if dimensions is not None:
             warnings.warn(
                 "The dimensions parameter is not supported by Tencent Hunyuan API. "
                 "Dimensions are fixed at 1024. The provided dimensions parameter will be ignored.",
                 UserWarning,
                 stacklevel=2,
             )
-            dimensions = None

src/pyseekdb/utils/embedding_functions/google_vertex_embedding_function.py (1)

242-243: Error message should reference client_kwargs for consistency.

The error message says "kwargs" but the config parameter is named "client_kwargs".

Proposed fix

         if not isinstance(kwargs, dict):
-            raise TypeError(f"kwargs must be a dictionary, but got {kwargs}")
+            raise TypeError(f"client_kwargs must be a dictionary, but got {kwargs}")

src/pyseekdb/utils/embedding_functions/cohere_embedding_function.py (1)

216-217: Error message should reference client_kwargs for consistency.

Same issue as in other embedding functions - the error message says "kwargs" but the config parameter is "client_kwargs".

Proposed fix

         if not isinstance(kwargs, dict):
-            raise TypeError(f"kwargs must be a dictionary, but got {kwargs}")
+            raise TypeError(f"client_kwargs must be a dictionary, but got {kwargs}")

src/pyseekdb/utils/embedding_functions/voyageai_embedding_function.py (1)

275-276: Error message should reference client_kwargs for consistency.

Same issue as in other embedding functions.

Proposed fix

         if not isinstance(kwargs, dict):
-            raise TypeError(f"kwargs must be a dictionary, but got {kwargs}")
+            raise TypeError(f"client_kwargs must be a dictionary, but got {kwargs}")

.github/workflows/ci.yml (1)

69-70: Fragile failure detection pattern.

The tail -n 1 | grep -q 'failed' approach can produce false negatives if:

• pytest crashes before producing summary output
• The summary format changes or doesn't end with "failed" on the last line
• Tests error (vs fail) may not contain "failed" in the last line

A more robust approach is to capture pytest's exit code directly:

♻️ Suggested improvement

-          uv run pytest tests/integration_tests/ -v --log-cli-level=${log_level} -k "not server and not embedded and not oceanbase" | tee pytest.log
-          tail -n 1 pytest.log | grep -q 'failed' && exit 1 || exit 0
+          uv run pytest tests/integration_tests/ -v --log-cli-level=${log_level} -k "not server and not embedded and not oceanbase" 2>&1 | tee pytest.log
+          exit ${PIPESTATUS[0]}

Note: This is an existing pattern throughout the workflow, so this is a low-priority suggestion that could be addressed separately across all steps.

src/pyseekdb/client/embedding_function.py (1)

517-519: Self return type on static method is unconventional.

Using Self as the return type for a @staticmethod is semantically incorrect. Self is designed for instance methods and @classmethod where it refers to the type of the current class instance. For a static method, it doesn't have access to the class context to determine Self.

Since DefaultEmbeddingFunction.build_from_config always returns DefaultEmbeddingFunction(), the return type should be the concrete class.

📝 Suggested fix

     `@staticmethod`
-    def build_from_config(_config: dict[str, Any]) -> Self:
+    def build_from_config(_config: dict[str, Any]) -> "DefaultEmbeddingFunction":
         return DefaultEmbeddingFunction()

src/pyseekdb/utils/embedding_functions/amazon_bedrock_embedding_function.py (1)

186-209: Sequential API calls may be slow for large document batches.

Each document triggers a separate API call. While this matches Bedrock's Titan embedding API constraints (single input per call), consider adding a note in the docstring about potential latency for large batches, or consider batching where supported by newer models.

tests/unit_tests/test_jina_embedding_function.py (1)

107-119: Consider using env_guard for consistency.

This test correctly uses try/finally for cleanup, but using env_guard(JINA_AI_API_KEY=None) would be more consistent with the persistence tests and cleaner.

tests/unit_tests/test_ollama_embedding_function.py (2)

54-55: Consider being more explicit about caught exceptions.

While catching all exceptions is acceptable for availability checks, consider catching more specific exceptions (e.g., openai.APIConnectionError, ConnectionError) or adding a brief comment explaining the intent. This helps distinguish intentional broad catching from accidental oversight.

💡 Optional: Add comment or narrow exception types

-    except Exception:
+    except Exception:  # Any failure means Ollama is not available
         return False

Or more specific:

-    except Exception:
+    except (ConnectionError, openai.APIConnectionError, openai.APIError):
         return False

---

67-77: Redundant environment checks in each test method.

The class-level @pytest.mark.skipif(not is_ollama_available(), ...) decorator already ensures tests only run when Ollama is available. Calling self.test_ollama_env() at the start of each test method is redundant and adds unnecessary overhead.

Consider removing these calls from individual test methods, keeping only test_ollama_env as a standalone sanity check if desired.

tests/unit_tests/test_tencent_hunyuan_embedding_function.py (2)

49-66: Redundant environment checks inside test methods.

Since the class is decorated with @pytest.mark.skipif that already checks both conditions, the self.test_hunyuan_env() calls inside each test method (lines 54, 71, 85, 97, etc.) are redundant. The class will only run if conditions are already met.

♻️ Suggested simplification

Remove the self.test_hunyuan_env() calls from individual test methods. The class-level skipif decorator already guarantees the prerequisites:

     def test_initialization_with_defaults(self):
         """Test TencentHunyuanEmbeddingFunction initialization with default values"""
         print("\n✅ Testing TencentHunyuanEmbeddingFunction initialization with defaults")
 
-        # Check if openai is available and env vars are set
-        self.test_hunyuan_env()
-
         ef = TencentHunyuanEmbeddingFunction()

Apply similar changes to other test methods that call self.test_hunyuan_env().

---

169-173: Consider suppressing expected warning.

Line 170 creates an instance with dimensions=512, which triggers a warning (as verified in test_initialization_with_dimensions_warning). Suppress it here to reduce test noise, consistent with test_get_config_with_dimensions (lines 316-318).

♻️ Suggested fix

         # Test that dimension is always 1024 regardless of initialization
-        ef2 = TencentHunyuanEmbeddingFunction(dimensions=512)  # Should warn but dimension still 1024
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            ef2 = TencentHunyuanEmbeddingFunction(dimensions=512)  # Should warn but dimension still 1024
         dim2 = ef2.dimension

tests/unit_tests/test_siliconflow_embedding_function.py (2)

42-46: Redundant environment checks within test methods.

Since the class is decorated with @pytest.mark.skipif that already verifies SILICONFLOW_API_KEY and is_openai_available(), the repeated calls to self.test_siliconflow_env() at the start of each test method (lines 53, 70, 94, 108, 120, 131, 155, 203, etc.) are redundant. These conditions are guaranteed to be true when any test in this class executes.

Consider removing the redundant self.test_siliconflow_env() calls from individual test methods to reduce test execution overhead.

---

90-102: Environment variable not cleaned up after test.

The custom environment variable set at line 98 is never cleaned up. If this test or subsequent tests fail, CUSTOM_SILICONFLOW_KEY remains in the environment, potentially affecting test isolation.

Consider using the env_guard context manager (already imported from test_utils) for consistency with the persistence tests.

♻️ Proposed fix

     def test_initialization_with_custom_api_key_env(self):
         """Test SiliconflowEmbeddingFunction initialization with custom API key env"""
         print("\n✅ Testing SiliconflowEmbeddingFunction initialization with custom API key env")

-        self.test_siliconflow_env()
-
         custom_key_env = "CUSTOM_SILICONFLOW_KEY"
-        if not os.environ.get(custom_key_env):
-            os.environ[custom_key_env] = os.environ.get("SILICONFLOW_API_KEY", "your-custom-key")
-
-        ef = SiliconflowEmbeddingFunction(model_name="BAAI/bge-large-zh-v1.5", api_key_env=custom_key_env)
-        assert ef.api_key_env == custom_key_env
-        print(f"   Custom API key env: {ef.api_key_env}")
+        with env_guard(**{custom_key_env: os.environ.get("SILICONFLOW_API_KEY", "test-key")}):
+            ef = SiliconflowEmbeddingFunction(model_name="BAAI/bge-large-zh-v1.5", api_key_env=custom_key_env)
+            assert ef.api_key_env == custom_key_env
+            print(f"   Custom API key env: {ef.api_key_env}")

tests/unit_tests/test_voyageai_embedding_function.py (1)

102-108: Avoid leaking CUSTOM_VOYAGE_KEY into the process environment.
Use env_guard (or restore the original value) so this test doesn’t affect later tests.

♻️ Proposed refactor

-        custom_key_env = "CUSTOM_VOYAGE_KEY"
-        if not os.environ.get(custom_key_env):
-            os.environ[custom_key_env] = os.environ.get("VOYAGE_API_KEY", "your-custom-key")
-
-        ef = VoyageaiEmbeddingFunction(model_name="voyage-4-large", api_key_env=custom_key_env)
-        assert ef.api_key_env == custom_key_env
-        print(f"   Custom API key env: {ef.api_key_env}")
+        custom_key_env = "CUSTOM_VOYAGE_KEY"
+        with env_guard(**{custom_key_env: os.environ.get("VOYAGE_API_KEY", "your-custom-key")}):
+            ef = VoyageaiEmbeddingFunction(model_name="voyage-4-large", api_key_env=custom_key_env)
+            assert ef.api_key_env == custom_key_env
+            print(f"   Custom API key env: {ef.api_key_env}")

src/pyseekdb/utils/embedding_functions/amazon_bedrock_embedding_function.py (1)

265-266: Passing None values in session_args to boto3.Session() may cause issues.

When session_args contains None values (e.g., {"region_name": None, "profile_name": None}), passing them to boto3.Session() may cause unexpected behavior compared to not passing them at all.

📝 Suggested fix (from past review)

         session_args = config.get("session_args")
-        session = boto3.Session(**session_args) if session_args else boto3.Session()
+        if session_args:
+            # Filter out None values to avoid passing them to boto3.Session
+            filtered_args = {k: v for k, v in session_args.items() if v is not None}
+            session = boto3.Session(**filtered_args) if filtered_args else boto3.Session()
+        else:
+            session = boto3.Session()

src/pyseekdb/utils/embedding_functions/jina_embedding_function.py (1)

196-199: Error message references "kwargs" but config key is "client_kwargs".

For consistency with the config structure and better debugging experience, the error message should reference the actual config key name.

📝 Suggested fix

         kwargs = config.get("client_kwargs", {})
         if not isinstance(kwargs, dict):
-            raise TypeError(f"kwargs must be a dictionary, but got {kwargs}")
+            raise TypeError(f"client_kwargs must be a dictionary, but got {kwargs}")

src/pyseekdb/utils/embedding_functions/cohere_embedding_function.py (1)

206-219: Error message references wrong parameter name.

The error message on line 212 says "kwargs must be a dictionary" but the config key being validated is client_kwargs. This inconsistency could confuse users debugging configuration issues.

🐛 Proposed fix

         kwargs = config.get("client_kwargs", {})
         if not isinstance(kwargs, dict):
-            raise TypeError(f"kwargs must be a dictionary, but got {kwargs}")
+            raise TypeError(f"client_kwargs must be a dictionary, but got {kwargs}")

src/pyseekdb/utils/embedding_functions/voyageai_embedding_function.py (1)

259-278: Error message references wrong parameter name.

Similar to the Cohere implementation, the error message on line 269 says "kwargs must be a dictionary" but the config key being validated is client_kwargs.

🐛 Proposed fix

         kwargs = config.get("client_kwargs", {})
         if not isinstance(kwargs, dict):
-            raise TypeError(f"kwargs must be a dictionary, but got {kwargs}")
+            raise TypeError(f"client_kwargs must be a dictionary, but got {kwargs}")

src/pyseekdb/utils/embedding_functions/google_vertex_embedding_function.py (1)

9-101: Missing dimension property for consistency.

Other embedding functions in this PR (Jina, Siliconflow, AmazonBedrock, Voyageai) implement a dimension property that returns the embedding vector size. This is missing from GoogleVertexEmbeddingFunction, which may cause issues if callers expect this property.

♻️ Suggested addition

`@property`
def dimension(self) -> int:
    """Get the dimension of embeddings produced by this function."""
    # textembedding-gecko produces 768-dimensional embeddings
    # Add known dimensions mapping or probe via API
    known_dims = {
        "textembedding-gecko": 768,
        "textembedding-gecko@003": 768,
        "text-embedding-004": 768,
    }
    if self.model_name in known_dims:
        return known_dims[self.model_name]
    
    # Fallback: probe via API
    test_embeddings = self(["dimension probing"])
    return len(test_embeddings[0])

src/pyseekdb/client/embedding_function.py (1)

631-660: Consider partial registration for optional embedding functions.

The current implementation uses a single try/except block for all imports. If any one optional dependency is missing, none of the embedding functions will be registered. This could be problematic if a user has only some dependencies installed (e.g., has voyageai but not boto3).

Consider importing and registering each embedding function individually to allow partial availability:

♻️ Suggested approach for partial registration

# Register each optional embedding function individually
optional_functions = [
    ("sentence_transformer", "SentenceTransformerEmbeddingFunction"),
    ("openai", "OpenAIEmbeddingFunction"),
    ("qwen", "QwenEmbeddingFunction"),
    ("siliconflow", "SiliconflowEmbeddingFunction"),
    ("tencent_hunyuan", "TencentHunyuanEmbeddingFunction"),
    ("ollama", "OllamaEmbeddingFunction"),
    ("voyageai", "VoyageaiEmbeddingFunction"),
    ("google_vertex", "GoogleVertexEmbeddingFunction"),
    ("cohere", "CohereEmbeddingFunction"),
    ("jina", "JinaEmbeddingFunction"),
    ("amazon_bedrock", "AmazonBedrockEmbeddingFunction"),
]

for registry_name, class_name in optional_functions:
    try:
        module = __import__(
            "pyseekdb.utils.embedding_functions",
            fromlist=[class_name]
        )
        cls = getattr(module, class_name)
        cls._registry[registry_name] = cls
    except (ImportError, AttributeError) as e:
        logger.debug(f"Optional embedding function {class_name} not available: {e}")

src/pyseekdb/utils/embedding_functions/amazon_bedrock_embedding_function.py (2)

243-267: Align client_kwargs naming and filter session_args None values.

Two small robustness fixes:

• The error message should reference client_kwargs for consistency.
• Filter None values before building a boto3.Session.

🛠️ Proposed fix

-        kwargs = config.get("client_kwargs", {})
-        if not isinstance(kwargs, dict):
-            raise TypeError(f"kwargs must be a dictionary, but got {kwargs}")
+        client_kwargs = config.get("client_kwargs", {})
+        if not isinstance(client_kwargs, dict):
+            raise TypeError(f"client_kwargs must be a dictionary, but got {client_kwargs}")
@@
-        session_args = config.get("session_args")
-        session = boto3.Session(**session_args) if session_args else boto3.Session()
+        session_args = config.get("session_args") or {}
+        if session_args:
+            session_args = {k: v for k, v in session_args.items() if v is not None}
+        session = boto3.Session(**session_args) if session_args else boto3.Session()
@@
-            **kwargs,
+            **client_kwargs,

Verification (boto3 Session args handling):

boto3 Session parameters None values behavior region_name profile_name

---

95-115: Config object is documented but will be rejected by validation, blocking a documented feature.

The docstring at line 92 explicitly documents the config: boto3 Config object as a valid kwarg, but the validation at lines 98–100 rejects it because botocore.config.Config is not a primitive type. Additionally, even if validation allowed it, storing Config in self._client_kwargs (line 111) would break serialization in get_config() since Config objects are not JSON-serializable.

To fix this, separate runtime kwargs (which may include Config) from serializable kwargs (which must remain primitive):

🐛 Suggested fix (allow runtime Config, keep persisted kwargs serializable)

-        for key, value in kwargs.items():
-            if not isinstance(value, (str, int, float, bool, list, dict, tuple)):
-                raise TypeError(f"Keyword argument {key} is not a primitive type: {type(value)!s}")
+        runtime_kwargs = kwargs.copy()
+        serializable_kwargs: dict[str, Any] = {}
+        try:
+            from botocore.config import Config as BotoConfig
+        except Exception:
+            BotoConfig = None
+        for key, value in kwargs.items():
+            if key == "config" and BotoConfig is not None and isinstance(value, BotoConfig):
+                # Allow at runtime, but don't persist
+                continue
+            if not isinstance(value, (str, int, float, bool, list, dict, tuple, type(None))):
+                raise TypeError(f"Keyword argument {key} is not a primitive type: {type(value)!s}")
+            serializable_kwargs[key] = value
@@
-        self._client_kwargs = kwargs
+        self._client_kwargs = serializable_kwargs
@@
-        self._client = session.client("bedrock-runtime", **kwargs)
+        self._client = session.client("bedrock-runtime", **runtime_kwargs)

src/pyseekdb/utils/embedding_functions/google_vertex_embedding_function.py (2)

45-45: Environment variable name uses "CHROMA" prefix.

The default api_key_env is "CHROMA_GOOGLE_VERTEX_API_KEY" which appears to be copied from Chroma's codebase. Consider using a more project-appropriate name like "GOOGLE_VERTEX_API_KEY" or "PYSEEKDB_GOOGLE_VERTEX_API_KEY".

---

107-122: build_from_config is stricter than the constructor.

The constructor has defaults for all parameters, but build_from_config raises an error if any are missing. This means configs saved with defaults cannot be restored if any field is None. Consider using defaults similar to the constructor.

src/pyseekdb/utils/embedding_functions/jina_embedding_function.py (1)

193-195: Error message references wrong parameter name.

The error message says "kwargs must be a dictionary" but the config parameter is client_kwargs. Update for consistency.

src/pyseekdb/utils/embedding_functions/cohere_embedding_function.py (1)

204-206: Error message references wrong parameter name.

The error message says "kwargs must be a dictionary" but the config parameter is client_kwargs. Update for consistency with the parameter name.

src/pyseekdb/utils/embedding_functions/voyageai_embedding_function.py (1)

163-199: Optional: normalize embeddings to plain lists for defensive robustness.

The Voyage AI API with output_dtype="float" (already set explicitly) returns List[List[float]] per official documentation, which matches the Embeddings type requirement. However, defensive normalization could guard against unexpected API behavior changes:

♻️ Proposed refinement

-        # Convert to list of lists (in case voyageai returns numpy arrays or other formats)
-        return embeddings
+        # Convert to list of lists (in case voyageai returns numpy arrays or other formats)
+        normalized = []
+        for embedding in embeddings:
+            if hasattr(embedding, "tolist"):
+                normalized.append(embedding.tolist())
+            else:
+                normalized.append(list(embedding))
+        return normalized

src/pyseekdb/utils/embedding_functions/cohere_embedding_function.py (1)

200-200: Inconsistent handling of missing model_name in build_from_config.

Unlike other embedding functions (e.g., JinaEmbeddingFunction, OpenAIEmbeddingFunction) that raise ValueError when model_name is missing, this implementation uses a default value. This inconsistency could lead to unexpected behavior if a config is partially corrupted. Consider aligning with other implementations by requiring model_name.

📝 Suggested fix for consistency

-        model_name = config.get("model_name", _DEFAULT_MODEL_NAME)
+        model_name = config.get("model_name")
+        if model_name is None:
+            raise ValueError("Missing required field 'model_name' in configuration")