Skip to content

Commit

Permalink
chore(embeddings): added tests for embeddings (#3183)
Browse files Browse the repository at this point in the history
# 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):
  • Loading branch information
StanGirard authored Sep 11, 2024
1 parent 7edc3a2 commit 758e87e
Showing 1 changed file with 38 additions and 0 deletions.
38 changes: 38 additions & 0 deletions backend/api/tests/settings/test_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
from unittest.mock import patch, MagicMock
from quivr_api.modules.dependencies import get_embedding_client
from langchain_community.embeddings.ollama import OllamaEmbeddings
from langchain_openai import AzureOpenAIEmbeddings

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

0 comments on commit 758e87e

Please sign in to comment.