sdk: implement promptAiSdkStream

Author: charleslien

backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts

@@ -140,14 +140,16 @@ export async function* promptAiSdkStream(
       return null
     }
     if (chunk.type === 'reasoning-delta') {
-      if (
-        (
-          params.providerOptions?.openrouter as
-            | OpenRouterProviderOptions
-            | undefined
-        )?.reasoning?.exclude
-      ) {
-        continue
+      for (const provider of ['openrouter', 'codebuff'] as const) {
+        if (
+          (
+            params.providerOptions?.[provider] as
+              | OpenRouterProviderOptions
+              | undefined
+          )?.reasoning?.exclude
+        ) {
+          continue
+        }
       }
       yield {
         type: 'reasoning',

common/src/types/contracts/database.ts

@@ -28,7 +28,7 @@ type AgentRun = {
 }
 export type AgentRunColumn = keyof AgentRun
 export type GetAgentRunFromIdInput<T extends AgentRunColumn> = {
-  agentRunId: string
+  runId: string
   userId: string
   fields: readonly T[]
 }

common/src/types/contracts/llm.ts

@@ -22,6 +22,7 @@ export type StreamChunk =
 
 export type PromptAiSdkStreamFn = (
   params: {
+    apiKey: string
     messages: Message[]
     clientSessionId: string
     fingerprintId: string

common/src/util/messages.ts

@@ -26,45 +26,35 @@ export function withCacheControl<
   if (!wrapper.providerOptions) {
     wrapper.providerOptions = {}
   }
-  if (!wrapper.providerOptions.anthropic) {
-    wrapper.providerOptions.anthropic = {}
-  }
-  wrapper.providerOptions.anthropic.cacheControl = { type: 'ephemeral' }
-  if (!wrapper.providerOptions.openrouter) {
-    wrapper.providerOptions.openrouter = {}
+
+  for (const provider of ['anthropic', 'openrouter', 'codebuff'] as const) {
+    if (!wrapper.providerOptions[provider]) {
+      wrapper.providerOptions[provider] = {}
+    }
+    wrapper.providerOptions[provider].cacheControl = { type: 'ephemeral' }
   }
-  wrapper.providerOptions.openrouter.cacheControl = { type: 'ephemeral' }
+
   return wrapper
 }
 
 export function withoutCacheControl<
   T extends { providerOptions?: ProviderMetadata },
 >(obj: T): T {
   const wrapper = cloneDeep(obj)
-  if (has(wrapper.providerOptions?.anthropic?.cacheControl, 'type')) {
-    delete wrapper.providerOptions?.anthropic?.cacheControl?.type
-  }
-  if (
-    Object.keys(wrapper.providerOptions?.anthropic?.cacheControl ?? {})
-      .length === 0
-  ) {
-    delete wrapper.providerOptions?.anthropic?.cacheControl
-  }
-  if (Object.keys(wrapper.providerOptions?.anthropic ?? {}).length === 0) {
-    delete wrapper.providerOptions?.anthropic
-  }
 
-  if (has(wrapper.providerOptions?.openrouter?.cacheControl, 'type')) {
-    delete wrapper.providerOptions?.openrouter?.cacheControl?.type
-  }
-  if (
-    Object.keys(wrapper.providerOptions?.openrouter?.cacheControl ?? {})
-      .length === 0
-  ) {
-    delete wrapper.providerOptions?.openrouter?.cacheControl
-  }
-  if (Object.keys(wrapper.providerOptions?.openrouter ?? {}).length === 0) {
-    delete wrapper.providerOptions?.openrouter
+  for (const provider of ['anthropic', 'openrouter', 'codebuff'] as const) {
+    if (has(wrapper.providerOptions?.[provider]?.cacheControl, 'type')) {
+      delete wrapper.providerOptions?.[provider]?.cacheControl?.type
+    }
+    if (
+      Object.keys(wrapper.providerOptions?.[provider]?.cacheControl ?? {})
+        .length === 0
+    ) {
+      delete wrapper.providerOptions?.[provider]?.cacheControl
+    }
+    if (Object.keys(wrapper.providerOptions?.[provider] ?? {}).length === 0) {
+      delete wrapper.providerOptions?.[provider]
+    }
   }
 
   if (Object.keys(wrapper.providerOptions ?? {}).length === 0) {

packages/agent-runtime/src/prompt-agent-stream.ts

@@ -1,5 +1,3 @@
-import { providerModelNames } from '@codebuff/common/old-constants'
-
 import { globalStopSequence } from './constants'
 
 import type { AgentTemplate } from './templates/types'
@@ -16,6 +14,7 @@ import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type { OpenRouterProviderOptions } from '@codebuff/internal/openrouter-ai-sdk'
 
 export const getAgentStreamFromTemplate = (params: {
+  apiKey: string
   clientSessionId: string
   fingerprintId: string
   userInputId: string
@@ -33,6 +32,7 @@ export const getAgentStreamFromTemplate = (params: {
   trackEvent: TrackEventFn
 }) => {
   const {
+    apiKey,
     clientSessionId,
     fingerprintId,
     userInputId,
@@ -57,6 +57,7 @@ export const getAgentStreamFromTemplate = (params: {
 
   const getStream = (messages: Message[]) => {
     const aiSdkStreamParams: ParamsOf<PromptAiSdkStreamFn> = {
+      apiKey,
       messages,
       model,
       stopSequences: [globalStopSequence],
@@ -78,18 +79,18 @@ export const getAgentStreamFromTemplate = (params: {
 
     // Add Gemini-specific options if needed
     const primaryModel = Array.isArray(model) ? model[0] : model
-    const provider =
-      providerModelNames[primaryModel as keyof typeof providerModelNames]
 
     if (!aiSdkStreamParams.providerOptions) {
       aiSdkStreamParams.providerOptions = {}
     }
-    if (!aiSdkStreamParams.providerOptions.openrouter) {
-      aiSdkStreamParams.providerOptions.openrouter = {}
+    for (const provider of ['openrouter', 'codebuff'] as const) {
+      if (!aiSdkStreamParams.providerOptions[provider]) {
+        aiSdkStreamParams.providerOptions[provider] = {}
+      }
+      ;(
+        aiSdkStreamParams.providerOptions[provider] as OpenRouterProviderOptions
+      ).reasoning = template.reasoningOptions
     }
-    ;(
-      aiSdkStreamParams.providerOptions.openrouter as OpenRouterProviderOptions
-    ).reasoning = template.reasoningOptions
 
     return promptAiSdkStream(aiSdkStreamParams)
   }

packages/agent-runtime/src/util/messages.ts

@@ -261,8 +261,9 @@ export function getMessagesSubset(params: {
 
   // Remove cache_control from all messages
   for (const message of messagesSubset) {
-    delete message.providerOptions?.anthropic?.cacheControl
-    delete message.providerOptions?.openrouter?.cacheControl
+    for (const provider of ['anthropic', 'openrouter', 'codebuff'] as const) {
+      delete message.providerOptions?.[provider]?.cacheControl
+    }
   }
 
   // Cache up to the last message!

sdk/src/impl/agent-runtime.ts

@@ -8,6 +8,7 @@ import {
   getUserInfoFromApiKey,
   startAgentRun,
 } from './database'
+import { promptAiSdkStream } from './llm'
 
 import type {
   AgentRuntimeDeps,
@@ -20,7 +21,6 @@ export function getAgentRuntimeImpl(params: {
   apiKey: string
 }): Omit<
   AgentRuntimeDeps & AgentRuntimeScopedDeps,
-  | 'promptAiSdkStream'
   | 'promptAiSdk'
   | 'promptAiSdkStructured'
   | 'handleStepsLogChunk'
@@ -48,7 +48,7 @@ export function getAgentRuntimeImpl(params: {
       }),
 
     // LLM
-    // promptAiSdkStream: PromptAiSdkStreamFn,
+    promptAiSdkStream,
     // promptAiSdk: PromptAiSdkFn,
     // promptAiSdkStructured: PromptAiSdkStructuredFn,