Add Completion Endpoint

src/index.ts

@@ -7,10 +7,14 @@ import { google } from '@ai-sdk/google'
 import { openai } from '@ai-sdk/openai'
 import { createOpenAICompatible } from '@ai-sdk/openai-compatible'
 import {
+  Output,
+  consumeStream,
+  pipeTextStreamToResponse,
   streamText,
   createUIMessageStream,
   pipeUIMessageStreamToResponse,
   convertToModelMessages,
+  toTextStream,
   toUIMessageStream,
 } from 'ai'
 
@@ -42,7 +46,11 @@ const maxOutputTokens = process.env.MAX_TOKENS
   ? Number.parseInt(process.env.MAX_TOKENS)
   : undefined
 console.log(`maxOutputTokens: ${maxOutputTokens}`)
-console.log(`INSTRUCTIONS: ${process.env.INSTRUCTIONS}`)
+console.log(
+  `CHAT_INSTRUCTIONS: ${process.env.INSTRUCTIONS || process.env.CHAT_INSTRUCTIONS}`, // NOSONAR
+)
+console.log(`COMPLETION_INSTRUCTIONS: ${process.env.COMPLETION_INSTRUCTIONS}`)
+console.log(`OBJECT_INSTRUCTIONS: ${process.env.OBJECT_INSTRUCTIONS}`)
 
 const model = getModel()
 console.log(`Loaded modelId: ${model.modelId}`)
@@ -61,7 +69,7 @@ app.listen(port, () => console.log(`Listening on PORT: ${port}`))
 
 // app.get('/app-health-check', (_req, res) => res.sendStatus(200))
 
-app.post('/', async (req: Request, res: Response) => {
+app.post(['/', '/chat'], async (req: Request, res: Response) => {
   // console.log('req.headers:', req.headers)
   // console.log('authorization:', req.headers.authorization)
   const { messages, system } = req.body
@@ -73,7 +81,7 @@ app.post('/', async (req: Request, res: Response) => {
       const result = streamText({
         model: model,
         messages: modelMessages,
-        system: system || process.env.INSTRUCTIONS,
+        system: system || process.env.INSTRUCTIONS || process.env.CHAT_INSTRUCTIONS,
         maxOutputTokens,
         providerOptions,
       })
@@ -84,6 +92,63 @@ app.post('/', async (req: Request, res: Response) => {
   pipeUIMessageStreamToResponse({ response: res, stream })
 })
 
+app.post('/completion', async (req: Request, res: Response) => {
+  const { prompt, system } = req.body
+  console.log('prompt:', prompt?.length, 'system:', system?.length)
+  const result = streamText({
+    model: model,
+    prompt,
+    system: system || process.env.COMPLETION_INSTRUCTIONS,
+    maxOutputTokens,
+    providerOptions,
+    onError(error) {
+      console.log('error:', error)
+    },
+    onEnd({ finalStep, finishReason, text, usage }) {
+      console.log('reasoning:', finalStep.reasoningText)
+      console.log('response:', text)
+      console.log('usage:', usage)
+      console.log('finishReason:', finishReason)
+    },
+  })
+  const stream = createUIMessageStream({
+    execute: ({ writer }) => {
+      writer.merge(toUIMessageStream({ stream: result.stream }))
+    },
+  })
+  pipeUIMessageStreamToResponse({
+    response: res,
+    stream,
+    consumeSseStream: consumeStream,
+  })
+})
+
+app.post('/object', async (req: Request, res: Response) => {
+  const { outputSchema, prompt, system } = req.body
+  console.log('prompt:', prompt?.length, 'system:', system?.length)
+  const result = streamText({
+    model: model,
+    prompt,
+    system: system || process.env.OBJECT_INSTRUCTIONS,
+    maxOutputTokens,
+    providerOptions,
+    output: outputSchema ? Output.object({ schema: outputSchema }) : Output.json(),
+    onError(error) {
+      console.log('error:', error)
+    },
+    onEnd({ finalStep, finishReason, text, usage }) {
+      console.log('reasoning:', finalStep.reasoningText)
+      console.log('response:', text)
+      console.log('usage:', usage)
+      console.log('finishReason:', finishReason)
+    },
+  })
+  pipeTextStreamToResponse({
+    response: res,
+    stream: toTextStream({ stream: result.stream }),
+  })
+})
+
 function corsCallback(
   origin: string | undefined,
   callback: (err: Error | null, origin?: boolean) => void,