115 lines
2.8 KiB
TypeScript
115 lines
2.8 KiB
TypeScript
import nock from "nock"
|
|
import { MockLLMResponseFn, MockLLMResponseOpts } from "."
|
|
import _ from "lodash"
|
|
import { ai } from "@budibase/pro"
|
|
|
|
let chatID = 1
|
|
const SPACE_REGEX = /\s+/g
|
|
|
|
interface Message {
|
|
role: string
|
|
content: string
|
|
}
|
|
|
|
interface Choice {
|
|
index: number
|
|
message: Message
|
|
logprobs: null
|
|
finish_reason: string
|
|
}
|
|
|
|
interface CompletionTokensDetails {
|
|
reasoning_tokens: number
|
|
accepted_prediction_tokens: number
|
|
rejected_prediction_tokens: number
|
|
}
|
|
|
|
interface Usage {
|
|
prompt_tokens: number
|
|
completion_tokens: number
|
|
total_tokens: number
|
|
completion_tokens_details: CompletionTokensDetails
|
|
}
|
|
|
|
interface ChatCompletionRequest {
|
|
messages: Message[]
|
|
model: string
|
|
}
|
|
|
|
interface ChatCompletionResponse {
|
|
id: string
|
|
object: string
|
|
created: number
|
|
model: string
|
|
system_fingerprint: string
|
|
choices: Choice[]
|
|
usage: Usage
|
|
}
|
|
|
|
export const mockChatGPTResponse: MockLLMResponseFn = (
|
|
answer: string | ((prompt: string) => string),
|
|
opts?: MockLLMResponseOpts
|
|
) => {
|
|
let body: any = undefined
|
|
|
|
if (opts?.format) {
|
|
body = _.matches({
|
|
response_format: ai.openai.parseResponseFormat(opts.format),
|
|
})
|
|
}
|
|
return nock(opts?.host || "https://api.openai.com")
|
|
.post("/v1/chat/completions", body)
|
|
.reply((uri: string, body: nock.Body) => {
|
|
const req = body as ChatCompletionRequest
|
|
const messages = req.messages
|
|
const prompt = messages[0].content
|
|
|
|
let content
|
|
if (typeof answer === "function") {
|
|
try {
|
|
content = answer(prompt)
|
|
} catch (e) {
|
|
return [500, "Internal Server Error"]
|
|
}
|
|
} else {
|
|
content = answer
|
|
}
|
|
|
|
chatID++
|
|
|
|
// We mock token usage because we use it to calculate Budibase AI quota
|
|
// usage when Budibase AI is enabled, and some tests assert against quota
|
|
// usage to make sure we're tracking correctly.
|
|
const prompt_tokens = messages[0].content.split(SPACE_REGEX).length
|
|
const completion_tokens = content.split(SPACE_REGEX).length
|
|
|
|
const response: ChatCompletionResponse = {
|
|
id: `chatcmpl-${chatID}`,
|
|
object: "chat.completion",
|
|
created: Math.floor(Date.now() / 1000),
|
|
model: req.model,
|
|
system_fingerprint: `fp_${chatID}`,
|
|
choices: [
|
|
{
|
|
index: 0,
|
|
message: { role: "assistant", content },
|
|
logprobs: null,
|
|
finish_reason: "stop",
|
|
},
|
|
],
|
|
usage: {
|
|
prompt_tokens,
|
|
completion_tokens,
|
|
total_tokens: prompt_tokens + completion_tokens,
|
|
completion_tokens_details: {
|
|
reasoning_tokens: 0,
|
|
accepted_prediction_tokens: 0,
|
|
rejected_prediction_tokens: 0,
|
|
},
|
|
},
|
|
}
|
|
return [200, response]
|
|
})
|
|
.persist()
|
|
}
|