Add cohere v2/rerank support (#8421) (#8605)

* Add cohere v2/rerank support (#8421)

* Support v2 endpoint cohere rerank

* Add tests and docs

* Make v1 default if old params used

* Update docs

* Update docs pt 2

* Update tests

* Add e2e test

* Clean up code

* Use inheritence for new config

* Fix linting issues (#8608)

* Fix cohere v2 failing test + linting (#8672)

* Fix test and unused imports

* Fix tests

* fix: fix linting errors

* test: handle tgai instability

* fix: skip service unavailable err

* test: print logs for unstable test

* test: skip unreliable tests

---------

Co-authored-by: vibhavbhat <vibhavb00@gmail.com>
This commit is contained in:
Krish Dholakia 2025-02-22 22:25:29 -08:00 committed by GitHub
parent c2aec21b4d
commit 09462ba80c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 257 additions and 40 deletions

View file

@ -1970,6 +1970,26 @@ def test_get_applied_guardrails(test_case):
# Assert
assert sorted(result) == sorted(test_case["expected"])
@pytest.mark.parametrize(
"endpoint, params, expected_bool",
[
("localhost:4000/v1/rerank", ["max_chunks_per_doc"], True),
("localhost:4000/v2/rerank", ["max_chunks_per_doc"], False),
("localhost:4000", ["max_chunks_per_doc"], True),
("localhost:4000/v1/rerank", ["max_tokens_per_doc"], True),
("localhost:4000/v2/rerank", ["max_tokens_per_doc"], False),
("localhost:4000", ["max_tokens_per_doc"], False),
("localhost:4000/v1/rerank", ["max_chunks_per_doc", "max_tokens_per_doc"], True),
("localhost:4000/v2/rerank", ["max_chunks_per_doc", "max_tokens_per_doc"], False),
("localhost:4000", ["max_chunks_per_doc", "max_tokens_per_doc"], False),
],
)
def test_should_use_cohere_v1_client(endpoint, params, expected_bool):
assert(litellm.utils.should_use_cohere_v1_client(endpoint, params) == expected_bool)
def test_add_openai_metadata():
from litellm.utils import add_openai_metadata