add prompt assistant

src/store/guidance.ts

@@ -0,0 +1,111 @@
+import { ref } from "vue"
+import { WS_BASE_URL } from "../utils/const"
+
+export interface GuidanceMessage {
+  role: "user" | "assistant"
+  content: string
+  id?: number
+}
+
+export const messages = ref<GuidanceMessage[]>([])
+export const connected = ref(false)
+export const streaming = ref(false)
+export const streamingContent = ref("")
+export const isReady = ref(false)
+export const isOpen = ref(false)
+export const initialPrompt = ref("")
+
+let ws: WebSocket | null = null
+let _taskId = 0
+
+function setupHandlers(socket: WebSocket) {
+  socket.onopen = () => {
+    connected.value = true
+    if (initialPrompt.value) {
+      sendToSocket(initialPrompt.value)
+    }
+  }
+
+  socket.onmessage = (event) => {
+    const data = JSON.parse(event.data)
+    if (data.type === "init") {
+      streaming.value = false
+      streamingContent.value = ""
+    } else if (data.type === "stream") {
+      streaming.value = true
+      streamingContent.value += data.content
+    } else if (data.type === "complete") {
+      streaming.value = false
+      const content = streamingContent.value.replace(/^\[READY\]\n?/, "")
+      messages.value.push({ role: "assistant", content, id: data.message_id })
+      streamingContent.value = ""
+      if (data.is_ready) isReady.value = true
+    } else if (data.type === "error") {
+      streaming.value = false
+      streamingContent.value = ""
+      messages.value.push({ role: "assistant", content: data.content })
+    }
+  }
+
+  socket.onclose = () => {
+    connected.value = false
+  }
+}
+
+function sendToSocket(content: string) {
+  if (!ws || ws.readyState !== WebSocket.OPEN) return
+  streaming.value = true
+  messages.value.push({ role: "user", content })
+  ws.send(JSON.stringify({ type: "message", content }))
+}
+
+export function openGuidance(taskId: number, prompt: string) {
+  if (!taskId) return
+
+  _taskId = taskId
+  initialPrompt.value = prompt
+  messages.value = []
+  isReady.value = false
+  streamingContent.value = ""
+  isOpen.value = true
+
+  if (ws) ws.close()
+  ws = new WebSocket(`${WS_BASE_URL}/ws/guidance/${taskId}/`)
+  setupHandlers(ws)
+}
+
+export function sendGuidance(content: string) {
+  sendToSocket(content)
+}
+
+export function stopGuidance() {
+  if (
+    messages.value.length > 0 &&
+    messages.value[messages.value.length - 1].role === "user"
+  ) {
+    messages.value.pop()
+  }
+  streaming.value = false
+  streamingContent.value = ""
+  if (_taskId) {
+    if (ws) ws.close()
+    ws = new WebSocket(`${WS_BASE_URL}/ws/guidance/${_taskId}/`)
+    initialPrompt.value = ""
+    setupHandlers(ws)
+  }
+}
+
+export function closeGuidance() {
+  if (ws) {
+    ws.close()
+    ws = null
+  }
+  isOpen.value = false
+  messages.value = []
+  isReady.value = false
+  connected.value = false
+  streaming.value = false
+  streamingContent.value = ""
+  initialPrompt.value = ""
+  _taskId = 0
+}

src/store/prompt.ts

@@ -15,6 +15,7 @@ export const conversationId = ref<string>("")
 export const connected = ref(false)
 export const streaming = ref(false)
 export const streamingContent = ref("")
+export const currentTaskId = ref(0)
 let _onCodeComplete:
   | ((
       code: { html: string | null; css: string | null; js: string | null },
@@ -31,6 +32,7 @@ let _currentTaskId = 0
 
 export function connectPrompt(taskId: number) {
   _currentTaskId = taskId
+  currentTaskId.value = taskId
   if (ws) ws.close()
 
   ws = new WebSocket(`${WS_BASE_URL}/ws/prompt/${taskId}/`)
@@ -99,6 +101,8 @@ export function disconnectPrompt() {
   connected.value = false
   streaming.value = false
   streamingContent.value = ""
+  currentTaskId.value = 0
+  _currentTaskId = 0
   _onCodeComplete = null
 }