use model_name in embedding

seehi · May 24, 2024 · 5f1ca3c · 5f1ca3c
1 parent db26406
commit 5f1ca3c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/metagpt/rag/factories/embedding.py b/metagpt/rag/factories/embedding.py
@@ -89,9 +89,9 @@ def _create_ollama(self) -> OllamaEmbedding:
         return OllamaEmbedding(**params)
 
     def _try_set_model_and_batch_size(self, params: dict):
-        """Set the model and embed_batch_size only when they are specified."""
+        """Set the model_name and embed_batch_size only when they are specified."""
         if config.embedding.model:
-            params["model"] = config.embedding.model
+            params["model_name"] = config.embedding.model
 
         if config.embedding.embed_batch_size:
             params["embed_batch_size"] = config.embedding.embed_batch_size

diff --git a/tests/metagpt/rag/factories/test_embedding.py b/tests/metagpt/rag/factories/test_embedding.py
@@ -66,7 +66,7 @@ def test_get_rag_embedding_default(self, mocker, mock_config):
 
     @pytest.mark.parametrize(
         "model, embed_batch_size, expected_params",
-        [("test_model", 100, {"model": "test_model", "embed_batch_size": 100}), (None, None, {})],
+        [("test_model", 100, {"model_name": "test_model", "embed_batch_size": 100}), (None, None, {})],
     )
     def test_try_set_model_and_batch_size(self, mock_config, model, embed_batch_size, expected_params):
         # Mock