gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-08-07 12:07:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(backend): integrate LiteLLM API Backend (#564)

Browse Source 
* File structure for supporting litellm

* more litellm support

* feat: Add CachedAPIBackend class and dynamic API backend retrieval function

* fix: update benchmark folder path and add default values for architecture and hyperparameters

* feat: add LiteLLMAPIBackend and DeprecBackend ; changed structure of the project ; with bus

* fix : deprec_backend

* feat: Add LiteLLMAPIBackend class and related features; update configuration and test cases.

* feat: Enhance LiteLLMAPIBackend with encoder support and dynamic argument handling;Enhance log Colors

* lint

* fix lint...

* fix: Lint

* fix:make auto-lint

* fix:test oai

* fix:redundant _abckend.py

* fix: Optimize LiteLLMAPIBackend on token counting functiona, and clean up unused code;add test on this function

* feat: Add LiteLLMSettings class and update model settings usage

* fix: Update LiteLLMSettings environment variable prefix and model configurations

* fix : gitignore

* test: Consolidate and relocate test files for litellm backend and oai

* fix : lint

* fix: lint

* auto lint

* lint

* LINT

* lint

* chore: remove deprecated backend configuration comments

* refactor: Remove unused functions and imports from deprec.py and llm_utils.py

* refactor: Move md5_hash function from deprec.py to llm_utils.py

* chore: Remove extra newline and add missing import in deprec.py

* lint

* refactor: Move md5_hash function to utils module

* lint

* lint

* lint

---------

Co-authored-by: Young <afe.young@gmail.com>
Co-authored-by: Yihua Chen <v-yihuachen@microsoft.com>
This commit is contained in:
炼金术师华华
2025-02-13 15:16:18 +08:00
committed by GitHub
parent f362a1618d
commit 97c1f7a021
18 changed files with 1216 additions and 954 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/oai/test_advanced.py
+162
View File
@@ -0,0 +1,162 @@
"""
We have implemented a basic version of litellm.
Not all features in the interface are included.
Therefore, the advanced tests will be placed in a separate file for easier testing of litellm.
"""
import json
import random
import unittest
from rdagent.oai.llm_utils import APIBackend
def _worker(system_prompt, user_prompt):
api = APIBackend()
return api.build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
class TestAdvanced(unittest.TestCase):
def test_chat_cache_multiprocess(self) -> None:
"""
Tests:
- Multi process, ask same question, enable cache
- 2 pass
- cache is not missed & same question get different answer.
"""
from rdagent.core.utils import LLM_CACHE_SEED_GEN, multiprocessing_wrapper
from rdagent.oai.llm_conf import LLM_SETTINGS
system_prompt = "You are a helpful assistant."
user_prompt = f"Give me {2} random country names, list {2} cities in each country, and introduce them"
origin_value = (
LLM_SETTINGS.use_auto_chat_cache_seed_gen,
LLM_SETTINGS.use_chat_cache,
LLM_SETTINGS.dump_chat_cache,
)
LLM_SETTINGS.use_chat_cache = True
LLM_SETTINGS.dump_chat_cache = True
LLM_SETTINGS.use_auto_chat_cache_seed_gen = True
func_calls = [(_worker, (system_prompt, user_prompt)) for _ in range(4)]
LLM_CACHE_SEED_GEN.set_seed(10)
responses1 = multiprocessing_wrapper(func_calls, n=4)
LLM_CACHE_SEED_GEN.set_seed(20)
responses2 = multiprocessing_wrapper(func_calls, n=4)
LLM_CACHE_SEED_GEN.set_seed(10)
responses3 = multiprocessing_wrapper(func_calls, n=4)
# Reset, for other tests
(
LLM_SETTINGS.use_auto_chat_cache_seed_gen,
LLM_SETTINGS.use_chat_cache,
LLM_SETTINGS.dump_chat_cache,
) = origin_value
for i in range(len(func_calls)):
assert (
responses1[i] != responses2[i] and responses1[i] == responses3[i]
), "Responses sequence should be determined by 'init_chat_cache_seed'"
for j in range(i + 1, len(func_calls)):
assert (
responses1[i] != responses1[j] and responses2[i] != responses2[j]
), "Same question should get different response when use_auto_chat_cache_seed_gen=True"
def test_chat_multi_round(self) -> None:
system_prompt = "You are a helpful assistant."
fruit_name = random.SystemRandom().choice(["apple", "banana", "orange", "grape", "watermelon"])
user_prompt_1 = (
f"I will tell you a name of fruit, please remember them and tell me later. "
f"The name is {fruit_name}. Once you remember it, please answer OK."
)
user_prompt_2 = "What is the name of the fruit I told you before?"
session = APIBackend().build_chat_session(session_system_prompt=system_prompt)
response_1 = session.build_chat_completion(user_prompt=user_prompt_1)
assert response_1 is not None
assert "ok" in response_1.lower()
response2 = session.build_chat_completion(user_prompt=user_prompt_2)
assert response2 is not None
def test_chat_cache(self) -> None:
"""
Tests:
- Single process, ask same question, enable cache
- 2 pass
- cache is not missed & same question get different answer.
"""
from rdagent.core.utils import LLM_CACHE_SEED_GEN
from rdagent.oai.llm_conf import LLM_SETTINGS
system_prompt = "You are a helpful assistant."
user_prompt = f"Give me {2} random country names, list {2} cities in each country, and introduce them"
origin_value = (
LLM_SETTINGS.use_auto_chat_cache_seed_gen,
LLM_SETTINGS.use_chat_cache,
LLM_SETTINGS.dump_chat_cache,
)
LLM_SETTINGS.use_chat_cache = True
LLM_SETTINGS.dump_chat_cache = True
LLM_SETTINGS.use_auto_chat_cache_seed_gen = True
LLM_CACHE_SEED_GEN.set_seed(10)
response1 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
response2 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
LLM_CACHE_SEED_GEN.set_seed(20)
response3 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
response4 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
LLM_CACHE_SEED_GEN.set_seed(10)
response5 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
response6 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
# Reset, for other tests
(
LLM_SETTINGS.use_auto_chat_cache_seed_gen,
LLM_SETTINGS.use_chat_cache,
LLM_SETTINGS.dump_chat_cache,
) = origin_value
assert (
response1 != response3 and response2 != response4
), "Responses sequence should be determined by 'init_chat_cache_seed'"
assert (
response1 == response5 and response2 == response6
), "Responses sequence should be determined by 'init_chat_cache_seed'"
assert (
response1 != response2 and response3 != response4 and response5 != response6
), "Same question should get different response when use_auto_chat_cache_seed_gen=True"
if __name__ == "__main__":
unittest.main()
test/oai/test_completion.py
+5 -143
View File
@@ -1,18 +1,9 @@
import json
import random
import unittest
from rdagent.oai.llm_utils import APIBackend
def _worker(system_prompt, user_prompt):
api = APIBackend()
return api.build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
class TestChatCompletion(unittest.TestCase):
def test_chat_completion(self) -> None:
system_prompt = "You are a helpful assistant."
@@ -36,141 +27,12 @@ class TestChatCompletion(unittest.TestCase):
assert isinstance(response, str)
json.loads(response)
def test_chat_multi_round(self) -> None:
def test_build_messages_and_calculate_token(self) -> None:
system_prompt = "You are a helpful assistant."
fruit_name = random.SystemRandom().choice(["apple", "banana", "orange", "grape", "watermelon"])
user_prompt_1 = (
f"I will tell you a name of fruit, please remember them and tell me later. "
f"The name is {fruit_name}. Once you remember it, please answer OK."
)
user_prompt_2 = "What is the name of the fruit I told you before?"
session = APIBackend().build_chat_session(session_system_prompt=system_prompt)
response_1 = session.build_chat_completion(user_prompt=user_prompt_1)
assert response_1 is not None
assert "ok" in response_1.lower()
response2 = session.build_chat_completion(user_prompt=user_prompt_2)
assert response2 is not None
def test_chat_cache(self) -> None:
"""
Tests:
- Single process, ask same question, enable cache
- 2 pass
- cache is not missed & same question get different answer.
"""
from rdagent.core.utils import LLM_CACHE_SEED_GEN
from rdagent.oai.llm_conf import LLM_SETTINGS
system_prompt = "You are a helpful assistant."
user_prompt = f"Give me {2} random country names, list {2} cities in each country, and introduce them"
origin_value = (
LLM_SETTINGS.use_auto_chat_cache_seed_gen,
LLM_SETTINGS.use_chat_cache,
LLM_SETTINGS.dump_chat_cache,
)
LLM_SETTINGS.use_chat_cache = True
LLM_SETTINGS.dump_chat_cache = True
LLM_SETTINGS.use_auto_chat_cache_seed_gen = True
LLM_CACHE_SEED_GEN.set_seed(10)
response1 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
response2 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
LLM_CACHE_SEED_GEN.set_seed(20)
response3 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
response4 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
LLM_CACHE_SEED_GEN.set_seed(10)
response5 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
response6 = APIBackend().build_messages_and_create_chat_completion(
system_prompt=system_prompt,
user_prompt=user_prompt,
)
# Reset, for other tests
(
LLM_SETTINGS.use_auto_chat_cache_seed_gen,
LLM_SETTINGS.use_chat_cache,
LLM_SETTINGS.dump_chat_cache,
) = origin_value
assert (
response1 != response3 and response2 != response4
), "Responses sequence should be determined by 'init_chat_cache_seed'"
assert (
response1 == response5 and response2 == response6
), "Responses sequence should be determined by 'init_chat_cache_seed'"
assert (
response1 != response2 and response3 != response4 and response5 != response6
), "Same question should get different response when use_auto_chat_cache_seed_gen=True"
def test_chat_cache_multiprocess(self) -> None:
"""
Tests:
- Multi process, ask same question, enable cache
- 2 pass
- cache is not missed & same question get different answer.
"""
from rdagent.core.utils import LLM_CACHE_SEED_GEN, multiprocessing_wrapper
from rdagent.oai.llm_conf import LLM_SETTINGS
system_prompt = "You are a helpful assistant."
user_prompt = f"Give me {2} random country names, list {2} cities in each country, and introduce them"
origin_value = (
LLM_SETTINGS.use_auto_chat_cache_seed_gen,
LLM_SETTINGS.use_chat_cache,
LLM_SETTINGS.dump_chat_cache,
)
LLM_SETTINGS.use_chat_cache = True
LLM_SETTINGS.dump_chat_cache = True
LLM_SETTINGS.use_auto_chat_cache_seed_gen = True
func_calls = [(_worker, (system_prompt, user_prompt)) for _ in range(4)]
LLM_CACHE_SEED_GEN.set_seed(10)
responses1 = multiprocessing_wrapper(func_calls, n=4)
LLM_CACHE_SEED_GEN.set_seed(20)
responses2 = multiprocessing_wrapper(func_calls, n=4)
LLM_CACHE_SEED_GEN.set_seed(10)
responses3 = multiprocessing_wrapper(func_calls, n=4)
# Reset, for other tests
(
LLM_SETTINGS.use_auto_chat_cache_seed_gen,
LLM_SETTINGS.use_chat_cache,
LLM_SETTINGS.dump_chat_cache,
) = origin_value
for i in range(len(func_calls)):
assert (
responses1[i] != responses2[i] and responses1[i] == responses3[i]
), "Responses sequence should be determined by 'init_chat_cache_seed'"
for j in range(i + 1, len(func_calls)):
assert (
responses1[i] != responses1[j] and responses2[i] != responses2[j]
), "Same question should get different response when use_auto_chat_cache_seed_gen=True"
user_prompt = "What is your name?"
token = APIBackend().build_messages_and_calculate_token(user_prompt=user_prompt, system_prompt=system_prompt)
assert token is not None
assert isinstance(token, int)
if __name__ == "__main__":
test/oai/test_embedding_and_similarity.py
+6
View File
@@ -13,6 +13,12 @@ class TestEmbedding(unittest.TestCase):
assert isinstance(emb, list)
assert len(emb) > 0
def test_embedding_list(self) -> None:
emb = APIBackend().create_embedding(["hello", "hi"])
assert emb is not None
assert isinstance(emb, list)
assert len(emb) == 2
def test_embedding_similarity(self) -> None:
similarity = calculate_embedding_distance_between_str_list(["Hello"], ["Hi"])[0][0]
assert similarity is not None