Skip to content

Commit

Permalink
Merge pull request #1599 from zhourunlai/tavily_sdk
Browse files Browse the repository at this point in the history 
feat: use tavily sdk
  • Loading branch information
@shakkernerd
shakkernerd authored Dec 31, 2024
2 parents 532a0ad + 4beabfa commit 5aab307
Show file tree
Hide file tree
Showing 4 changed files with 163 additions and 152 deletions.
155 changes: 78 additions & 77 deletions packages/core/package.json
View file
Original file line number Diff line number Diff line change
@@ -1,79 +1,80 @@
{
"name": "@elizaos/core",
"version": "0.1.7-alpha.2",
"description": "",
"main": "dist/index.js",
"type": "module",
"types": "dist/index.d.ts",
"scripts": {
"build": "tsup --format esm --dts",
"lint": "eslint --fix --cache .",
"watch": "tsc --watch",
"dev": "tsup --format esm --dts --watch",
"build:docs": "cd docs && pnpm run build",
"test": "vitest run",
"test:coverage": "vitest run --coverage",
"test:watch": "vitest"
},
"author": "",
"license": "MIT",
"devDependencies": {
"@eslint/js": "9.16.0",
"@rollup/plugin-commonjs": "25.0.8",
"@rollup/plugin-json": "6.1.0",
"@rollup/plugin-node-resolve": "15.3.0",
"@rollup/plugin-replace": "5.0.7",
"@rollup/plugin-terser": "0.1.0",
"@rollup/plugin-typescript": "11.1.6",
"@solana/web3.js": "1.95.8",
"@types/fluent-ffmpeg": "2.1.27",
"@types/jest": "29.5.14",
"@types/mocha": "10.0.10",
"@types/node": "22.8.4",
"@types/pdfjs-dist": "2.10.378",
"@types/tar": "6.1.13",
"@types/wav-encoder": "1.3.3",
"@typescript-eslint/eslint-plugin": "8.16.0",
"@typescript-eslint/parser": "8.16.0",
"@vitest/coverage-v8": "2.1.5",
"dotenv": "16.4.5",
"jest": "29.7.0",
"lint-staged": "15.2.10",
"nodemon": "3.1.7",
"pm2": "5.4.3",
"rimraf": "6.0.1",
"rollup": "2.79.2",
"ts-jest": "29.2.5",
"ts-node": "10.9.2",
"tslib": "2.8.1",
"tsup": "8.3.5",
"typescript": "5.6.3"
},
"dependencies": {
"@ai-sdk/anthropic": "0.0.56",
"@ai-sdk/google": "0.0.55",
"@ai-sdk/google-vertex": "0.0.43",
"@ai-sdk/groq": "0.0.3",
"@ai-sdk/openai": "1.0.5",
"@anthropic-ai/sdk": "0.30.1",
"@fal-ai/client": "1.2.0",
"@types/uuid": "10.0.0",
"ai": "3.4.33",
"anthropic-vertex-ai": "1.0.2",
"fastembed": "1.14.1",
"fastestsmallesttextencoderdecoder": "1.0.22",
"gaxios": "6.7.1",
"glob": "11.0.0",
"handlebars": "^4.7.8",
"js-sha1": "0.7.0",
"js-tiktoken": "1.0.15",
"langchain": "0.3.6",
"ollama-ai-provider": "0.16.1",
"openai": "4.73.0",
"tinyld": "1.3.4",
"together-ai": "0.7.0",
"unique-names-generator": "4.7.1",
"uuid": "11.0.3",
"zod": "3.23.8"
}
"name": "@elizaos/core",
"version": "0.1.7-alpha.2",
"description": "",
"main": "dist/index.js",
"type": "module",
"types": "dist/index.d.ts",
"scripts": {
"build": "tsup --format esm --dts",
"lint": "eslint --fix --cache .",
"watch": "tsc --watch",
"dev": "tsup --format esm --dts --watch",
"build:docs": "cd docs && pnpm run build",
"test": "vitest run",
"test:coverage": "vitest run --coverage",
"test:watch": "vitest"
},
"author": "",
"license": "MIT",
"devDependencies": {
"@eslint/js": "9.16.0",
"@rollup/plugin-commonjs": "25.0.8",
"@rollup/plugin-json": "6.1.0",
"@rollup/plugin-node-resolve": "15.3.0",
"@rollup/plugin-replace": "5.0.7",
"@rollup/plugin-terser": "0.1.0",
"@rollup/plugin-typescript": "11.1.6",
"@solana/web3.js": "1.95.8",
"@tavily/core": "^0.0.2",
"@types/fluent-ffmpeg": "2.1.27",
"@types/jest": "29.5.14",
"@types/mocha": "10.0.10",
"@types/node": "22.8.4",
"@types/pdfjs-dist": "2.10.378",
"@types/tar": "6.1.13",
"@types/wav-encoder": "1.3.3",
"@typescript-eslint/eslint-plugin": "8.16.0",
"@typescript-eslint/parser": "8.16.0",
"@vitest/coverage-v8": "2.1.5",
"dotenv": "16.4.5",
"jest": "29.7.0",
"lint-staged": "15.2.10",
"nodemon": "3.1.7",
"pm2": "5.4.3",
"rimraf": "6.0.1",
"rollup": "2.79.2",
"ts-jest": "29.2.5",
"ts-node": "10.9.2",
"tslib": "2.8.1",
"tsup": "8.3.5",
"typescript": "5.6.3"
},
"dependencies": {
"@ai-sdk/anthropic": "0.0.56",
"@ai-sdk/google": "0.0.55",
"@ai-sdk/google-vertex": "0.0.43",
"@ai-sdk/groq": "0.0.3",
"@ai-sdk/openai": "1.0.5",
"@anthropic-ai/sdk": "0.30.1",
"@fal-ai/client": "1.2.0",
"@types/uuid": "10.0.0",
"ai": "3.4.33",
"anthropic-vertex-ai": "1.0.2",
"fastembed": "1.14.1",
"fastestsmallesttextencoderdecoder": "1.0.22",
"gaxios": "6.7.1",
"glob": "11.0.0",
"handlebars": "^4.7.8",
"js-sha1": "0.7.0",
"js-tiktoken": "1.0.15",
"langchain": "0.3.6",
"ollama-ai-provider": "0.16.1",
"openai": "4.73.0",
"tinyld": "1.3.4",
"together-ai": "0.7.0",
"unique-names-generator": "4.7.1",
"uuid": "11.0.3",
"zod": "3.23.8"
}
}
132 changes: 62 additions & 70 deletions packages/core/src/generation.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ import {
TelemetrySettings,
} from "./types.ts";
import { fal } from "@fal-ai/client";
import { tavily } from "@tavily/core";

/**
* Send a message to the model for a text generateText - receive a string back and parse how you'd like
Expand Down Expand Up @@ -982,33 +983,35 @@ export const generateImage = async (
});

const apiKey =
runtime.imageModelProvider === runtime.modelProvider
? runtime.token
: (() => {
// First try to match the specific provider
switch (runtime.imageModelProvider) {
case ModelProviderName.HEURIST:
return runtime.getSetting("HEURIST_API_KEY");
case ModelProviderName.TOGETHER:
return runtime.getSetting("TOGETHER_API_KEY");
case ModelProviderName.FAL:
return runtime.getSetting("FAL_API_KEY");
case ModelProviderName.OPENAI:
return runtime.getSetting("OPENAI_API_KEY");
case ModelProviderName.VENICE:
return runtime.getSetting("VENICE_API_KEY");
case ModelProviderName.LIVEPEER:
return runtime.getSetting("LIVEPEER_GATEWAY_URL");
default:
// If no specific match, try the fallback chain
return (runtime.getSetting("HEURIST_API_KEY") ??
runtime.getSetting("TOGETHER_API_KEY") ??
runtime.getSetting("FAL_API_KEY") ??
runtime.getSetting("OPENAI_API_KEY") ??
runtime.getSetting("VENICE_API_KEY"))??
runtime.getSetting("LIVEPEER_GATEWAY_URL");
}
})();
runtime.imageModelProvider === runtime.modelProvider
? runtime.token
: (() => {
// First try to match the specific provider
switch (runtime.imageModelProvider) {
case ModelProviderName.HEURIST:
return runtime.getSetting("HEURIST_API_KEY");
case ModelProviderName.TOGETHER:
return runtime.getSetting("TOGETHER_API_KEY");
case ModelProviderName.FAL:
return runtime.getSetting("FAL_API_KEY");
case ModelProviderName.OPENAI:
return runtime.getSetting("OPENAI_API_KEY");
case ModelProviderName.VENICE:
return runtime.getSetting("VENICE_API_KEY");
case ModelProviderName.LIVEPEER:
return runtime.getSetting("LIVEPEER_GATEWAY_URL");
default:
// If no specific match, try the fallback chain
return (
runtime.getSetting("HEURIST_API_KEY") ??
runtime.getSetting("TOGETHER_API_KEY") ??
runtime.getSetting("FAL_API_KEY") ??
runtime.getSetting("OPENAI_API_KEY") ??
runtime.getSetting("VENICE_API_KEY") ??
runtime.getSetting("LIVEPEER_GATEWAY_URL")
);
}
})();
try {
if (runtime.imageModelProvider === ModelProviderName.HEURIST) {
const response = await fetch(
Expand Down Expand Up @@ -1198,28 +1201,31 @@ export const generateImage = async (
});

return { success: true, data: base64s };

} else if (runtime.imageModelProvider === ModelProviderName.LIVEPEER) {
if (!apiKey) {
throw new Error("Livepeer Gateway is not defined");
}
try {
const baseUrl = new URL(apiKey);
if (!baseUrl.protocol.startsWith('http')) {
if (!baseUrl.protocol.startsWith("http")) {
throw new Error("Invalid Livepeer Gateway URL protocol");
}
const response = await fetch(`${baseUrl.toString()}text-to-image`, {
method: "POST",
headers: {
"Content-Type": "application/json"
},
body: JSON.stringify({
model_id: data.modelId || "ByteDance/SDXL-Lightning",
prompt: data.prompt,
width: data.width || 1024,
height: data.height || 1024
})
});
const response = await fetch(
`${baseUrl.toString()}text-to-image`,
{
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
model_id:
data.modelId || "ByteDance/SDXL-Lightning",
prompt: data.prompt,
width: data.width || 1024,
height: data.height || 1024,
}),
}
);
const result = await response.json();
if (!result.images?.length) {
throw new Error("No images generated");
Expand All @@ -1241,19 +1247,19 @@ export const generateImage = async (
}
const blob = await imageResponse.blob();
const arrayBuffer = await blob.arrayBuffer();
const base64 = Buffer.from(arrayBuffer).toString("base64");
const base64 =
Buffer.from(arrayBuffer).toString("base64");
return `data:image/jpeg;base64,${base64}`;
})
);
return {
success: true,
data: base64Images
data: base64Images,
};
} catch (error) {
console.error(error);
return { success: false, error: error };
}

} else {
let targetSize = `${data.width}x${data.height}`;
if (
Expand Down Expand Up @@ -1316,34 +1322,20 @@ export const generateWebSearch = async (
query: string,
runtime: IAgentRuntime
): Promise<SearchResponse> => {
const apiUrl = "https://api.tavily.com/search";
const apiKey = runtime.getSetting("TAVILY_API_KEY");

try {
const response = await fetch(apiUrl, {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
api_key: apiKey,
query,
include_answer: true,
max_results: 3, // 5 (default)
topic: "general", // "general"(default) "news"
search_depth: "basic", // "basic"(default) "advanced"
include_images: false, // false (default) true
}),
});

if (!response.ok) {
throw new elizaLogger.error(
`HTTP error! status: ${response.status}`
);
const apiKey = runtime.getSetting("TAVILY_API_KEY") as string;
if (!apiKey) {
throw new Error("TAVILY_API_KEY is not set");
}

const data: SearchResponse = await response.json();
return data;
const tvly = tavily({ apiKey });
const response = await tvly.search(query, {
includeAnswer: true,
maxResults: 3, // 5 (default)
topic: "general", // "general"(default) "news"
searchDepth: "basic", // "basic"(default) "advanced"
includeImages: false, // false (default) true
});
return response;
} catch (error) {
elizaLogger.error("Error:", error);
}
Expand Down
15 changes: 10 additions & 5 deletions packages/core/src/types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1255,21 +1255,26 @@ export interface IAwsS3Service extends Service {
generateSignedUrl(fileName: string, expiresIn: number): Promise<string>;
}

export type SearchImage = {
url: string;
description?: string;
};

export type SearchResult = {
title: string;
url: string;
content: string;
rawContent?: string;
score: number;
raw_content: string | null;
publishedDate?: string;
};

export type SearchResponse = {
answer?: string;
query: string;
follow_up_questions: string[] | null;
answer: string | null;
images: string[];
responseTime: number;
images: SearchImage[];
results: SearchResult[];
response_time: number;
};

export enum ServiceType {
Expand Down
Loading

0 comments on commit 5aab307

Please sign in to comment.