From 20ab429d2c3a2db9b1ff7cdc9fd71532fac583a1 Mon Sep 17 00:00:00 2001
From: YangQianli92 <108046369+YangQianli92@users.noreply.github.com>
Date: Mon, 15 Apr 2024 14:59:07 +0800
Subject: [PATCH] Add files via upload

---
 metagpt/rag/schema.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/metagpt/rag/schema.py b/metagpt/rag/schema.py
index c00486c82..1036c9dad 100644
--- a/metagpt/rag/schema.py
+++ b/metagpt/rag/schema.py
@@ -119,6 +119,18 @@ class ColbertRerankConfig(BaseRankerConfig):
     keep_retrieval_score: bool = Field(default=False, description="Whether to keep the retrieval score in metadata.")
 
 
+class CohereRerankConfig(BaseRankerConfig):
+    model: str = Field(default="rerank-english-v3.0")
+    top_n: int = Field(default=2)
+    api_key: str = Field(default="YOUR_COHERE_API")
+
+
+class FlagEmbeddingConfig(BaseRankerConfig):
+    model: str = Field(default="BAAI/bge-reranker-large", description="BAAI Reranker model name.")
+    top_n: int = Field(default=2, description="Number of nodes to return sorted by score.")
+    use_fp16: bool = Field(default=True, description="Whether to use fp16 for inference.")
+
+
 class ObjectRankerConfig(BaseRankerConfig):
     field_name: str = Field(..., description="field name of the object, field's value must can be compared.")
     order: Literal["desc", "asc"] = Field(default="desc", description="the direction of order.")