From 6a61fd38c3d9ac9e6739869205b667eb36a06d92 Mon Sep 17 00:00:00 2001
From: Anupam Singh <96116740+ianupamsingh@users.noreply.github.com>
Date: Fri, 7 Jul 2023 17:56:43 +0530
Subject: [PATCH] feat: Add config based template to query function (#177)

---
 README.md                        |  4 +++-
 embedchain/config/QueryConfig.py | 32 ++++++++++++++++++++++++++++++--
 embedchain/embedchain.py         | 12 +++++-------
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index dd752283..500c044c 100644
--- a/README.md
+++ b/README.md
@@ -317,7 +317,9 @@ This section describes all possible config options.
 
 #### **Query Config**
 
-*coming soon*
+|option|description|type|default|
+|---|---|---|---|
+|template|custom template for prompt|Template|Template("Use the following pieces of context to answer the query at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. \$context Query: $query Helpful Answer:")|
 
 #### **Chat Config**
 
diff --git a/embedchain/config/QueryConfig.py b/embedchain/config/QueryConfig.py
index ad69c11d..f40c4a30 100644
--- a/embedchain/config/QueryConfig.py
+++ b/embedchain/config/QueryConfig.py
@@ -1,8 +1,36 @@
 from embedchain.config.BaseConfig import BaseConfig
+from string import Template
+import re
+
+
+DEFAULT_PROMPT_TEMPLATE = Template("""
+  Use the following pieces of context to answer the query at the end.
+  If you don't know the answer, just say that you don't know, don't try to make up an answer.
+
+  $context
+  
+  Query: $query
+  
+  Helpful Answer:
+""")
+query_re = re.compile(r"\$\{*query\}*")
+context_re = re.compile(r"\$\{*context\}*")
+
 
 class QueryConfig(BaseConfig):
     """
     Config for the `query` method.
     """
-    def __init__(self):
-        pass
+    def __init__(self, template: Template = None):
+        """
+        Initializes the QueryConfig instance.
+
+        :param template: Optional. The `Template` instance to use as a template for prompt.
+        :raises ValueError: If the template is not valid as template should contain $context and $query
+        """
+        if template is None:
+            template = DEFAULT_PROMPT_TEMPLATE
+        if not (re.search(query_re, template.template) \
+            and re.search(context_re, template.template)):
+            raise ValueError("`template` should have `query` and `context` keys")
+        self.template = template
diff --git a/embedchain/embedchain.py b/embedchain/embedchain.py
index ea8892ac..ab70c62f 100644
--- a/embedchain/embedchain.py
+++ b/embedchain/embedchain.py
@@ -1,5 +1,6 @@
 import openai
 import os
+from string import Template
 
 from chromadb.utils import embedding_functions
 from dotenv import load_dotenv
@@ -192,19 +193,16 @@ class EmbedChain:
             content = ""
         return content
 
-    def generate_prompt(self, input_query, context):
+    def generate_prompt(self, input_query, context, template: Template = None):
         """
         Generates a prompt based on the given query and context, ready to be passed to an LLM
 
         :param input_query: The query to use.
         :param context: Similar documents to the query used as context.
+        :param template: Optional. The `Template` instance to use as a template for prompt.
         :return: The prompt
         """
-        prompt = f"""Use the following pieces of context to answer the query at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
-        {context}
-        Query: {input_query}
-        Helpful Answer:
-        """
+        prompt = template.substitute(context = context, query = input_query)
         return prompt
 
     def get_answer_from_llm(self, prompt):
@@ -232,7 +230,7 @@ class EmbedChain:
         if config is None:
             config = QueryConfig()
         context = self.retrieve_from_database(input_query)
-        prompt = self.generate_prompt(input_query, context)
+        prompt = self.generate_prompt(input_query, context, config.template)
         answer = self.get_answer_from_llm(prompt)
         return answer