Skip to content

Commit 34a8200

Browse files
john0isaacjohn0isaac
authored
Merge pull request #84 from Azure-Samples/add-thought-process
add thought process
2 parents 3f9d941 + b787931 commit 34a8200

File tree

5 files changed

+116
-16
lines changed

5 files changed

+116
-16
lines changed

‎src/quartapp/approaches/rag.py

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -88,14 +88,20 @@ async def run(
8888
context_prompt_template = ChatPromptTemplate.from_template(CONTEXT_PROMPT)
8989
self._chat.temperature = temperature
9090
context_chain = context_prompt_template | self._chat
91-
91+
documents_list: list[Document] = []
9292
if data_points:
9393
# Perform RAG search
9494
response = await context_chain.ainvoke(
9595
{"context": [dp.to_dict() for dp in data_points], "input": rephrased_question.content}
9696
)
97-
98-
return vector_context, str(response.content)
97+
for document in vector_context:
98+
documents_list.append(
99+
Document(page_content=document.page_content, metadata={"source": document.metadata["source"]})
100+
)
101+
formatted_response = (
102+
f'{{"response": "{response.content}", "rephrased_response": "{rephrased_question.content}"}}'
103+
)
104+
return documents_list, str(formatted_response)
99105

100106
# Perform RAG search with no context
101107
response = await context_chain.ainvoke({"context": [], "input": rephrased_question.content})
@@ -127,14 +133,18 @@ async def run_stream(
127133
context_prompt_template = ChatPromptTemplate.from_template(CONTEXT_PROMPT)
128134
self._chat.temperature = temperature
129135
context_chain = context_prompt_template | self._chat
136+
documents_list: list[Document] = []
130137

131138
if data_points:
132139
# Perform RAG search
133140
response = context_chain.astream(
134141
{"context": [dp.to_dict() for dp in data_points], "input": rephrased_question.content}
135142
)
136-
137-
return vector_context, response
143+
for document in vector_context:
144+
documents_list.append(
145+
Document(page_content=document.page_content, metadata={"source": document.metadata["source"]})
146+
)
147+
return documents_list, response
138148

139149
# Perform RAG search with no context
140150
response = context_chain.astream({"context": [], "input": rephrased_question.content})

‎src/quartapp/approaches/vector.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,13 @@ async def run(
1212
search_type="similarity", search_kwargs={"k": limit, "score_threshold": score_threshold}
1313
)
1414
vector_response = await retriever.ainvoke(query)
15+
documents_list: list[Document] = []
16+
1517
if vector_response:
16-
return vector_response, vector_response[0].page_content
18+
for document in vector_response:
19+
documents_list.append(
20+
Document(page_content=document.page_content, metadata={"source": document.metadata["source"]})
21+
)
22+
if documents_list:
23+
return documents_list, documents_list[0].page_content
1724
return [], ""

‎src/quartapp/config.py

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,9 @@ async def run_keyword(
3939
"""
4040

4141
context: Context = await self.get_context(keyword_response)
42-
42+
context.thoughts.insert(0, Thought(description=answer, title="Cosmos Text Search Top Result"))
43+
context.thoughts.insert(0, Thought(description=str(keyword_response), title="Cosmos Text Search Result"))
44+
context.thoughts.insert(0, Thought(description=messages[-1]["content"], title="Cosmos Text Search Query"))
4345
message: Message = Message(content=message_content, role=AIChatRoles.ASSISTANT)
4446

4547
await self.add_to_cosmos(
@@ -75,6 +77,9 @@ async def run_vector(
7577
"""
7678

7779
context: Context = await self.get_context(vector_response)
80+
context.thoughts.insert(0, Thought(description=answer, title="Cosmos Vector Search Top Result"))
81+
context.thoughts.insert(0, Thought(description=str(vector_response), title="Cosmos Vector Search Result"))
82+
context.thoughts.insert(0, Thought(description=messages[-1]["content"], title="Cosmos Vector Search Query"))
7883
message: Message = Message(content=message_content, role=AIChatRoles.ASSISTANT)
7984

8085
await self.add_to_cosmos(
@@ -90,6 +95,7 @@ async def run_rag(
9095
self, session_state: str | None, messages: list, temperature: float, limit: int, score_threshold: float
9196
) -> RetrievalResponse:
9297
rag_response, answer = await self.setup.rag.run(messages, temperature, limit, score_threshold)
98+
json_answer = json.loads(answer)
9399

94100
new_session_state: str = session_state if session_state else str(uuid4())
95101

@@ -98,7 +104,7 @@ async def run_rag(
98104
return RetrievalResponse(
99105
sessionState=new_session_state,
100106
context=Context([DataPoint()], [Thought()]),
101-
message=Message(content=answer, role=AIChatRoles.ASSISTANT),
107+
message=Message(content=json_answer.get("response"), role=AIChatRoles.ASSISTANT),
102108
)
103109
else:
104110
return RetrievalResponse(
@@ -108,7 +114,17 @@ async def run_rag(
108114
)
109115

110116
context: Context = await self.get_context(rag_response)
111-
message: Message = Message(content=answer, role=AIChatRoles.ASSISTANT)
117+
context.thoughts.insert(
118+
0, Thought(description=json_answer.get("response"), title="Cosmos RAG OpenAI Rephrased Response")
119+
)
120+
context.thoughts.insert(
121+
0, Thought(description=str(rag_response), title="Cosmos RAG Search Vector Search Result")
122+
)
123+
context.thoughts.insert(
124+
0, Thought(description=json_answer.get("rephrased_response"), title="Cosmos RAG OpenAI Rephrased Query")
125+
)
126+
context.thoughts.insert(0, Thought(description=messages[-1]["content"], title="Cosmos RAG Query"))
127+
message: Message = Message(content=json_answer.get("response"), role=AIChatRoles.ASSISTANT)
112128

113129
await self.add_to_cosmos(
114130
old_messages=messages,
@@ -127,6 +143,10 @@ async def run_rag_stream(
127143
new_session_state: str = session_state if session_state else str(uuid4())
128144

129145
context: Context = await self.get_context(rag_response)
146+
context.thoughts.insert(
147+
0, Thought(description=str(rag_response), title="Cosmos RAG Search Vector Search Result")
148+
)
149+
context.thoughts.insert(0, Thought(description=messages[-1]["content"], title="Cosmos RAG Query"))
130150

131151
yield RetrievalResponseDelta(context=context, sessionState=new_session_state)
132152

‎tests/conftest.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,8 @@ def approaches_base_mock():
7878
# Mock Vector Store
7979
mock_vector_store = MagicMock()
8080
mock_document = Document(
81-
page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}'
81+
page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}',
82+
metadata={"source": "test"},
8283
)
8384
mock_vector_store.as_retriever.return_value.ainvoke = AsyncMock(return_value=[mock_document])
8485

‎tests/test_approaches.py

Lines changed: 68 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,12 @@ async def test_vector_run(vector_mock):
5858
"""Test the Vector class run method."""
5959
result = vector_mock.run([{"content": "test"}], 0.0, 0, 0.0)
6060
assert await result == (
61-
[Document(page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}')],
61+
[
62+
Document(
63+
metadata={"source": "test"},
64+
page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}',
65+
)
66+
],
6267
'{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}',
6368
)
6469

@@ -79,8 +84,13 @@ async def test_rag_run(rag_mock):
7984
"""Test the RAG class run method."""
8085
result = rag_mock.run([{"content": "test"}], 0.0, 0, 0.0)
8186
assert await result == (
82-
[Document(page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}')],
83-
"content",
87+
[
88+
Document(
89+
metadata={"source": "test"},
90+
page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}',
91+
)
92+
],
93+
'{"response": "content", "rephrased_response": "content"}',
8494
)
8595

8696

@@ -120,7 +130,23 @@ async def test_app_config_run_keyword(app_config_mock):
120130
collection="collection_name",
121131
)
122132
],
123-
thoughts=[Thought(title="Source", description="test")],
133+
thoughts=[
134+
Thought(
135+
title="Cosmos Text Search Query",
136+
description="test",
137+
),
138+
Thought(
139+
title="Cosmos Text Search Result",
140+
description="[Document(metadata={'source': 'test'}, "
141+
'page_content=\'{"name": "test", "description": "test", '
142+
'"price": "5.0USD", "category": "test"}\')]',
143+
),
144+
Thought(
145+
title="Cosmos Text Search Top Result",
146+
description='{"name": "test", "description": "test", "price": ' '"5.0USD", "category": "test"}',
147+
),
148+
Thought(title="Source", description="test"),
149+
],
124150
),
125151
message=Message(
126152
content="\n Name: test\n Description: test\n Price: 5.0USD\n"
@@ -146,7 +172,23 @@ async def test_app_config_run_vector(app_config_mock):
146172
collection="collection_name",
147173
)
148174
],
149-
thoughts=[Thought(title="Source", description=None)],
175+
thoughts=[
176+
Thought(
177+
title="Cosmos Vector Search Query",
178+
description="test",
179+
),
180+
Thought(
181+
title="Cosmos Vector Search Result",
182+
description="[Document(metadata={'source': 'test'}, "
183+
'page_content=\'{"name": "test", "description": "test", '
184+
'"price": "5.0USD", "category": "test"}\')]',
185+
),
186+
Thought(
187+
title="Cosmos Vector Search Top Result",
188+
description='{"name": "test", "description": "test", "price": ' '"5.0USD", "category": "test"}',
189+
),
190+
Thought(title="Source", description="test"),
191+
],
150192
),
151193
message=Message(
152194
content="\n Name: test\n Description: test\n Price: 5.0USD\n"
@@ -172,7 +214,27 @@ async def test_app_config_run_rag(app_config_mock):
172214
collection="collection_name",
173215
)
174216
],
175-
thoughts=[Thought(title="Source", description=None)],
217+
thoughts=[
218+
Thought(
219+
title="Cosmos RAG Query",
220+
description="test",
221+
),
222+
Thought(
223+
title="Cosmos RAG OpenAI Rephrased Query",
224+
description="content",
225+
),
226+
Thought(
227+
title="Cosmos RAG Search Vector Search Result",
228+
description="[Document(metadata={'source': 'test'}, "
229+
'page_content=\'{"name": "test", "description": "test", '
230+
'"price": "5.0USD", "category": "test"}\')]',
231+
),
232+
Thought(
233+
title="Cosmos RAG OpenAI Rephrased Response",
234+
description="content",
235+
),
236+
Thought(title="Source", description="test"),
237+
],
176238
),
177239
message=Message(content="content", role=AIChatRoles.ASSISTANT),
178240
sessionState="test",

0 commit comments

Comments
 (0)