google-genai: Fix TaskType for embedding queries when using vector store retriever

[pranav-kural]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

When using the GoogleGenerativeAIEmbeddings class from @langchain/google-genai to create an instance for embedding model to use with a vector store, we can optionally specify a TaskType.

For the task of generating embeddings for documents of a knowledge base (to store in the vector store), this task would ideally be RETRIEVAL_DOCUMENT. Example below:

const embeddingModel = new GoogleGenerativeAIEmbeddings({
	apiKey: process.env.GOOGLE_GENAI_API_KEY,
	model: 'text-embedding-004',
	taskType: TaskType.RETRIEVAL_DOCUMENT,
});

const vectorStore = new MemoryVectorStore(embeddingModel);

However, when using the vector store retriever, the query being sent should ideally be embedded with the TaskType of RETRIEVAL_QUERY.

Currently, this is not the case.

_embedQueryContent currently calls _convertToContent only by supplying it the query text. In the object returned by _convertToContent, the taskType property is set to this.taskType. If the TaskType was provided when the instance of GoogleGenerativeAIEmbeddings was created, then that task type gets used, which might be inaccurate. For example, if the instance was configured to use task type TaskType.RETRIEVAL_DOCUMENT when created (as in code above), it will use this task type, even though its actually supposed to be TaskType.RETRIEVAL_QUERY when embedding query text.

File: https://github.com/langchain-ai/langchainjs/blob/main/libs/langchain-google-genai/src/embeddings.ts

  protected async _embedQueryContent(text: string): Promise<number[]> {
    const req = this._convertToContent(text);
    const res = await this.client.embedContent(req);
    return res.embedding.values ?? [];
  }

  private _convertToContent(text: string): EmbedContentRequest {
    const cleanedText = this.stripNewLines ? text.replace(/\n/g, " ") : text;
    return {
      content: { role: "user", parts: [{ text: cleanedText }] },
      taskType: this.taskType,
      title: this.title,
    };
  }

Motivation

Without this feature/fix being implemented, full potential benefit of specifying the TaskType parameter for a Google Generative model can't be observed. Instead, it may lead instances of using (knowingly or unknowingly) the wrong value for the TaskType.

Proposal (If applicable)

I would like to suggest the below solution.

Updating the _convertToContent method to accept an additional optional argument for taskType, and updating the _embedQueryContent method to send RETRIEVAL_QUERY as the task type by default to the _convertToContent method.

File: https://github.com/langchain-ai/langchainjs/blob/main/libs/langchain-google-genai/src/embeddings.ts

protected async _embedQueryContent(text: string): Promise<number[]> {
    const req = this._convertToContent(text, TaskType.RETRIEVAL_QUERY);
    const res = await this.client.embedContent(req);
    return res.embedding.values ?? [];
  }

private _convertToContent(text: string, taskType?: TaskType): EmbedContentRequest {
    const cleanedText = this.stripNewLines ? text.replace(/\n/g, " ") : text;
    return {
      content: { role: "user", parts: [{ text: cleanedText }] },
      taskType: taskType ?? this.taskType,
      title: this.title,
    };
  }

There is a workaround to this for now (but its a bit tricky to manage):

// embedding model
const embeddingModel = new GoogleGenerativeAIEmbeddings({
	apiKey: process.env.GOOGLE_GENAI_API_KEY,
	model: 'text-embedding-004'.name,
	taskType: TaskType.RETRIEVAL_DOCUMENT,
});

// create a vector store instance
const vectorStore = new MemoryVectorStore(embeddingModel);

// add documents to the vector store
await vectorStore.addDocuments(splitDocs);

// switch the embedding model to the task type of retrieval query, since it will be used for that purpose now
// remember to switch the model back to the task of retrieval document if need to add more documents later
vectorStore.embeddings = new GoogleGenerativeAIEmbeddings({
	apiKey: process.env.GOOGLE_GENAI_API_KEY,
	model: 'text-embedding-004',
	taskType: TaskType.RETRIEVAL_QUERY,
});
		
// return the vector store retriever
return vectorStore.asRetriever(retrievalOptions);