Compare commits

...

1 Commits

Author SHA1 Message Date
google-labs-jules[bot] 68529cadc2 chore: update to google/gen-ai fixing #1934 2025-05-23 04:08:18 +00:00
5 changed files with 260 additions and 172 deletions
+1 -3
View File
@@ -38,8 +38,6 @@
"@llamaindex/env": "workspace:*"
},
"dependencies": {
"@google-cloud/vertexai": "1.9.0",
"@google/genai": "^0.12.0",
"@google/generative-ai": "0.24.0"
"@google/genai": "^1.0.0"
}
}
@@ -15,20 +15,23 @@ export class GeminiEmbedding extends BaseEmbedding {
model: GEMINI_EMBEDDING_MODEL;
session: GeminiSession;
constructor(init?: Partial<GeminiEmbedding>) {
apiKey?: string; // Added for clarity, though session handles it
constructor(init?: Partial<GeminiEmbedding> & { apiKey?: string }) {
super();
this.model = init?.model ?? GEMINI_EMBEDDING_MODEL.EMBEDDING_001;
this.apiKey = init?.apiKey; // Store apiKey if provided
this.session =
init?.session ??
(GeminiSessionStore.get({
backend: GEMINI_BACKENDS.GOOGLE,
}) as GeminiSession);
(GeminiSessionStore.get(
{ backend: GEMINI_BACKENDS.GOOGLE, apiKey: this.apiKey }, // Use stored or passed apiKey
{ model: this.model }, // Pass modelParams
) as GeminiSession);
}
private async getEmbedding(prompt: string): Promise<number[]> {
const client = this.session.getGenerativeModel({
model: this.model,
});
// getGenerativeModel no longer takes arguments
const client = this.session.getGenerativeModel();
const result = await client.embedContent(prompt);
return result.embedding.values;
}
+178 -105
View File
@@ -8,9 +8,8 @@ import {
type StartChatParams as GoogleStartChatParams,
type GenerateContentStreamResult as GoogleStreamGenerateContentResult,
type SafetySetting,
} from "@google/generative-ai";
} from "@google/genai";
import type { StartChatParams as VertexStartChatParams } from "@google-cloud/vertexai";
import { wrapLLMEvent } from "@llamaindex/core/decorator";
import type {
CompletionResponse,
@@ -105,77 +104,113 @@ export type GeminiConfig = Partial<typeof DEFAULT_GEMINI_PARAMS> & {
voiceName?: GeminiVoiceName;
};
type StartChatParams = GoogleStartChatParams & VertexStartChatParams;
type StartChatParams = GoogleStartChatParams;
/**
* Gemini Session to manage the connection to the Gemini API
*/
export class GeminiSession implements IGeminiSession {
private gemini: GoogleGenerativeAI;
private generativeModel: GoogleGenerativeModel; // Added to store the model
constructor(options: GoogleGeminiSessionOptions) {
if (!options.apiKey) {
options.apiKey = getEnv("GOOGLE_API_KEY")!;
constructor(options: GoogleGeminiSessionOptions, modelParams: GoogleModelParams, requestOpts?: GoogleRequestOptions) {
if (options.backend === GEMINI_BACKENDS.VERTEX) {
// Vertex AI initialization
const project = options.project ?? getEnv("GOOGLE_VERTEX_PROJECT_ID");
const location = options.location ?? getEnv("GOOGLE_VERTEX_LOCATION");
if (!project || !location) {
throw new Error(
"Set GOOGLE_VERTEX_PROJECT_ID and GOOGLE_VERTEX_LOCATION env variables for Vertex AI",
);
}
this.gemini = new GoogleGenerativeAI({
vertexai: true,
project,
location,
apiVersion: options.apiVersion,
});
} else {
// Gemini API initialization
if (!options.apiKey) {
options.apiKey = getEnv("GOOGLE_API_KEY")!;
}
if (!options.apiKey) {
throw new Error("Set Google API Key in GOOGLE_API_KEY env variable");
}
this.gemini = new GoogleGenerativeAI({
apiKey: options.apiKey,
apiVersion: options.apiVersion,
});
}
if (!options.apiKey) {
throw new Error("Set Google API Key in GOOGLE_API_KEY env variable");
}
this.gemini = new GoogleGenerativeAI(options.apiKey);
}
getGenerativeModel(
metadata: GoogleModelParams,
requestOpts?: GoogleRequestOptions,
): GoogleGenerativeModel {
return this.gemini.getGenerativeModel(
// Store the generative model
this.generativeModel = this.gemini.getGenerativeModel(
{
safetySettings: metadata.safetySettings ?? DEFAULT_SAFETY_SETTINGS,
...metadata,
safetySettings: modelParams.safetySettings ?? DEFAULT_SAFETY_SETTINGS,
...modelParams,
},
requestOpts,
);
}
getGenerativeModel(
_metadata: GoogleModelParams, // metadata is now part of constructor
_requestOpts?: GoogleRequestOptions, // requestOpts is now part of constructor
): GoogleGenerativeModel {
return this.generativeModel;
}
getResponseText(response: EnhancedGenerateContentResponse): string {
return response.text();
// Updated to access text as a property
return response.text ?? "";
}
getToolsFromResponse(
response: EnhancedGenerateContentResponse,
): ToolCall[] | undefined {
return response.functionCalls()?.map(
(call: FunctionCall) =>
({
name: call.name,
input: call.args,
id: randomUUID(),
}) as ToolCall,
);
// Updated to access functionCalls as a property
const fc = response.functionCalls;
if (fc) {
return fc.map( // Assuming fc is an array
(call: FunctionCall) =>
({
name: call.name,
input: call.args, // Ensure call.args is the correct field for input
id: randomUUID(),
}) as ToolCall,
);
}
return undefined;
}
async *getChatStream(
result: GoogleStreamGenerateContentResult,
result: GoogleStreamGenerateContentResult, // This will likely change to an AsyncIterableIterator
): GeminiChatStreamResponse {
yield* streamConverter(result.stream, (response) => {
// The new SDK returns an async iterator directly from generateContentStream
// This method needs to adapt to that
for await (const response of result.stream) { // Assuming result.stream is the async iterator
const tools = this.getToolsFromResponse(response);
const options: ToolCallLLMMessageOptions = tools?.length
? { toolCall: tools }
: {};
return {
yield {
delta: this.getResponseText(response),
raw: response,
options,
};
});
}
}
getCompletionStream(
result: GoogleStreamGenerateContentResult,
async *getCompletionStream( // Changed to async generator
result: GoogleStreamGenerateContentResult, // This will likely change to an AsyncIterableIterator
): AsyncIterable<CompletionResponse> {
return streamConverter(result.stream, (response) => ({
text: this.getResponseText(response),
raw: response,
}));
// The new SDK returns an async iterator directly from generateContentStream
// This method needs to adapt to that
for await (const response of result.stream) { // Assuming result.stream is the async iterator
yield {
text: this.getResponseText(response),
raw: response,
};
}
}
}
@@ -186,42 +221,48 @@ export class GeminiSessionStore {
static sessions: Array<{
session: IGeminiSession;
options: GeminiSessionOptions;
// Store modelParams and requestOpts used for this session
modelParams: GoogleModelParams;
requestOpts?: GoogleRequestOptions;
}> = [];
private static getSessionId(options: GeminiSessionOptions): string {
if (options.backend === GEMINI_BACKENDS.GOOGLE) {
return options?.apiKey ?? "";
if (options.backend === GEMINI_BACKENDS.VERTEX) {
return `${options.project}-${options.location}-${options.apiVersion ?? ""}`;
}
return "";
return `${options.apiKey ?? ""}-${options.apiVersion ?? ""}`;
}
private static sessionMatched(
o1: GeminiSessionOptions,
o2: GeminiSessionOptions,
s1: { options: GeminiSessionOptions; modelParams: GoogleModelParams; requestOpts?: GoogleRequestOptions },
options: GeminiSessionOptions,
modelParams: GoogleModelParams,
requestOpts?: GoogleRequestOptions,
): boolean {
// #TODO: check if the session is matched
// Q: should we check the requestOptions?
// A: wait for confirmation from author
// Also check if modelParams and requestOpts match, if necessary for session reuse.
// For simplicity, this example only matches on options.
// A more robust solution might involve deep comparison of modelParams and requestOpts.
return (
GeminiSessionStore.getSessionId(o1) ===
GeminiSessionStore.getSessionId(o2)
GeminiSessionStore.getSessionId(s1.options) ===
GeminiSessionStore.getSessionId(options) &&
JSON.stringify(s1.modelParams) === JSON.stringify(modelParams) && // Basic comparison
JSON.stringify(s1.requestOpts) === JSON.stringify(requestOpts) // Basic comparison
);
}
static get(
options: GeminiSessionOptions = { backend: GEMINI_BACKENDS.GOOGLE },
modelParams: GoogleModelParams, // Pass modelParams
requestOpts?: GoogleRequestOptions, // Pass requestOpts
): IGeminiSession {
let session = this.sessions.find((session) =>
this.sessionMatched(session.options, options),
)?.session;
if (session) return session;
let sessionEntry = this.sessions.find((entry) =>
this.sessionMatched(entry, options, modelParams, requestOpts),
);
if (sessionEntry) return sessionEntry.session;
if (options.backend === GEMINI_BACKENDS.VERTEX) {
throw Error("No Session");
} else {
session = new GeminiSession(options);
}
this.sessions.push({ session, options });
return session;
const newSession = new GeminiSession(options, modelParams, requestOpts);
this.sessions.push({ session: newSession, options, modelParams, requestOpts });
return newSession;
}
}
@@ -245,10 +286,20 @@ export class Gemini extends ToolCallLLM<GeminiAdditionalChatOptions> {
this.temperature = init?.temperature ?? 0.1;
this.topP = init?.topP ?? 1;
this.maxTokens = init?.maxTokens ?? undefined;
this.session = init?.session ?? GeminiSessionStore.get();
this.#requestOptions = init?.requestOptions ?? undefined;
this.#requestOptions = init?.requestOptions ?? undefined; // Store requestOptions
this.safetySettings = init?.safetySettings ?? DEFAULT_SAFETY_SETTINGS;
this.apiKey = init?.apiKey ?? getEnv("GOOGLE_API_KEY");
// Pass model metadata and request options to GeminiSessionStore.get
// The session is now initialized with modelParams and requestOpts
this.session = init?.session ?? GeminiSessionStore.get(
{
apiKey: this.apiKey,
backend: this.apiKey ? GEMINI_BACKENDS.GOOGLE : GEMINI_BACKENDS.VERTEX, // Infer backend
// Potentially add project/location if known, or handle in GeminiSession constructor
},
{ model: this.model, temperature: this.temperature, topP: this.topP, maxTokens: this.maxTokens, safetySettings: this.safetySettings },
this.#requestOptions
);
this.voiceName = init?.voiceName ?? undefined;
}
@@ -280,45 +331,54 @@ export class Gemini extends ToolCallLLM<GeminiAdditionalChatOptions> {
};
}
private async createStartChatParams(
private async createGenerationParams( // Renamed and updated
params: GeminiChatParamsNonStreaming | GeminiChatParamsStreaming,
) {
const context = await getChatContext(params);
const common = {
history: context.history,
safetySettings: this.safetySettings as SafetySetting[],
const { messages, tools } = params; // Assuming params.messages is the new way to pass chat history
const lastMessage = messages[messages.length - 1];
const history = messages.slice(0, -1).map(GeminiHelper.messageToGeminiContent);
const generationConfig = {
candidateCount: 1, // Default or make configurable
stopSequences: undefined, // Default or make configurable
maxOutputTokens: this.maxTokens,
temperature: this.temperature,
topP: this.topP,
// topK: undefined, // Default or make configurable
};
return params.tools?.length
const safetySettings = this.safetySettings;
const toolParams = tools?.length
? {
...common,
// only if non-empty tools list
tools: [
{
functionDeclarations: params.tools.map(
functionDeclarations: tools.map(
mapBaseToolToGeminiFunctionDeclaration,
),
},
],
safetySettings: this.safetySettings,
}
: common;
: {};
return {
contents: [...history, await GeminiHelper.messageContentToGeminiParts(lastMessage.content, lastMessage.role)],
generationConfig,
safetySettings,
...toolParams,
};
}
protected async nonStreamChat(
params: GeminiChatParamsNonStreaming,
): Promise<GeminiChatNonStreamResponse> {
const context = await getChatContext(params);
const client = this.session.getGenerativeModel(
this.metadata,
this.#requestOptions,
);
const chat = client.startChat(
(await this.createStartChatParams(params)) as StartChatParams,
);
const { response } = await chat.sendMessage(context.message);
const topCandidate = response.candidates![0]!;
const generationParams = await this.createGenerationParams(params);
const client = this.session.getGenerativeModel();
// Use generateContent from the new SDK
const result = await client.generateContent(generationParams);
const response = result.response; // Response from generateContent
const topCandidate = response.candidates![0]!;
const tools = this.session.getToolsFromResponse(response);
const options: ToolCallLLMMessageOptions = tools?.length
? { toolCall: tools }
@@ -339,16 +399,30 @@ export class Gemini extends ToolCallLLM<GeminiAdditionalChatOptions> {
protected async *streamChat(
params: GeminiChatParamsStreaming,
): GeminiChatStreamResponse {
const context = await getChatContext(params);
const client = this.session.getGenerativeModel(
this.metadata,
this.#requestOptions,
);
const chat = client.startChat(
(await this.createStartChatParams(params)) as StartChatParams,
);
const result = await chat.sendMessageStream(context.message);
yield* this.session.getChatStream(result);
const generationParams = await this.createGenerationParams(params);
const client = this.session.getGenerativeModel();
// Use generateContentStream from the new SDK
const result = await client.generateContentStream(generationParams);
// getChatStream in GeminiSession expects a GoogleStreamGenerateContentResult,
// which has a `stream` property that is an async iterator.
// The result from generateContentStream is directly the async iterator over chunks.
// So we need to adapt how we call getChatStream or how getChatStream is implemented.
// Option 1: Adapt getChatStream to take the raw stream
// For now, let's assume getChatStream is updated to handle the direct stream from generateContentStream
// This means `result` itself is the async iterator over the chunks (EnhancedGenerateContentResponse)
// However, the current `getChatStream` expects `result.stream`.
// Let's create a compatible object for now, or refactor `getChatStream` later.
// Create a compatible structure for the current getChatStream
const compatibleResult: GoogleStreamGenerateContentResult = {
stream: result, // result is already the async iterator
response: Promise.resolve({} as EnhancedGenerateContentResponse), // Dummy promise, not used by current getChatStream logic
};
yield* this.session.getChatStream(compatibleResult);
}
chat(params: GeminiChatParamsStreaming): Promise<GeminiChatStreamResponse>;
@@ -373,24 +447,23 @@ export class Gemini extends ToolCallLLM<GeminiAdditionalChatOptions> {
params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
): Promise<CompletionResponse | AsyncIterable<CompletionResponse>> {
const { prompt, stream } = params;
const client = this.session.getGenerativeModel(
this.metadata,
this.#requestOptions,
// getGenerativeModel no longer takes arguments
const client = this.session.getGenerativeModel();
const content = getPartsText(
await GeminiHelper.messageContentToGeminiParts({ content: prompt }),
);
if (stream) {
const result = await client.generateContentStream(
getPartsText(
await GeminiHelper.messageContentToGeminiParts({ content: prompt }),
),
);
// generateContentStream is now directly on the model (client)
const result = await client.generateContentStream({ contents: [{ parts: [{ text: content }] }] });
// getCompletionStream in GeminiSession expects an AsyncIterable
// result from generateContentStream should be directly usable if it's an AsyncIterableIterator<EnhancedGenerateContentResponse>
// or if result.stream is the correct iterator
return this.session.getCompletionStream(result);
}
const result = await client.generateContent(
getPartsText(
await GeminiHelper.messageContentToGeminiParts({ content: prompt }),
),
);
// generateContent is now directly on the model (client)
const result = await client.generateContent({ contents: [{ parts: [{ text: content }] }] });
return {
text: this.session.getResponseText(result.response),
raw: result.response,
@@ -81,7 +81,7 @@ export class GoogleStudio extends ToolCallLLM<GoogleAdditionalChatOptions> {
temperature: number;
topP: number;
maxTokens?: number | undefined;
topK?: number;
topK?: number; // Keep an eye on topK if it's used by the SDK's GenerationConfig
constructor({
temperature,
+70 -56
View File
@@ -1,12 +1,13 @@
import {
type GenerateContentResponse,
GoogleGenerativeAI,
GenerativeModel as GoogleGenerativeModel,
type EnhancedGenerateContentResponse,
type FunctionCall,
type ModelParams as GoogleModelParams,
type RequestOptions as GoogleRequestOptions, // Added for consistency if needed
type GenerateContentStreamResult as GoogleStreamGenerateContentResult,
type SafetySetting,
VertexAI,
GenerativeModel as VertexGenerativeModel,
GenerativeModelPreview as VertexGenerativeModelPreview,
type ModelParams as VertexModelParams,
type StreamGenerateContentResult as VertexStreamGenerateContentResult,
} from "@google-cloud/vertexai";
} from "@google/genai";
import type {
GeminiChatStreamResponse,
@@ -14,7 +15,7 @@ import type {
VertexGeminiSessionOptions,
} from "./types.js";
import type { FunctionCall } from "@google/generative-ai";
// ToolCall related imports are already here from @llamaindex/core/llms
import type {
CompletionResponse,
ToolCall,
@@ -22,7 +23,7 @@ import type {
} from "@llamaindex/core/llms";
import { streamConverter } from "@llamaindex/core/utils";
import { getEnv, randomUUID } from "@llamaindex/env";
import { DEFAULT_SAFETY_SETTINGS, getFunctionCalls, getText } from "./utils.js";
import { DEFAULT_SAFETY_SETTINGS } from "./utils.js"; // getText and getFunctionCalls might be replaced or inlined
/* To use Google's Vertex AI backend, it doesn't use api key authentication.
*
@@ -38,80 +39,93 @@ import { DEFAULT_SAFETY_SETTINGS, getFunctionCalls, getText } from "./utils.js";
* */
export class GeminiVertexSession implements IGeminiSession {
private vertex: VertexAI;
private preview: boolean = false;
private gemini: GoogleGenerativeAI;
private generativeModel: GoogleGenerativeModel;
constructor(options?: Partial<VertexGeminiSessionOptions>) {
const project = options?.project ?? getEnv("GOOGLE_VERTEX_PROJECT");
constructor(options?: VertexGeminiSessionOptions, modelParams?: GoogleModelParams, requestOpts?: GoogleRequestOptions) {
const project = options?.project ?? getEnv("GOOGLE_VERTEX_PROJECT_ID") ?? 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",
"Set Google Vertex project and location in GOOGLE_VERTEX_PROJECT_ID/GOOGLE_VERTEX_PROJECT and GOOGLE_VERTEX_LOCATION env variables",
);
}
this.vertex = new VertexAI({
...options,
this.gemini = new GoogleGenerativeAI({
vertexai: true,
project,
location,
apiVersion: options?.apiVersion, // Pass apiVersion if provided
});
this.preview = options?.preview ?? false;
}
getGenerativeModel(
metadata: VertexModelParams,
): VertexGenerativeModelPreview | VertexGenerativeModel {
const safetySettings = metadata.safetySettings ?? DEFAULT_SAFETY_SETTINGS;
if (this.preview) {
return this.vertex.preview.getGenerativeModel({
safetySettings: safetySettings as SafetySetting[],
...metadata,
});
if (!modelParams || !modelParams.model) {
throw new Error("Model parameters (modelParams) with a model name are required.");
}
return this.vertex.getGenerativeModel({
safetySettings: safetySettings as SafetySetting[],
...metadata,
});
}
getResponseText(response: GenerateContentResponse): string {
return getText(response);
}
getToolsFromResponse(
response: GenerateContentResponse,
): ToolCall[] | undefined {
return getFunctionCalls(response)?.map(
(call: FunctionCall) =>
({
name: call.name,
input: call.args,
id: randomUUID(),
}) as ToolCall,
this.generativeModel = this.gemini.getGenerativeModel(
{
safetySettings: modelParams.safetySettings ?? DEFAULT_SAFETY_SETTINGS,
...modelParams,
},
requestOpts,
);
}
getGenerativeModel(
_metadata?: GoogleModelParams, // No longer takes metadata here, it's part of constructor
_requestOpts?: GoogleRequestOptions,
): GoogleGenerativeModel {
return this.generativeModel;
}
getResponseText(response: EnhancedGenerateContentResponse): string {
// Align with base.ts GeminiSession
return response.text ?? "";
}
getToolsFromResponse(
response: EnhancedGenerateContentResponse,
): ToolCall[] | undefined {
// Align with base.ts GeminiSession
const fc = response.functionCalls;
if (fc) {
return fc.map(
(call: FunctionCall) =>
({
name: call.name,
input: call.args,
id: randomUUID(),
}) as ToolCall,
);
}
return undefined;
}
async *getChatStream(
result: VertexStreamGenerateContentResult,
result: GoogleStreamGenerateContentResult, // Use GoogleStreamGenerateContentResult from @google/genai
): GeminiChatStreamResponse {
yield* streamConverter(result.stream, (response) => {
// Align with base.ts GeminiSession
for await (const response of result.stream) {
const tools = this.getToolsFromResponse(response);
const options: ToolCallLLMMessageOptions = tools?.length
? { toolCall: tools }
: {};
return {
yield {
delta: this.getResponseText(response),
raw: response,
options,
};
});
}
}
getCompletionStream(
result: VertexStreamGenerateContentResult,
async *getCompletionStream( // Changed to async generator and use GoogleStreamGenerateContentResult
result: GoogleStreamGenerateContentResult,
): AsyncIterable<CompletionResponse> {
return streamConverter(result.stream, (response) => ({
text: this.getResponseText(response),
raw: response,
}));
// Align with base.ts GeminiSession
for await (const response of result.stream) {
yield {
text: this.getResponseText(response),
raw: response,
};
}
}
}