add promptAiSdkStructured to AgentRuntimeDeps

Author: charleslien

backend/src/impl/agent-runtime.ts

@@ -2,18 +2,23 @@ import { addAgentStep, finishAgentRun, startAgentRun } from '../agent-run'
 import {
   promptAiSdk,
   promptAiSdkStream,
+  promptAiSdkStructured,
 } from '../llm-apis/vercel-ai-sdk/ai-sdk'
 import { logger } from '../util/logger'
 
 import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
 
 export const BACKEND_AGENT_RUNTIME_IMPL: AgentRuntimeDeps = Object.freeze({
-  logger,
-
+  // Database
   startAgentRun,
   finishAgentRun,
   addAgentStep,
 
+  // LLM
   promptAiSdkStream,
   promptAiSdk,
+  promptAiSdkStructured,
+
+  // Other
+  logger,
 })

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

@@ -20,10 +20,10 @@ import type { Model, OpenAIModel } from '@codebuff/common/old-constants'
 import type {
   PromptAiSdkFn,
   PromptAiSdkStreamFn,
+  PromptAiSdkStructuredInput,
+  PromptAiSdkStructuredOutput,
 } from '@codebuff/common/types/contracts/llm'
-import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsOf } from '@codebuff/common/types/function-params'
-import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type {
   OpenRouterProviderOptions,
   OpenRouterUsageAccounting,
@@ -287,24 +287,9 @@ export const promptAiSdk = async function (
 }
 
 // Copied over exactly from promptAiSdk but with a schema
-export const promptAiSdkStructured = async function <T>(params: {
-  messages: Message[]
-  schema: z.ZodType<T>
-  clientSessionId: string
-  fingerprintId: string
-  userInputId: string
-  model: Model
-  userId: string | undefined
-  maxTokens?: number
-  temperature?: number
-  timeout?: number
-  chargeUser?: boolean
-  agentId?: string
-  onCostCalculated?: (credits: number) => Promise<void>
-  includeCacheControl?: boolean
-  maxRetries?: number
-  logger: Logger
-}): Promise<T> {
+export const promptAiSdkStructured = async function <T>(
+  params: PromptAiSdkStructuredInput<T>,
+): PromptAiSdkStructuredOutput<T> {
   const { logger } = params
 
   if (!checkLiveUserInput(params)) {

common/src/testing/impl/agent-runtime.ts

@@ -21,6 +21,9 @@ export const TEST_AGENT_RUNTIME_IMPL: AgentRuntimeDeps = Object.freeze({
   promptAiSdk: async function () {
     throw new Error('promptAiSdk not implemented in test runtime')
   },
+  promptAiSdkStructured: async function () {
+    throw new Error('promptAiSdkStructured not implemented in test runtime')
+  },
 
   // Other
   logger: testLogger,

common/src/types/contracts/agent-runtime.ts

@@ -3,7 +3,11 @@ import type {
   FinishAgentRunFn,
   StartAgentRunFn,
 } from './database'
-import type { PromptAiSdkFn, PromptAiSdkStreamFn } from './llm'
+import type {
+  PromptAiSdkFn,
+  PromptAiSdkStreamFn,
+  PromptAiSdkStructuredFn,
+} from './llm'
 import type { Logger } from './logger'
 
 export type AgentRuntimeDeps = {
@@ -15,6 +19,7 @@ export type AgentRuntimeDeps = {
   // LLM
   promptAiSdkStream: PromptAiSdkStreamFn
   promptAiSdk: PromptAiSdkFn
+  promptAiSdkStructured: PromptAiSdkStructuredFn
 
   // Other
   logger: Logger

common/src/types/contracts/llm.ts

@@ -3,6 +3,7 @@ import type { Logger } from './logger'
 import type { Model } from '../../old-constants'
 import type { Message } from '../messages/codebuff-message'
 import type { generateText, streamText } from 'ai'
+import type z from 'zod/v4'
 
 export type StreamChunk =
   | {
@@ -49,3 +50,26 @@ export type PromptAiSdkFn = (
     logger: Logger
   } & ParamsExcluding<typeof generateText, 'model' | 'messages'>,
 ) => Promise<string>
+
+export type PromptAiSdkStructuredInput<T> = {
+  messages: Message[]
+  schema: z.ZodType<T>
+  clientSessionId: string
+  fingerprintId: string
+  userInputId: string
+  model: Model
+  userId: string | undefined
+  maxTokens?: number
+  temperature?: number
+  timeout?: number
+  chargeUser?: boolean
+  agentId?: string
+  onCostCalculated?: (credits: number) => Promise<void>
+  includeCacheControl?: boolean
+  maxRetries?: number
+  logger: Logger
+}
+export type PromptAiSdkStructuredOutput<T> = Promise<T>
+export type PromptAiSdkStructuredFn = <T>(
+  params: PromptAiSdkStructuredInput<T>,
+) => PromptAiSdkStructuredOutput<T>

evals/impl/agent-runtime.ts

@@ -13,4 +13,7 @@ export const EVALS_AGENT_RUNTIME_IMPL: AgentRuntimeDeps = Object.freeze({
   promptAiSdk: async function () {
     throw new Error('promptAiSdk not implemented in eval runtime')
   },
+  promptAiSdkStructured: async function () {
+    throw new Error('promptAiSdkStructured not implemented in eval runtime')
+  },
 })