rocket-ai
Version:
Simple AI Client that lets you access different LLMs in a unified way.
161 lines (124 loc) • 3.78 kB
Markdown
# RocketAI
Simple AI client to interface with different LLMs in a unified way.
## Installation
```bash
npm install rocket-ai
# or
yarn add rocket-ai
```
## API Keys
You need to set the api keys as environment variables. Those environment variables need to be present, even if they are just an empty string:
```bash
OPENAI_API_KEY=
ANTHROPIC_API_KEY=
GOOGLE_API_KEY=
FIREWORKS_API_KEY=
```
## Usage
### Client
```typescript jsx
import {AiClient} from 'rocket-ai';
const aiClient = new AiClient();
// invoke
const response = await aiClient.invoke({
model: 'gpt-4o',
message: [{role: "user", content: "tell me a joke"}],
systemPrompt: 'You are a helpful assistant.',
});
// stream
const stream = await aiClient.stream({
model: 'gpt-4o',
message: [{role: "user", content: "tell me a joke"}],
systemPrompt: 'You are a helpful assistant.',
});
for await (const chunk of stream) {
console.log(response);
}
// generate image
const image = await aiClient.generateImage({
model: 'dall-e3',
prompt: 'a painting of a flower vase',
size: '1024x1024',
n: 1
})
// generate speech
const speech = await aiClient.generateSpeech({
model: 'tts-1',
text: 'Hello, how are you doing today?',
voice: 'alloy',
})
// structured output, works for stream an invoke
import {z} from 'zod';
const schema = z.object({
joke: z.string(),
})
const structuredOutput = aiClient
.withStructuredOutput(schema)
.invoke({
model: 'gpt-4o',
message: [{role: "user", content: "tell me a joke"}],
systemPrompt: 'You are a helpful assistant.',
});
```
### Agent
To use the agent, you need to have at least one tool. Create a sample tool that fetches weather data:
```typescript jsx
//Weather API Tool
import {z} from "zod";
import {tool} from "rocket-ai"
export const weatherApi = tool(
async ({city, country}: { city: string, country?: string }): Promise<any> => {
if (!city) {
throw new Error('City is required');
}
const key = "xxxxxx";
const response = await fetch(`https://api.openweathermap.org/data/2.5/weather?q=${city},${country}&appid=${key}`)
const res = await response.json()
return JSON.stringify(res, null, 2)
},
{
name: "weatherApi",
description: "Get weather data for a city.",
queryFormat: z.object({
city: z.string(),
country: z.string(),
}),
}
)
```
Now you can use the agent to interact with the tool:
```typescript jsx
import {Agent} from ".rocket-ai";
import {weatherApi} from "./weather-api";
(async () => {
const systemInstruction = "You run in a loop of Thought, Action, PAUSE, Observation.\n" +
"At the end of the loop you output an Answer\n" +
"Strictly follow the provided response format.\n" +
"Use Thought to describe your thoughts about the question you have been asked.\n" +
"Use Action to run one of the actions available to you\n" +
"Observation will be the result of running those actions.\n";
const agent = new Agent(systemInstruction);
agent.registerTools([weatherApi]);
const response = await agent.executeTask("how is the weather in berlin??");
console.log(response);
})();
```
## Available Models
All models are available through the `AiModelType` enum:
- `gpt-4o`
- `gpt-4o-mini`
- `dall-e3`
- `tts-1`
- `o1-preview`
- `o1-mini`
- `claude-3-5-sonnet-latest`
- `claude-3-5-haiku-latest`
- `gemini-2.0-flash-latest`
- `gemini-1.5-flash-latest`
- `gemini-1.5-pro-latest`
- `fw-llama-3-3`
- `fw-deepseek-v3`
- `fw-deepseek-r1`
> note: `fw-llama-3-3`, `fw-deepseek-v3` and `fw-deepseek-r1` are used from the [Fireworks AI API](https://fireworks.ai/).
## License
MIT