feat(model): Support deploy rerank model (#1522)

dbgpt/core/__init__.py

@@ -7,7 +7,7 @@ from dbgpt.core.interface.cache import (  # noqa: F401
     CachePolicy,
     CacheValue,
 )
-from dbgpt.core.interface.embeddings import Embeddings  # noqa: F401
+from dbgpt.core.interface.embeddings import Embeddings, RerankEmbeddings  # noqa: F401
 from dbgpt.core.interface.knowledge import Chunk, Document  # noqa: F401
 from dbgpt.core.interface.llm import (  # noqa: F401
     DefaultMessageConverter,
@@ -106,6 +106,7 @@ __ALL__ = [
     "QuerySpec",
     "StorageError",
     "Embeddings",
+    "RerankEmbeddings",
     "Chunk",
     "Document",
 ]

dbgpt/core/interface/embeddings.py

@@ -1,9 +1,24 @@
 """Interface for embedding models."""
+
 import asyncio
 from abc import ABC, abstractmethod
 from typing import List
 
 
+class RerankEmbeddings(ABC):
+    """Interface for rerank models."""
+
+    @abstractmethod
+    def predict(self, query: str, candidates: List[str]) -> List[float]:
+        """Predict the scores of the candidates."""
+
+    async def apredict(self, query: str, candidates: List[str]) -> List[float]:
+        """Asynchronously predict the scores of the candidates."""
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self.predict, query, candidates
+        )
+
+
 class Embeddings(ABC):
     """Interface for embedding models.
 

dbgpt/core/schema/api.py

@@ -166,6 +166,23 @@ class EmbeddingsResponse(BaseModel):
     usage: UsageInfo = Field(..., description="Usage info")
 
 
+class RelevanceRequest(BaseModel):
+    """Relevance request entity."""
+
+    model: str = Field(..., description="Rerank model name")
+    query: str = Field(..., description="Query text")
+    documents: List[str] = Field(..., description="Document texts")
+
+
+class RelevanceResponse(BaseModel):
+    """Relevance response entity."""
+
+    object: str = Field("list", description="Object type")
+    model: str = Field(..., description="Rerank model name")
+    data: List[float] = Field(..., description="Data list, relevance scores")
+    usage: UsageInfo = Field(..., description="Usage info")
+
+
 class ModelPermission(BaseModel):
     """Model permission entity."""