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