51 lines
1.5 KiB
TypeScript
51 lines
1.5 KiB
TypeScript
import { streamText } from 'ai'
|
|
import { z } from 'zod'
|
|
import { zodToJsonSchema } from 'zod-to-json-schema'
|
|
|
|
import { languagePrompt, systemPrompt } from './prompt'
|
|
import { useAiModel } from '~/composables/useAiProvider'
|
|
|
|
type PartialFeedback = DeepPartial<z.infer<typeof feedbackTypeSchema>>
|
|
|
|
export const feedbackTypeSchema = z.object({
|
|
questions: z.array(z.string()),
|
|
})
|
|
|
|
export function generateFeedback({
|
|
query,
|
|
language,
|
|
numQuestions = 3,
|
|
}: {
|
|
query: string
|
|
language: string
|
|
numQuestions?: number
|
|
}) {
|
|
const schema = z.object({
|
|
questions: z
|
|
.array(z.string())
|
|
.describe(`Follow up questions to clarify the research direction`),
|
|
})
|
|
const jsonSchema = JSON.stringify(zodToJsonSchema(schema))
|
|
const prompt = [
|
|
`Given the following query from the user, ask ${numQuestions} follow up questions to clarify the research direction. Return a maximum of ${numQuestions} questions, but feel free to return less if the original query is clear: <query>${query}</query>`,
|
|
`You MUST respond in JSON matching this JSON schema: ${jsonSchema}`,
|
|
languagePrompt(language),
|
|
].join('\n\n')
|
|
|
|
const stream = streamText({
|
|
model: useAiModel(),
|
|
system: systemPrompt(),
|
|
prompt,
|
|
onError({ error }) {
|
|
console.error(`generateFeedback`, error)
|
|
throw error instanceof Error ? error : new Error(String(error))
|
|
},
|
|
})
|
|
|
|
return parseStreamingJson(
|
|
stream.fullStream,
|
|
feedbackTypeSchema,
|
|
(value: PartialFeedback) => !!value.questions && value.questions.length > 0,
|
|
)
|
|
}
|