Skip to content

Commit 4edb682

Browse files
authored
Merge pull request #54 from whyhow-ai/37-bug-graceful-exit-error-message-for-wrong-or-incomplete-qdrant-env-settings
37 bug graceful exit error message for wrong or incomplete qdrant env settings
2 parents 4ecfca9 + f4984fb commit 4edb682

File tree

7 files changed

+202
-16
lines changed

7 files changed

+202
-16
lines changed

backend/CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1111

1212
- Added support for queries without source data in vector database
1313
- Graceful failure of triple export when no chunks are found
14+
- Tested Qdrant vector database service
1415

1516
### Changed
1617

backend/src/app/schemas/query_api.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from typing import Any, List, Optional, Union
44

5-
from pydantic import BaseModel
5+
from pydantic import BaseModel, ConfigDict
66

77
from app.models.query_core import Chunk, FormatType, Rule
88

@@ -23,10 +23,7 @@ class QueryRequestSchema(BaseModel):
2323
document_id: str
2424
prompt: QueryPromptSchema
2525

26-
class Config:
27-
"""Pydantic configuration."""
28-
29-
extra = "allow"
26+
model_config = ConfigDict(extra="allow")
3027

3128

3229
class VectorResponseSchema(BaseModel):

backend/src/app/services/vector_db/base.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,15 +82,20 @@ async def ensure_collection_exists(self) -> None:
8282
async def get_embeddings(
8383
self, texts: Union[str, List[str]]
8484
) -> List[List[float]]:
85-
"""Get embeddings for the given text(s) using the LLM service."""
85+
"""Get embeddings for the given text(s) using the embedding service."""
8686
if isinstance(texts, str):
8787
texts = [texts]
8888
return await self.embedding_service.get_embeddings(texts)
8989

90+
async def get_single_embedding(self, text: str) -> List[float]:
91+
"""Get a single embedding for the given text."""
92+
embeddings = await self.get_embeddings(text)
93+
return embeddings[0]
94+
9095
async def prepare_chunks(
9196
self, document_id: str, chunks: List[Document]
9297
) -> List[Dict[str, Any]]:
93-
"""Prepare chunks for insertion into the Milvus database."""
98+
"""Prepare chunks for insertion into the vector database."""
9499
logger.info(f"Preparing {len(chunks)} chunks")
95100

96101
# Clean the chunks

backend/src/app/services/vector_db/milvus_service.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -155,9 +155,8 @@ async def vector_search(
155155
# Search for each query
156156
for query in queries:
157157
logger.info("Generating embedding.")
158-
159-
# Embed the query
160-
embedded_query = await self.get_embeddings(query)
158+
# Use get_single_embedding but wrap result in list for Milvus
159+
embedded_query = [await self.get_single_embedding(query)]
161160

162161
logger.info("Searching...")
163162

@@ -323,7 +322,7 @@ def count_keywords(text: str, keywords: List[str]) -> int:
323322
)
324323

325324
# Embed the query
326-
embedded_query = await self.get_embeddings(query)
325+
embedded_query = [await self.get_single_embedding(query)]
327326

328327
try:
329328
# First, let's check if there are any vectors for this document_id

backend/src/app/services/vector_db/qdrant_service.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ async def vector_search(
7575

7676
for query in queries:
7777
logger.info("Generating embedding.")
78-
embedded_query = await self.get_embeddings(query)
78+
embedded_query = await self.get_single_embedding(query)
7979
logger.info("Searching...")
8080

8181
query_response = self.client.query_points(
@@ -162,7 +162,7 @@ def count_keywords(text: str, keywords: List[str]) -> int:
162162
reverse=True,
163163
)
164164

165-
embedded_query = await self.get_embeddings(query)
165+
embedded_query = await self.get_single_embedding(query)
166166
logger.info("Running semantic similarity search.")
167167

168168
semantic_response = self.client.query_points(
@@ -194,8 +194,6 @@ def count_keywords(text: str, keywords: List[str]) -> int:
194194
combined_chunks, key=lambda chunk: chunk["chunk_number"]
195195
)
196196

197-
# Optionally, for each chunk, retrieve neighbouring chunks to ensure full context is retrieved
198-
199197
# Eliminate duplicate chunks
200198
seen_chunks = set()
201199
formatted_output = []

backend/tests/test_service_vector_db_milvus.py

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ def __init__(self, embedding_service, llm_service, settings):
1313
self.embedding_service = embedding_service
1414
self.llm_service = llm_service
1515
self.settings = settings
16-
self.client = Mock() # Use regular Mock instead of AsyncMock
16+
self.client = Mock()
1717

1818
# Set up synchronous return values
1919
self.client.has_collection.return_value = True
@@ -36,12 +36,17 @@ async def upsert_vectors(self, vectors):
3636
}
3737

3838
async def vector_search(self, queries, document_id):
39+
# Mock using get_single_embedding
40+
for query in queries:
41+
_ = await self.get_single_embedding(query)
3942
return VectorResponseSchema(message="success", chunks=[])
4043

4144
async def keyword_search(self, query, document_id, keywords):
4245
return VectorResponseSchema(message="success", chunks=[])
4346

4447
async def hybrid_search(self, query, document_id, rules):
48+
# Mock using get_single_embedding
49+
_ = await self.get_single_embedding(query)
4550
return VectorResponseSchema(
4651
message="Query processed successfully.", chunks=[]
4752
)
@@ -107,3 +112,27 @@ async def test_delete_document(vector_db_service):
107112

108113
assert result["status"] == "success"
109114
assert result["message"] == "Document deleted successfully."
115+
116+
117+
@pytest.mark.asyncio
118+
async def test_get_single_embedding(vector_db_service):
119+
# Reset the mock before the test
120+
vector_db_service.embedding_service.get_embeddings.reset_mock()
121+
122+
# Mock the embedding service to return a known value
123+
vector_db_service.embedding_service.get_embeddings.return_value = [
124+
[0.1, 0.2, 0.3]
125+
]
126+
127+
# Test getting a single embedding
128+
result = await vector_db_service.get_single_embedding("test text")
129+
130+
# Verify the result
131+
assert isinstance(result, list)
132+
assert len(result) == 3 # Length of our mock embedding
133+
assert result == [0.1, 0.2, 0.3]
134+
135+
# Verify the embedding service was called correctly
136+
vector_db_service.embedding_service.get_embeddings.assert_called_once_with(
137+
["test text"]
138+
)
Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
from unittest.mock import Mock, patch
2+
3+
import pytest
4+
5+
from app.schemas.query_api import VectorResponseSchema
6+
from app.services.vector_db.qdrant_service import QdrantService
7+
8+
9+
@pytest.fixture
10+
def mock_qdrant_client():
11+
with patch("app.services.vector_db.qdrant_service.QdrantClient") as mock:
12+
client = Mock()
13+
# Set up mock responses
14+
client.collection_exists.return_value = True
15+
client.upsert.return_value = None
16+
17+
# Use a simple Mock instead of Qdrant models
18+
response_mock = Mock()
19+
response_mock.points = [
20+
Mock(
21+
payload={
22+
"text": "test text",
23+
"page_number": 1,
24+
"chunk_number": 1,
25+
"document_id": "test_doc",
26+
}
27+
)
28+
]
29+
client.query_points.return_value = response_mock
30+
31+
client.delete.return_value = None
32+
mock.return_value = client
33+
yield client
34+
35+
36+
@pytest.fixture
37+
def qdrant_service(
38+
mock_embeddings_service,
39+
mock_llm_service,
40+
test_settings,
41+
mock_qdrant_client,
42+
):
43+
service = QdrantService(
44+
embedding_service=mock_embeddings_service,
45+
llm_service=mock_llm_service,
46+
settings=test_settings,
47+
)
48+
# Override the client with our mock
49+
service.client = mock_qdrant_client
50+
return service
51+
52+
53+
@pytest.mark.asyncio
54+
async def test_ensure_collection_exists(qdrant_service):
55+
await qdrant_service.ensure_collection_exists()
56+
assert qdrant_service.client.collection_exists.called
57+
58+
59+
@pytest.mark.asyncio
60+
async def test_upsert_vectors(qdrant_service):
61+
vectors = [
62+
{
63+
"id": "1",
64+
"vector": [0.1, 0.2],
65+
"text": "test",
66+
"page_number": 1,
67+
"chunk_number": 1,
68+
"document_id": "doc1",
69+
}
70+
]
71+
72+
result = await qdrant_service.upsert_vectors(vectors)
73+
74+
assert "message" in result
75+
assert qdrant_service.client.upsert.called
76+
77+
78+
@pytest.mark.asyncio
79+
async def test_vector_search(qdrant_service, mock_embeddings_service):
80+
mock_embeddings_service.get_embeddings.return_value = [[0.1, 0.2]]
81+
82+
result = await qdrant_service.vector_search(["test query"], "test_doc")
83+
84+
assert isinstance(result, VectorResponseSchema)
85+
assert result.message == "Query processed successfully."
86+
assert qdrant_service.client.query_points.called
87+
88+
89+
@pytest.mark.asyncio
90+
async def test_hybrid_search(qdrant_service, mock_embeddings_service):
91+
mock_embeddings_service.get_embeddings.return_value = [[0.1, 0.2]]
92+
93+
with patch.object(
94+
qdrant_service,
95+
"extract_keywords",
96+
return_value=["keyword1", "keyword2"],
97+
):
98+
result = await qdrant_service.hybrid_search(
99+
"test query", "test_doc", []
100+
)
101+
102+
assert isinstance(result, VectorResponseSchema)
103+
assert result.message == "Query processed successfully."
104+
assert qdrant_service.client.query_points.called
105+
106+
107+
@pytest.mark.asyncio
108+
async def test_decomposed_search(qdrant_service, mock_llm_service):
109+
mock_llm_service.decompose_query.return_value = {
110+
"sub-queries": ["query1", "query2"]
111+
}
112+
113+
result = await qdrant_service.decomposed_search(
114+
"test query", "test_doc", []
115+
)
116+
117+
assert "sub_queries" in result
118+
assert "chunks" in result
119+
120+
121+
@pytest.mark.asyncio
122+
async def test_delete_document(qdrant_service):
123+
result = await qdrant_service.delete_document("test_doc")
124+
125+
assert result["status"] == "success"
126+
assert result["message"] == "Document deleted successfully."
127+
assert qdrant_service.client.delete.called
128+
129+
130+
@pytest.mark.asyncio
131+
async def test_keyword_search_not_implemented(qdrant_service):
132+
with pytest.raises(NotImplementedError):
133+
await qdrant_service.keyword_search("query", "doc_id", ["keyword"])
134+
135+
136+
@pytest.mark.asyncio
137+
async def test_get_single_embedding(qdrant_service):
138+
# Reset the mock before the test
139+
qdrant_service.embedding_service.get_embeddings.reset_mock()
140+
141+
# Mock the embedding service to return a known value
142+
qdrant_service.embedding_service.get_embeddings.return_value = [
143+
[0.1, 0.2, 0.3]
144+
]
145+
146+
# Test getting a single embedding
147+
result = await qdrant_service.get_single_embedding("test text")
148+
149+
# Verify the result
150+
assert isinstance(result, list)
151+
assert len(result) == 3
152+
assert result == [0.1, 0.2, 0.3]
153+
154+
# Verify the embedding service was called correctly
155+
qdrant_service.embedding_service.get_embeddings.assert_called_once_with(
156+
["test text"]
157+
)

0 commit comments

Comments
 (0)