Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: atoma provider #2082

Open
wants to merge 23 commits into
base: develop
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
d785426
feat: added atoma-sdk dependency
francis2tm Jan 9, 2025
2c908d1
feat: atoma-sdk integrated
francis2tm Jan 9, 2025
36ab70d
fix: update handleAtoma response structure to include dynamic finishR…
francis2tm Jan 9, 2025
88cd7be
feat: add new model classes for embedding and image processing
francis2tm Jan 14, 2025
2a88cdb
feat: extend handleProvider function to include runtime, context, and…
francis2tm Jan 14, 2025
1d8bf86
feat: add Atoma SDK configuration to .env.example for Bearer Auth token
francis2tm Jan 15, 2025
0480f65
feat: integrate Atoma model support with configuration updates and SD…
francis2tm Jan 15, 2025
f490729
refactor: remove unused model configurations from .env.example and mo…
francis2tm Jan 15, 2025
57054df
Merge remote-tracking branch 'upstream/main' into feat/atoma-provider
francis2tm Jan 15, 2025
10d8144
chore: update lockfile
francis2tm Jan 15, 2025
f19eade
chore: update lockfile
francis2tm Jan 15, 2025
28358b6
chore: update lockfile
francis2tm Jan 15, 2025
bfef39b
refactor: restructure Atoma model configuration for improved clarity …
francis2tm Jan 16, 2025
ef2f9a1
chore: restore original formatting
francis2tm Jan 16, 2025
6c8475c
chore: restore original formatting
francis2tm Jan 16, 2025
ed17dd2
chore: restore original formatting
francis2tm Jan 16, 2025
cb20d0e
chore: restore original formatting
francis2tm Jan 16, 2025
58e66e3
chore: revert types.ts
francis2tm Jan 16, 2025
af9acd8
Merge remote-tracking branch 'upstream/develop' into feat/atoma-provider
francis2tm Jan 16, 2025
be5aa19
refactor: improve code formatting and readability in generation.ts
francis2tm Jan 16, 2025
3c04b92
refactor: remove Atoma model generation function from generation.ts
francis2tm Jan 16, 2025
1697270
chore: revert deps
francis2tm Jan 16, 2025
c3f8617
chore: remove atoma-sdk
francis2tm Jan 16, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions packages/core/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@
"@types/uuid": "10.0.0",
"ai": "3.4.33",
"anthropic-vertex-ai": "1.0.2",
"atoma-sdk": "github:atoma-network/atoma-sdk-typescript",
"fastembed": "1.14.1",
"fastestsmallesttextencoderdecoder": "1.0.22",
"gaxios": "6.7.1",
Expand Down
168 changes: 140 additions & 28 deletions packages/core/src/generation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ import {
} from "./types.ts";
import { fal } from "@fal-ai/client";
import { tavily } from "@tavily/core";
import { AtomaSDK } from "atoma-sdk";

type Tool = CoreTool<any, any>;
type StepResult = AIStepResult<any>;
Expand Down Expand Up @@ -1632,45 +1633,32 @@ interface ProviderOptions {
export async function handleProvider(
options: ProviderOptions
): Promise<GenerateObjectResult<unknown>> {
const { provider, runtime, context, modelClass } = options;
const { provider } = options;
francis2tm marked this conversation as resolved.
Show resolved Hide resolved

switch (provider) {
case ModelProviderName.OPENAI:
case ModelProviderName.ETERNALAI:
case ModelProviderName.ALI_BAILIAN:
case ModelProviderName.VOLENGINE:
case ModelProviderName.LLAMACLOUD:
case ModelProviderName.TOGETHER:
case ModelProviderName.NANOGPT:
case ModelProviderName.AKASH_CHAT_API:
return await handleOpenAI(options);
return handleOpenAI(options);
case ModelProviderName.ANTHROPIC:
case ModelProviderName.CLAUDE_VERTEX:
return await handleAnthropic(options);
return handleAnthropic(options);
case ModelProviderName.GROK:
return await handleGrok(options);
return handleGrok(options);
case ModelProviderName.GROQ:
return await handleGroq(options);
case ModelProviderName.LLAMALOCAL:
return await generateObjectDeprecated({
runtime,
context,
modelClass,
});
return handleGroq(options);
case ModelProviderName.GOOGLE:
return await handleGoogle(options);
return handleGoogle(options);
case ModelProviderName.REDPILL:
return await handleRedPill(options);
return handleRedPill(options);
case ModelProviderName.OPENROUTER:
return await handleOpenRouter(options);
return handleOpenRouter(options);
case ModelProviderName.OLLAMA:
return await handleOllama(options);
default: {
const errorMessage = `Unsupported provider: ${provider}`;
elizaLogger.error(errorMessage);
throw new Error(errorMessage);
}
return handleOllama(options);
case ModelProviderName.ATOMA:
return handleAtoma(options);
default:
throw new Error(`Unsupported provider: ${provider}`);
}
}

/**
* Handles object generation for OpenAI.
*
Expand Down Expand Up @@ -1886,6 +1874,130 @@ async function handleOllama({
});
}

/**
* Handles object generation for Atoma models.
*
* @param {ProviderOptions} options - Options specific to Atoma.
* @returns {Promise<GenerateObjectResult<unknown>>} - A promise that resolves to generated objects.
*/
async function handleAtoma({
model,
apiKey,
schema,
schemaName: _schemaName,
schemaDescription: _schemaDescription,
mode: _mode,
modelOptions: _modelOptions,
context,
runtime: _runtime,
provider: _provider,
modelClass: _modelClass,
}: ProviderOptions): Promise<GenerateObjectResult<unknown>> {
const atomaSDK = new AtomaSDK({
bearerAuth: apiKey,
});

try {
const result = await atomaSDK.chat.create({
messages: [
{
content: context,
role: "user",
},
],
model: model,
});

const completion = result.choices[0].message.content;

if (schema) {
// For schema-based generation, we'll parse the completion as JSON
// and validate it against the schema
const parsedCompletion = JSON.parse(completion);
const validatedObject = schema.parse(parsedCompletion);

return {
object: validatedObject,
finishReason: (result.choices[0].finishReason ?? "stop") as
| "stop"
| "length"
| "error"
| "tool-calls",
usage: {
promptTokens: result.usage?.promptTokens ?? 0,
completionTokens: result.usage?.completionTokens ?? 0,
totalTokens: result.usage?.totalTokens ?? 0,
},
warnings: undefined,
request: {},
response: {
id: result.id,
timestamp: new Date(result.created * 1000), // Convert Unix timestamp to Date
modelId: result.model,
},
logprobs: undefined,
experimental_providerMetadata: result.systemFingerprint
? {
metadata: {
systemFingerprint: result.systemFingerprint,
},
}
: undefined,
toJsonResponse(init?: ResponseInit) {
return new Response(
JSON.stringify({ object: validatedObject }),
{
...init,
headers: {
...init?.headers,
"content-type":
"application/json; charset=utf-8",
},
}
);
},
};
}

return {
object: completion,
finishReason: (result.choices[0].finishReason ?? "stop") as
| "stop"
| "length"
| "error"
| "tool-calls",
usage: {
promptTokens: result.usage?.promptTokens ?? 0,
completionTokens: result.usage?.completionTokens ?? 0,
totalTokens: result.usage?.totalTokens ?? 0,
},
warnings: undefined,
request: {},
response: {
id: result.id,
timestamp: new Date(result.created * 1000), // Convert Unix timestamp to Date
modelId: result.model,
},
logprobs: undefined,
experimental_providerMetadata: result.systemFingerprint
? { metadata: { systemFingerprint: result.systemFingerprint } }
: undefined,
toJsonResponse(init?: ResponseInit) {
return new Response(JSON.stringify({ object: completion }), {
...init,
headers: {
...init?.headers,
"content-type": "application/json; charset=utf-8",
},
});
},
};
} catch (error) {
console.error("Error in Atoma handler:", error);
throw error;
}
}

// Add type definition for Together AI response
interface TogetherAIImageResponse {
data: Array<{
Expand Down
45 changes: 34 additions & 11 deletions packages/core/src/models.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@ export const models: Models = {
[ModelClass.SMALL]: settings.SMALL_OPENAI_MODEL || "gpt-4o-mini",
[ModelClass.MEDIUM]: settings.MEDIUM_OPENAI_MODEL || "gpt-4o",
[ModelClass.LARGE]: settings.LARGE_OPENAI_MODEL || "gpt-4o",
[ModelClass.EMBEDDING]: settings.EMBEDDING_OPENAI_MODEL || "text-embedding-3-small",
[ModelClass.EMBEDDING]:
settings.EMBEDDING_OPENAI_MODEL || "text-embedding-3-small",
[ModelClass.IMAGE]: settings.IMAGE_OPENAI_MODEL || "dall-e-3",
},
},
Expand Down Expand Up @@ -55,9 +56,12 @@ export const models: Models = {
},
endpoint: "https://api.anthropic.com/v1",
model: {
[ModelClass.SMALL]: settings.SMALL_ANTHROPIC_MODEL || "claude-3-haiku-20240307",
[ModelClass.MEDIUM]: settings.MEDIUM_ANTHROPIC_MODEL || "claude-3-5-sonnet-20241022",
[ModelClass.LARGE]: settings.LARGE_ANTHROPIC_MODEL || "claude-3-5-sonnet-20241022",
[ModelClass.SMALL]:
settings.SMALL_ANTHROPIC_MODEL || "claude-3-haiku-20240307",
[ModelClass.MEDIUM]:
settings.MEDIUM_ANTHROPIC_MODEL || "claude-3-5-sonnet-20241022",
[ModelClass.LARGE]:
settings.LARGE_ANTHROPIC_MODEL || "claude-3-5-sonnet-20241022",
},
},
[ModelProviderName.CLAUDE_VERTEX]: {
Expand Down Expand Up @@ -90,7 +94,8 @@ export const models: Models = {
[ModelClass.SMALL]: settings.SMALL_GROK_MODEL || "grok-2-1212",
[ModelClass.MEDIUM]: settings.MEDIUM_GROK_MODEL || "grok-2-1212",
[ModelClass.LARGE]: settings.LARGE_GROK_MODEL || "grok-2-1212",
[ModelClass.EMBEDDING]: settings.EMBEDDING_GROK_MODEL || "grok-2-1212", // not sure about this one
[ModelClass.EMBEDDING]:
settings.EMBEDDING_GROK_MODEL || "grok-2-1212", // not sure about this one
},
},
[ModelProviderName.GROQ]: {
Expand Down Expand Up @@ -395,7 +400,9 @@ export const models: Models = {
},
},
[ModelProviderName.VOLENGINE]: {
endpoint: settings.VOLENGINE_API_URL || "https://open.volcengineapi.com/api/v3/",
endpoint:
settings.VOLENGINE_API_URL ||
"https://open.volcengineapi.com/api/v3/",
settings: {
stop: [],
maxInputTokens: 128000,
Expand All @@ -418,8 +425,7 @@ export const models: Models = {
settings.VOLENGINE_MODEL ||
"doubao-pro-256k",
[ModelClass.EMBEDDING]:
settings.VOLENGINE_EMBEDDING_MODEL ||
"doubao-embedding",
settings.VOLENGINE_EMBEDDING_MODEL || "doubao-embedding",
},
},
[ModelProviderName.NANOGPT]: {
Expand All @@ -436,7 +442,7 @@ export const models: Models = {
[ModelClass.SMALL]: settings.SMALL_NANOGPT_MODEL || "gpt-4o-mini",
[ModelClass.MEDIUM]: settings.MEDIUM_NANOGPT_MODEL || "gpt-4o",
[ModelClass.LARGE]: settings.LARGE_NANOGPT_MODEL || "gpt-4o",
}
},
},
[ModelProviderName.HYPERBOLIC]: {
endpoint: "https://api.hyperbolic.xyz/v1",
Expand Down Expand Up @@ -472,7 +478,8 @@ export const models: Models = {
},
model: {
[ModelClass.SMALL]: settings.SMALL_VENICE_MODEL || "llama-3.3-70b",
[ModelClass.MEDIUM]: settings.MEDIUM_VENICE_MODEL || "llama-3.3-70b",
[ModelClass.MEDIUM]:
settings.MEDIUM_VENICE_MODEL || "llama-3.3-70b",
[ModelClass.LARGE]: settings.LARGE_VENICE_MODEL || "llama-3.1-405b",
[ModelClass.IMAGE]: settings.IMAGE_VENICE_MODEL || "fluently-xl",
},
Expand Down Expand Up @@ -511,7 +518,23 @@ export const models: Models = {
[ModelClass.MEDIUM]: "",
[ModelClass.LARGE]: "",
[ModelClass.EMBEDDING]: "",
[ModelClass.IMAGE]: settings.LIVEPEER_IMAGE_MODEL || "ByteDance/SDXL-Lightning",
[ModelClass.IMAGE]:
settings.LIVEPEER_IMAGE_MODEL || "ByteDance/SDXL-Lightning",
},
},
[ModelProviderName.ATOMA]: {
settings: {
stop: [],
maxInputTokens: 128000,
maxOutputTokens: 8192,
frequency_penalty: 0.0,
presence_penalty: 0.0,
temperature: 0.7,
},
model: {
[ModelClass.SMALL]: "meta-llama/Llama-3.3-70B-Instruct",
[ModelClass.MEDIUM]: "meta-llama/Llama-3.3-70B-Instruct",
[ModelClass.LARGE]: "meta-llama/Llama-3.3-70B-Instruct",
francis2tm marked this conversation as resolved.
Show resolved Hide resolved
},
},
};
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,7 @@ export type Models = {
[ModelProviderName.VENICE]: Model;
[ModelProviderName.AKASH_CHAT_API]: Model;
[ModelProviderName.LIVEPEER]: Model;
[ModelProviderName.ATOMA]: Model;
};

/**
Expand Down Expand Up @@ -243,6 +244,7 @@ export enum ModelProviderName {
VENICE = "venice",
AKASH_CHAT_API = "akash_chat_api",
LIVEPEER = "livepeer",
ATOMA = "atoma",
}

/**
Expand Down
Loading
Loading