Skip to content

Commit ccb97ff

Browse files
Simplify vectorstore factory to only support Pydantic ModelConfig objects
Remove dictionary support from create_embeddings() function and update tests to use proper ModelConfig instances instead of dictionaries. This improves type safety and consistency across the codebase. Co-authored-by: Andrew Brookins <[email protected]>
1 parent ddd747d commit ccb97ff

File tree

2 files changed

+14
-12
lines changed

2 files changed

+14
-12
lines changed

agent_memory_server/vectorstore_factory.py

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -46,15 +46,8 @@ def create_embeddings() -> Embeddings:
4646
An Embeddings instance
4747
"""
4848
embedding_config = settings.embedding_model_config
49-
if embedding_config:
50-
# Handle both dict (for tests with mocked config) and ModelConfig (actual config)
51-
provider = (
52-
embedding_config.get("provider", "openai")
53-
if hasattr(embedding_config, 'get')
54-
else embedding_config.provider
55-
)
56-
else:
57-
provider = "openai"
49+
# Only support ModelConfig objects
50+
provider = embedding_config.provider if embedding_config else "openai"
5851

5952
if provider == "openai":
6053
try:

tests/integration/test_vectorstore_factory_integration.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import pytest
1111
from langchain_core.embeddings import Embeddings
1212

13+
from agent_memory_server.config import ModelConfig, ModelProvider
1314
from agent_memory_server.vectorstore_factory import (
1415
_import_and_call_factory,
1516
create_embeddings,
@@ -89,8 +90,13 @@ class TestEmbeddingsCreation:
8990
def test_create_openai_embeddings(self, mock_settings):
9091
"""Test OpenAI embeddings creation."""
9192

92-
# Configure mock settings
93-
mock_settings.embedding_model_config = {"provider": "openai"}
93+
# Configure mock settings with ModelConfig object
94+
mock_settings.embedding_model_config = ModelConfig(
95+
provider=ModelProvider.OPENAI,
96+
name="text-embedding-3-small",
97+
max_tokens=8191,
98+
embedding_dimensions=1536,
99+
)
94100
mock_settings.embedding_model = "text-embedding-3-small"
95101
mock_settings.openai_api_key = "test-key"
96102

@@ -107,7 +113,10 @@ def test_create_openai_embeddings(self, mock_settings):
107113
def test_create_embeddings_unsupported_provider(self, mock_settings):
108114
"""Test embeddings creation with unsupported provider."""
109115

110-
mock_settings.embedding_model_config = {"provider": "unsupported"}
116+
# Create a mock model config with unsupported provider
117+
mock_config = Mock()
118+
mock_config.provider = "unsupported" # Set directly as string, bypassing enum validation
119+
mock_settings.embedding_model_config = mock_config
111120

112121
with pytest.raises(ValueError, match="Unsupported embedding provider"):
113122
create_embeddings()

0 commit comments

Comments
 (0)