Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

updated tests to reflect changes to the rag pipeline #381

Merged
merged 1 commit into from
Oct 20, 2024
Merged

updated tests to reflect changes to the rag pipeline #381

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions App_Function_Libraries/RAG/RAG_Library_2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -343,8 +343,7 @@ def perform_vector_search(query: str, relevant_media_ids: List[str] = None, top_
vector_results = []
try:
for collection in all_collections:
k = top_k
collection_results = vector_search(collection.name, query, k)
collection_results = vector_search(collection.name, query, k=top_k)
filtered_results = [
result for result in collection_results
if relevant_media_ids is None or result['metadata'].get('media_id') in relevant_media_ids
Expand Down
22 changes: 11 additions & 11 deletions Tests/RAG/test_RAG_Library_2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ def test_perform_vector_search_with_relevant_media_ids(self, mock_chroma_client,
mock_chroma_client.list_collections.assert_called_once()

# Assert vector_search was called with correct arguments
mock_vector_search.assert_called_once_with('collection1', query, k=5)
mock_vector_search.assert_called_once_with('collection1', query, k=10)

@patch('App_Function_Libraries.RAG.RAG_Library_2.vector_search')
@patch('App_Function_Libraries.RAG.RAG_Library_2.chroma_client')
Expand Down Expand Up @@ -152,7 +152,7 @@ def test_perform_vector_search_without_relevant_media_ids(self, mock_chroma_clie
mock_chroma_client.list_collections.assert_called_once()

# Assert vector_search was called with correct arguments
mock_vector_search.assert_called_once_with('collection1', query, k=5)
mock_vector_search.assert_called_once_with('collection1', query, k=10)

@patch('App_Function_Libraries.RAG.RAG_Library_2.search_db')
def test_perform_full_text_search_with_relevant_media_ids(self, mock_search_db):
Expand All @@ -171,7 +171,7 @@ def test_perform_full_text_search_with_relevant_media_ids(self, mock_search_db):
relevant_media_ids = [1, 3]

# Call the function
result = perform_full_text_search(query, relevant_media_ids)
result = perform_full_text_search(query, relevant_media_ids, fts_top_k=10)

# Expected to filter out id 2
expected = [
Expand All @@ -182,7 +182,7 @@ def test_perform_full_text_search_with_relevant_media_ids(self, mock_search_db):

# Assert search_db was called with correct arguments
mock_search_db.assert_called_once_with(
query, ['content'], '', page=1, results_per_page=5)
query, ['content'], '', page=1, results_per_page=10)

@patch('App_Function_Libraries.RAG.RAG_Library_2.search_db')
def test_perform_full_text_search_without_relevant_media_ids(self, mock_search_db):
Expand Down Expand Up @@ -211,7 +211,7 @@ def test_perform_full_text_search_without_relevant_media_ids(self, mock_search_d

# Assert search_db was called with correct arguments
mock_search_db.assert_called_once_with(
query, ['content'], '', page=1, results_per_page=5)
query, ['content'], '', page=1, results_per_page=10)

@patch('App_Function_Libraries.RAG.RAG_Library_2.search_db')
def test_perform_full_text_search_empty_results(self, mock_search_db):
Expand All @@ -234,7 +234,7 @@ def test_perform_full_text_search_empty_results(self, mock_search_db):

# Assert search_db was called with correct arguments
mock_search_db.assert_called_once_with(
query, ['content'], '', page=1, results_per_page=5)
query, ['content'], '', page=1, results_per_page=10)

@patch('App_Function_Libraries.RAG.RAG_Library_2.fetch_keywords_for_media')
@patch('App_Function_Libraries.RAG.RAG_Library_2.logging')
Expand Down Expand Up @@ -344,7 +344,7 @@ def test_perform_full_text_search_case_insensitive_filtering(self, mock_search_d

# Assert search_db was called with correct arguments
mock_search_db.assert_called_once_with(
query, ['content'], '', page=1, results_per_page=5)
query, ['content'], '', page=1, results_per_page=10)

@patch('App_Function_Libraries.RAG.RAG_Library_2.search_db')
def test_perform_full_text_search_multiple_pages(self, mock_search_db):
Expand Down Expand Up @@ -380,7 +380,7 @@ def test_perform_full_text_search_multiple_pages(self, mock_search_db):

# Assert search_db was called with correct arguments
mock_search_db.assert_called_once_with(
query, ['content'], '', page=1, results_per_page=5)
query, ['content'], '', page=1, results_per_page=10)

@patch('App_Function_Libraries.RAG.RAG_Library_2.chroma_client')
@patch('App_Function_Libraries.RAG.RAG_Library_2.vector_search')
Expand Down Expand Up @@ -429,8 +429,8 @@ def vector_search_side_effect(collection_name, query, k):
mock_chroma_client.list_collections.assert_called_once()

# Assert vector_search was called twice with correct arguments
mock_vector_search.assert_any_call('collection1', query, k=5)
mock_vector_search.assert_any_call('collection2', query, k=5)
mock_vector_search.assert_any_call('collection1', query, k=10)
mock_vector_search.assert_any_call('collection2', query, k=10)
self.assertEqual(mock_vector_search.call_count, 2)

@patch('App_Function_Libraries.RAG.RAG_Library_2.search_db')
Expand Down Expand Up @@ -462,7 +462,7 @@ def test_perform_full_text_search_partial_matches(self, mock_search_db):

# Assert search_db was called with correct arguments
mock_search_db.assert_called_once_with(
query, ['content'], '', page=1, results_per_page=5)
query, ['content'], '', page=1, results_per_page=10)


if __name__ == '__main__':
Expand Down
Loading