From 47b371333e7368d1c3ca9c2c9586c61c7ebd3516 Mon Sep 17 00:00:00 2001
From: Juan <juan@tymit.com>
Date: Sat, 9 Dec 2023 12:12:32 -0300
Subject: [PATCH] feat: Context Template from settings for QueryDocs Allowing
 the text values "context_template"

---
 private_gpt/server/chat/chat_service.py |  9 +++++++++
 private_gpt/settings/settings.py        | 11 +++++++++++
 settings.yaml                           |  5 ++++-
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/private_gpt/server/chat/chat_service.py b/private_gpt/server/chat/chat_service.py
index ffdb3f90..77afe114 100644
--- a/private_gpt/server/chat/chat_service.py
+++ b/private_gpt/server/chat/chat_service.py
@@ -19,6 +19,9 @@ from private_gpt.components.vector_store.vector_store_component import (
 )
 from private_gpt.open_ai.extensions.context_filter import ContextFilter
 from private_gpt.server.chunks.chunks_service import Chunk
+from private_gpt.settings.settings import settings
+
+DEFAULT_CONTEXT_TEMPLATE = settings().rag.default_context_template
 
 
 class Completion(BaseModel):
@@ -97,6 +100,7 @@ class ChatService:
         system_prompt: str | None = None,
         use_context: bool = False,
         context_filter: ContextFilter | None = None,
+        context_template: str | None = DEFAULT_CONTEXT_TEMPLATE,
     ) -> BaseChatEngine:
         if use_context:
             vector_index_retriever = self.vector_store_component.get_retriever(
@@ -109,6 +113,7 @@ class ChatService:
                 node_postprocessors=[
                     MetadataReplacementPostProcessor(target_metadata_key="window"),
                 ],
+                context_template=context_template,
             )
         else:
             return SimpleChatEngine.from_defaults(
@@ -121,6 +126,7 @@ class ChatService:
         messages: list[ChatMessage],
         use_context: bool = False,
         context_filter: ContextFilter | None = None,
+        context_template: str | None = None,
     ) -> CompletionGen:
         chat_engine_input = ChatEngineInput.from_messages(messages)
         last_message = (
@@ -141,6 +147,7 @@ class ChatService:
             system_prompt=system_prompt,
             use_context=use_context,
             context_filter=context_filter,
+            context_template=context_template,
         )
         streaming_response = chat_engine.stream_chat(
             message=last_message if last_message is not None else "",
@@ -157,6 +164,7 @@ class ChatService:
         messages: list[ChatMessage],
         use_context: bool = False,
         context_filter: ContextFilter | None = None,
+        context_template: str | None = None,
     ) -> Completion:
         chat_engine_input = ChatEngineInput.from_messages(messages)
         last_message = (
@@ -177,6 +185,7 @@ class ChatService:
             system_prompt=system_prompt,
             use_context=use_context,
             context_filter=context_filter,
+            context_template=context_template,
         )
         wrapped_response = chat_engine.chat(
             message=last_message if last_message is not None else "",
diff --git a/private_gpt/settings/settings.py b/private_gpt/settings/settings.py
index 8b03f611..ccf6b177 100644
--- a/private_gpt/settings/settings.py
+++ b/private_gpt/settings/settings.py
@@ -217,6 +217,16 @@ class QdrantSettings(BaseModel):
     )
 
 
+class RagSettings(BaseModel):
+    default_context_template: str | None = Field(
+        None,
+        description=(
+            "The default context template to use for the chat engine when using RAG. "
+            "If none is given - use the default system prompt (from the llama_index). "
+        ),
+    )
+
+
 class Settings(BaseModel):
     server: ServerSettings
     data: DataSettings
@@ -228,6 +238,7 @@ class Settings(BaseModel):
     openai: OpenAISettings
     vectorstore: VectorstoreSettings
     qdrant: QdrantSettings | None = None
+    rag: RagSettings
 
 
 """
diff --git a/settings.yaml b/settings.yaml
index af51a7f7..066c0064 100644
--- a/settings.yaml
+++ b/settings.yaml
@@ -56,4 +56,7 @@ sagemaker:
 
 openai:
   api_key: ${OPENAI_API_KEY:}
-  model: gpt-3.5-turbo
\ No newline at end of file
+  model: gpt-3.5-turbo
+
+rag:
+  default_context_template: "Context information is below.\n--------------------\n{context_str}\n--------------------\n"