fix: generate structured objects and images with NEAR AI (#3644)

* fix: nearai generate object and image model * fix: require structured output if schema is specified * feat: update warning msg * style: update comments * style: update comments * fix: env typo
elizaOS · Feb 23, 2025 · a003d14 · a003d14
1 parent b6d944c
commit a003d14
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 0 deletions.
diff --git a/.env.example b/.env.example
@@ -366,6 +366,7 @@ NEARAI_MODEL=
 SMALL_NEARAI_MODEL=  # Default: fireworks::accounts/fireworks/models/llama-v3p2-3b-instruct
 MEDIUM_NEARAI_MODEL= # Default: fireworks::accounts/fireworks/models/llama-v3p1-70b-instruct
 LARGE_NEARAI_MODEL=  # Default: fireworks::accounts/fireworks/models/llama-v3p1-405b-instruct
+IMAGE_NEARAI_MODEL=   # Default: fireworks::accounts/fireworks/models/playground-v2-5-1024px-aesthetic
 
 # Remaining Provider Configurations
 GOOGLE_GENERATIVE_AI_API_KEY= # Gemini API key

diff --git a/packages/core/src/generation.ts b/packages/core/src/generation.ts
@@ -53,6 +53,9 @@ import { fal } from "@fal-ai/client";
 
 import BigNumber from "bignumber.js";
 import { createPublicClient, http } from "viem";
+import fs from "fs";
+import os from "os";
+import path from "path";
 
 type Tool = CoreTool<any, any>;
 type StepResult = AIStepResult<any>;
@@ -1734,6 +1737,15 @@ export const generateImage = async (
                           return runtime.getSetting("LIVEPEER_GATEWAY_URL");
                       case ModelProviderName.SECRETAI:
                           return runtime.getSetting("SECRET_AI_API_KEY");
+                      case ModelProviderName.NEARAI:
+                          try {
+                            // Read auth config from ~/.nearai/config.json if it exists
+                            const config = JSON.parse(fs.readFileSync(path.join(os.homedir(), '.nearai/config.json'), 'utf8'));
+                            return JSON.stringify(config?.auth);
+                          } catch (e) {
+                            elizaLogger.warn(`Error loading NEAR AI config. The environment variable NEARAI_API_KEY will be used. ${e}`);
+                          }
+                          return runtime.getSetting("NEARAI_API_KEY");
                       default:
                           // If no specific match, try the fallback chain
                           return (
@@ -2279,6 +2291,8 @@ export async function handleProvider(
             return await handleLivepeer(options);
         case ModelProviderName.SECRETAI:
             return await handleSecretAi(options);
+        case ModelProviderName.NEARAI:
+            return await handleNearAi(options);
         default: {
             const errorMessage = `Unsupported provider: ${provider}`;
             elizaLogger.error(errorMessage);
@@ -2661,6 +2675,34 @@ async function handleSecretAi({
     });
 }
 
+/**
+ * Handles object generation for NEAR AI models.
+ *
+ * @param {ProviderOptions} options - Options specific to NEAR AI.
+ * @returns {Promise<GenerateObjectResult<unknown>>} - A promise that resolves to generated objects.
+ */
+async function handleNearAi({
+    model,
+    apiKey,
+    schema,
+    schemaName,
+    schemaDescription,
+    mode = "json",
+    modelOptions,
+}: ProviderOptions): Promise<GenerateObjectResult<unknown>> {
+    const nearai = createOpenAI({ apiKey, baseURL: models.nearai.endpoint });
+    // Require structured output if schema is provided
+    const settings = schema ? { structuredOutputs: true } : undefined;
+    return await aiGenerateObject({
+        model: nearai.languageModel(model, settings),
+        schema,
+        schemaName,
+        schemaDescription,
+        mode,
+        ...modelOptions,
+    });
+}
+
 // Add type definition for Together AI response
 interface TogetherAIImageResponse {
     data: Array<{

diff --git a/packages/core/src/models.ts b/packages/core/src/models.ts
@@ -1216,6 +1216,9 @@ export const models: Models = {
                 maxOutputTokens: 8192,
                 temperature: 0.6,
             },
+            [ModelClass.IMAGE]: {
+                name: settings.IMAGE_NEARAI_MODEL || "fireworks::accounts/fireworks/models/playground-v2-5-1024px-aesthetic",
+            },
         },
     },
 };