From cf99dce9405a120d9e778d80e818d2492b96656a Mon Sep 17 00:00:00 2001
From: cachho <admin@ch-webdev.com>
Date: Fri, 23 Jun 2023 18:47:52 +0200
Subject: [PATCH] Refactor query endpoint into 3 parts (#42)

Query endpoint now consists of 3 sub functions
- get data from db
- get prompt
- get answer from the data retrieved above by passing to LLM
---
 embedchain/embedchain.py | 47 ++++++++++++++++++++++++++++++----------
 1 file changed, 35 insertions(+), 12 deletions(-)

diff --git a/embedchain/embedchain.py b/embedchain/embedchain.py
index 1cd01ea3..0f12b230 100644
--- a/embedchain/embedchain.py
+++ b/embedchain/embedchain.py
@@ -163,8 +163,39 @@ class EmbedChain:
             top_p=1,
         )
         return response["choices"][0]["message"]["content"]
+    
+    def retrieve_from_database(self, input_query):
+        """
+        Queries the vector database based on the given input query.
+        Gets relevant doc based on the query
 
-    def get_answer_from_llm(self, query, context):
+        :param input_query: The query to use.
+        :return: The content of the document that matched your query.
+        """
+        result = self.collection.query(
+            query_texts=[input_query,],
+            n_results=1,
+        )
+        result_formatted = self._format_result(result)
+        content = result_formatted[0][0].page_content
+        return content
+    
+    def generate_prompt(self, input_query, context):
+        """
+        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.
+        :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:
+        """
+        return prompt
+
+    def get_answer_from_llm(self, prompt):
         """
         Gets an answer based on the given query and context by passing it
         to an LLM.
@@ -173,11 +204,6 @@ class EmbedChain:
         :param context: Similar documents to the query used as context.
         :return: The answer.
         """
-        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: {query}
-        Helpful Answer:
-        """
         answer = self.get_openai_answer(prompt)
         return answer
 
@@ -190,12 +216,9 @@ class EmbedChain:
         :param input_query: The query to use.
         :return: The answer to the query.
         """
-        result = self.collection.query(
-            query_texts=[input_query,],
-            n_results=1,
-        )
-        result_formatted = self._format_result(result)
-        answer = self.get_answer_from_llm(input_query, result_formatted[0][0].page_content)
+        context = self.retrieve_from_database(input_query)
+        prompt = self.generate_prompt(input_query, context)
+        answer = self.get_answer_from_llm(prompt)
         return answer