refactor vector search to chat widget

src/fn/rag-chat.ts

@@ -0,0 +1,45 @@
+import { createServerFn } from "@tanstack/react-start";
+import { adminMiddleware } from "~/lib/auth";
+import { z } from "zod";
+import { ragChatUseCase } from "~/use-cases/rag-chat";
+
+const videoSourceSchema = z.object({
+  segmentId: z.number(),
+  segmentTitle: z.string(),
+  segmentSlug: z.string(),
+  moduleTitle: z.string(),
+  chunkText: z.string(),
+  similarity: z.number(),
+});
+
+const MAX_MESSAGE_CONTENT_LENGTH = 2000;
+
+const conversationMessageSchema = z.object({
+  id: z.string(),
+  role: z.enum(["user", "assistant"]),
+  content: z
+    .string()
+    .max(MAX_MESSAGE_CONTENT_LENGTH, "Message content too long"),
+  timestamp: z.string(),
+  sources: z.array(videoSourceSchema).optional(),
+});
+
+const ragChatInputSchema = z.object({
+  userMessage: z
+    .string()
+    .min(1, "Message cannot be empty")
+    .max(MAX_MESSAGE_CONTENT_LENGTH, "Message too long"),
+  conversationHistory: z
+    .array(conversationMessageSchema)
+    .max(20, "Too many messages in history"),
+});
+
+export const ragChatFn = createServerFn({ method: "POST" })
+  .middleware([adminMiddleware])
+  .inputValidator(ragChatInputSchema)
+  .handler(async ({ data }) => {
+    return ragChatUseCase({
+      userMessage: data.userMessage,
+      conversationHistory: data.conversationHistory,
+    });
+  });

src/hooks/use-rag-chat.ts

@@ -0,0 +1,204 @@
+import { useState, useCallback, useEffect, useRef } from "react";
+import { useMutation } from "@tanstack/react-query";
+import { z } from "zod";
+import { ragChatFn } from "~/fn/rag-chat";
+import type { VideoSource, ConversationMessage } from "~/use-cases/rag-chat";
+
+export type { VideoSource, ConversationMessage };
+
+const STORAGE_KEY = "rag-chat-history";
+const SOURCES_STORAGE_KEY = "rag-chat-sources";
+const MAX_MESSAGE_LENGTH = 2000;
+const MAX_HISTORY_LENGTH = 20000;
+
+const videoSourceSchema = z.object({
+  segmentId: z.number(),
+  segmentTitle: z.string(),
+  segmentSlug: z.string(),
+  moduleTitle: z.string(),
+  chunkText: z.string(),
+  similarity: z.number(),
+});
+
+const conversationMessageSchema = z.object({
+  id: z.string(),
+  role: z.enum(["user", "assistant"]),
+  content: z.string(),
+  timestamp: z.string(),
+  sources: z.array(videoSourceSchema).optional(),
+});
+
+const conversationMessagesSchema = z.array(conversationMessageSchema);
+const videoSourcesSchema = z.array(videoSourceSchema);
+
+function loadFromStorage<T>(
+  key: string,
+  fallback: T,
+  schema: z.ZodType<T>
+): T {
+  if (typeof window === "undefined") return fallback;
+  try {
+    const stored = sessionStorage.getItem(key);
+    if (stored) {
+      const parsed = JSON.parse(stored);
+      const result = schema.safeParse(parsed);
+      if (result.success) {
+        return result.data;
+      }
+      console.error(`[RAG Chat] Validation failed for ${key}:`, result.error);
+    }
+  } catch (error) {
+    console.error(`[RAG Chat] Failed to load from sessionStorage:`, error);
+  }
+  return fallback;
+}
+
+function saveToStorage<T>(key: string, value: T): void {
+  if (typeof window === "undefined") return;
+  try {
+    sessionStorage.setItem(key, JSON.stringify(value));
+  } catch (error) {
+    console.error(`[RAG Chat] Failed to save to sessionStorage:`, error);
+  }
+}
+
+function clearStorage(): void {
+  if (typeof window === "undefined") return;
+  try {
+    sessionStorage.removeItem(STORAGE_KEY);
+    sessionStorage.removeItem(SOURCES_STORAGE_KEY);
+  } catch (error) {
+    console.error(`[RAG Chat] Failed to clear sessionStorage:`, error);
+  }
+}
+
+function trimConversationHistory(
+  history: ConversationMessage[],
+  maxLength: number
+): ConversationMessage[] {
+  if (history.length === 0) return history;
+
+  const truncateMessage = (msg: ConversationMessage): ConversationMessage => ({
+    ...msg,
+    content: msg.content.slice(0, MAX_MESSAGE_LENGTH),
+    sources: undefined,
+  });
+
+  let trimmed = history.map(truncateMessage);
+  let serialized = JSON.stringify(trimmed);
+
+  while (serialized.length > maxLength && trimmed.length > 0) {
+    trimmed = trimmed.slice(1);
+    serialized = JSON.stringify(trimmed);
+  }
+
+  return trimmed;
+}
+
+export interface UseRagChatReturn {
+  messages: ConversationMessage[];
+  isLoading: boolean;
+  error: Error | null;
+  sendMessage: (content: string) => Promise<void>;
+  clearChat: () => void;
+  currentSources: VideoSource[];
+}
+
+export function useRagChat(): UseRagChatReturn {
+  const [messages, setMessages] = useState<ConversationMessage[]>(() =>
+    loadFromStorage<ConversationMessage[]>(STORAGE_KEY, [], conversationMessagesSchema)
+  );
+  const [currentSources, setCurrentSources] = useState<VideoSource[]>(() =>
+    loadFromStorage<VideoSource[]>(SOURCES_STORAGE_KEY, [], videoSourcesSchema)
+  );
+  const messagesRef = useRef<ConversationMessage[]>(messages);
+
+  useEffect(() => {
+    messagesRef.current = messages;
+  }, [messages]);
+
+  useEffect(() => {
+    saveToStorage(STORAGE_KEY, messages);
+  }, [messages]);
+
+  useEffect(() => {
+    saveToStorage(SOURCES_STORAGE_KEY, currentSources);
+  }, [currentSources]);
+
+  const mutation = useMutation({
+    mutationFn: async (userMessage: string) => {
+      const trimmedHistory = trimConversationHistory(
+        messagesRef.current,
+        MAX_HISTORY_LENGTH
+      );
+      const result = await ragChatFn({
+        data: {
+          userMessage,
+          conversationHistory: trimmedHistory,
+        },
+      });
+      return result;
+    },
+    onMutate: (userMessage) => {
+      const userMsg: ConversationMessage = {
+        id: crypto.randomUUID(),
+        role: "user",
+        content: userMessage,
+        timestamp: new Date().toISOString(),
+      };
+      setMessages((prev) => [...prev, userMsg]);
+      setCurrentSources([]);
+    },
+    onSuccess: (result) => {
+      const assistantMsg: ConversationMessage = {
+        id: crypto.randomUUID(),
+        role: "assistant",
+        content: result.response,
+        timestamp: new Date().toISOString(),
+        sources: result.sources,
+      };
+      setMessages((prev) => [...prev, assistantMsg]);
+      setCurrentSources(result.sources);
+    },
+    onError: (error) => {
+      setMessages((prev) => {
+        const lastMessage = prev[prev.length - 1];
+        if (lastMessage?.role === "user") {
+          return prev.slice(0, -1);
+        }
+        return prev;
+      });
+      console.error("[RAG Chat] Error:", error);
+    },
+  });
+
+  const sendMessage = useCallback(
+    async (content: string) => {
+      const trimmed = content.trim();
+      if (!trimmed || mutation.isPending) return;
+      if (trimmed.length > MAX_MESSAGE_LENGTH) {
+        throw new Error(
+          `Message exceeds maximum length of ${MAX_MESSAGE_LENGTH} characters`
+        );
+      }
+      await mutation.mutateAsync(trimmed);
+    },
+    [mutation]
+  );
+
+  const clearChat = useCallback(() => {
+    setMessages([]);
+    setCurrentSources([]);
+    clearStorage();
+    mutation.reset();
+  }, [mutation]);
+
+  return {
+    messages,
+    isLoading: mutation.isPending,
+    error: mutation.error,
+    sendMessage,
+    clearChat,
+    currentSources,
+  };
+}

src/lib/openai-chat.ts

@@ -0,0 +1,146 @@
+import OpenAI from "openai";
+import { env } from "~/utils/env";
+
+const openai = new OpenAI({
+  apiKey: env.OPENAI_API_KEY,
+});
+
+const CHAT_MODEL = "gpt-4o";
+const MAX_RETRIES = 3;
+const INITIAL_RETRY_DELAY_MS = 1000;
+
+export interface ChatMessage {
+  role: "system" | "user" | "assistant";
+  content: string;
+}
+
+export interface ChatCompletionOptions {
+  model?: "gpt-4o-mini" | "gpt-4o";
+  temperature?: number;
+  maxTokens?: number;
+}
+
+export interface ChatCompletionResult {
+  content: string;
+  usage?: {
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+  };
+}
+
+class ChatCompletionError extends Error {
+  constructor(
+    message: string,
+    public readonly code?: string,
+    public readonly status?: number,
+    public readonly context?: Record<string, unknown>
+  ) {
+    super(message);
+    this.name = "ChatCompletionError";
+  }
+}
+
+async function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+async function withRetry<T>(
+  fn: () => Promise<T>,
+  context: Record<string, unknown>
+): Promise<T> {
+  let lastError: Error | undefined;
+
+  for (let attempt = 0; attempt < MAX_RETRIES; attempt++) {
+    try {
+      return await fn();
+    } catch (error) {
+      lastError = error instanceof Error ? error : new Error(String(error));
+
+      const isRetryable =
+        error instanceof OpenAI.APIError &&
+        (error.status === 429 ||
+          error.status === 500 ||
+          error.status === 502 ||
+          error.status === 503);
+
+      if (!isRetryable || attempt === MAX_RETRIES - 1) {
+        break;
+      }
+
+      const delay = INITIAL_RETRY_DELAY_MS * Math.pow(2, attempt);
+      console.warn(
+        `Chat completion API call failed (attempt ${attempt + 1}/${MAX_RETRIES}), retrying in ${delay}ms...`,
+        { error: lastError.message, ...context }
+      );
+      await sleep(delay);
+    }
+  }
+
+  if (lastError instanceof OpenAI.APIError) {
+    throw new ChatCompletionError(
+      `OpenAI API error: ${lastError.message}`,
+      lastError.code ?? undefined,
+      lastError.status,
+      context
+    );
+  }
+
+  throw new ChatCompletionError(
+    `Chat completion failed: ${lastError?.message ?? "Unknown error"}`,
+    undefined,
+    undefined,
+    context
+  );
+}
+
+export async function createChatCompletion(
+  messages: ChatMessage[],
+  options?: ChatCompletionOptions
+): Promise<ChatCompletionResult> {
+  if (!messages || !Array.isArray(messages) || messages.length === 0) {
+    throw new ChatCompletionError(
+      "Messages must be a non-empty array",
+      undefined,
+      undefined,
+      { messagesLength: messages?.length ?? 0 }
+    );
+  }
+
+  const model = options?.model ?? CHAT_MODEL;
+  const temperature = options?.temperature ?? 0.7;
+  const maxTokens = options?.maxTokens ?? 2048;
+
+  return withRetry(
+    async () => {
+      const response = await openai.chat.completions.create({
+        model,
+        messages,
+        temperature,
+        max_tokens: maxTokens,
+      });
+
+      const choice = response.choices[0];
+      if (!choice?.message?.content) {
+        throw new ChatCompletionError(
+          "Invalid API response: missing message content",
+          undefined,
+          undefined,
+          { choicesLength: response.choices.length }
+        );
+      }
+
+      return {
+        content: choice.message.content,
+        usage: response.usage
+          ? {
+              promptTokens: response.usage.prompt_tokens,
+              completionTokens: response.usage.completion_tokens,
+              totalTokens: response.usage.total_tokens,
+            }
+          : undefined,
+      };
+    },
+    { model, messagesCount: messages.length }
+  );
+}