Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit cc5fe19

Browse files
add tests and fix db_session issue
1 parent 29692ee commit cc5fe19

File tree

3 files changed

+127
-0
lines changed

3 files changed

+127
-0
lines changed

‎src/backend/fastapi_app/rag_advanced.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,11 @@ async def run_stream(
183183
stream=True,
184184
)
185185

186+
# Forcefully Close the database session before yielding the response
187+
# Yielding keeps the connection open while streaming the response until the end
188+
# The connection closes when it returns back to the context manger in the dependencies
189+
await self.searcher.db_session.close()
190+
186191
yield RetrievalResponse(
187192
message=Message(content="", role="assistant"),
188193
context=RAGContext(

‎src/backend/fastapi_app/rag_simple.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,11 @@ async def run_stream(
201201
stream=True,
202202
)
203203

204+
# Forcefully Close the database session before yielding the response
205+
# Yielding keeps the connection open while streaming the response until the end
206+
# The connection closes when it returns back to the context manger in the dependencies
207+
await self.searcher.db_session.close()
208+
204209
yield RetrievalResponse(
205210
message=Message(content="", role="assistant"),
206211
context=RAGContext(

‎tests/test_api_routes.py

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,123 @@ async def test_simple_chat_flow(test_client):
227227
assert response_data["session_state"] is None
228228

229229

230+
@pytest.mark.asyncio
231+
async def test_simple_chat_streaming_flow(test_client):
232+
"""test the simple chat streaming flow route with hybrid retrieval mode"""
233+
response = test_client.post(
234+
"/chat/stream",
235+
json={
236+
"context": {
237+
"overrides": {"top": 1, "use_advanced_flow": False, "retrieval_mode": "hybrid", "temperature": 0.3}
238+
},
239+
"messages": [{"content": "What is the capital of France?", "role": "user"}],
240+
},
241+
)
242+
response_data = response.content.split(b"\n")
243+
assert response.status_code == 200
244+
assert response.headers["Content-Type"] == "application/x-ndjson"
245+
assert response_data[0] == (
246+
b'{"message": {"content": "", "role": "assistant"}, "context": {"data_points":'
247+
+ b' {"1": {"id": 1, "type": "Footwear", "brand": "Daybird", "name": "Wanderer B'
248+
+ b'lack Hiking Boots", "description": "Daybird\'s Wanderer Hiking Boots in s'
249+
+ b"leek black are perfect for all your outdoor adventures. These boots are made"
250+
+ b" with a waterproof leather upper and a durable rubber sole for superior trac"
251+
+ b"tion. With their cushioned insole and padded collar, these boots will keep y"
252+
+ b'ou comfortable all day long.", "price": 109.99}}, "thoughts": [{"title": "Se'
253+
+ b'arch query for database", "description": "What is the capital of France?", "'
254+
+ b'props": {"top": 1, "vector_search": true, "text_search": true}}, {"title": "'
255+
+ b'Search results", "description": [{"id": 1, "type": "Footwear", "brand": "Day'
256+
+ b'bird", "name": "Wanderer Black Hiking Boots", "description": "Daybird\'s '
257+
+ b"Wanderer Hiking Boots in sleek black are perfect for all your outdoor advent"
258+
+ b"ures. These boots are made with a waterproof leather upper and a durable rub"
259+
+ b"ber sole for superior traction. With their cushioned insole and padded colla"
260+
+ b'r, these boots will keep you comfortable all day long.", "price": 109.99}], '
261+
+ b'"props": {}}, {"title": "Prompt to generate answer", "description": ["{\''
262+
+ b"role': 'system', 'content': \\\"Assistant helps customers with questio"
263+
+ b"ns about products.\\\\nRespond as if you are a salesperson helping a custo"
264+
+ b"mer in a store. Do NOT respond with tables.\\\\nAnswer ONLY with the produ"
265+
+ b"ct details listed in the products.\\\\nIf there isn't enough information b"
266+
+ b"elow, say you don't know.\\\\nDo not generate answers that don't use the s"
267+
+ b"ources below.\\\\nEach product has an ID in brackets followed by colon and"
268+
+ b" the product details.\\\\nAlways include the product ID for each product y"
269+
+ b"ou use in the response.\\\\nUse square brackets to reference the source, f"
270+
+ b"or example [52].\\\\nDon't combine citations, list each product separately"
271+
+ b", for example [27][51].\\\"}\", \"{'role': 'user', 'content': \\\"What is "
272+
+ b"the capital of France?\\\\n\\\\nSources:\\\\n[1]:Name:Wanderer Black Hikin"
273+
+ b"g Boots Description:Daybird's Wanderer Hiking Boots in sleek black are perfe"
274+
+ b"ct for all your outdoor adventures. These boots are made with a waterproof l"
275+
+ b"eather upper and a durable rubber sole for superior traction. With their cus"
276+
+ b"hioned insole and padded collar, these boots will keep you comfortable all d"
277+
+ b'ay long. Price:109.99 Brand:Daybird Type:Footwear\\\\n\\\\n\\"}"], "props'
278+
+ b'": {"model": "gpt-35-turbo", "deployment": "gpt-35-turbo"}}], "followup_ques'
279+
+ b'tions": null}, "session_state": null}'
280+
)
281+
282+
283+
@pytest.mark.asyncio
284+
async def test_advanved_chat_streaming_flow(test_client):
285+
"""test the advanced chat streaming flow route with hybrid retrieval mode"""
286+
response = test_client.post(
287+
"/chat/stream",
288+
json={
289+
"context": {
290+
"overrides": {"top": 1, "use_advanced_flow": True, "retrieval_mode": "hybrid", "temperature": 0.3}
291+
},
292+
"messages": [{"content": "What is the capital of France?", "role": "user"}],
293+
},
294+
)
295+
response_data = response.content.split(b"\n")
296+
assert response.status_code == 200
297+
assert response.headers["Content-Type"] == "application/x-ndjson"
298+
assert response_data[0] == (
299+
b'{"message": {"content": "", "role": "assistant"}, "context": {"data_points":'
300+
+ b' {"1": {"id": 1, "type": "Footwear", "brand": "Daybird", "name": "Wanderer B'
301+
+ b'lack Hiking Boots", "description": "Daybird\'s Wanderer Hiking Boots in s'
302+
+ b'leek black are perfect for all your outdoor adventures. These boots are made'
303+
+ b' with a waterproof leather upper and a durable rubber sole for superior trac'
304+
+ b'tion. With their cushioned insole and padded collar, these boots will keep y'
305+
+ b'ou comfortable all day long.", "price": 109.99}}, "thoughts": [{"title": "Pr'
306+
+ b'ompt to generate search arguments", "description": ["{\'role\': \'system\', '
307+
+ b"'content': 'Below is a history of the conversation so far, and a new questio"
308+
+ b'n asked by the user that needs to be answered by searching database rows'
309+
+ b'.\\\\nYou have access to an Azure PostgreSQL database with an items table '
310+
+ b'that has columns for title, description, brand, price, and type.\\\\nGener'
311+
+ b'ate a search query based on the conversation and the new question.\\\\nIf '
312+
+ b'the question is not in English, translate the question to English before gen'
313+
+ b'erating the search query.\\\\nIf you cannot generate a search query, retur'
314+
+ b'n the original user question.\\\\nDO NOT return anything besides the query'
315+
+ b'.\'}", "{\'role\': \'user\', \'content\': \'What is the capital of Franc'
316+
+ b'e?\'}"], "props": {"model": "gpt-35-turbo", "deployment": "gpt-35-turbo"}'
317+
+ b'}, {"title": "Search using generated search arguments", "description": "The '
318+
+ b'capital of France is Paris. [Benefit_Options-2.pdf].", "props": {"top": 1, "'
319+
+ b'vector_search": true, "text_search": true, "filters": []}}, {"title": "Searc'
320+
+ b'h results", "description": [{"id": 1, "type": "Footwear", "brand": "Daybird"'
321+
+ b', "name": "Wanderer Black Hiking Boots", "description": "Daybird\'s Wande'
322+
+ b'rer Hiking Boots in sleek black are perfect for all your outdoor adventures.'
323+
+ b' These boots are made with a waterproof leather upper and a durable rubber s'
324+
+ b'ole for superior traction. With their cushioned insole and padded collar, th'
325+
+ b'ese boots will keep you comfortable all day long.", "price": 109.99}], "prop'
326+
+ b's": {}}, {"title": "Prompt to generate answer", "description": ["{\'role\''
327+
+ b': \'system\', \'content\': \\"Assistant helps customers with questions ab'
328+
+ b'out products.\\\\nRespond as if you are a salesperson helping a customer i'
329+
+ b'n a store. Do NOT respond with tables.\\\\nAnswer ONLY with the product de'
330+
+ b"tails listed in the products.\\\\nIf there isn't enough information below,"
331+
+ b" say you don't know.\\\\nDo not generate answers that don't use the source"
332+
+ b's below.\\\\nEach product has an ID in brackets followed by colon and the '
333+
+ b'product details.\\\\nAlways include the product ID for each product you us'
334+
+ b'e in the response.\\\\nUse square brackets to reference the source, for ex'
335+
+ b"ample [52].\\\\nDon't combine citations, list each product separately, for"
336+
+ b' example [27][51].\\"}", "{\'role\': \'user\', \'content\': \\"What is the c'
337+
+ b'apital of France?\\\\n\\\\nSources:\\\\n[1]:Name:Wanderer Black Hiking Boo'
338+
+ b"ts Description:Daybird's Wanderer Hiking Boots in sleek black are perfect fo"
339+
+ b'r all your outdoor adventures. These boots are made with a waterproof leathe'
340+
+ b'r upper and a durable rubber sole for superior traction. With their cushione'
341+
+ b'd insole and padded collar, these boots will keep you comfortable all day lo'
342+
+ b'ng. Price:109.99 Brand:Daybird Type:Footwear\\\\n\\\\n\\"}"], "props": {"'
343+
+ b'model": "gpt-35-turbo", "deployment": "gpt-35-turbo"}}], "followup_questions'
344+
+ b'": null}, "session_state": null}'
345+
)
346+
230347
@pytest.mark.asyncio
231348
async def test_advanced_chat_flow(test_client):
232349
"""test the advanced chat flow route with hybrid retrieval mode"""

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /