feat: sanitize prompts (#10)

stacklok · Dec 17, 2024 · 5b9dbdd · 5b9dbdd
1 parent 7b70afe
commit 5b9dbdd
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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;
+  }
+}