mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-14 17:03:29 +03:00
282fa0e3f8
# Description Please include a summary of the changes and the related issue. Please also include relevant motivation and context. ## Checklist before requesting a review Please delete options that are not relevant. - [ ] My code follows the style guidelines of this project - [ ] I have performed a self-review of my code - [ ] I have commented hard-to-understand areas - [ ] I have ideally added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [ ] Any dependent changes have been merged ## Screenshots (if appropriate):
47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
from langchain_community.embeddings.ollama import OllamaEmbeddings
|
|
from langchain_openai import AzureOpenAIEmbeddings
|
|
from quivr_api.modules.dependencies import get_embedding_client
|
|
|
|
|
|
def test_ollama_embedding():
|
|
with patch("quivr_api.modules.dependencies.settings") as mock_settings:
|
|
mock_settings.ollama_api_base_url = "http://ollama.example.com"
|
|
mock_settings.azure_openai_embeddings_url = None
|
|
|
|
embedding_client = get_embedding_client()
|
|
|
|
assert isinstance(embedding_client, OllamaEmbeddings)
|
|
assert embedding_client.base_url == "http://ollama.example.com"
|
|
|
|
|
|
def test_azure_embedding():
|
|
with patch("quivr_api.modules.dependencies.settings") as mock_settings:
|
|
mock_settings.ollama_api_base_url = None
|
|
mock_settings.azure_openai_embeddings_url = "https://quivr-test.openai.azure.com/openai/deployments/embedding/embeddings?api-version=2023-05-15"
|
|
|
|
embedding_client = get_embedding_client()
|
|
|
|
assert isinstance(embedding_client, AzureOpenAIEmbeddings)
|
|
assert embedding_client.azure_endpoint == "https://quivr-test.openai.azure.com"
|
|
|
|
|
|
def test_openai_embedding():
|
|
with (
|
|
patch("quivr_api.modules.dependencies.settings") as mock_settings,
|
|
patch(
|
|
"quivr_api.modules.dependencies.OpenAIEmbeddings"
|
|
) as mock_openai_embeddings,
|
|
):
|
|
mock_settings.ollama_api_base_url = None
|
|
mock_settings.azure_openai_embeddings_url = None
|
|
|
|
# Create a mock instance for OpenAIEmbeddings
|
|
mock_openai_instance = MagicMock()
|
|
mock_openai_embeddings.return_value = mock_openai_instance
|
|
|
|
embedding_client = get_embedding_client()
|
|
|
|
assert embedding_client == mock_openai_instance
|