From 4ebf8e8814b7d728e8ffe32ef55c6b49191d7188 Mon Sep 17 00:00:00 2001
From: Wesley Stewart <wstewart3@gmail.com>
Date: Tue, 12 Mar 2024 19:32:23 +0000
Subject: [PATCH] Added postprocessor checks

---
 private_gpt/server/chat/chat_service.py | 36 +++++++++++++++++++++++++
 private_gpt/settings/settings.py        | 13 ++++++++-
 settings.yaml                           |  3 +++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/private_gpt/server/chat/chat_service.py b/private_gpt/server/chat/chat_service.py
index 5369200b..09dd1a2e 100644
--- a/private_gpt/server/chat/chat_service.py
+++ b/private_gpt/server/chat/chat_service.py
@@ -8,6 +8,10 @@ from llama_index.core.chat_engine.types import (
 from llama_index.core.indices import VectorStoreIndex
 from llama_index.core.indices.postprocessor import MetadataReplacementPostProcessor
 from llama_index.core.llms import ChatMessage, MessageRole
+from llama_index.core.postprocessor import (
+    KeywordNodePostprocessor,
+    SimilarityPostprocessor,
+)
 from llama_index.core.storage import StorageContext
 from llama_index.core.types import TokenGen
 from pydantic import BaseModel
@@ -20,6 +24,7 @@ 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
 
 
 class Completion(BaseModel):
@@ -102,12 +107,43 @@ class ChatService:
             vector_index_retriever = self.vector_store_component.get_retriever(
                 index=self.index, context_filter=context_filter
             )
+            # Initialize node_postporcessors
+            node_postprocessors_tmp: list[
+                MetadataReplacementPostProcessor
+                | SimilarityPostprocessor
+                | KeywordNodePostprocessor
+            ]
+            node_postprocessors_tmp = [
+                MetadataReplacementPostProcessor(target_metadata_key="window"),
+            ]
+            # If similarity value is set, use it.  If not, dont add it
+            if settings().llm.similarity_value is not None:
+                node_postprocessors_tmp.append(
+                    SimilarityPostprocessor(
+                        similarity_cutoff=settings().llm.similarity_value
+                    )
+                )
+
+            # If similarity value is set, use it.  If not, dont add it
+            # if settings().llm.keywords_include is not empty or
+            # settings().llm.keywords_exclude is not empty
+            if settings().llm.keywords_include or settings().llm.keywords_exclude:
+                node_postprocessors_tmp.append(
+                    KeywordNodePostprocessor(
+                        required_keywords=settings().llm.keywords_include,
+                        exclude_keywords=settings().llm.keywords_exclude,
+                    )
+                )
+
             return ContextChatEngine.from_defaults(
                 system_prompt=system_prompt,
                 retriever=vector_index_retriever,
                 llm=self.llm_component.llm,  # Takes no effect at the moment
                 node_postprocessors=[
                     MetadataReplacementPostProcessor(target_metadata_key="window"),
+                    SimilarityPostprocessor(
+                        similarity_cutoff=settings().llm.similarity_value
+                    ),
                 ],
             )
         else:
diff --git a/private_gpt/settings/settings.py b/private_gpt/settings/settings.py
index 62af3f34..b345c738 100644
--- a/private_gpt/settings/settings.py
+++ b/private_gpt/settings/settings.py
@@ -102,6 +102,18 @@ class LLMSettings(BaseModel):
         0.1,
         description="The temperature of the model. Increasing the temperature will make the model answer more creatively. A value of 0.1 would be more factual.",
     )
+    similarity_value: float = Field(
+        None,
+        description="If set, any documents retrieved from the RAG must meet a certain score. Acceptable values are between 0 and 1.",
+    )
+    keywords_include: list[str] = Field(
+        [],
+        description="If set, any documents retrieved from the RAG Must include this keyword",
+    )
+    keywords_exclude: list[str] = Field(
+        [],
+        description="If set, any documents retrieved from the RAG Must exclude this keyword",
+    )
 
 
 class VectorstoreSettings(BaseModel):
@@ -122,7 +134,6 @@ class LlamaCPPSettings(BaseModel):
             "`llama2` is the historic behaviour. `default` might work better with your custom models."
         ),
     )
-
     tfs_z: float = Field(
         1.0,
         description="Tail free sampling is used to reduce the impact of less probable tokens from the output. A higher value (e.g., 2.0) will reduce the impact more, while a value of 1.0 disables this setting.",
diff --git a/settings.yaml b/settings.yaml
index a9a676bd..6ee5d35c 100644
--- a/settings.yaml
+++ b/settings.yaml
@@ -40,6 +40,9 @@ llm:
   max_new_tokens: 512
   context_window: 3900
   temperature: 0.1      # The temperature of the model. Increasing the temperature will make the model answer more creatively. A value of 0.1 would be more factual. (Default: 0.1)
+  #similarity_value: 0.4    #If set, any documents retrieved from the RAG must meet a certain score. Acceptable values are between 0 and 1.
+  #keywords_include: ["Apples","Bananas"] #Optional - requires spacy package: pip install spacy
+  #keywords_exclude: ["Pears","Mangos"]   #Optional - requires spacy package: pip install spacy
 
 llamacpp:
   prompt_style: "mistral"