135 lines
3.4 KiB
TypeScript
135 lines
3.4 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
|
|
id?: number // assistant message backend pk (for deletion)
|
|
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("")
|
|
let _onCodeComplete:
|
|
| ((
|
|
code: { html: string | null; css: string | null; js: string | null },
|
|
messageId: number
|
|
) => void)
|
|
| null = null
|
|
|
|
export function setOnCodeComplete(fn: typeof _onCodeComplete) {
|
|
_onCodeComplete = fn
|
|
}
|
|
|
|
let ws: WebSocket | null = null
|
|
let _currentTaskId = 0
|
|
|
|
export function connectPrompt(taskId: number) {
|
|
_currentTaskId = taskId
|
|
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") {
|
|
streaming.value = false
|
|
streamingContent.value = ""
|
|
const alreadyLoaded = conversationId.value === data.conversation_id
|
|
conversationId.value = data.conversation_id
|
|
if (!alreadyLoaded) {
|
|
messages.value = data.messages || []
|
|
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
|
|
messages.value.push({
|
|
role: "assistant",
|
|
content: streamingContent.value,
|
|
id: data.message_id,
|
|
code: data.code,
|
|
})
|
|
streamingContent.value = ""
|
|
if (data.code) {
|
|
applyCode(data.code)
|
|
if (_onCodeComplete) {
|
|
_onCodeComplete(data.code, data.message_id)
|
|
}
|
|
}
|
|
} 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, model: string = "") {
|
|
if (!ws || ws.readyState !== WebSocket.OPEN) return
|
|
streaming.value = true
|
|
messages.value.push({ role: "user", content })
|
|
ws.send(JSON.stringify({ type: "message", content, model }))
|
|
}
|
|
|
|
export function stopPrompt() {
|
|
if (
|
|
messages.value.length > 0 &&
|
|
messages.value[messages.value.length - 1].role === "user"
|
|
) {
|
|
messages.value.pop()
|
|
}
|
|
streaming.value = false
|
|
streamingContent.value = ""
|
|
if (_currentTaskId) {
|
|
connectPrompt(_currentTaskId)
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|