Add reranking support

src/browser.ts

@@ -21,6 +21,8 @@ import type {
   ProgressResponse,
   PullRequest,
   PushRequest,
+  RerankRequest,
+  RerankResponse,
   ShowRequest,
   ShowResponse,
   StatusResponse,
@@ -299,6 +301,21 @@ async encodeImage(image: Uint8Array | string): Promise<string> {
     return (await response.json()) as EmbeddingsResponse
   }
 
+  /**
+   * Rerank the documents and return at most top_n of them in rank order with scores.
+   * @param request {RerankRequest} - The request object.
+   * @returns {Promise<RerankResponse>} - The response object.
+   */
+  async rerank(request: RerankRequest): Promise<RerankResponse> {
+    if (request.top_n === undefined) {
+      request.top_n = request.documents.length;
+    }
+    const response = await utils.post(this.fetch, `${this.config.host}/api/rerank`, {
+      ...request,
+    })
+    return (await response.json()) as RerankResponse
+  }
+
   /**
    * Lists the running models on the server
    * @returns {Promise<ListResponse>} - The response object.

src/interfaces.ts

@@ -159,6 +159,13 @@ export interface EmbeddingsRequest {
   options?: Partial<Options>
 }
 
+export interface RerankRequest {
+  model: string
+  query: string
+  documents: string[]
+  top_n?: number
+}
+
 // response types
 
 export interface GenerateResponse {
@@ -199,6 +206,17 @@ export interface EmbeddingsResponse {
   embedding: number[]
 }
 
+export interface RerankResponse {
+  model: string
+  results: {
+    document: string
+    relevance_score: number
+  }[]
+  total_duration: number
+  load_duration: number
+  prompt_eval_count: number
+}
+
 export interface ProgressResponse {
   status: string
   digest: string