Skip to content

Commit

Permalink
feat: watsonx.ai llm compatibility
Browse files Browse the repository at this point in the history
  • Loading branch information
grahamwhiteuk committed Sep 17, 2024
1 parent 18b1eff commit 2dd4adc
Showing 2 changed files with 19 additions and 3 deletions.
4 changes: 4 additions & 0 deletions .env.template
Original file line number Diff line number Diff line change
@@ -6,6 +6,10 @@ BEE_FRAMEWORK_LOG_SINGLE_LINE="false"
# For BAM LLM Adapter
# GENAI_API_KEY=

# For Watsonx LLM Adapter
# WATSONX_API_KEY=
# WATSONX_PROJECT_ID=

# For OpenAI LLM Adapter
# OPENAI_API_KEY=

18 changes: 15 additions & 3 deletions src/index.ts
Original file line number Diff line number Diff line change
@@ -23,6 +23,7 @@ import { TokenMemory } from "bee-agent-framework/memory/tokenMemory";
import { Logger } from "bee-agent-framework/logger/logger";
import { OllamaChatLLM } from "bee-agent-framework/adapters/ollama/chat";
import { OpenAIChatLLM } from "bee-agent-framework/adapters/openai/chat";
import { WatsonXChatLLM } from "bee-agent-framework/adapters/watsonx/chat";

// core tools
import { DuckDuckGoSearchTool } from "bee-agent-framework/tools/search/duckDuckGoSearch";
@@ -40,14 +41,25 @@ const logger = new Logger({ name: "app", level: "trace" });

async function runBeeAgent() {
// use BAM if GENAI_API_KEY env var is defined
// else use Watsonx if WATSONX_API_KEY and WATSONX_PROJECT_ID env vars are defined
// else use OpenAI if OPENAI_API_KEY env var is defined
// else use Ollama
const llm =
process.env.GENAI_API_KEY !== undefined
? BAMChatLLM.fromPreset("meta-llama/llama-3-1-70b-instruct")
: process.env.OPENAI_API_KEY !== undefined
? new OpenAIChatLLM()
: new OllamaChatLLM({ modelId: "llama3.1" });
: process.env.WATSONX_API_KEY !== undefined && process.env.WATSONX_PROJECT_ID !== undefined
? WatsonXChatLLM.fromPreset("meta-llama/llama-3-1-70b-instruct", {
apiKey: process.env.WATSONX_API_KEY, //pragma: allowlist secret
projectId: process.env.WATSONX_PROJECT_ID,
parameters: {
decoding_method: "greedy",
min_new_tokens: 5,
max_new_tokens: 50,
},
})
: process.env.OPENAI_API_KEY !== undefined
? new OpenAIChatLLM()
: new OllamaChatLLM({ modelId: "llama3.1" });

const agent = new BeeAgent({
llm,

0 comments on commit 2dd4adc

Please sign in to comment.