forked from phoenix/litellm-mirror
fac3b2ee42
* fix: fix type-checking errors * fix: fix additional type-checking errors * fix: additional type-checking error fixes * fix: fix additional type-checking errors * fix: additional type-check fixes * fix: fix all type-checking errors + add pyright to ci/cd * fix: fix incorrect import * ci(config.yml): use mypy on ci/cd * fix: fix type-checking errors in utils.py * fix: fix all type-checking errors on main.py * fix: fix mypy linting errors * fix(anthropic/cost_calculator.py): fix linting errors * fix: fix mypy linting errors * fix: fix linting errors
37 lines
962 B
Python
37 lines
962 B
Python
"""
|
|
LiteLLM Follows the cohere API format for the re rank API
|
|
https://docs.cohere.com/reference/rerank
|
|
|
|
"""
|
|
|
|
from typing import List, Optional, Union
|
|
|
|
from pydantic import BaseModel, PrivateAttr
|
|
|
|
|
|
class RerankRequest(BaseModel):
|
|
model: str
|
|
query: str
|
|
top_n: Optional[int] = None
|
|
documents: List[Union[str, dict]]
|
|
rank_fields: Optional[List[str]] = None
|
|
return_documents: Optional[bool] = None
|
|
max_chunks_per_doc: Optional[int] = None
|
|
|
|
|
|
class RerankResponse(BaseModel):
|
|
id: str
|
|
results: List[dict] # Contains index and relevance_score
|
|
meta: dict # Contains api_version and billed_units
|
|
|
|
# Define private attributes using PrivateAttr
|
|
_hidden_params: dict = PrivateAttr(default_factory=dict)
|
|
|
|
def __getitem__(self, key):
|
|
return self.__dict__[key]
|
|
|
|
def get(self, key, default=None):
|
|
return self.__dict__.get(key, default)
|
|
|
|
def __contains__(self, key):
|
|
return key in self.__dict__
|