mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 00:16:23 +02:00
0a2ce47a88
Base Service (trustgraph-base/trustgraph/base/embeddings_service.py): - Changed on_request to use request.texts FastEmbed Processor (trustgraph-flow/trustgraph/embeddings/fastembed/processor.py): - on_embeddings(texts, model=None) now processes full batch efficiently - Returns [[v.tolist()] for v in vecs] - list of vector sets Ollama Processor (trustgraph-flow/trustgraph/embeddings/ollama/processor.py): - on_embeddings(texts, model=None) passes list directly to Ollama - Returns [[embedding] for embedding in embeds.embeddings] EmbeddingsClient (trustgraph-base/trustgraph/base/embeddings_client.py): - embed(texts, timeout=300) accepts list of texts Tests Updated: - test_fastembed_dynamic_model.py - 4 tests updated for new interface - test_ollama_dynamic_model.py - 4 tests updated for new interface Updated CLI, SDK and APIs
167 lines
7.2 KiB
Python
167 lines
7.2 KiB
Python
"""
|
|
Unit tests for Ollama dynamic model loading
|
|
|
|
Tests the dynamic model selection functionality for Ollama embeddings service.
|
|
Since Ollama is server-side, no model caching is needed on the client side.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, patch, AsyncMock
|
|
from unittest import IsolatedAsyncioTestCase
|
|
from trustgraph.embeddings.ollama.processor import Processor
|
|
|
|
|
|
class TestOllamaDynamicModelLoading(IsolatedAsyncioTestCase):
|
|
"""Test Ollama dynamic model selection"""
|
|
|
|
@patch('trustgraph.embeddings.ollama.processor.Client')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
|
@patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__')
|
|
async def test_client_initialized_with_host(self, mock_embeddings_init, mock_async_init, mock_client_class):
|
|
"""Test that Ollama client is initialized with correct host"""
|
|
# Arrange
|
|
mock_ollama_client = Mock()
|
|
mock_response = Mock()
|
|
mock_response.embeddings = [[0.1, 0.2, 0.3, 0.4, 0.5]]
|
|
mock_ollama_client.embed.return_value = mock_response
|
|
mock_client_class.return_value = mock_ollama_client
|
|
mock_async_init.return_value = None
|
|
mock_embeddings_init.return_value = None
|
|
|
|
# Act
|
|
processor = Processor(id="test", concurrency=1, model="test-model",
|
|
ollama="http://localhost:11434", taskgroup=AsyncMock())
|
|
|
|
# Assert
|
|
mock_client_class.assert_called_once_with(host="http://localhost:11434")
|
|
assert processor.default_model == "test-model"
|
|
|
|
@patch('trustgraph.embeddings.ollama.processor.Client')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
|
@patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__')
|
|
async def test_on_embeddings_uses_default_model(self, mock_embeddings_init, mock_async_init, mock_client_class):
|
|
"""Test that on_embeddings uses default model when no model specified"""
|
|
# Arrange
|
|
mock_ollama_client = Mock()
|
|
mock_response = Mock()
|
|
mock_response.embeddings = [[0.1, 0.2, 0.3, 0.4, 0.5]]
|
|
mock_ollama_client.embed.return_value = mock_response
|
|
mock_client_class.return_value = mock_ollama_client
|
|
mock_async_init.return_value = None
|
|
mock_embeddings_init.return_value = None
|
|
|
|
processor = Processor(id="test", concurrency=1, model="test-model", taskgroup=AsyncMock())
|
|
|
|
# Act
|
|
result = await processor.on_embeddings(["test text"])
|
|
|
|
# Assert
|
|
mock_ollama_client.embed.assert_called_once_with(
|
|
model="test-model",
|
|
input=["test text"]
|
|
)
|
|
assert result == [[[0.1, 0.2, 0.3, 0.4, 0.5]]]
|
|
|
|
@patch('trustgraph.embeddings.ollama.processor.Client')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
|
@patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__')
|
|
async def test_on_embeddings_uses_specified_model(self, mock_embeddings_init, mock_async_init, mock_client_class):
|
|
"""Test that on_embeddings uses specified model when provided"""
|
|
# Arrange
|
|
mock_ollama_client = Mock()
|
|
mock_response = Mock()
|
|
mock_response.embeddings = [[0.1, 0.2, 0.3, 0.4, 0.5]]
|
|
mock_ollama_client.embed.return_value = mock_response
|
|
mock_client_class.return_value = mock_ollama_client
|
|
mock_async_init.return_value = None
|
|
mock_embeddings_init.return_value = None
|
|
|
|
processor = Processor(id="test", concurrency=1, model="test-model", taskgroup=AsyncMock())
|
|
|
|
# Act
|
|
result = await processor.on_embeddings(["test text"], model="custom-model")
|
|
|
|
# Assert
|
|
mock_ollama_client.embed.assert_called_once_with(
|
|
model="custom-model",
|
|
input=["test text"]
|
|
)
|
|
assert result == [[[0.1, 0.2, 0.3, 0.4, 0.5]]]
|
|
|
|
@patch('trustgraph.embeddings.ollama.processor.Client')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
|
@patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__')
|
|
async def test_multiple_model_switches(self, mock_embeddings_init, mock_async_init, mock_client_class):
|
|
"""Test switching between multiple models"""
|
|
# Arrange
|
|
mock_ollama_client = Mock()
|
|
mock_response = Mock()
|
|
mock_response.embeddings = [[0.1, 0.2, 0.3, 0.4, 0.5]]
|
|
mock_ollama_client.embed.return_value = mock_response
|
|
mock_client_class.return_value = mock_ollama_client
|
|
mock_async_init.return_value = None
|
|
mock_embeddings_init.return_value = None
|
|
|
|
processor = Processor(id="test", concurrency=1, model="test-model", taskgroup=AsyncMock())
|
|
|
|
# Act - switch between different models
|
|
await processor.on_embeddings(["text1"], model="model-a")
|
|
await processor.on_embeddings(["text2"], model="model-b")
|
|
await processor.on_embeddings(["text3"], model="model-a")
|
|
await processor.on_embeddings(["text4"]) # Use default
|
|
|
|
# Assert
|
|
calls = mock_ollama_client.embed.call_args_list
|
|
assert len(calls) == 4
|
|
assert calls[0][1]['model'] == "model-a"
|
|
assert calls[1][1]['model'] == "model-b"
|
|
assert calls[2][1]['model'] == "model-a"
|
|
assert calls[3][1]['model'] == "test-model" # Default
|
|
|
|
@patch('trustgraph.embeddings.ollama.processor.Client')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
|
@patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__')
|
|
async def test_none_model_uses_default(self, mock_embeddings_init, mock_async_init, mock_client_class):
|
|
"""Test that None model parameter falls back to default"""
|
|
# Arrange
|
|
mock_ollama_client = Mock()
|
|
mock_response = Mock()
|
|
mock_response.embeddings = [[0.1, 0.2, 0.3, 0.4, 0.5]]
|
|
mock_ollama_client.embed.return_value = mock_response
|
|
mock_client_class.return_value = mock_ollama_client
|
|
mock_async_init.return_value = None
|
|
mock_embeddings_init.return_value = None
|
|
|
|
processor = Processor(id="test", concurrency=1, model="test-model", taskgroup=AsyncMock())
|
|
|
|
# Act
|
|
result = await processor.on_embeddings(["test text"], model=None)
|
|
|
|
# Assert
|
|
mock_ollama_client.embed.assert_called_once_with(
|
|
model="test-model",
|
|
input=["test text"]
|
|
)
|
|
|
|
@patch('trustgraph.embeddings.ollama.processor.Client')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
|
@patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__')
|
|
async def test_initialization_without_model_uses_default(self, mock_embeddings_init, mock_async_init, mock_client_class):
|
|
"""Test initialization without model parameter uses module default"""
|
|
# Arrange
|
|
mock_ollama_client = Mock()
|
|
mock_client_class.return_value = mock_ollama_client
|
|
mock_async_init.return_value = None
|
|
mock_embeddings_init.return_value = None
|
|
|
|
# Act
|
|
processor = Processor(id="test-embeddings", concurrency=1, taskgroup=AsyncMock())
|
|
|
|
# Assert
|
|
# Should use default_model from module
|
|
expected_default = "mxbai-embed-large"
|
|
assert processor.default_model == expected_default
|
|
|
|
|
|
if __name__ == '__main__':
|
|
pytest.main([__file__])
|