Add sources to completions APIs and UI (#1206)

private_gpt/server/chat/chat_service.py

@@ -1,13 +1,14 @@
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any
-
 from injector import inject, singleton
 from llama_index import ServiceContext, StorageContext, VectorStoreIndex
 from llama_index.chat_engine import ContextChatEngine
+from llama_index.chat_engine.types import (
+    BaseChatEngine,
+)
 from llama_index.indices.postprocessor import MetadataReplacementPostProcessor
 from llama_index.llm_predictor.utils import stream_chat_response_to_tokens
 from llama_index.llms import ChatMessage
 from llama_index.types import TokenGen
+from pydantic import BaseModel
 
 from private_gpt.components.embedding.embedding_component import EmbeddingComponent
 from private_gpt.components.llm.llm_component import LLMComponent
@@ -16,12 +17,17 @@ from private_gpt.components.vector_store.vector_store_component import (
     VectorStoreComponent,
 )
 from private_gpt.open_ai.extensions.context_filter import ContextFilter
+from private_gpt.server.chunks.chunks_service import Chunk
 
-if TYPE_CHECKING:
-    from llama_index.chat_engine.types import (
-        AgentChatResponse,
-        StreamingAgentChatResponse,
-    )
+
+class Completion(BaseModel):
+    response: str
+    sources: list[Chunk] | None = None
+
+
+class CompletionGen(BaseModel):
+    response: TokenGen
+    sources: list[Chunk] | None = None
 
 
 @singleton
@@ -51,66 +57,64 @@ class ChatService:
             show_progress=True,
         )
 
-    def _chat_with_contex(
-        self,
-        message: str,
-        context_filter: ContextFilter | None = None,
-        chat_history: Sequence[ChatMessage] | None = None,
-        streaming: bool = False,
-    ) -> Any:
+    def _chat_engine(
+        self, context_filter: ContextFilter | None = None
+    ) -> BaseChatEngine:
         vector_index_retriever = self.vector_store_component.get_retriever(
             index=self.index, context_filter=context_filter
         )
-        chat_engine = ContextChatEngine.from_defaults(
+        return ContextChatEngine.from_defaults(
             retriever=vector_index_retriever,
             service_context=self.service_context,
             node_postprocessors=[
                 MetadataReplacementPostProcessor(target_metadata_key="window"),
             ],
         )
-        if streaming:
-            result = chat_engine.stream_chat(message, chat_history)
-        else:
-            result = chat_engine.chat(message, chat_history)
-        return result
 
     def stream_chat(
         self,
         messages: list[ChatMessage],
         use_context: bool = False,
         context_filter: ContextFilter | None = None,
-    ) -> TokenGen:
+    ) -> CompletionGen:
         if use_context:
             last_message = messages[-1].content
-            response: StreamingAgentChatResponse = self._chat_with_contex(
+            chat_engine = self._chat_engine(context_filter=context_filter)
+            streaming_response = chat_engine.stream_chat(
                 message=last_message if last_message is not None else "",
                 chat_history=messages[:-1],
-                context_filter=context_filter,
-                streaming=True,
             )
-            response_gen = response.response_gen
+            sources = [
+                Chunk.from_node(node) for node in streaming_response.source_nodes
+            ]
+            completion_gen = CompletionGen(
+                response=streaming_response.response_gen, sources=sources
+            )
         else:
             stream = self.llm_service.llm.stream_chat(messages)
-            response_gen = stream_chat_response_to_tokens(stream)
-        return response_gen
+            completion_gen = CompletionGen(
+                response=stream_chat_response_to_tokens(stream)
+            )
+        return completion_gen
 
     def chat(
         self,
         messages: list[ChatMessage],
         use_context: bool = False,
         context_filter: ContextFilter | None = None,
-    ) -> str:
+    ) -> Completion:
         if use_context:
             last_message = messages[-1].content
-            wrapped_response: AgentChatResponse = self._chat_with_contex(
+            chat_engine = self._chat_engine(context_filter=context_filter)
+            wrapped_response = chat_engine.chat(
                 message=last_message if last_message is not None else "",
                 chat_history=messages[:-1],
-                context_filter=context_filter,
-                streaming=False,
             )
-            response = wrapped_response.response
+            sources = [Chunk.from_node(node) for node in wrapped_response.source_nodes]
+            completion = Completion(response=wrapped_response.response, sources=sources)
         else:
             chat_response = self.llm_service.llm.chat(messages)
             response_content = chat_response.message.content
             response = response_content if response_content is not None else ""
-        return response
+            completion = Completion(response=response)
+        return completion