Merge pull request #335 from tsubasakong/main

feat: Add Heurist API Integration as New Model Provider
elizaOS · Nov 15, 2024 · 4d1e66c · 4d1e66c
2 parents c4e8601 + 2e3d507
commit 4d1e66c
Show file tree

Hide file tree

Showing 6 changed files with 90 additions and 25,474 deletions.
diff --git a/.env.example b/.env.example
@@ -55,6 +55,9 @@ LARGE_OLLAMA_MODEL=         #default hermes3:70b
 # For asking Claude stuff
 ANTHROPIC_API_KEY=
 
+# Heurist API 
+HEURIST_API_KEY=
+
 WALLET_PRIVATE_KEY=EXAMPLE_WALLET_PRIVATE_KEY
 WALLET_PUBLIC_KEY=EXAMPLE_WALLET_PUBLIC_KEY
 

diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts
@@ -150,6 +150,11 @@ export function getTokenForProvider(
                 character.settings?.secrets?.GROK_API_KEY ||
                 settings.GROK_API_KEY
             );
+        case ModelProviderName.HEURIST:
+            return (
+                character.settings?.secrets?.HEURIST_API_KEY ||
+                settings.HEURIST_API_KEY
+            );
     }
 }
 

diff --git a/packages/core/src/generation.ts b/packages/core/src/generation.ts
@@ -312,6 +312,25 @@ export async function generateText({
                 console.debug("Received response from Ollama model.");
                 break;
 
+            case ModelProviderName.HEURIST: {
+                elizaLogger.debug("Initializing Heurist model.");
+                const heurist = createOpenAI({ apiKey: apiKey, baseURL: endpoint });
+
+                const { text: heuristResponse } = await aiGenerateText({
+                    model: heurist.languageModel(model),
+                    prompt: context,
+                    system: runtime.character.system ?? settings.SYSTEM_PROMPT ?? undefined,
+                    temperature: temperature,
+                    maxTokens: max_response_length,
+                    frequencyPenalty: frequency_penalty,
+                    presencePenalty: presence_penalty,
+                });
+
+                response = heuristResponse;
+                elizaLogger.debug("Received response from Heurist model.");
+                break;
+            }
+
             default: {
                 const errorMessage = `Unsupported provider: ${provider}`;
                 elizaLogger.error(errorMessage);
@@ -681,6 +700,12 @@ export const generateImage = async (
         width: number;
         height: number;
         count?: number;
+        negativePrompt?: string;
+        numIterations?: number;
+        guidanceScale?: number;
+        seed?: number;
+        modelId?: string;
+        jobId?: string;
     },
     runtime: IAgentRuntime
 ): Promise<{
@@ -696,14 +721,40 @@ export const generateImage = async (
 
     const model = getModel(runtime.character.modelProvider, ModelClass.IMAGE);
     const modelSettings = models[runtime.character.modelProvider].imageSettings;
-    // some fallbacks for backwards compat, should remove in the future
-    const apiKey =
-        runtime.token ??
-        runtime.getSetting("TOGETHER_API_KEY") ??
-        runtime.getSetting("OPENAI_API_KEY");
-
+    const apiKey = runtime.token ?? runtime.getSetting("HEURIST_API_KEY") ??  runtime.getSetting("TOGETHER_API_KEY") ?? runtime.getSetting("OPENAI_API_KEY");    
     try {
-        if (runtime.character.modelProvider === ModelProviderName.LLAMACLOUD) {
+        if (runtime.character.modelProvider === ModelProviderName.HEURIST) {
+            const response = await fetch('http://sequencer.heurist.xyz/submit_job', {
+                method: 'POST',
+                headers: {
+                    'Authorization': `Bearer ${apiKey}`,
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify({
+                    job_id: data.jobId || crypto.randomUUID(),
+                    model_input: {
+                        SD: {
+                            prompt: data.prompt,
+                            neg_prompt: data.negativePrompt,
+                            num_iterations: data.numIterations || 20,
+                            width: data.width || 512,
+                            height: data.height || 512,
+                            guidance_scale: data.guidanceScale,
+                            seed: data.seed || -1,
+                        }
+                    },
+                    model_id: data.modelId || 'PepeXL', // Default to SD 1.5 if not specified
+                })
+            });
+
+            if (!response.ok) {
+                throw new Error(`Heurist image generation failed: ${response.statusText}`);
+            }
+
+            const result = await response.json();
+            return { success: true, data: [result.url] };
+        }
+        else if (runtime.character.modelProvider === ModelProviderName.LLAMACLOUD) {
             const together = new Together({ apiKey: apiKey as string });
             const response = await together.images.create({
                 model: "black-forest-labs/FLUX.1-schnell",

diff --git a/packages/core/src/models.ts b/packages/core/src/models.ts
@@ -217,8 +217,30 @@ const models: Models = {
                 settings.OLLAMA_EMBEDDING_MODEL || "mxbai-embed-large",
         },
     },
+    [ModelProviderName.HEURIST]: {
+        settings: {
+            stop: [],
+            maxInputTokens: 128000,
+            maxOutputTokens: 8192,
+            repetition_penalty: 0.0,
+            temperature: 0.7,
+        },
+        imageSettings: {
+            steps: 20,
+        },
+        endpoint: "https://llm-gateway.heurist.xyz",
+        model: {
+            [ModelClass.SMALL]: "meta-llama/llama-3-70b-instruct",
+            [ModelClass.MEDIUM]: "meta-llama/llama-3-70b-instruct",
+            [ModelClass.LARGE]: "meta-llama/llama-3.1-405b-instruct",
+            [ModelClass.EMBEDDING]: "" , //Add later,
+            [ModelClass.IMAGE]: "PepeXL",
+        },
+    }
 };
 
+
+
 export function getModel(provider: ModelProviderName, type: ModelClass) {
     return models[provider].model[type];
 }

diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts
@@ -114,6 +114,7 @@ export type Models = {
     [ModelProviderName.REDPILL]: Model;
     [ModelProviderName.OPENROUTER]: Model;
     [ModelProviderName.OLLAMA]: Model;
+    [ModelProviderName.HEURIST]: Model;
 };
 
 export enum ModelProviderName {
@@ -128,6 +129,7 @@ export enum ModelProviderName {
     REDPILL = "redpill",
     OPENROUTER = "openrouter",
     OLLAMA = "ollama",
+    HEURIST = "heurist",
 }
 
 /**