Skip to content

Commit

Permalink
支持谷歌 Gemini 模型 (#7)
Browse files Browse the repository at this point in the history
  • Loading branch information
zhengxs2018 authored Dec 16, 2023
1 parent c215120 commit e83bd50
Show file tree
Hide file tree
Showing 16 changed files with 542 additions and 4 deletions.
1 change: 1 addition & 0 deletions .cspell.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ words:
- longcontext
- openai
- openpose
- permissioning
- qwen
- qwenai
- sdxl
Expand Down
6 changes: 6 additions & 0 deletions .env.example
Original file line number Diff line number Diff line change
Expand Up @@ -26,3 +26,9 @@ VYRO_API_KEY="vk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
# See https://api.minimax.chat/user-center/basic-information/interface-key
MINIMAX_API_ORG="xxxxxxxx"
MINIMAX_API_KEY="eyJhxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"

# Google Gemini AI
# Documentation:
# https://ai.google.dev/tutorials/ai-studio_quickstart
# https://makersuite.google.com/app/apikey
GEMINI_API_KEY="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "@zhengxs/ai",
"version": "0.4.0",
"version": "0.5.0",
"description": "llm sdk",
"keywords": [
"ai",
Expand Down
14 changes: 14 additions & 0 deletions samples/gemini/chat/chat.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
import { GeminiAI } from '../../../src';

const ai = new GeminiAI();

async function main() {
const chatCompletion = await ai.chat.completions.create({
model: 'gemini-pro',
messages: [{ role: 'user', content: 'Say this is a test' }],
});

console.log(chatCompletion.choices[0].message.content);
}

main();
17 changes: 17 additions & 0 deletions samples/gemini/chat/stream.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
import { GeminiAI } from '../../../src';

const client = new GeminiAI();

async function main() {
const stream = await client.chat.completions.create({
stream: true,
model: 'gemini-pro',
messages: [{ role: 'user', content: 'Say this is a test' }],
});

for await (const chunk of stream) {
console.log(chunk.choices[0]?.delta?.content || '');
}
}

main();
13 changes: 13 additions & 0 deletions samples/gemini/models/list.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
import { GeminiAI } from '../../../src';

const ai = new GeminiAI();

async function main() {
const list = await ai.models.list();

for await (const model of list) {
console.log(model);
}
}

main();
11 changes: 11 additions & 0 deletions samples/gemini/models/retrieve.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
import { GeminiAI } from '../../../src';

const ai = new GeminiAI();

async function main() {
const model = await ai.models.retrieve('gemini-pro');

console.log(model);
}

main();
73 changes: 73 additions & 0 deletions src/gemini/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
import type { Agent } from 'node:http';

import { APIClient, type DefaultQuery, type Fetch, type FinalRequestOptions, type Headers } from 'openai/core';

import * as API from './resources';

const BASE_URL = 'https://generativelanguage.googleapis.com/v1';

export interface GeminiAIOptions {
baseURL?: string;
apiKey?: string;
timeout?: number | undefined;
httpAgent?: Agent;
fetch?: Fetch | undefined;
defaultHeaders?: Headers;
defaultQuery?: DefaultQuery;
}

export class GeminiAI extends APIClient {
apiKey: string;

private _options: GeminiAIOptions;

constructor(options: GeminiAIOptions = {}) {
const {
apiKey = process.env.GEMINI_API_KEY || '',
baseURL = process.env.GEMINI_BASE_URL || BASE_URL,
timeout = 30000,
fetch = globalThis.fetch,
httpAgent = undefined,
...rest
} = options;

super({
baseURL,
timeout,
fetch,
httpAgent,
...rest,
});

this._options = options;

this.apiKey = apiKey;
}

chat = new API.Chat(this);

models = new API.Models(this);

protected override defaultHeaders(opts: FinalRequestOptions): Headers {
return {
...super.defaultHeaders(opts),
...this._options.defaultHeaders,
};
}

protected override defaultQuery(): DefaultQuery | undefined {
return {
...this._options.defaultQuery,
key: this.apiKey,
};
}
}

export namespace GeminiAI {
export type ChatModel = API.ChatModel;
export type ChatCompletionCreateParams = API.ChatCompletionCreateParams;
export type ChatCompletionCreateParamsStreaming = API.ChatCompletionCreateParamsStreaming;
export type ChatCompletionCreateParamsNonStreaming = API.ChatCompletionCreateParamsNonStreaming;
}

export default GeminiAI;
9 changes: 9 additions & 0 deletions src/gemini/resource.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
import { GeminiAI } from './index';

export class APIResource {
protected _client: GeminiAI;

constructor(client: GeminiAI) {
this._client = client;
}
}
6 changes: 6 additions & 0 deletions src/gemini/resources/chat/chat.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
import { APIResource } from '../../resource';
import { Completions } from './completions';

export class Chat extends APIResource {
completions = new Completions(this._client);
}
Loading

0 comments on commit e83bd50

Please sign in to comment.