From e4bc793be6b734c5f2dd570d46f5ca199070a2c4 Mon Sep 17 00:00:00 2001
From: Giuseppe Scuglia <peppescg@gmail.com>
Date: Tue, 17 Dec 2024 19:55:34 +0100
Subject: [PATCH] feat: sanitize prompts

---
 src/components/Chat.tsx       |  6 +++++-
 src/components/PromptList.tsx | 11 ++++++++--
 src/lib/utils.ts              | 40 +++++++++++++++++++++++++++--------
 3 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/src/components/Chat.tsx b/src/components/Chat.tsx
index bde08b9e..29f59ba3 100644
--- a/src/components/Chat.tsx
+++ b/src/components/Chat.tsx
@@ -8,6 +8,7 @@ import { useParams } from "react-router-dom";
 import { usePromptsStore } from "@/hooks/usePromptsStore";
 import { Markdown } from "./Markdown";
 import { useEffect } from "react";
+import { sanitizeQuestionPrompt } from "@/lib/utils";
 
 export function Chat() {
   const { id } = useParams();
@@ -31,7 +32,10 @@ export function Chat() {
               <ChatBubbleAvatar fallback="User" className="w-14" />
               <ChatBubbleMessage variant="sent" className="bg-zinc-700">
                 <Markdown className="text-gray-300">
-                  {question?.message}
+                  {sanitizeQuestionPrompt({
+                    question: question?.message,
+                    answer: answer?.message,
+                  })}
                 </Markdown>
               </ChatBubbleMessage>
             </ChatBubble>
diff --git a/src/components/PromptList.tsx b/src/components/PromptList.tsx
index 224a662f..9b7c0291 100644
--- a/src/components/PromptList.tsx
+++ b/src/components/PromptList.tsx
@@ -3,6 +3,7 @@ import { Link } from "react-router-dom";
 import {
   extractTitleFromMessage,
   groupPromptsByRelativeDate,
+  sanitizeQuestionPrompt,
 } from "@/lib/utils";
 import { usePromptsStore } from "@/hooks/usePromptsStore";
 import clsx from "clsx";
@@ -28,8 +29,14 @@ export function PromptList({ prompts }: { prompts: Prompt[] }) {
                   )}
                 >
                   {extractTitleFromMessage(
-                    prompt.question_answers?.[0].question.message ??
-                      `Prompt ${prompt.conversation_timestamp}`
+                    prompt.question_answers?.[0].question.message
+                      ? sanitizeQuestionPrompt({
+                          question:
+                            prompt.question_answers?.[0].question.message,
+                          answer:
+                            prompt.question_answers?.[0]?.answer?.message ?? "",
+                        })
+                      : `Prompt ${prompt.conversation_timestamp}`
                   )}
                 </Link>
               </li>
diff --git a/src/lib/utils.ts b/src/lib/utils.ts
index 06e97732..502f353f 100644
--- a/src/lib/utils.ts
+++ b/src/lib/utils.ts
@@ -13,18 +13,22 @@ export function cn(...inputs: ClassValue[]) {
 }
 
 export function extractTitleFromMessage(message: string) {
-  const regex = /^(.*)```[\s\S]*?```(.*)$/s;
-  const match = message.match(regex);
+  try {
+    const regex = /^(.*)```[\s\S]*?```(.*)$/s;
+    const match = message.match(regex);
 
-  if (match) {
-    const beforeMarkdown = match[1].trim();
-    const afterMarkdown = match[2].trim();
+    if (match) {
+      const beforeMarkdown = match[1].trim();
+      const afterMarkdown = match[2].trim();
 
-    const title = beforeMarkdown || afterMarkdown;
-    return title;
-  }
+      const title = beforeMarkdown || afterMarkdown;
+      return title;
+    }
 
-  return message.trim();
+    return message.trim();
+  } catch {
+    return message.trim();
+  }
 }
 
 function getGroup(differenceInMs: number, promptDate: Date): string {
@@ -108,3 +112,21 @@ export function getMaliciousPackages() {
 
   return chartData;
 }
+
+export function sanitizeQuestionPrompt({
+  question,
+  answer,
+}: {
+  question: string;
+  answer: string;
+}) {
+  try {
+    if (answer) {
+      return question.split("Query:").pop() ?? "";
+    }
+
+    return question;
+  } catch {
+    return question;
+  }
+}