Merge pull request #84 from Azure-Samples/add-thought-process

add thought process

src/quartapp/approaches/rag.py

@@ -88,14 +88,20 @@ async def run(
         context_prompt_template = ChatPromptTemplate.from_template(CONTEXT_PROMPT)
         self._chat.temperature = temperature
         context_chain = context_prompt_template | self._chat
-
+        documents_list: list[Document] = []
         if data_points:
             # Perform RAG search
             response = await context_chain.ainvoke(
                 {"context": [dp.to_dict() for dp in data_points], "input": rephrased_question.content}
             )
-
-            return vector_context, str(response.content)
+            for document in vector_context:
+                documents_list.append(
+                    Document(page_content=document.page_content, metadata={"source": document.metadata["source"]})
+                )
+            formatted_response = (
+                f'{{"response": "{response.content}", "rephrased_response": "{rephrased_question.content}"}}'
+            )
+            return documents_list, str(formatted_response)
 
         # Perform RAG search with no context
         response = await context_chain.ainvoke({"context": [], "input": rephrased_question.content})
@@ -127,14 +133,18 @@ async def run_stream(
         context_prompt_template = ChatPromptTemplate.from_template(CONTEXT_PROMPT)
         self._chat.temperature = temperature
         context_chain = context_prompt_template | self._chat
+        documents_list: list[Document] = []
 
         if data_points:
             # Perform RAG search
             response = context_chain.astream(
                 {"context": [dp.to_dict() for dp in data_points], "input": rephrased_question.content}
             )
-
-            return vector_context, response
+            for document in vector_context:
+                documents_list.append(
+                    Document(page_content=document.page_content, metadata={"source": document.metadata["source"]})
+                )
+            return documents_list, response
 
         # Perform RAG search with no context
         response = context_chain.astream({"context": [], "input": rephrased_question.content})

src/quartapp/approaches/vector.py

@@ -12,6 +12,13 @@ async def run(
             search_type="similarity", search_kwargs={"k": limit, "score_threshold": score_threshold}
         )
         vector_response = await retriever.ainvoke(query)
+        documents_list: list[Document] = []
+
         if vector_response:
-            return vector_response, vector_response[0].page_content
+            for document in vector_response:
+                documents_list.append(
+                    Document(page_content=document.page_content, metadata={"source": document.metadata["source"]})
+                )
+            if documents_list:
+                return documents_list, documents_list[0].page_content
         return [], ""

src/quartapp/config.py

@@ -39,7 +39,9 @@ async def run_keyword(
         """
 
         context: Context = await self.get_context(keyword_response)
-
+        context.thoughts.insert(0, Thought(description=answer, title="Cosmos Text Search Top Result"))
+        context.thoughts.insert(0, Thought(description=str(keyword_response), title="Cosmos Text Search Result"))
+        context.thoughts.insert(0, Thought(description=messages[-1]["content"], title="Cosmos Text Search Query"))
         message: Message = Message(content=message_content, role=AIChatRoles.ASSISTANT)
 
         await self.add_to_cosmos(
@@ -75,6 +77,9 @@ async def run_vector(
         """
 
         context: Context = await self.get_context(vector_response)
+        context.thoughts.insert(0, Thought(description=answer, title="Cosmos Vector Search Top Result"))
+        context.thoughts.insert(0, Thought(description=str(vector_response), title="Cosmos Vector Search Result"))
+        context.thoughts.insert(0, Thought(description=messages[-1]["content"], title="Cosmos Vector Search Query"))
         message: Message = Message(content=message_content, role=AIChatRoles.ASSISTANT)
 
         await self.add_to_cosmos(
@@ -90,6 +95,7 @@ async def run_rag(
         self, session_state: str | None, messages: list, temperature: float, limit: int, score_threshold: float
     ) -> RetrievalResponse:
         rag_response, answer = await self.setup.rag.run(messages, temperature, limit, score_threshold)
+        json_answer = json.loads(answer)
 
         new_session_state: str = session_state if session_state else str(uuid4())
 
@@ -98,7 +104,7 @@ async def run_rag(
                 return RetrievalResponse(
                     sessionState=new_session_state,
                     context=Context([DataPoint()], [Thought()]),
-                    message=Message(content=answer, role=AIChatRoles.ASSISTANT),
+                    message=Message(content=json_answer.get("response"), role=AIChatRoles.ASSISTANT),
                 )
             else:
                 return RetrievalResponse(
@@ -108,7 +114,17 @@ async def run_rag(
                 )
 
         context: Context = await self.get_context(rag_response)
-        message: Message = Message(content=answer, role=AIChatRoles.ASSISTANT)
+        context.thoughts.insert(
+            0, Thought(description=json_answer.get("response"), title="Cosmos RAG OpenAI Rephrased Response")
+        )
+        context.thoughts.insert(
+            0, Thought(description=str(rag_response), title="Cosmos RAG Search Vector Search Result")
+        )
+        context.thoughts.insert(
+            0, Thought(description=json_answer.get("rephrased_response"), title="Cosmos RAG OpenAI Rephrased Query")
+        )
+        context.thoughts.insert(0, Thought(description=messages[-1]["content"], title="Cosmos RAG Query"))
+        message: Message = Message(content=json_answer.get("response"), role=AIChatRoles.ASSISTANT)
 
         await self.add_to_cosmos(
             old_messages=messages,
@@ -127,6 +143,10 @@ async def run_rag_stream(
         new_session_state: str = session_state if session_state else str(uuid4())
 
         context: Context = await self.get_context(rag_response)
+        context.thoughts.insert(
+            0, Thought(description=str(rag_response), title="Cosmos RAG Search Vector Search Result")
+        )
+        context.thoughts.insert(0, Thought(description=messages[-1]["content"], title="Cosmos RAG Query"))
 
         yield RetrievalResponseDelta(context=context, sessionState=new_session_state)
 

tests/conftest.py

@@ -78,7 +78,8 @@ def approaches_base_mock():
     # Mock Vector Store
     mock_vector_store = MagicMock()
     mock_document = Document(
-        page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}'
+        page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}',
+        metadata={"source": "test"},
     )
     mock_vector_store.as_retriever.return_value.ainvoke = AsyncMock(return_value=[mock_document])
 

tests/test_approaches.py

@@ -58,7 +58,12 @@ async def test_vector_run(vector_mock):
     """Test the Vector class run method."""
     result = vector_mock.run([{"content": "test"}], 0.0, 0, 0.0)
     assert await result == (
-        [Document(page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}')],
+        [
+            Document(
+                metadata={"source": "test"},
+                page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}',
+            )
+        ],
         '{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}',
     )
 
@@ -79,8 +84,13 @@ async def test_rag_run(rag_mock):
     """Test the RAG class run method."""
     result = rag_mock.run([{"content": "test"}], 0.0, 0, 0.0)
     assert await result == (
-        [Document(page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}')],
-        "content",
+        [
+            Document(
+                metadata={"source": "test"},
+                page_content='{"name": "test", "description": "test", "price": "5.0USD", "category": "test"}',
+            )
+        ],
+        '{"response": "content", "rephrased_response": "content"}',
     )
 
 
@@ -120,7 +130,23 @@ async def test_app_config_run_keyword(app_config_mock):
                     collection="collection_name",
                 )
             ],
-            thoughts=[Thought(title="Source", description="test")],
+            thoughts=[
+                Thought(
+                    title="Cosmos Text Search Query",
+                    description="test",
+                ),
+                Thought(
+                    title="Cosmos Text Search Result",
+                    description="[Document(metadata={'source': 'test'}, "
+                    'page_content=\'{"name": "test", "description": "test", '
+                    '"price": "5.0USD", "category": "test"}\')]',
+                ),
+                Thought(
+                    title="Cosmos Text Search Top Result",
+                    description='{"name": "test", "description": "test", "price": ' '"5.0USD", "category": "test"}',
+                ),
+                Thought(title="Source", description="test"),
+            ],
         ),
         message=Message(
             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):
                     collection="collection_name",
                 )
             ],
-            thoughts=[Thought(title="Source", description=None)],
+            thoughts=[
+                Thought(
+                    title="Cosmos Vector Search Query",
+                    description="test",
+                ),
+                Thought(
+                    title="Cosmos Vector Search Result",
+                    description="[Document(metadata={'source': 'test'}, "
+                    'page_content=\'{"name": "test", "description": "test", '
+                    '"price": "5.0USD", "category": "test"}\')]',
+                ),
+                Thought(
+                    title="Cosmos Vector Search Top Result",
+                    description='{"name": "test", "description": "test", "price": ' '"5.0USD", "category": "test"}',
+                ),
+                Thought(title="Source", description="test"),
+            ],
         ),
         message=Message(
             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):
                     collection="collection_name",
                 )
             ],
-            thoughts=[Thought(title="Source", description=None)],
+            thoughts=[
+                Thought(
+                    title="Cosmos RAG Query",
+                    description="test",
+                ),
+                Thought(
+                    title="Cosmos RAG OpenAI Rephrased Query",
+                    description="content",
+                ),
+                Thought(
+                    title="Cosmos RAG Search Vector Search Result",
+                    description="[Document(metadata={'source': 'test'}, "
+                    'page_content=\'{"name": "test", "description": "test", '
+                    '"price": "5.0USD", "category": "test"}\')]',
+                ),
+                Thought(
+                    title="Cosmos RAG OpenAI Rephrased Response",
+                    description="content",
+                ),
+                Thought(title="Source", description="test"),
+            ],
         ),
         message=Message(content="content", role=AIChatRoles.ASSISTANT),
         sessionState="test",