apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-09 13:22:10 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Fix embeddings integration

Browse source
This commit is contained in:
Cyber MacGeddon 2026-03-08 19:38:08 +00:00
parent ea95effbf2
commit 22ba73e680
6 changed files with 31 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

10
tests/integration/test_document_rag_integration.py
View file

@ -27,9 +27,13 @@ class TestDocumentRagIntegration:
def mock_embeddings_client(self):
"""Mock embeddings client that returns realistic vector embeddings"""
client = AsyncMock()
# New batch format: [[[vectors_for_text1], ...]]
# One text input returns one vector set containing two vectors
client.embed.return_value = [
[0.1, 0.2, 0.3, 0.4, 0.5], # Realistic 5-dimensional embedding
[0.6, 0.7, 0.8, 0.9, 1.0] # Second embedding for testing
[
[0.1, 0.2, 0.3, 0.4, 0.5], # First vector for text
[0.6, 0.7, 0.8, 0.9, 1.0] # Second vector for text
]
]
return client
@ -90,7 +94,7 @@ class TestDocumentRagIntegration:
)
# Assert - Verify service coordination
mock_embeddings_client.embed.assert_called_once_with(query)
mock_embeddings_client.embed.assert_called_once_with([query])
mock_doc_embeddings_client.query.assert_called_once_with(
[[0.1, 0.2, 0.3, 0.4, 0.5], [0.6, 0.7, 0.8, 0.9, 1.0]],

3
tests/integration/test_document_rag_streaming_integration.py
View file

@ -30,7 +30,8 @@ class TestDocumentRagStreaming:
def mock_embeddings_client(self):
"""Mock embeddings client"""
client = AsyncMock()
client.embed.return_value = [[0.1, 0.2, 0.3, 0.4, 0.5]]
# New batch format: [[[vectors_for_text1]]]
client.embed.return_value = [[[0.1, 0.2, 0.3, 0.4, 0.5]]]
return client
@pytest.fixture

10
tests/integration/test_graph_rag_integration.py
View file

@ -21,8 +21,12 @@ class TestGraphRagIntegration:
def mock_embeddings_client(self):
"""Mock embeddings client that returns realistic vector embeddings"""
client = AsyncMock()
# New batch format: [[[vectors_for_text1], ...]]
# One text input returns one vector set containing one vector
client.embed.return_value = [
[0.1, 0.2, 0.3, 0.4, 0.5], # Realistic 5-dimensional embedding
[
[0.1, 0.2, 0.3, 0.4, 0.5], # Vector for text
]
]
return client
@ -120,8 +124,8 @@ class TestGraphRagIntegration:
# Assert - Verify service coordination
# 1. Should compute embeddings for query
mock_embeddings_client.embed.assert_called_once_with(query)
# 1. Should compute embeddings for query (now expects list of texts)
mock_embeddings_client.embed.assert_called_once_with([query])
# 2. Should query graph embeddings to find relevant entities
mock_graph_embeddings_client.query.assert_called_once()

3
tests/integration/test_graph_rag_streaming_integration.py
View file

@ -24,7 +24,8 @@ class TestGraphRagStreaming:
def mock_embeddings_client(self):
"""Mock embeddings client"""
client = AsyncMock()
client.embed.return_value = [[0.1, 0.2, 0.3, 0.4, 0.5]]
# New batch format: [[[vectors_for_text1]]]
client.embed.return_value = [[[0.1, 0.2, 0.3, 0.4, 0.5]]]
return client
@pytest.fixture

4
tests/integration/test_rag_streaming_protocol.py
View file

@ -18,7 +18,7 @@ class TestGraphRagStreamingProtocol:
def mock_embeddings_client(self):
"""Mock embeddings client"""
client = AsyncMock()
client.embed.return_value = [[0.1, 0.2, 0.3]]
client.embed.return_value = [[[0.1, 0.2, 0.3]]]
return client
@pytest.fixture
@ -197,7 +197,7 @@ class TestDocumentRagStreamingProtocol:
def mock_embeddings_client(self):
"""Mock embeddings client"""
client = AsyncMock()
client.embed.return_value = [[0.1, 0.2, 0.3]]
client.embed.return_value = [[[0.1, 0.2, 0.3]]]
return client
@pytest.fixture

20
tests/unit/test_retrieval/test_document_rag.py
View file

@ -132,9 +132,10 @@ class TestQuery:
mock_embeddings_client = AsyncMock()
mock_rag.embeddings_client = mock_embeddings_client
# Mock the embed method to return test vectors
# Mock the embed method to return test vectors in batch format
# New format: [[[vectors_for_text1]]] - returns first text's vector set
expected_vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
mock_embeddings_client.embed.return_value = expected_vectors
mock_embeddings_client.embed.return_value = [expected_vectors]
# Initialize Query
query = Query(
@ -148,10 +149,10 @@ class TestQuery:
test_query = "What documents are relevant?"
result = await query.get_vector(test_query)
# Verify embeddings client was called correctly
mock_embeddings_client.embed.assert_called_once_with(test_query)
# Verify embeddings client was called correctly (now expects list)
mock_embeddings_client.embed.assert_called_once_with([test_query])
# Verify result matches expected vectors
# Verify result matches expected vectors (extracted from batch)
assert result == expected_vectors
@pytest.mark.asyncio
@ -170,8 +171,9 @@ class TestQuery:
mock_rag.fetch_chunk = mock_fetch
# Mock the embedding and document query responses
# New batch format: [[[vectors]]] - get_vector extracts [0]
test_vectors = [[0.1, 0.2, 0.3]]
mock_embeddings_client.embed.return_value = test_vectors
mock_embeddings_client.embed.return_value = [test_vectors]
# Mock document embeddings returns chunk_ids
test_chunk_ids = ["doc/c1", "doc/c2"]
@ -190,10 +192,10 @@ class TestQuery:
test_query = "Find relevant documents"
result = await query.get_docs(test_query)
# Verify embeddings client was called
mock_embeddings_client.embed.assert_called_once_with(test_query)
# Verify embeddings client was called (now expects list)
mock_embeddings_client.embed.assert_called_once_with([test_query])
# Verify doc embeddings client was called correctly
# Verify doc embeddings client was called correctly (with extracted vectors)
mock_doc_embeddings_client.query.assert_called_once_with(
test_vectors,
limit=15,