Merge pull request #72 from eren23/context-restarter-branch

a flag to restart the context mid conversation
geeks-of-data · Apr 8, 2023 · b3f6646 · b3f6646
2 parents a5aac49 + 6ca2a9a
commit b3f6646
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/knowledgegpt/__init__.py b/knowledgegpt/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.0.6b"
+__version__ = "0.0.7b"
 
 from .extractors.yt_subs_extractor import YTSubsExtractor
 from .extractors.yt_audio_extractor import YoutubeAudioExtractor

diff --git a/knowledgegpt/extractors/base_extractor.py b/knowledgegpt/extractors/base_extractor.py
@@ -63,7 +63,7 @@ def load_embeddings_indexes(self):
         if self.df is None:
             self.df = pd.read_csv(self.index_path + "/df.csv")
 
-    def extract(self, query, max_tokens, load_index=False) -> tuple[str, str, list]:
+    def extract(self, query, max_tokens, load_index=False, context_restarter=False) -> tuple[str, str, list]:
         """
         param query: Query to answer
         param max_tokens: Maximum number of tokens to generate
@@ -85,7 +85,8 @@ def extract(self, query, max_tokens, load_index=False) -> tuple[str, str, list]:
             messages=self.messages,
             max_tokens=max_tokens,
             index_type=self.index_type,
-            prompt_template=self.prompt_template
+            prompt_template=self.prompt_template,
+            context_restarter=context_restarter
         )
         if not self.verbose:
             print("all_done!")

diff --git a/knowledgegpt/utils/utils_completion.py b/knowledgegpt/utils/utils_completion.py
@@ -37,7 +37,8 @@ def answer_query_with_context(
         messages: list = None,
         index_type: str = "basic",
         max_tokens=1000,
-        prompt_template=None
+        prompt_template=None,
+        context_restarter: bool = False
 ) -> str:
     """
     Answer a query using the provided context.
@@ -55,7 +56,7 @@ def answer_query_with_context(
 
     """
 
-    if len(messages) < 3 or not is_turbo:
+    if len(messages) < 3 or not is_turbo or context_restarter:
         prompt = construct_prompt(
             verbose=verbose,
             question=query,