diff --git a/.gitignore b/.gitignore
index 423ab4d6f10..71e424acf56 100644
--- a/.gitignore
+++ b/.gitignore
@@ -38,4 +38,3 @@ characters/
packages/core/src/providers/cache
packages/core/src/providers/cache/*
cache/*
-packages/core/cache/*
\ No newline at end of file
diff --git a/package.json b/package.json
index 808745a33bf..0310a66e308 100644
--- a/package.json
+++ b/package.json
@@ -35,5 +35,6 @@
"ollama-ai-provider": "^0.16.1",
"optional": "^0.1.4",
"sharp": "^0.33.5"
- }
+ },
+ "packageManager": "pnpm@9.12.3+sha512.cce0f9de9c5a7c95bef944169cc5dfe8741abfb145078c0d508b868056848a87c81e626246cb60967cbd7fd29a6c062ef73ff840d96b3c86c40ac92cf4a813ee"
}
diff --git a/packages/core/package.json b/packages/core/package.json
index 2e4a63deab2..0c3312ba291 100644
--- a/packages/core/package.json
+++ b/packages/core/package.json
@@ -26,6 +26,7 @@
"@rollup/plugin-typescript": "11.1.6",
"@types/fluent-ffmpeg": "2.1.27",
"@types/jest": "29.5.14",
+ "@types/mocha": "^10.0.9",
"@types/node": "22.8.4",
"@types/pdfjs-dist": "^2.10.378",
"@types/tar": "6.1.13",
diff --git a/packages/core/src/generation.ts b/packages/core/src/generation.ts
index 07909847b21..6c9e1d59477 100644
--- a/packages/core/src/generation.ts
+++ b/packages/core/src/generation.ts
@@ -2,8 +2,12 @@ import { createAnthropic } from "@ai-sdk/anthropic";
import { createGroq } from "@ai-sdk/groq";
import { createOpenAI } from "@ai-sdk/openai";
import { getModel } from "./models.ts";
+import {
+ generateText as aiGenerateText,
+ generateObject as aiGenerateObject,
+ GenerateObjectResult,
+} from "ai";
import { IImageDescriptionService, ModelClass, Service } from "./types.ts";
-import { generateText as aiGenerateText } from "ai";
import { Buffer } from "buffer";
import { createOllama } from "ollama-ai-provider";
import OpenAI from "openai";
@@ -26,6 +30,7 @@ import {
ModelProviderName,
ServiceType,
} from "./types.ts";
+import { ZodSchema } from "zod";
/**
* Send a message to the model for a text generateText - receive a string back and parse how you'd like
@@ -748,3 +753,368 @@ export const generateCaption = async (
description: resp.description.trim(),
};
};
+
+/**
+ * Configuration options for generating objects with a model.
+ */
+export interface GenerationOptions {
+ runtime: IAgentRuntime;
+ context: string;
+ modelClass: ModelClass;
+ schema?: ZodSchema;
+ schemaName?: string;
+ schemaDescription?: string;
+ stop?: string[];
+ mode?: "auto" | "json" | "tool";
+ experimental_providerMetadata?: Record<string, unknown>;
+}
+
+/**
+ * Base settings for model generation.
+ */
+interface ModelSettings {
+ prompt: string;
+ temperature: number;
+ maxTokens: number;
+ frequencyPenalty: number;
+ presencePenalty: number;
+ stop?: string[];
+}
+
+/**
+ * Generates structured objects from a prompt using specified AI models and configuration options.
+ *
+ * @param {GenerationOptions} options - Configuration options for generating objects.
+ * @returns {Promise<any[]>} - A promise that resolves to an array of generated objects.
+ * @throws {Error} - Throws an error if the provider is unsupported or if generation fails.
+ */
+export const generateObjectV2 = async ({
+ runtime,
+ context,
+ modelClass,
+ schema,
+ schemaName,
+ schemaDescription,
+ stop,
+ mode = "json",
+}: GenerationOptions): Promise<GenerateObjectResult<unknown>> => {
+ if (!context) {
+ const errorMessage = "generateObject context is empty";
+ console.error(errorMessage);
+ throw new Error(errorMessage);
+ }
+
+ const provider = runtime.modelProvider;
+ const model = models[provider].model[modelClass];
+ const temperature = models[provider].settings.temperature;
+ const frequency_penalty = models[provider].settings.frequency_penalty;
+ const presence_penalty = models[provider].settings.presence_penalty;
+ const max_context_length = models[provider].settings.maxInputTokens;
+ const max_response_length = models[provider].settings.maxOutputTokens;
+ const apiKey = runtime.token;
+
+ try {
+ context = await trimTokens(context, max_context_length, modelClass);
+
+ const modelOptions: ModelSettings = {
+ prompt: context,
+ temperature,
+ maxTokens: max_response_length,
+ frequencyPenalty: frequency_penalty,
+ presencePenalty: presence_penalty,
+ stop: stop || models[provider].settings.stop,
+ };
+
+ const response = await handleProvider({
+ provider,
+ model,
+ apiKey,
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ modelOptions,
+ runtime,
+ context,
+ modelClass,
+ });
+
+ return response;
+ } catch (error) {
+ console.error("Error in generateObject:", error);
+ throw error;
+ }
+};
+
+/**
+ * Interface for provider-specific generation options.
+ */
+interface ProviderOptions {
+ runtime: IAgentRuntime;
+ provider: ModelProviderName;
+ model: any;
+ apiKey: string;
+ schema?: ZodSchema;
+ schemaName?: string;
+ schemaDescription?: string;
+ mode?: "auto" | "json" | "tool";
+ experimental_providerMetadata?: Record<string, unknown>;
+ modelOptions: ModelSettings;
+ modelClass: string;
+ context: string;
+}
+
+/**
+ * Handles AI generation based on the specified provider.
+ *
+ * @param {ProviderOptions} options - Configuration options specific to the provider.
+ * @returns {Promise<any[]>} - A promise that resolves to an array of generated objects.
+ */
+export async function handleProvider(
+ options: ProviderOptions
+): Promise<GenerateObjectResult<unknown>> {
+ const { provider, runtime, context, modelClass } = options;
+ switch (provider) {
+ case ModelProviderName.OPENAI:
+ case ModelProviderName.LLAMACLOUD:
+ return await handleOpenAI(options);
+ case ModelProviderName.ANTHROPIC:
+ return await handleAnthropic(options);
+ case ModelProviderName.GROK:
+ return await handleGrok(options);
+ case ModelProviderName.GROQ:
+ return await handleGroq(options);
+ case ModelProviderName.LLAMALOCAL:
+ return await generateObject({
+ runtime,
+ context,
+ modelClass,
+ });
+ case ModelProviderName.GOOGLE:
+ return await handleGoogle(options);
+ case ModelProviderName.REDPILL:
+ return await handleRedPill(options);
+ case ModelProviderName.OPENROUTER:
+ return await handleOpenRouter(options);
+ case ModelProviderName.OLLAMA:
+ return await handleOllama(options);
+ default: {
+ const errorMessage = `Unsupported provider: ${provider}`;
+ elizaLogger.error(errorMessage);
+ throw new Error(errorMessage);
+ }
+ }
+}
+/**
+ * Handles object generation for OpenAI.
+ *
+ * @param {ProviderOptions} options - Options specific to OpenAI.
+ * @returns {Promise<GenerateObjectResult<unknown>>} - A promise that resolves to generated objects.
+ */
+async function handleOpenAI({
+ model,
+ apiKey,
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ modelOptions,
+}: ProviderOptions): Promise<GenerateObjectResult<unknown>> {
+ const openai = createOpenAI({ apiKey });
+ return await aiGenerateObject({
+ model: openai.languageModel(model),
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ ...modelOptions,
+ });
+}
+
+/**
+ * Handles object generation for Anthropic models.
+ *
+ * @param {ProviderOptions} options - Options specific to Anthropic.
+ * @returns {Promise<GenerateObjectResult<unknown>>} - A promise that resolves to generated objects.
+ */
+async function handleAnthropic({
+ model,
+ apiKey,
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ modelOptions,
+}: ProviderOptions): Promise<GenerateObjectResult<unknown>> {
+ const anthropic = createAnthropic({ apiKey });
+ return await aiGenerateObject({
+ model: anthropic.languageModel(model),
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ ...modelOptions,
+ });
+}
+
+/**
+ * Handles object generation for Grok models.
+ *
+ * @param {ProviderOptions} options - Options specific to Grok.
+ * @returns {Promise<GenerateObjectResult<unknown>>} - A promise that resolves to generated objects.
+ */
+async function handleGrok({
+ model,
+ apiKey,
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ modelOptions,
+}: ProviderOptions): Promise<GenerateObjectResult<unknown>> {
+ const grok = createOpenAI({ apiKey, baseURL: models.grok.endpoint });
+ return await aiGenerateObject({
+ model: grok.languageModel(model, { parallelToolCalls: false }),
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ ...modelOptions,
+ });
+}
+
+/**
+ * Handles object generation for Groq models.
+ *
+ * @param {ProviderOptions} options - Options specific to Groq.
+ * @returns {Promise<GenerateObjectResult<unknown>>} - A promise that resolves to generated objects.
+ */
+async function handleGroq({
+ model,
+ apiKey,
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ modelOptions,
+}: ProviderOptions): Promise<GenerateObjectResult<unknown>> {
+ const groq = createGroq({ apiKey });
+ return await aiGenerateObject({
+ model: groq.languageModel(model),
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ ...modelOptions,
+ });
+}
+
+/**
+ * Handles object generation for Google models.
+ *
+ * @param {ProviderOptions} options - Options specific to Google.
+ * @returns {Promise<GenerateObjectResult<unknown>>} - A promise that resolves to generated objects.
+ */
+async function handleGoogle({
+ model,
+ apiKey: _apiKey,
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ modelOptions,
+}: ProviderOptions): Promise<GenerateObjectResult<unknown>> {
+ const google = createGoogleGenerativeAI();
+ return await aiGenerateObject({
+ model: google(model),
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ ...modelOptions,
+ });
+}
+
+/**
+ * Handles object generation for Redpill models.
+ *
+ * @param {ProviderOptions} options - Options specific to Redpill.
+ * @returns {Promise<GenerateObjectResult<unknown>>} - A promise that resolves to generated objects.
+ */
+async function handleRedPill({
+ model,
+ apiKey,
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ modelOptions,
+}: ProviderOptions): Promise<GenerateObjectResult<unknown>> {
+ const redPill = createOpenAI({ apiKey, baseURL: models.redpill.endpoint });
+ return await aiGenerateObject({
+ model: redPill.languageModel(model),
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ ...modelOptions,
+ });
+}
+
+/**
+ * Handles object generation for OpenRouter models.
+ *
+ * @param {ProviderOptions} options - Options specific to OpenRouter.
+ * @returns {Promise<GenerateObjectResult<unknown>>} - A promise that resolves to generated objects.
+ */
+async function handleOpenRouter({
+ model,
+ apiKey,
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ modelOptions,
+}: ProviderOptions): Promise<GenerateObjectResult<unknown>> {
+ const openRouter = createOpenAI({
+ apiKey,
+ baseURL: models.openrouter.endpoint,
+ });
+ return await aiGenerateObject({
+ model: openRouter.languageModel(model),
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ ...modelOptions,
+ });
+}
+
+/**
+ * Handles object generation for Ollama models.
+ *
+ * @param {ProviderOptions} options - Options specific to Ollama.
+ * @returns {Promise<GenerateObjectResult<unknown>>} - A promise that resolves to generated objects.
+ */
+async function handleOllama({
+ model,
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ modelOptions,
+ provider,
+}: ProviderOptions): Promise<GenerateObjectResult<unknown>> {
+ const ollamaProvider = createOllama({
+ baseURL: models[provider].endpoint + "/api",
+ });
+ const ollama = ollamaProvider(model);
+ return await aiGenerateObject({
+ model: ollama,
+ schema,
+ schemaName,
+ schemaDescription,
+ mode,
+ ...modelOptions,
+ });
+}
diff --git a/packages/core/src/tests/generation.test.ts b/packages/core/src/tests/generation.test.ts
new file mode 100644
index 00000000000..b5230259254
--- /dev/null
+++ b/packages/core/src/tests/generation.test.ts
@@ -0,0 +1,166 @@
+import {
+ generateObject,
+ GenerationOptions,
+ trimTokens,
+ handleProvider,
+} from "../generation";
+import { createRuntime } from "../test_resources/createRuntime";
+import { ModelProviderName, ModelClass } from "../types";
+import { ZodSchema } from "zod";
+import dotenv from "dotenv";
+
+dotenv.config({ path: ".dev.vars" });
+
+describe("generateObject", () => {
+ let runtime;
+
+ beforeAll(async () => {
+ // Create runtime with a mock environment
+ const setup = await createRuntime({
+ env: process.env as Record<string, string>,
+ });
+ runtime = setup.runtime;
+ });
+
+ test("should throw an error when context is empty", async () => {
+ const options: GenerationOptions = {
+ runtime,
+ context: "",
+ modelClass: ModelClass.SMALL,
+ };
+
+ await expect(generateObject(options)).rejects.toThrow(
+ "generateObject context is empty"
+ );
+ });
+
+ test("should handle supported provider calls", async () => {
+ // Mock provider and trimTokens response
+ const context = "Test prompt for generation";
+ const provider = ModelProviderName.OPENAI;
+ const schema: ZodSchema = ZodSchema.any(); // Replace with a valid schema if needed
+
+ runtime.modelProvider = provider;
+
+ (trimTokens as jest.Mock).mockResolvedValue(context);
+ (handleProvider as jest.Mock).mockResolvedValue([
+ { response: "Generated text" },
+ ]);
+
+ const options: GenerationOptions = {
+ runtime,
+ context,
+ modelClass: ModelClass.SMALL,
+ schema,
+ schemaName: "TestSchema",
+ schemaDescription: "A schema for testing purposes",
+ mode: "json",
+ };
+
+ const result = await generateObject(options);
+
+ expect(trimTokens).toHaveBeenCalledWith(
+ context,
+ expect.any(Number),
+ ModelClass.SMALL
+ );
+ expect(handleProvider).toHaveBeenCalledWith(
+ expect.objectContaining({
+ provider,
+ model: expect.anything(),
+ schema,
+ schemaName: "TestSchema",
+ schemaDescription: "A schema for testing purposes",
+ })
+ );
+ expect(result).toEqual([{ response: "Generated text" }]);
+ });
+
+ test("should throw an error for unsupported provider", async () => {
+ runtime.modelProvider = "unsupportedProvider" as ModelProviderName;
+
+ const options: GenerationOptions = {
+ runtime,
+ context: "This should fail",
+ modelClass: ModelClass.SMALL,
+ };
+
+ await expect(generateObject(options)).rejects.toThrow(
+ "Unsupported provider"
+ );
+ });
+});
+
+describe("handleProvider", () => {
+ let runtime;
+
+ beforeAll(async () => {
+ const setup = await createRuntime({
+ env: process.env as Record<string, string>,
+ });
+ runtime = setup.runtime;
+ });
+
+ test("should handle OpenAI provider call", async () => {
+ const options = {
+ runtime,
+ provider: ModelProviderName.OPENAI,
+ model: "text-davinci-003",
+ apiKey: "testApiKey",
+ schema: ZodSchema.any(),
+ schemaName: "TestSchema",
+ schemaDescription: "A test schema",
+ mode: "json",
+ modelOptions: {
+ prompt: "Test prompt",
+ temperature: 0.7,
+ maxTokens: 100,
+ frequencyPenalty: 0,
+ presencePenalty: 0,
+ },
+ modelClass: ModelClass.SMALL,
+ context: "This is a test context",
+ };
+
+ (handleOpenAI as jest.Mock).mockResolvedValue([
+ { response: "Generated by OpenAI" },
+ ]);
+
+ const result = await handleProvider(options);
+
+ expect(handleOpenAI).toHaveBeenCalledWith(
+ expect.objectContaining({
+ model: "text-davinci-003",
+ apiKey: "testApiKey",
+ schemaName: "TestSchema",
+ })
+ );
+ expect(result).toEqual([{ response: "Generated by OpenAI" }]);
+ });
+
+ test("should throw error on unsupported provider in handleProvider", async () => {
+ const options = {
+ runtime,
+ provider: "unsupportedProvider" as ModelProviderName,
+ model: "unsupportedModel",
+ apiKey: "testApiKey",
+ schema: ZodSchema.any(),
+ schemaName: "UnsupportedSchema",
+ schemaDescription: "This should fail",
+ mode: "json",
+ modelOptions: {
+ prompt: "Test unsupported provider",
+ temperature: 0.7,
+ maxTokens: 100,
+ frequencyPenalty: 0,
+ presencePenalty: 0,
+ },
+ modelClass: ModelClass.SMALL,
+ context: "This is an unsupported provider context",
+ };
+
+ await expect(handleProvider(options)).rejects.toThrow(
+ "Unsupported provider"
+ );
+ });
+});
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index d64fef99e65..27df3956edc 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -459,6 +459,9 @@ importers:
'@types/jest':
specifier: 29.5.14
version: 29.5.14
+ '@types/mocha':
+ specifier: ^10.0.9
+ version: 10.0.9
'@types/node':
specifier: 22.8.4
version: 22.8.4
@@ -512,7 +515,7 @@ importers:
version: 2.79.2
ts-jest:
specifier: 29.2.5
- version: 29.2.5(@babel/core@7.26.0)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.26.0))(jest@29.7.0(@types/node@22.8.4)(ts-node@10.9.2(@types/node@22.8.4)(typescript@5.6.3)))(typescript@5.6.3)
+ version: 29.2.5(@babel/core@7.26.0)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.26.0))(esbuild@0.24.0)(jest@29.7.0(@types/node@22.8.4)(ts-node@10.9.2(@types/node@22.8.4)(typescript@5.6.3)))(typescript@5.6.3)
ts-node:
specifier: 10.9.2
version: 10.9.2(@types/node@22.8.4)(typescript@5.6.3)
@@ -3776,6 +3779,9 @@ packages:
'@types/minimist@1.2.5':
resolution: {integrity: sha512-hov8bUuiLiyFPGyFPE1lwWhmzYbirOXQNNo40+y3zow8aFVTeyn3VWL0VFFfdNddA8S4Vf0Tc062rzyNr7Paag==}
+ '@types/mocha@10.0.9':
+ resolution: {integrity: sha512-sicdRoWtYevwxjOHNMPTl3vSfJM6oyW8o1wXeI7uww6b6xHg8eBznQDNSGBCDJmsE8UMxP05JgZRtsKbTqt//Q==}
+
'@types/ms@0.7.34':
resolution: {integrity: sha512-nG96G3Wp6acyAgJqGasjODb+acrI7KltPiRxzHPXnP3NgI28bpQDRv53olbqGXbfcgF5aiiHmO3xpwEpS5Ld9g==}
@@ -15989,6 +15995,8 @@ snapshots:
'@types/minimist@1.2.5': {}
+ '@types/mocha@10.0.9': {}
+
'@types/ms@0.7.34': {}
'@types/node-fetch@2.6.12':
@@ -24636,7 +24644,7 @@ snapshots:
ts-interface-checker@0.1.13: {}
- ts-jest@29.2.5(@babel/core@7.26.0)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.26.0))(jest@29.7.0(@types/node@22.8.4)(ts-node@10.9.2(@types/node@22.8.4)(typescript@5.6.3)))(typescript@5.6.3):
+ ts-jest@29.2.5(@babel/core@7.26.0)(@jest/transform@29.7.0)(@jest/types@29.6.3)(babel-jest@29.7.0(@babel/core@7.26.0))(esbuild@0.24.0)(jest@29.7.0(@types/node@22.8.4)(ts-node@10.9.2(@types/node@22.8.4)(typescript@5.6.3)))(typescript@5.6.3):
dependencies:
bs-logger: 0.2.6
ejs: 3.1.10
@@ -24654,6 +24662,7 @@ snapshots:
'@jest/transform': 29.7.0
'@jest/types': 29.6.3
babel-jest: 29.7.0(@babel/core@7.26.0)
+ esbuild: 0.24.0
ts-mixer@6.0.4: {}