"""Tests for embedding functionality."""
from preloop.models.crud import crud_embedding_model, crud_issue_embedding, crud_comment
from preloop.models.models.issue import IssueEmbedding
from preloop.models.models.comment import Comment
from preloop.models.models import Issue # Ensure Issue is imported
from sqlalchemy.orm import Session
import pytest
def test_create_embedding_model(db_session):
"""Test creating a new embedding model."""
model_data = {
"name": "test-embeddings",
"provider": "openai",
"version": "text-embedding-ada-002",
"dimensions": 1536,
"is_active": True,
"meta_data": {},
}
model = crud_embedding_model.create(db_session, obj_in=model_data)
assert model.name == model_data["name"]
assert model.provider == model_data["provider"]
assert model.dimensions == model_data["dimensions"]
assert model.is_active is True
# Verify it was saved in the database
db_model = crud_embedding_model.get(db_session, id=model.id)
assert db_model is not None
assert db_model.id == model.id
def test_create_issue_embedding(
db_session, create_issue, create_embedding_model, mocker
):
"""Test creating an issue embedding."""
# Create test issue and embedding model
issue = create_issue()
model = create_embedding_model()
# Mock the embedding vector generation
mock_generate_embedding = mocker.patch(
"preloop.models.crud.crud_issue_embedding._generate_embedding_vector"
)
mock_generate_embedding.return_value = [0.1] * model.dimensions
# Create embedding for the issue
result = crud_issue_embedding.create_embeddings(
db_session, issue_id=issue.id, force_update=False
)
# Check result contains the model name with success status
assert model.name in result
assert "created" in result[model.name] or "exists" in result[model.name]
# Verify embedding was created in the database
embeddings = crud_issue_embedding.get_for_issue(db_session, issue_id=issue.id)
assert model.name in embeddings
# Verify the embedding has the correct dimension
embedding = embeddings[model.name]
assert len(embedding.embedding) == model.dimensions
def test_create_comment_embedding(
db_session, create_issue, create_comment, create_embedding_model, mocker
):
"""Test creating an embedding for a specific comment."""
issue = create_issue()
comment = create_comment(
issue_id=str(issue.id), external_id="901", body="This is a test comment."
)
model = create_embedding_model()
# Mock the embedding vector generation
mock_generate_embedding = mocker.patch(
"preloop.models.crud.crud_issue_embedding._generate_embedding_vector"
)
mock_generate_embedding.return_value = [0.1] * model.dimensions
# Create embedding for the comment
result = crud_issue_embedding.create_embeddings(
db_session, issue_id=issue.id, comment_id=comment.id, force_update=False
)
assert model.name in result
assert (
f"created_for_comment {comment.id}" in result[model.name]
or "exists" in result[model.name]
)
# Verify embedding was created in the database for the comment
comment_embeddings = crud_issue_embedding.get_for_comment(
db_session, comment_id=comment.id
)
assert model.name in comment_embeddings
embedding = comment_embeddings[model.name]
assert embedding.comment_id == comment.id
assert embedding.issue_id == issue.id
assert len(embedding.embedding) == model.dimensions
# Verify it's not mistaken for an issue content embedding
issue_content_embeddings = crud_issue_embedding.get_for_issue_content(
db_session, issue_id=issue.id
)
assert model.name not in issue_content_embeddings
def test_get_issue_content_embedding_distinct_from_comment(
db_session, create_issue, create_comment, create_embedding_model, mocker
):
"""Test that issue content embeddings are distinct from comment embeddings."""
issue = create_issue()
comment = create_comment(
issue_id=str(issue.id), external_id="901", body="Another test comment."
)
model = create_embedding_model()
# Mock the embedding vector generation
mock_generate_embedding = mocker.patch(
"preloop.models.crud.crud_issue_embedding._generate_embedding_vector"
)
mock_generate_embedding.return_value = [0.1] * model.dimensions
# Create embedding for issue content
crud_issue_embedding.create_embeddings(db_session, issue_id=issue.id)
# Create embedding for comment
crud_issue_embedding.create_embeddings(
db_session, issue_id=issue.id, comment_id=comment.id
)
# Get issue content embeddings
issue_embeddings = crud_issue_embedding.get_for_issue_content(
db_session, issue_id=issue.id
)
assert model.name in issue_embeddings
assert issue_embeddings[model.name].comment_id is None
# Get comment embeddings
comment_embeddings = crud_issue_embedding.get_for_comment(
db_session, comment_id=comment.id
)
assert model.name in comment_embeddings
assert comment_embeddings[model.name].comment_id == comment.id
def test_delete_comment_cascades_embeddings(
db_session, create_issue, create_comment, create_embedding_model, mocker
):
"""Test that deleting a comment also deletes its associated embeddings."""
issue = create_issue()
comment = create_comment(
issue_id=str(issue.id), external_id="901", body="Comment to be deleted."
)
model = create_embedding_model()
# Mock the private method that generates embeddings
mock_generate_embedding = mocker.patch(
"preloop.models.crud.crud_issue_embedding._generate_embedding_vector"
)
mock_generate_embedding.return_value = [0.1] * model.dimensions
# Create embedding for the comment
crud_issue_embedding.create_embeddings(
db_session, issue_id=issue.id, comment_id=comment.id
)
comment_embeddings_before_delete = crud_issue_embedding.get_for_comment(
db_session, comment_id=comment.id
)
assert model.name in comment_embeddings_before_delete
embedding_id = comment_embeddings_before_delete[model.name].id
# Delete the comment
crud_comment.delete(db_session, id=comment.id)
# Verify the comment is deleted
deleted_comment = db_session.get(Comment, comment.id)
assert deleted_comment is None
# Verify the embedding associated with the comment is also deleted
retrieved_embedding = db_session.get(IssueEmbedding, embedding_id)
assert retrieved_embedding is None
# Verify using the CRUD method too
comment_embeddings_after_delete = crud_issue_embedding.get_for_comment(
db_session, comment_id=comment.id
)
assert model.name not in comment_embeddings_after_delete
def test_similarity_search(
db_session, create_issue, create_embedding_model, create_tracker
):
"""Test similarity search functionality."""
# Create embedding model (defaulting to 1536 dimensions from conftest fixture)
model = create_embedding_model()
# Create a single tracker instance using the injected factory fixture
# The `create_tracker` argument in the function signature provides the factory function.
tracker = create_tracker() # Call the factory to create the tracker
# Create test issues using the *same* tracker instance
issue1 = create_issue(title="Bug in login feature", tracker=tracker)
issue2 = create_issue(title="Feature request: dashboard", tracker=tracker)
issue3 = create_issue(title="Login page crashes on mobile", tracker=tracker)
# Create fixed embeddings for testing similarity
default_dimension = 1536 # Should match the default in EmbeddingModel or conftest
embeddings_data = [
([1.0, 0.0, 0.0, 0.0], default_dimension), # issue1
([0.0, 1.0, 0.0, 0.0], default_dimension), # issue2
([0.8, 0.2, 0.0, 0.0], default_dimension), # issue3 (more similar to issue1)
]
embeddings = []
for vec, dim in embeddings_data:
embeddings.append(vec + [0.0] * (dim - len(vec)))
# Manually create embeddings
for issue, emb in zip([issue1, issue2, issue3], embeddings, strict=False):
embedding = IssueEmbedding(
issue_id=issue.id,
embedding_model_id=model.id,
embedding=emb,
meta_data={"test": True},
)
db_session.add(embedding)
db_session.commit()
# Search with a vector similar to issue1
query_vector_short = [0.9, 0.1, 0.0, 0.0]
query_vector = query_vector_short + [0.0] * (
default_dimension - len(query_vector_short)
)
# Get the tracker_id from one of the created issues
tracker_id = tracker.id # Use the id from the shared tracker created above
results = crud_issue_embedding.similarity_search(
db_session,
model_id=model.id,
query_vector=query_vector,
distance_type="cosine",
tracker_ids=[tracker_id], # Pass the tracker_id
)
# Check we got results in the right order (issue1, issue3, issue2)
assert len(results) == 3
assert results[0][0].id == issue1.id # Most similar
assert results[1][0].id == issue3.id # Second most similar
assert results[2][0].id == issue2.id # Least similar
def test_similarity_search_with_embedding_type(
db_session: Session,
create_issue,
create_embedding_model,
create_tracker,
create_comment,
):
"""Test similarity search with embedding_type filtering."""
# Use default 1536 dimensions for the model
model = create_embedding_model(name="test_model_type_filter")
default_dimension = (
1536 # model.dimensions would also work if fixture guarantees it
)
tracker = create_tracker(name="test_tracker_type_filter")
issue1 = create_issue(title="Issue One For Type Filter", tracker_id=tracker.id)
issue2 = create_issue(title="Issue Two For Type Filter", tracker_id=tracker.id)
# create_comment should return the Comment object and handle db_session internally or take it as arg
comment1 = create_comment(
issue_id=str(issue1.id), external_id="901", body="A comment on Issue One"
)
embedding_vector_issue1_short = [1.0, 0.0, 0.0]
embedding_vector_issue2_short = [0.0, 1.0, 0.0]
embedding_vector_comment1_short = [
0.0,
0.0,
1.0,
] # Distinct embedding for the comment
embedding_vector_issue1 = embedding_vector_issue1_short + [0.0] * (
default_dimension - len(embedding_vector_issue1_short)
)
embedding_vector_issue2 = embedding_vector_issue2_short + [0.0] * (
default_dimension - len(embedding_vector_issue2_short)
)
embedding_vector_comment1 = embedding_vector_comment1_short + [0.0] * (
default_dimension - len(embedding_vector_comment1_short)
)
# Create embeddings directly for precise testing
embedding_for_issue1 = IssueEmbedding(
issue_id=issue1.id,
embedding_model_id=model.id,
embedding=embedding_vector_issue1,
comment_id=None, # Explicitly for an issue embedding
)
embedding_for_issue2 = IssueEmbedding(
issue_id=issue2.id,
embedding_model_id=model.id,
embedding=embedding_vector_issue2,
comment_id=None, # Explicitly for an issue embedding
)
embedding_for_comment1 = IssueEmbedding(
issue_id=comment1.issue_id, # Link to the parent issue of the comment
comment_id=comment1.id,
embedding_model_id=model.id,
embedding=embedding_vector_comment1,
)
db_session.add_all(
[embedding_for_issue1, embedding_for_issue2, embedding_for_comment1]
)
db_session.commit()
# Test with embedding_type="issue"
results_issue_type_query_issue1_vec = crud_issue_embedding.similarity_search(
db_session,
model_id=model.id,
query_vector=embedding_vector_issue1,
embedding_type="issue",
limit=5,
)
assert len(results_issue_type_query_issue1_vec) == 2
assert results_issue_type_query_issue1_vec[0][0].id == issue1.id
assert isinstance(results_issue_type_query_issue1_vec[0][0], Issue)
assert isinstance(results_issue_type_query_issue1_vec[1][0], Issue)
results_issue_type_query_comment1_vec = crud_issue_embedding.similarity_search(
db_session,
model_id=model.id,
query_vector=embedding_vector_comment1,
embedding_type="issue",
limit=5,
)
assert len(results_issue_type_query_comment1_vec) == 2
assert isinstance(results_issue_type_query_comment1_vec[0][0], Issue)
assert isinstance(results_issue_type_query_comment1_vec[1][0], Issue)
# Test with embedding_type="comment"
results_comment_type_query_comment1_vec = crud_issue_embedding.similarity_search(
db_session,
model_id=model.id,
query_vector=embedding_vector_comment1,
embedding_type="comment",
limit=5,
)
assert len(results_comment_type_query_comment1_vec) == 1
assert (
results_comment_type_query_comment1_vec[0][0].id == comment1.id
) # Should return the Comment object itself
assert isinstance(results_comment_type_query_comment1_vec[0][0], Comment)
results_comment_type_query_issue1_vec = crud_issue_embedding.similarity_search(
db_session,
model_id=model.id,
query_vector=embedding_vector_issue1,
embedding_type="comment",
limit=5,
)
assert len(results_comment_type_query_issue1_vec) == 1
assert (
results_comment_type_query_issue1_vec[0][0].id == comment1.id
) # The only comment available
assert isinstance(results_comment_type_query_issue1_vec[0][0], Comment)
# Test with embedding_type=None (mixed results), querying with issue1's vector
results_none_type_query_issue1_vec = crud_issue_embedding.similarity_search(
db_session,
model_id=model.id,
query_vector=embedding_vector_issue1,
embedding_type=None,
limit=5,
)
assert len(results_none_type_query_issue1_vec) == 3
assert results_none_type_query_issue1_vec[0][0].id == issue1.id
assert isinstance(results_none_type_query_issue1_vec[0][0], Issue)
returned_ids_and_types_issue_query = {
(item.id, type(item)) for item, score in results_none_type_query_issue1_vec
}
expected_ids_and_types_issue_query = {
(issue1.id, Issue),
(issue2.id, Issue),
(comment1.id, Comment),
}
assert returned_ids_and_types_issue_query == expected_ids_and_types_issue_query
# Test with embedding_type=None (mixed results), querying with comment1's vector
results_none_type_query_comment1_vec = crud_issue_embedding.similarity_search(
db_session,
model_id=model.id,
query_vector=embedding_vector_comment1,
embedding_type=None,
limit=5,
)
assert len(results_none_type_query_comment1_vec) == 3
assert results_none_type_query_comment1_vec[0][0].id == comment1.id
assert isinstance(results_none_type_query_comment1_vec[0][0], Comment)
returned_ids_and_types_comment_query = {
(item.id, type(item)) for item, score in results_none_type_query_comment1_vec
}
expected_ids_and_types_comment_query = {
(comment1.id, Comment),
(issue1.id, Issue),
(issue2.id, Issue),
}
assert returned_ids_and_types_comment_query == expected_ids_and_types_comment_query
# Test with an invalid embedding_type (should raise ValueError)
with pytest.raises(
ValueError, match="Unsupported embedding_type: invalid_type_string"
):
crud_issue_embedding.similarity_search(
db_session,
model_id=model.id,
query_vector=embedding_vector_issue1,
embedding_type="invalid_type_string",
limit=5,
)
with pytest.raises(ValueError, match="Unsupported embedding_type: another_invalid"):
crud_issue_embedding.similarity_search(
db_session,
model_id=model.id,
query_vector=embedding_vector_comment1,
embedding_type="another_invalid",
limit=5,
)
