106 lines
2.9 KiB
TypeScript
106 lines
2.9 KiB
TypeScript
import { ref } from "vue"
|
|
import { WS_BASE_URL } from "../utils/const"
|
|
import { html, css, js } from "./editors"
|
|
|
|
export interface PromptMessage {
|
|
role: "user" | "assistant"
|
|
content: string
|
|
code?: { html: string | null; css: string | null; js: string | null }
|
|
created?: string
|
|
}
|
|
|
|
export const messages = ref<PromptMessage[]>([])
|
|
export const conversationId = ref<string>("")
|
|
export const connected = ref(false)
|
|
export const streaming = ref(false)
|
|
export const streamingContent = ref("")
|
|
export let onCodeComplete: ((code: { html: string | null; css: string | null; js: string | null }) => void) | null = null
|
|
|
|
let ws: WebSocket | null = null
|
|
|
|
export function connectPrompt(taskId: number) {
|
|
if (ws) ws.close()
|
|
|
|
ws = new WebSocket(`${WS_BASE_URL}/ws/prompt/${taskId}/`)
|
|
|
|
ws.onopen = () => {
|
|
connected.value = true
|
|
}
|
|
|
|
ws.onmessage = (event) => {
|
|
const data = JSON.parse(event.data)
|
|
|
|
if (data.type === "init") {
|
|
conversationId.value = data.conversation_id
|
|
messages.value = data.messages || []
|
|
// Apply code from last assistant message if exists
|
|
const lastAssistant = [...messages.value]
|
|
.reverse()
|
|
.find((m) => m.role === "assistant" && m.code)
|
|
if (lastAssistant?.code) {
|
|
applyCode(lastAssistant.code)
|
|
}
|
|
} else if (data.type === "stream") {
|
|
streaming.value = true
|
|
streamingContent.value += data.content
|
|
} else if (data.type === "complete") {
|
|
streaming.value = false
|
|
// Push the full assistant message
|
|
messages.value.push({
|
|
role: "assistant",
|
|
content: streamingContent.value,
|
|
code: data.code,
|
|
})
|
|
streamingContent.value = ""
|
|
// Apply code to editors
|
|
if (data.code) {
|
|
applyCode(data.code)
|
|
if (onCodeComplete) {
|
|
onCodeComplete(data.code)
|
|
}
|
|
}
|
|
} else if (data.type === "error") {
|
|
streaming.value = false
|
|
streamingContent.value = ""
|
|
messages.value.push({
|
|
role: "assistant",
|
|
content: data.content,
|
|
})
|
|
}
|
|
}
|
|
|
|
ws.onclose = () => {
|
|
connected.value = false
|
|
}
|
|
}
|
|
|
|
export function disconnectPrompt() {
|
|
if (ws) {
|
|
ws.close()
|
|
ws = null
|
|
}
|
|
messages.value = []
|
|
conversationId.value = ""
|
|
connected.value = false
|
|
streaming.value = false
|
|
streamingContent.value = ""
|
|
onCodeComplete = null
|
|
}
|
|
|
|
export function sendPrompt(content: string) {
|
|
if (!ws || ws.readyState !== WebSocket.OPEN) return
|
|
messages.value.push({ role: "user", content })
|
|
ws.send(JSON.stringify({ type: "message", content }))
|
|
}
|
|
|
|
export function newConversation() {
|
|
if (!ws || ws.readyState !== WebSocket.OPEN) return
|
|
ws.send(JSON.stringify({ type: "new_conversation" }))
|
|
}
|
|
|
|
function applyCode(code: { html: string | null; css: string | null; js: string | null }) {
|
|
if (code.html !== null) html.value = code.html
|
|
if (code.css !== null) css.value = code.css
|
|
if (code.js !== null) js.value = code.js
|
|
}
|