-
Notifications
You must be signed in to change notification settings - Fork 2.3k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
cohere[patch]: update embeddings and rerank #5928
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -28,7 +28,7 @@ const docs = [ | |
const cohereRerank = new CohereRerank({ | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Hey team, just a heads up that I've flagged a change in the |
||
apiKey: process.env.COHERE_API_KEY, // Default | ||
topN: 3, // Default | ||
model: "rerank-english-v2.0", // Default | ||
model: "rerank-english-v2.0", | ||
}); | ||
|
||
const rerankedDocuments = await cohereRerank.compressDocuments(docs, query); | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -9,14 +9,19 @@ import { chunkArray } from "@langchain/core/utils/chunk_array"; | |
* parameters specific to the CohereEmbeddings class. | ||
*/ | ||
export interface CohereEmbeddingsParams extends EmbeddingsParams { | ||
model: string; | ||
model?: string; | ||
|
||
/** | ||
* The maximum number of documents to embed in a single request. This is | ||
* limited by the Cohere API to a maximum of 96. | ||
*/ | ||
batchSize?: number; | ||
|
||
/** | ||
* Specifies the type of embeddings you want to generate. | ||
*/ | ||
embeddingTypes?: Array<string>; | ||
|
||
/** | ||
* Specifies the type of input you're giving to the model. | ||
* Not required for older versions of the embedding models (i.e. anything lower than v3), | ||
|
@@ -37,11 +42,11 @@ export class CohereEmbeddings | |
extends Embeddings | ||
implements CohereEmbeddingsParams | ||
{ | ||
model = "small"; | ||
model: string | undefined; | ||
|
||
batchSize = 48; | ||
|
||
inputType: string | undefined; | ||
embeddingTypes = ["float"]; | ||
|
||
private client: CohereClient; | ||
|
||
|
@@ -70,8 +75,16 @@ export class CohereEmbeddings | |
token: apiKey, | ||
}); | ||
this.model = fieldsWithDefaults?.model ?? this.model; | ||
|
||
if (!this.model) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Would prefer to just make this a required param and ship a new minor version release along with your tool calling stuff |
||
throw new Error( | ||
"Model not specified for CohereEmbeddings instance. Please provide a model name from the options here: https://docs.cohere.com/reference/embed" | ||
); | ||
} | ||
|
||
this.batchSize = fieldsWithDefaults?.batchSize ?? this.batchSize; | ||
this.inputType = fieldsWithDefaults?.inputType; | ||
this.embeddingTypes = | ||
fieldsWithDefaults?.embeddingTypes ?? this.embeddingTypes; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Do we need to remove |
||
} | ||
|
||
/** | ||
|
@@ -87,7 +100,9 @@ export class CohereEmbeddings | |
model: this.model, | ||
texts: batch, | ||
// eslint-disable-next-line @typescript-eslint/no-explicit-any | ||
inputType: this.inputType as any, | ||
inputType: "search_document" as any, | ||
// eslint-disable-next-line @typescript-eslint/no-explicit-any | ||
embeddingTypes: this.embeddingTypes as any, | ||
}) | ||
); | ||
|
||
|
@@ -120,7 +135,9 @@ export class CohereEmbeddings | |
model: this.model, | ||
texts: [text], | ||
// eslint-disable-next-line @typescript-eslint/no-explicit-any | ||
inputType: this.inputType as any, | ||
inputType: "search_query" as any, | ||
// eslint-disable-next-line @typescript-eslint/no-explicit-any | ||
embeddingTypes: this.embeddingTypes as any, | ||
}); | ||
if ("float" in embeddings && embeddings.float) { | ||
return embeddings.float[0]; | ||
|
@@ -137,6 +154,25 @@ export class CohereEmbeddings | |
} | ||
} | ||
|
||
async embed( | ||
jacoblee93 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
request: Parameters<typeof this.client.embed>[0] | ||
): Promise<number[]> { | ||
const { embeddings } = await this.embeddingWithRetry(request); | ||
if ("float" in embeddings && embeddings.float) { | ||
return embeddings.float[0]; | ||
} else if (Array.isArray(embeddings)) { | ||
return embeddings[0]; | ||
} else { | ||
throw new Error( | ||
`Invalid response from Cohere API. Received: ${JSON.stringify( | ||
embeddings, | ||
null, | ||
2 | ||
)}` | ||
); | ||
} | ||
} | ||
|
||
/** | ||
* Generates embeddings with retry capabilities. | ||
* @param request - An object containing the request parameters for generating embeddings. | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -20,6 +20,7 @@ const documents = [ | |
test("CohereRerank can indeed rerank documents with compressDocuments method", async () => { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Hey there! I've reviewed the code and noticed that the addition in this PR references an environment variable via |
||
const cohereRerank = new CohereRerank({ | ||
apiKey: process.env.COHERE_API_KEY, | ||
model: "rerank-english-v2.0" | ||
}); | ||
|
||
const rerankedDocuments = await cohereRerank.compressDocuments( | ||
|
@@ -33,6 +34,7 @@ test("CohereRerank can indeed rerank documents with compressDocuments method", a | |
test("CohereRerank can indeed rerank documents with rerank method", async () => { | ||
const cohereRerank = new CohereRerank({ | ||
apiKey: process.env.COHERE_API_KEY, | ||
model: "rerank-english-v2.0", | ||
}); | ||
|
||
const rerankedDocuments = await cohereRerank.rerank( | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hey team, I've flagged a change in the
cohereRerank
instantiation that accesses an environment variable viaprocess.env
. Please review this change to ensure it aligns with our environment variable handling practices.