From 5681069495bbea25bb47d536531026e649da2b6d Mon Sep 17 00:00:00 2001
From: Muspi Merol <wjj123345567@icloud.com>
Date: Wed, 17 Apr 2024 13:52:34 +0800
Subject: [PATCH] feat: content filtering

---
 src/routes/run.py   | 5 +++++
 src/utils/config.py | 3 +++
 2 files changed, 8 insertions(+)

diff --git a/src/routes/run.py b/src/routes/run.py
index 5490c376..23d0e91b 100644
--- a/src/routes/run.py
+++ b/src/routes/run.py
@@ -103,6 +103,11 @@ async def make_stream():
 
 @run_router.put(f"{env.base}/single/{{template}}")
 async def step_run(data: ChainInput, node: Node = Depends(get_node)):
+    for msg in data.messages:
+        for string in env.banned_substrings:
+            if string in msg.content:
+                return PlainTextResponse(env.banned_response)
+
     async def make_stream():
         last = ""
         async for c in node.astream(data.context, find_llm(data.model).generate, **data.config):
diff --git a/src/utils/config.py b/src/utils/config.py
index 96a29bf5..e320a059 100644
--- a/src/utils/config.py
+++ b/src/utils/config.py
@@ -18,6 +18,9 @@ class Config(BaseSettings):
 
     base_path: str = ""
 
+    banned_substrings: list[str] = []
+    banned_response: str = "Sorry, I can't help with that."
+
     @property
     def base(self):
         return f"/{self.base_path}" if self.base_path else ""