From 3a82334762d647f6a79e9e8cf8c45399301631ef Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 27 Aug 2024 15:55:43 -0700 Subject: [PATCH] add basic cohere rerank --- litellm/__init__.py | 1 + litellm/rerank_api/types.py | 20 ++++++++++++++++++++ 2 files changed, 21 insertions(+) create mode 100644 litellm/rerank_api/types.py diff --git a/litellm/__init__.py b/litellm/__init__.py index a627061cfe1..581db4fcb88 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -934,6 +934,7 @@ from .proxy.proxy_cli import run_server from .router import Router from .assistants.main import * from .batches.main import * +from .rerank_api.main import * from .fine_tuning.main import * from .files.main import * from .scheduler import * diff --git a/litellm/rerank_api/types.py b/litellm/rerank_api/types.py new file mode 100644 index 00000000000..9d53cf278cb --- /dev/null +++ b/litellm/rerank_api/types.py @@ -0,0 +1,20 @@ +""" +LiteLLM Follows the cohere API format for the re rank API +https://docs.cohere.com/reference/rerank + +""" + +from pydantic import BaseModel + + +class RerankRequest(BaseModel): + model: str + query: str + top_n: int + documents: list[str] + + +class RerankResponse(BaseModel): + id: str + results: list[dict] # Contains index and relevance_score + meta: dict # Contains api_version and billed_units