Skip to content
Snippets Groups Projects
Unverified Commit 37525df5 authored by Parham Saidi's avatar Parham Saidi Committed by GitHub
Browse files

feat: Gemini Access via Vertex AI (#838)

parent a1f24753
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 878 additions and 8 deletions
.changeset/fuzzy-coins-breathe.md 0 → 100644
+ 6
− 0
View file @ 37525df5
---
"llamaindex": patch
"@llamaindex/examples": patch
---
Added support for accessing Gemini via Vertex AI
apps/docs/docs/modules/llms/available_llms/gemini.md
+ 30
− 0
View file @ 37525df5
......@@ -10,6 +10,36 @@ Settings.llm = new Gemini({
});
```
### Usage with Vertex AI
To use Gemini via Vertex AI you can use `GeminiVertexSession`.
GeminiVertexSession accepts the env variables: `GOOGLE_VERTEX_LOCATION` and `GOOGLE_VERTEX_PROJECT`
```ts
import { Gemini, GEMINI_MODEL, GeminiVertexSession } from "llamaindex";
const gemini = new Gemini({
model: GEMINI_MODEL.GEMINI_PRO,
session: new GeminiVertexSession({
location: "us-central1", // optional if provided by GOOGLE_VERTEX_LOCATION env variable
project: "project1", // optional if provided by GOOGLE_VERTEX_PROJECT env variable
googleAuthOptions: {...}, // optional, but useful for production. It accepts all values from `GoogleAuthOptions`
}),
});
```
[GoogleAuthOptions](https://github.com/googleapis/google-auth-library-nodejs/blob/main/src/auth/googleauth.ts)
To authenticate for local development:
```bash
npm install @google-cloud/vertexai
gcloud auth application-default login
```
To authenticate for production you'll have to use a [service account](https://cloud.google.com/docs/authentication/). `googleAuthOptions` has `credentials` which might be useful for you.
## Load and index documents
For this example, we will use a single document. In a real-world scenario, you would have multiple documents to index.
......
examples/gemini/chatVertex.ts 0 → 100644
+ 19
− 0
View file @ 37525df5
import { Gemini, GEMINI_MODEL, GeminiVertexSession } from "llamaindex";
(async () => {
const gemini = new Gemini({
model: GEMINI_MODEL.GEMINI_PRO,
session: new GeminiVertexSession(),
});
const result = await gemini.chat({
messages: [
{ content: "You want to talk in rhymes.", role: "system" },
{
content:
"How much wood would a woodchuck chuck if a woodchuck could chuck wood?",
role: "user",
},
],
});
console.log(result);
})();
packages/autotool/examples/02_nextjs/next.config.mjs
+ 2
− 1
View file @ 37525df5
import { withNext } from "@llamaindex/autotool/next";
import withLlamaIndex from "llamaindex/next";
/** @type {import('next').NextConfig} */
const nextConfig = {};
export default withNext(nextConfig);
export default withLlamaIndex(withNext(nextConfig));
packages/core/package.json
+ 1
− 0
View file @ 37525df5
......@@ -24,6 +24,7 @@
"@anthropic-ai/sdk": "^0.20.9",
"@aws-crypto/sha256-js": "^5.2.0",
"@datastax/astra-db-ts": "^1.1.0",
"@google-cloud/vertexai": "^1.1.0",
"@google/generative-ai": "^0.11.0",
"@grpc/grpc-js": "^1.10.7",
"@huggingface/inference": "^2.6.7",
......
packages/core/src/embeddings/GeminiEmbedding.ts
+ 9
− 3
View file @ 37525df5
import { GeminiSessionStore, type GeminiSession } from "../llm/gemini.js";
import { GeminiSession, GeminiSessionStore } from "../llm/gemini/base.js";
import { GEMINI_BACKENDS } from "../llm/gemini/types.js";
import { BaseEmbedding } from "./types.js";
export enum GEMINI_EMBEDDING_MODEL {
......@@ -8,6 +9,7 @@ export enum GEMINI_EMBEDDING_MODEL {
/**
* GeminiEmbedding is an alias for Gemini that implements the BaseEmbedding interface.
* Note: Vertex SDK currently does not support embeddings
*/
export class GeminiEmbedding extends BaseEmbedding {
model: GEMINI_EMBEDDING_MODEL;
......@@ -16,11 +18,15 @@ export class GeminiEmbedding extends BaseEmbedding {
constructor(init?: Partial<GeminiEmbedding>) {
super();
this.model = init?.model ?? GEMINI_EMBEDDING_MODEL.EMBEDDING_001;
this.session = init?.session ?? GeminiSessionStore.get();
this.session =
init?.session ??
(GeminiSessionStore.get({
backend: GEMINI_BACKENDS.GOOGLE,
}) as GeminiSession);
}
private async getEmbedding(prompt: string): Promise<number[]> {
const client = this.session.gemini.getGenerativeModel({
const client = this.session.getGenerativeModel({
model: this.model,
});
const result = await client.embedContent(prompt);
......
packages/core/src/index.ts
+ 3
− 0
View file @ 37525df5
......@@ -10,3 +10,6 @@ export {
HuggingFaceEmbedding,
HuggingFaceEmbeddingModelType,
} from "./embeddings/HuggingFaceEmbedding.js";
export { type VertexGeminiSessionOptions } from "./llm/gemini/types.js";
export { GeminiVertexSession } from "./llm/gemini/vertex.js";
packages/core/src/llm/gemini.ts → packages/core/src/llm/gemini/base.ts
+ 248
− 0
View file @ 37525df5
import {
ChatSession,
GoogleGenerativeAI,
type Content as GeminiMessageContent,
type Part,
GenerativeModel as GoogleGenerativeModel,
type EnhancedGenerateContentResponse,
type ModelParams as GoogleModelParams,
type GenerateContentStreamResult as GoogleStreamGenerateContentResult,
} from "@google/generative-ai";
import { getEnv } from "@llamaindex/env";
import { ToolCallLLM } from "./base.js";
import { ToolCallLLM } from "../base.js";
import type {
ChatMessage,
ChatResponse,
ChatResponseChunk,
CompletionResponse,
LLMChatParamsNonStreaming,
LLMChatParamsStreaming,
LLMCompletionParamsNonStreaming,
LLMCompletionParamsStreaming,
LLMMetadata,
MessageContent,
MessageContentImageDetail,
MessageContentTextDetail,
MessageType,
ToolCallLLMMessageOptions,
} from "../types.js";
import { streamConverter, wrapLLMEvent } from "../utils.js";
import {
GEMINI_BACKENDS,
GEMINI_MODEL,
type GeminiAdditionalChatOptions,
type GeminiChatNonStreamResponse,
type GeminiChatParamsNonStreaming,
type GeminiChatParamsStreaming,
type GeminiChatStreamResponse,
type GeminiMessageRole,
type GeminiModelInfo,
type GeminiSessionOptions,
type GoogleGeminiSessionOptions,
type IGeminiSession,
} from "./types.js";
import { streamConverter, wrapLLMEvent } from "./utils.js";
// Session and Model Type Definitions
type GeminiSessionOptions = {
apiKey?: string;
};
export enum GEMINI_MODEL {
GEMINI_PRO = "gemini-pro",
GEMINI_PRO_VISION = "gemini-pro-vision",
GEMINI_PRO_LATEST = "gemini-1.5-pro-latest",
}
export interface GeminiModelInfo {
contextWindow: number;
}
import { GeminiHelper, getChatContext, getPartsText } from "./utils.js";
export const GEMINI_MODEL_INFO_MAP: Record<GEMINI_MODEL, GeminiModelInfo> = {
[GEMINI_MODEL.GEMINI_PRO]: { contextWindow: 30720 },
......@@ -58,38 +50,16 @@ const DEFAULT_GEMINI_PARAMS = {
};
export type GeminiConfig = Partial<typeof DEFAULT_GEMINI_PARAMS> & {
session?: GeminiSession;
session?: IGeminiSession;
};
/// Chat Type Definitions
type GeminiMessageRole = "user" | "model";
export type GeminiAdditionalChatOptions = {};
export type GeminiChatParamsStreaming = LLMChatParamsStreaming<
GeminiAdditionalChatOptions,
ToolCallLLMMessageOptions
>;
export type GeminiChatStreamResponse = AsyncIterable<
ChatResponseChunk<ToolCallLLMMessageOptions>
>;
export type GeminiChatParamsNonStreaming = LLMChatParamsNonStreaming<
GeminiAdditionalChatOptions,
ToolCallLLMMessageOptions
>;
export type GeminiChatNonStreamResponse =
ChatResponse<ToolCallLLMMessageOptions>;
/**
* Gemini Session to manage the connection to the Gemini API
*/
export class GeminiSession {
gemini: GoogleGenerativeAI;
export class GeminiSession implements IGeminiSession {
private gemini: GoogleGenerativeAI;
constructor(options: GeminiSessionOptions) {
constructor(options: GoogleGeminiSessionOptions) {
if (!options.apiKey) {
options.apiKey = getEnv("GOOGLE_API_KEY");
}
......@@ -98,6 +68,32 @@ export class GeminiSession {
}
this.gemini = new GoogleGenerativeAI(options.apiKey);
}
getGenerativeModel(metadata: GoogleModelParams): GoogleGenerativeModel {
return this.gemini.getGenerativeModel(metadata);
}
getResponseText(response: EnhancedGenerateContentResponse): string {
return response.text();
}
async *getChatStream(
result: GoogleStreamGenerateContentResult,
): GeminiChatStreamResponse {
yield* streamConverter(result.stream, (response) => ({
delta: this.getResponseText(response),
raw: response,
}));
}
getCompletionStream(
result: GoogleStreamGenerateContentResult,
): AsyncIterable<CompletionResponse> {
return streamConverter(result.stream, (response) => ({
text: this.getResponseText(response),
raw: response,
}));
}
}
/**
......@@ -105,109 +101,43 @@ export class GeminiSession {
*/
export class GeminiSessionStore {
static sessions: Array<{
session: GeminiSession;
session: IGeminiSession;
options: GeminiSessionOptions;
}> = [];
private static getSessionId(options: GeminiSessionOptions): string {
if (options.backend === GEMINI_BACKENDS.GOOGLE)
return options?.apiKey ?? "";
return "";
}
private static sessionMatched(
o1: GeminiSessionOptions,
o2: GeminiSessionOptions,
): boolean {
return o1.apiKey === o2.apiKey;
return (
GeminiSessionStore.getSessionId(o1) ===
GeminiSessionStore.getSessionId(o2)
);
}
static get(options: GeminiSessionOptions = {}): GeminiSession {
static get(
options: GeminiSessionOptions = { backend: GEMINI_BACKENDS.GOOGLE },
): IGeminiSession {
let session = this.sessions.find((session) =>
this.sessionMatched(session.options, options),
)?.session;
if (!session) {
if (session) return session;
if (options.backend === GEMINI_BACKENDS.VERTEX) {
throw Error("No Session");
} else {
session = new GeminiSession(options);
this.sessions.push({ session, options });
}
this.sessions.push({ session, options });
return session;
}
}
/**
* Helper class providing utility functions for Gemini
*/
class GeminiHelper {
// Gemini only has user and model roles. Put the rest in user role.
public static readonly ROLES_TO_GEMINI: Record<
MessageType,
GeminiMessageRole
> = {
user: "user",
system: "user",
assistant: "model",
memory: "user",
};
public static readonly ROLES_FROM_GEMINI: Record<
GeminiMessageRole,
MessageType
> = {
user: "user",
model: "assistant",
};
public static mergeNeighboringSameRoleMessages(
messages: GeminiMessageContent[],
): GeminiMessageContent[] {
return messages.reduce(
(
result: GeminiMessageContent[],
current: GeminiMessageContent,
index: number,
) => {
if (index > 0 && messages[index - 1].role === current.role) {
result[result.length - 1].parts = [
...result[result.length - 1].parts,
...current.parts,
];
} else {
result.push(current);
}
return result;
},
[],
);
}
public static messageContentToGeminiParts(content: MessageContent): Part[] {
if (typeof content === "string") {
return [{ text: content }];
}
const parts: Part[] = [];
const imageContents = content.filter(
(i) => i.type === "image_url",
) as MessageContentImageDetail[];
parts.push(
...imageContents.map((i) => ({
fileData: {
mimeType: i.type,
fileUri: i.image_url.url,
},
})),
);
const textContents = content.filter(
(i) => i.type === "text",
) as MessageContentTextDetail[];
parts.push(...textContents.map((t) => ({ text: t.text })));
return parts;
}
public static chatMessageToGemini(
message: ChatMessage,
): GeminiMessageContent {
return {
role: GeminiHelper.ROLES_TO_GEMINI[message.role],
parts: GeminiHelper.messageContentToGeminiParts(message.content),
};
}
}
/**
* ToolCallLLM for Gemini
*/
......@@ -216,7 +146,7 @@ export class Gemini extends ToolCallLLM<GeminiAdditionalChatOptions> {
temperature: number;
topP: number;
maxTokens?: number;
session: GeminiSession;
session: IGeminiSession;
constructor(init?: GeminiConfig) {
super();
......@@ -242,40 +172,21 @@ export class Gemini extends ToolCallLLM<GeminiAdditionalChatOptions> {
};
}
private prepareChat(
params: GeminiChatParamsStreaming | GeminiChatParamsNonStreaming,
): {
chat: ChatSession;
messageContent: Part[];
} {
const messages = GeminiHelper.mergeNeighboringSameRoleMessages(
params.messages.map(GeminiHelper.chatMessageToGemini),
);
const history = messages.slice(0, -1);
const client = this.session.gemini.getGenerativeModel(this.metadata);
const chat = client.startChat({
history,
});
return {
chat,
messageContent: messages[messages.length - 1].parts,
};
}
protected async nonStreamChat(
params: GeminiChatParamsNonStreaming,
): Promise<GeminiChatNonStreamResponse> {
const { chat, messageContent } = this.prepareChat(params);
const result = await chat.sendMessage(messageContent);
const { response } = result;
const context = getChatContext(params);
const client = this.session.getGenerativeModel(this.metadata);
const chat = client.startChat({
history: context.history,
});
const { response } = await chat.sendMessage(context.message);
const topCandidate = response.candidates![0];
return {
raw: response,
message: {
content: response.text(),
content: this.session.getResponseText(response),
role: GeminiHelper.ROLES_FROM_GEMINI[
topCandidate.content.role as GeminiMessageRole
],
......@@ -286,12 +197,13 @@ export class Gemini extends ToolCallLLM<GeminiAdditionalChatOptions> {
protected async *streamChat(
params: GeminiChatParamsStreaming,
): GeminiChatStreamResponse {
const { chat, messageContent } = this.prepareChat(params);
const result = await chat.sendMessageStream(messageContent);
yield* streamConverter(result.stream, (response) => ({
delta: response.text(),
raw: response,
}));
const context = getChatContext(params);
const client = this.session.getGenerativeModel(this.metadata);
const chat = client.startChat({
history: context.history,
});
const result = await chat.sendMessageStream(context.message);
yield* this.session.getChatStream(result);
}
chat(params: GeminiChatParamsStreaming): Promise<GeminiChatStreamResponse>;
......@@ -316,25 +228,20 @@ export class Gemini extends ToolCallLLM<GeminiAdditionalChatOptions> {
params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
): Promise<CompletionResponse | AsyncIterable<CompletionResponse>> {
const { prompt, stream } = params;
const client = this.session.gemini.getGenerativeModel(this.metadata);
const client = this.session.getGenerativeModel(this.metadata);
if (stream) {
const result = await client.generateContentStream(
GeminiHelper.messageContentToGeminiParts(prompt),
getPartsText(GeminiHelper.messageContentToGeminiParts(prompt)),
);
return streamConverter(result.stream, (response) => {
return {
text: response.text(),
raw: response,
};
});
return this.session.getCompletionStream(result);
}
const result = await client.generateContent(
GeminiHelper.messageContentToGeminiParts(prompt),
getPartsText(GeminiHelper.messageContentToGeminiParts(prompt)),
);
return {
text: result.response.text(),
text: this.session.getResponseText(result.response),
raw: result.response,
};
}
......
packages/core/src/llm/gemini/types.ts 0 → 100644
+ 111
− 0
View file @ 37525df5
import {
GenerativeModel as GoogleGenerativeModel,
type EnhancedGenerateContentResponse,
type Content as GeminiMessageContent,
type FileDataPart as GoogleFileDataPart,
type InlineDataPart as GoogleInlineFileDataPart,
type ModelParams as GoogleModelParams,
type Part as GooglePart,
type GenerateContentStreamResult as GoogleStreamGenerateContentResult,
} from "@google/generative-ai";
import {
GenerativeModel as VertexGenerativeModel,
GenerativeModelPreview as VertexGenerativeModelPreview,
type GenerateContentResponse,
type FileDataPart as VertexFileDataPart,
type VertexInit,
type InlineDataPart as VertexInlineFileDataPart,
type ModelParams as VertexModelParams,
type Part as VertexPart,
type StreamGenerateContentResult as VertexStreamGenerateContentResult,
} from "@google-cloud/vertexai";
import type {
ChatResponse,
ChatResponseChunk,
CompletionResponse,
LLMChatParamsNonStreaming,
LLMChatParamsStreaming,
ToolCallLLMMessageOptions,
} from "../types.js";
export enum GEMINI_BACKENDS {
GOOGLE = "google",
VERTEX = "vertex",
}
export type GoogleGeminiSessionOptions = {
apiKey?: string;
};
export type VertexGeminiSessionOptions = {
preview?: boolean;
} & VertexInit;
export type GeminiSessionOptions =
| (GoogleGeminiSessionOptions & { backend: GEMINI_BACKENDS.GOOGLE })
| (VertexGeminiSessionOptions & { backend: GEMINI_BACKENDS.VERTEX });
export enum GEMINI_MODEL {
GEMINI_PRO = "gemini-pro",
GEMINI_PRO_VISION = "gemini-pro-vision",
GEMINI_PRO_LATEST = "gemini-1.5-pro-latest",
}
export interface GeminiModelInfo {
contextWindow: number;
}
export type Part = GooglePart | VertexPart;
export type FileDataPart = GoogleFileDataPart | VertexFileDataPart;
export type InlineDataPart =
| GoogleInlineFileDataPart
| VertexInlineFileDataPart;
export type ModelParams = GoogleModelParams | VertexModelParams;
export type GenerativeModel =
| VertexGenerativeModelPreview
| VertexGenerativeModel
| GoogleGenerativeModel;
export type ChatContext = { message: Part[]; history: GeminiMessageContent[] };
export type GeminiMessageRole = "user" | "model";
export type GeminiAdditionalChatOptions = {};
export type GeminiChatParamsStreaming = LLMChatParamsStreaming<
GeminiAdditionalChatOptions,
ToolCallLLMMessageOptions
>;
export type GeminiChatStreamResponse = AsyncIterable<
ChatResponseChunk<ToolCallLLMMessageOptions>
>;
export type GeminiChatParamsNonStreaming = LLMChatParamsNonStreaming<
GeminiAdditionalChatOptions,
ToolCallLLMMessageOptions
>;
export type GeminiChatNonStreamResponse =
ChatResponse<ToolCallLLMMessageOptions>;
export interface IGeminiSession {
getGenerativeModel(metadata: ModelParams): GenerativeModel;
getResponseText(
response: EnhancedGenerateContentResponse | GenerateContentResponse,
): string;
getCompletionStream(
result:
| GoogleStreamGenerateContentResult
| VertexStreamGenerateContentResult,
): AsyncIterable<CompletionResponse>;
getChatStream(
result:
| GoogleStreamGenerateContentResult
| VertexStreamGenerateContentResult,
): GeminiChatStreamResponse;
}
packages/core/src/llm/gemini/utils.ts 0 → 100644
+ 197
− 0
View file @ 37525df5
import { type Content as GeminiMessageContent } from "@google/generative-ai";
import { type GenerateContentResponse } from "@google-cloud/vertexai";
import type {
ChatMessage,
MessageContent,
MessageContentImageDetail,
MessageContentTextDetail,
MessageType,
} from "../types.js";
import { extractDataUrlComponents } from "../utils.js";
import type {
ChatContext,
FileDataPart,
GeminiChatParamsNonStreaming,
GeminiChatParamsStreaming,
GeminiMessageRole,
InlineDataPart,
Part,
} from "./types.js";
const FILE_EXT_MIME_TYPES: { [key: string]: string } = {
png: "image/png",
jpeg: "image/jpeg",
jpg: "image/jpeg",
webp: "image/webp",
heic: "image/heic",
heif: "image/heif",
};
const ACCEPTED_IMAGE_MIME_TYPES = Object.values(FILE_EXT_MIME_TYPES);
const getFileURLExtension = (url: string): string | null => {
const pathname = new URL(url).pathname;
const parts = pathname.split(".");
return parts.length > 1 ? parts.pop()?.toLowerCase() || null : null;
};
const getFileURLMimeType = (url: string): string | null => {
const ext = getFileURLExtension(url);
return ext ? FILE_EXT_MIME_TYPES[ext] || null : null;
};
const getImageParts = (
message: MessageContentImageDetail,
): InlineDataPart | FileDataPart => {
if (message.image_url.url.startsWith("data:")) {
const { mimeType, base64: data } = extractDataUrlComponents(
message.image_url.url,
);
if (!mimeType || !ACCEPTED_IMAGE_MIME_TYPES.includes(mimeType))
throw new Error(
`Gemini only accepts the following mimeTypes: ${ACCEPTED_IMAGE_MIME_TYPES.join("\n")}`,
);
return {
inlineData: {
mimeType,
data,
},
};
}
const mimeType = getFileURLMimeType(message.image_url.url);
if (!mimeType || !ACCEPTED_IMAGE_MIME_TYPES.includes(mimeType))
throw new Error(
`Gemini only accepts the following mimeTypes: ${ACCEPTED_IMAGE_MIME_TYPES.join("\n")}`,
);
return {
fileData: { mimeType, fileUri: message.image_url.url },
};
};
export const getPartsText = (parts: Part[]): string => {
const textStrings = [];
if (parts.length) {
for (const part of parts) {
if (part.text) {
textStrings.push(part.text);
}
}
}
if (textStrings.length > 0) {
return textStrings.join("");
} else {
return "";
}
};
/**
* Returns all text found in all parts of first candidate.
*/
export const getText = (response: GenerateContentResponse): string => {
if (response.candidates?.[0].content?.parts) {
return getPartsText(response.candidates?.[0].content?.parts);
}
return "";
};
export const cleanParts = (
message: GeminiMessageContent,
): GeminiMessageContent => {
return {
...message,
parts: message.parts.filter((part) => part.text?.trim()),
};
};
export const getChatContext = (
params: GeminiChatParamsStreaming | GeminiChatParamsNonStreaming,
): ChatContext => {
// Gemini doesn't allow:
// 1. Consecutive messages from the same role
// 2. Parts that have empty text
const messages = GeminiHelper.mergeNeighboringSameRoleMessages(
params.messages.map(GeminiHelper.chatMessageToGemini),
).map(cleanParts);
const history = messages.slice(0, -1);
const message = messages[messages.length - 1].parts;
return {
history,
message,
};
};
/**
* Helper class providing utility functions for Gemini
*/
export class GeminiHelper {
// Gemini only has user and model roles. Put the rest in user role.
public static readonly ROLES_TO_GEMINI: Record<
MessageType,
GeminiMessageRole
> = {
user: "user",
system: "user",
assistant: "model",
memory: "user",
};
public static readonly ROLES_FROM_GEMINI: Record<
GeminiMessageRole,
MessageType
> = {
user: "user",
model: "assistant",
};
public static mergeNeighboringSameRoleMessages(
messages: GeminiMessageContent[],
): GeminiMessageContent[] {
return messages.reduce(
(
result: GeminiMessageContent[],
current: GeminiMessageContent,
index: number,
) => {
if (index > 0 && messages[index - 1].role === current.role) {
result[result.length - 1].parts = [
...result[result.length - 1].parts,
...current.parts,
];
} else {
result.push(current);
}
return result;
},
[],
);
}
public static messageContentToGeminiParts(content: MessageContent): Part[] {
if (typeof content === "string") {
return [{ text: content }];
}
const parts: Part[] = [];
const imageContents = content.filter(
(i) => i.type === "image_url",
) as MessageContentImageDetail[];
parts.push(...imageContents.map(getImageParts));
const textContents = content.filter(
(i) => i.type === "text",
) as MessageContentTextDetail[];
parts.push(...textContents.map((t) => ({ text: t.text })));
return parts;
}
public static chatMessageToGemini(
message: ChatMessage,
): GeminiMessageContent {
return {
role: GeminiHelper.ROLES_TO_GEMINI[message.role],
parts: GeminiHelper.messageContentToGeminiParts(message.content),
};
}
}
packages/core/src/llm/gemini/vertex.ts 0 → 100644
+ 81
− 0
View file @ 37525df5
import {
VertexAI,
GenerativeModel as VertexGenerativeModel,
GenerativeModelPreview as VertexGenerativeModelPreview,
type GenerateContentResponse,
type ModelParams as VertexModelParams,
type StreamGenerateContentResult as VertexStreamGenerateContentResult,
} from "@google-cloud/vertexai";
import type {
GeminiChatStreamResponse,
IGeminiSession,
VertexGeminiSessionOptions,
} from "./types.js";
import { getEnv } from "@llamaindex/env";
import type { CompletionResponse } from "../types.js";
import { streamConverter } from "../utils.js";
import { getText } from "./utils.js";
/* To use Google's Vertex AI backend, it doesn't use api key authentication.
*
* To authenticate for local development:
*
* ```
* npm install @google-cloud/vertexai
* gcloud auth application-default login
* ```
* For production the prefered method is via a service account, more
* details: https://cloud.google.com/docs/authentication/
*
* */
export class GeminiVertexSession implements IGeminiSession {
private vertex: VertexAI;
private preview: boolean = false;
constructor(options?: Partial<VertexGeminiSessionOptions>) {
const project = options?.project ?? getEnv("GOOGLE_VERTEX_PROJECT");
const location = options?.location ?? getEnv("GOOGLE_VERTEX_LOCATION");
if (!project || !location) {
throw new Error(
"Set Google Vertex project and location in GOOGLE_VERTEX_PROJECT and GOOGLE_VERTEX_LOCATION env variables",
);
}
this.vertex = new VertexAI({
...options,
project,
location,
});
this.preview = options?.preview ?? false;
}
getGenerativeModel(
metadata: VertexModelParams,
): VertexGenerativeModelPreview | VertexGenerativeModel {
if (this.preview) return this.vertex.preview.getGenerativeModel(metadata);
return this.vertex.getGenerativeModel(metadata);
}
getResponseText(response: GenerateContentResponse): string {
return getText(response);
}
async *getChatStream(
result: VertexStreamGenerateContentResult,
): GeminiChatStreamResponse {
yield* streamConverter(result.stream, (response) => ({
delta: this.getResponseText(response),
raw: response,
}));
}
getCompletionStream(
result: VertexStreamGenerateContentResult,
): AsyncIterable<CompletionResponse> {
return streamConverter(result.stream, (response) => ({
text: this.getResponseText(response),
raw: response,
}));
}
}
packages/core/src/llm/index.ts
+ 7
− 1
View file @ 37525df5
......@@ -5,7 +5,13 @@ export {
Anthropic,
} from "./anthropic.js";
export { FireworksLLM } from "./fireworks.js";
export { GEMINI_MODEL, Gemini, GeminiSession } from "./gemini.js";
export { Gemini, GeminiSession } from "./gemini/base.js";
export {
GEMINI_MODEL,
type GoogleGeminiSessionOptions,
} from "./gemini/types.js";
export { Groq } from "./groq.js";
export { HuggingFaceInferenceAPI, HuggingFaceLLM } from "./huggingface.js";
export {
......
packages/core/src/llm/utils.ts
+ 21
− 0
View file @ 37525df5
......@@ -77,6 +77,27 @@ export function extractText(message: MessageContent): string {
}
}
export const extractDataUrlComponents = (
dataUrl: string,
): {
mimeType: string;
base64: string;
} => {
const parts = dataUrl.split(";base64,");
if (parts.length !== 2 || !parts[0].startsWith("data:")) {
throw new Error("Invalid data URL");
}
const mimeType = parts[0].slice(5);
const base64 = parts[1];
return {
mimeType,
base64,
};
};
/**
* @internal
*/
......
packages/core/src/next.ts
+ 1
− 0
View file @ 37525df5
......@@ -28,6 +28,7 @@ export default function withLlamaIndex(config: any) {
...webpackConfig.resolve.alias,
sharp$: false,
"onnxruntime-node$": false,
"@google-cloud/vertexai": false,
};
return webpackConfig;
};
......
pnpm-lock.yaml
+ 142
− 3
View file @ 37525df5
......@@ -157,7 +157,7 @@ importers:
version: link:../packages/core
mongodb:
specifier: ^6.6.1
version: 6.6.1
version: 6.6.1(gcp-metadata@6.1.0(encoding@0.1.13))
pathe:
specifier: ^1.1.2
version: 1.1.2
......@@ -343,6 +343,9 @@ importers:
'@datastax/astra-db-ts':
specifier: ^1.1.0
version: 1.1.0
'@google-cloud/vertexai':
specifier: ^1.1.0
version: 1.1.0(encoding@0.1.13)
'@google/generative-ai':
specifier: ^0.11.0
version: 0.11.0
......@@ -411,7 +414,7 @@ importers:
version: 2.0.0
mongodb:
specifier: ^6.6.1
version: 6.6.1
version: 6.6.1(gcp-metadata@6.1.0(encoding@0.1.13))
notion-md-crawler:
specifier: ^1.0.0
version: 1.0.0(encoding@0.1.13)
......@@ -2106,6 +2109,10 @@ packages:
'@fastify/deepmerge@1.3.0':
resolution: {integrity: sha512-J8TOSBq3SoZbDhM9+R/u77hP93gz/rajSA+K2kGyijPpORPWUXHUpTaleoj+92As0S9uPRP7Oi8IqMf0u+ro6A==}
 
'@google-cloud/vertexai@1.1.0':
resolution: {integrity: sha512-hfwfdlVpJ+kM6o2b5UFfPnweBcz8tgHAFRswnqUKYqLJsvKU0DDD0Z2/YKoHyAUoPJAv20qg6KlC3msNeUKUiw==}
engines: {node: '>=18.0.0'}
'@google/generative-ai@0.11.0':
resolution: {integrity: sha512-nyvrIltnq4RRkUj4FBeDPebWHB1oeALK75aa0xdatEriH05lzBJD6t91afRm5ldtQ5Txu6UwPLuHu5jM1CjtQg==}
engines: {node: '>=18.0.0'}
......@@ -3521,6 +3528,10 @@ packages:
resolution: {integrity: sha512-4B/qKCfeE/ODUaAUpSwfzazo5x29WD4r3vXiWsB7I2mSDAihwEqKO+g8GELZUQSSAo5e1XTYh3ZVfLyxBc12nA==}
engines: {node: '>= 10.0.0'}
 
agent-base@7.1.1:
resolution: {integrity: sha512-H0TSyFNDMomMNJQBn8wFV5YC/2eJ+VXECwOadZJT554xP6cODZHPX3H9QMQECxvrgiSOP1pHjy1sMWQVYJOUOA==}
engines: {node: '>= 14'}
agentkeepalive@4.5.0:
resolution: {integrity: sha512-5GG/5IbQQpC9FpkRGsSvZI5QYeSCzlJHdpBQntCsuTOxhKD8lqKhrleg2Yi7yvMIf82Ycmmqln9U8V9qwEiJew==}
engines: {node: '>= 8.0.0'}
......@@ -3831,6 +3842,9 @@ packages:
big.js@5.2.2:
resolution: {integrity: sha512-vyL2OymJxmarO8gxMr0mhChsO9QGwhynfuu4+MHTAW6czfq9humCB7rKpUjDd9YUiDPU4mzpyupFSvOClAwbmQ==}
 
bignumber.js@9.1.2:
resolution: {integrity: sha512-2/mKyZH9K85bzOEfhXDBFZTGd1CTs+5IHpeFQo9luiBG7hghdC851Pj2WAhb6E3R6b9tZj/XKhbg4fum+Kepug==}
bin-check@4.1.0:
resolution: {integrity: sha512-b6weQyEUKsDGFlACWSIOfveEnImkJyK/FGW6FAG42loyoquvjdtOIqO6yBFzHyqyVVhNgNkQxxx09SFLK28YnA==}
engines: {node: '>=4'}
......@@ -3906,6 +3920,9 @@ packages:
resolution: {integrity: sha512-w2IquM5mYzYZv6rs3uN2DZTOBe2a0zXLj53TGDqwF4l6Sz/XsISrisXOJihArF9+BZ6Cq/GjVht7Sjfmri7ytQ==}
engines: {node: '>=16.20.1'}
 
buffer-equal-constant-time@1.0.1:
resolution: {integrity: sha512-zRpUiDwd/xk6ADqPMATG8vc9VPrkck7T07OIx0gnjmJAnHnTVXNQG3vfvWNuiZIkwu9KrKdA1iJKfsfTVxE6NA==}
buffer-from@1.1.2:
resolution: {integrity: sha512-E+XQCRwSbaaiChtv6k6Dwgc+bx+Bs6vuKJHHl5kox/BaKbhiXzqQOwK4cO22yElGp2OCmjwVhT3HmxgyPGnJfQ==}
 
......@@ -4778,6 +4795,9 @@ packages:
eastasianwidth@0.2.0:
resolution: {integrity: sha512-I88TYZWc9XiYHRQ4/3c5rjjfgkjhLyW2luGIheGERbNQ6OY7yTybanSpDXZa8y7VUP9YmDcYa+eyq4ca7iLqWA==}
 
ecdsa-sig-formatter@1.0.11:
resolution: {integrity: sha512-nagl3RYrbNv6kQkeJIpt6NJZy8twLB/2vtz6yN9Z4vRKHN4/QZJIEbqohALSgwKdnksuY3k5Addp5lg8sVoVcQ==}
ee-first@1.1.1:
resolution: {integrity: sha512-WMwm9LhRUo+WUaRN+vRuETqG89IgZphVSNkdFgeb6sS/E4OrDIN7t48CAewSHXc6C8lefD8KKfr5vY61brQlow==}
 
......@@ -5406,6 +5426,14 @@ packages:
functions-have-names@1.2.3:
resolution: {integrity: sha512-xckBUXyTIqT97tq2x2AMb+g163b5JFysYk0x4qxNFwbfQkmNZoiRHb6sPzI9/QV33WeuvVYBUIiD4NzNIyqaRQ==}
 
gaxios@6.5.0:
resolution: {integrity: sha512-R9QGdv8j4/dlNoQbX3hSaK/S0rkMijqjVvW3YM06CoBdbU/VdKd159j4hePpng0KuE6Lh6JJ7UdmVGJZFcAG1w==}
engines: {node: '>=14'}
gcp-metadata@6.1.0:
resolution: {integrity: sha512-Jh/AIwwgaxan+7ZUUmRLCjtchyDiqh4KjBJ5tW3plBZb5iL/BPcso8A5DlzeD9qlw0duCamnNdpFjxwaT0KyKg==}
engines: {node: '>=14'}
generic-pool@3.9.0:
resolution: {integrity: sha512-hymDOu5B53XvN4QT9dBmZxPX4CWhBPPLguTZ9MMFeFa/Kg0xWVfylOVNlJji/E7yTZWFd/q9GO5TxDLq156D7g==}
engines: {node: '>= 4'}
......@@ -5541,6 +5569,10 @@ packages:
engines: {node: '>=0.6.0'}
hasBin: true
 
google-auth-library@9.10.0:
resolution: {integrity: sha512-ol+oSa5NbcGdDqA+gZ3G3mev59OHBZksBTxY/tYwjtcp1H/scAFwJfSQU9/1RALoyZ7FslNbke8j4i3ipwlyuQ==}
engines: {node: '>=14'}
gopd@1.0.1:
resolution: {integrity: sha512-d65bNlIadxvpb/A2abVdlqKqV563juRnZ1Wtk6s1sIR8uNsXR70xqIzVqxVf1eTqDunwT2MkczEeaezCKTZhwA==}
 
......@@ -5568,6 +5600,10 @@ packages:
resolution: {integrity: sha512-5v6yZd4JK3eMI3FqqCouswVqwugaA9r4dNZB1wwcmrD02QkV5H0y7XBQW8QwQqEaZY1pM9aqORSORhJRdNK44Q==}
engines: {node: '>=6.0'}
 
gtoken@7.1.0:
resolution: {integrity: sha512-pCcEwRi+TKpMlxAQObHDQ56KawURgyAf6jtIY046fJ5tIv3zDe/LEIubckAO8fj6JnAxLdmWkUfNyulQ2iKdEw==}
engines: {node: '>=14.0.0'}
guid-typescript@1.0.9:
resolution: {integrity: sha512-Y8T4vYhEfwJOTbouREvG+3XDsjr8E3kIr7uf+JZ0BYloFsttiHU0WfvANVsR7TxNUJa/WpCnw/Ino/p+DeBhBQ==}
 
......@@ -5763,6 +5799,10 @@ packages:
resolution: {integrity: sha512-V5nVw1PAOgfI3Lmeaj2Exmeg7fenjhRUgz1lPSezy1CuhPYbgQtbQj4jZfEAEMlaL+vupsvhjqCyjzob0yxsmQ==}
engines: {node: '>=10.19.0'}
 
https-proxy-agent@7.0.4:
resolution: {integrity: sha512-wlwpilI7YdjSkWaQ/7omYBMTliDcmCN8OLihO6I9B86g06lMyAoqgoDpV0XqoaPOKj+0DIdAvnsWfyAAhmimcg==}
engines: {node: '>= 14'}
human-id@1.0.2:
resolution: {integrity: sha512-UNopramDEhHJD+VR+ehk8rOslwSfByxPIZyJRfV739NDhN5LF1fa1MqnzKm2lGTQRjNrjK19Q5fhkgIfjlVUKw==}
 
......@@ -6250,6 +6290,9 @@ packages:
engines: {node: '>=4'}
hasBin: true
 
json-bigint@1.0.0:
resolution: {integrity: sha512-SiPv/8VpZuWbvLSMtTDU8hEfrZWg/mH/nV/b4o0CYbSxu1UIQPLdwKOCIyLQX+VIPO5vrLX3i8qtqFyhdPSUSQ==}
json-buffer@3.0.1:
resolution: {integrity: sha512-4bV5BfR2mqfQTJm+V5tPPdf+ZpuhiIvTuAB5g8kcrXOZpTT/QwwVRWBywX1ozr6lEuPdbHxwaJlm9G6mI2sfSQ==}
 
......@@ -6304,6 +6347,12 @@ packages:
jszip@3.10.1:
resolution: {integrity: sha512-xXDvecyTpGLrqFrvkrUSoxxfJI5AH7U8zxxtVclpsUtMCq4JQ290LY8AW5c7Ggnr/Y/oK+bQMbqK2qmtk3pN4g==}
 
jwa@2.0.0:
resolution: {integrity: sha512-jrZ2Qx916EA+fq9cEAeCROWPTfCwi1IVHqT2tapuqLEVVDKFDENFw1oL+MwrTvH6msKxsd1YTDVw6uKEcsrLEA==}
jws@4.0.0:
resolution: {integrity: sha512-KDncfTmOZoOMTFG4mBlG0qUIOlc03fmzH+ru6RgYVZhPkyiy/92Owlt/8UEN+a4TXR1FQetfIpJE8ApdvdVxTg==}
keyv@4.5.4:
resolution: {integrity: sha512-oxVHkHR/EJf2CNXnWxRLW6mg7JyCCUcG0DtEGmL2ctUo1PNTin1PUil+r/+4r5MpVgC/fn1kjsx7mjSujKqIpw==}
 
......@@ -9488,6 +9537,10 @@ packages:
resolution: {integrity: sha512-+NYs2QeMWy+GWFOEm9xnn6HCDp0l7QBD7ml8zLUmJ+93Q5NF0NocErnwkTkXVFNiX3/fpC6afS8Dhb/gz7R7eg==}
hasBin: true
 
uuid@9.0.1:
resolution: {integrity: sha512-b+1eJOlsR9K8HJpow9Ok3fiWOWSIcIzXodvv0rQjVoOVNpWMpxf1wZNpt4y9h10odCNrqnYp1OBzRktckBe3sA==}
hasBin: true
uuidv7@0.6.3:
resolution: {integrity: sha512-zV3eW2NlXTsun/aJ7AixxZjH/byQcH/r3J99MI0dDEkU2cJIBJxhEWUHDTpOaLPRNhebPZoeHuykYREkI9HafA==}
hasBin: true
......@@ -12021,6 +12074,13 @@ snapshots:
 
'@fastify/deepmerge@1.3.0': {}
 
'@google-cloud/vertexai@1.1.0(encoding@0.1.13)':
dependencies:
google-auth-library: 9.10.0(encoding@0.1.13)
transitivePeerDependencies:
- encoding
- supports-color
'@google/generative-ai@0.11.0': {}
 
'@grpc/grpc-js@1.10.7':
......@@ -13486,6 +13546,12 @@ snapshots:
 
address@1.2.2: {}
 
agent-base@7.1.1:
dependencies:
debug: 4.3.4
transitivePeerDependencies:
- supports-color
agentkeepalive@4.5.0:
dependencies:
humanize-ms: 1.2.1
......@@ -13848,6 +13914,8 @@ snapshots:
 
big.js@5.2.2: {}
 
bignumber.js@9.1.2: {}
bin-check@4.1.0:
dependencies:
execa: 0.7.0
......@@ -13954,6 +14022,8 @@ snapshots:
 
bson@6.7.0: {}
 
buffer-equal-constant-time@1.0.1: {}
buffer-from@1.1.2: {}
 
buffer@5.7.1:
......@@ -14871,6 +14941,10 @@ snapshots:
 
eastasianwidth@0.2.0: {}
 
ecdsa-sig-formatter@1.0.11:
dependencies:
safe-buffer: 5.2.1
ee-first@1.1.1: {}
 
electron-to-chromium@1.4.730: {}
......@@ -15809,6 +15883,25 @@ snapshots:
 
functions-have-names@1.2.3: {}
 
gaxios@6.5.0(encoding@0.1.13):
dependencies:
extend: 3.0.2
https-proxy-agent: 7.0.4
is-stream: 2.0.1
node-fetch: 2.7.0(encoding@0.1.13)
uuid: 9.0.1
transitivePeerDependencies:
- encoding
- supports-color
gcp-metadata@6.1.0(encoding@0.1.13):
dependencies:
gaxios: 6.5.0(encoding@0.1.13)
json-bigint: 1.0.0
transitivePeerDependencies:
- encoding
- supports-color
generic-pool@3.9.0: {}
 
gensync@1.0.0-beta.2: {}
......@@ -15961,6 +16054,18 @@ snapshots:
dependencies:
minimist: 1.2.8
 
google-auth-library@9.10.0(encoding@0.1.13):
dependencies:
base64-js: 1.5.1
ecdsa-sig-formatter: 1.0.11
gaxios: 6.5.0(encoding@0.1.13)
gcp-metadata: 6.1.0(encoding@0.1.13)
gtoken: 7.1.0(encoding@0.1.13)
jws: 4.0.0
transitivePeerDependencies:
- encoding
- supports-color
gopd@1.0.1:
dependencies:
get-intrinsic: 1.2.4
......@@ -16008,6 +16113,14 @@ snapshots:
section-matter: 1.0.0
strip-bom-string: 1.0.0
 
gtoken@7.1.0(encoding@0.1.13):
dependencies:
gaxios: 6.5.0(encoding@0.1.13)
jws: 4.0.0
transitivePeerDependencies:
- encoding
- supports-color
guid-typescript@1.0.9: {}
 
gzip-size@6.0.0:
......@@ -16289,6 +16402,13 @@ snapshots:
quick-lru: 5.1.1
resolve-alpn: 1.2.1
 
https-proxy-agent@7.0.4:
dependencies:
agent-base: 7.1.1
debug: 4.3.4
transitivePeerDependencies:
- supports-color
human-id@1.0.2: {}
 
human-signals@2.1.0: {}
......@@ -16699,6 +16819,10 @@ snapshots:
 
jsesc@2.5.2: {}
 
json-bigint@1.0.0:
dependencies:
bignumber.js: 9.1.2
json-buffer@3.0.1: {}
 
json-parse-even-better-errors@2.3.1: {}
......@@ -16759,6 +16883,17 @@ snapshots:
readable-stream: 2.3.8
setimmediate: 1.0.5
 
jwa@2.0.0:
dependencies:
buffer-equal-constant-time: 1.0.1
ecdsa-sig-formatter: 1.0.11
safe-buffer: 5.2.1
jws@4.0.0:
dependencies:
jwa: 2.0.0
safe-buffer: 5.2.1
keyv@4.5.4:
dependencies:
json-buffer: 3.0.1
......@@ -17658,11 +17793,13 @@ snapshots:
'@types/whatwg-url': 11.0.4
whatwg-url: 13.0.0
 
mongodb@6.6.1:
mongodb@6.6.1(gcp-metadata@6.1.0(encoding@0.1.13)):
dependencies:
'@mongodb-js/saslprep': 1.1.6
bson: 6.7.0
mongodb-connection-string-url: 3.0.0
optionalDependencies:
gcp-metadata: 6.1.0(encoding@0.1.13)
 
mrmime@2.0.0: {}
 
......@@ -20574,6 +20711,8 @@ snapshots:
 
uuid@8.3.2: {}
 
uuid@9.0.1: {}
uuidv7@0.6.3: {}
 
v8-compile-cache-lib@3.0.1: {}
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

:feet::heart: Strive to be the person your dogs believe you are :heart::feet: