|
| 1 | +import { ApolloLink } from "@apollo/client"; |
| 2 | +import { type LanguageModel, generateObject } from "ai"; |
| 3 | +import { AIAdapter } from "@apollo/client-ai"; |
| 4 | +import { isFormattedExecutionResult } from "@apollo/client/utilities"; |
| 5 | + |
| 6 | +namespace VercelAIAdapter { |
| 7 | + export interface Options extends AIAdapter.Options { |
| 8 | + model: LanguageModel; |
| 9 | + } |
| 10 | +} |
| 11 | + |
| 12 | +type GenerateObjectOptions = { |
| 13 | + model: LanguageModel; |
| 14 | + mode: "json"; |
| 15 | + prompt: string; |
| 16 | + system?: string; |
| 17 | + output: "no-schema"; |
| 18 | +}; |
| 19 | + |
| 20 | +export class VercelAIAdapter extends AIAdapter { |
| 21 | + public model: LanguageModel; |
| 22 | + |
| 23 | + constructor(options: VercelAIAdapter.Options) { |
| 24 | + super(options); |
| 25 | + |
| 26 | + this.model = options.model; |
| 27 | + } |
| 28 | + |
| 29 | + public async generateResponseForOperation( |
| 30 | + operation: ApolloLink.Operation, |
| 31 | + prompt: string |
| 32 | + ): Promise<AIAdapter.Result> { |
| 33 | + const promptOptions: GenerateObjectOptions = { |
| 34 | + mode: "json", |
| 35 | + model: this.model, |
| 36 | + prompt: this.createPrompt(operation, prompt), |
| 37 | + system: this.systemPrompt, |
| 38 | + output: "no-schema", |
| 39 | + }; |
| 40 | + |
| 41 | + return generateObject(promptOptions).then( |
| 42 | + ({ object: result, finishReason, usage, warnings }) => { |
| 43 | + if (!result || typeof result !== "object") { |
| 44 | + return { data: null }; |
| 45 | + } |
| 46 | + // Type guard to ensure result is a valid FormattedExecutionResult |
| 47 | + if (isFormattedExecutionResult(result)) { |
| 48 | + return result; |
| 49 | + } |
| 50 | + // Fallback: wrap in data property if not a valid execution result |
| 51 | + return { data: result }; |
| 52 | + } |
| 53 | + ); |
| 54 | + } |
| 55 | +} |
0 commit comments