You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
66 lines
2.1 KiB
66 lines
2.1 KiB
import importlib.util
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from types import ModuleType
|
|
|
|
# Load Conversation without importing the full swarms package
|
|
CONV_PATH = (
|
|
Path(__file__).resolve().parents[2]
|
|
/ "swarms"
|
|
/ "structs"
|
|
/ "conversation.py"
|
|
)
|
|
spec = importlib.util.spec_from_file_location(
|
|
"swarms.structs.conversation", CONV_PATH
|
|
)
|
|
conversation = importlib.util.module_from_spec(spec)
|
|
sys.modules.setdefault("swarms", ModuleType("swarms"))
|
|
sys.modules.setdefault("swarms.structs", ModuleType("swarms.structs"))
|
|
base_module = ModuleType("swarms.structs.base_structure")
|
|
base_module.BaseStructure = object
|
|
sys.modules["swarms.structs.base_structure"] = base_module
|
|
utils_root = ModuleType("swarms.utils")
|
|
any_to_str_mod = ModuleType("swarms.utils.any_to_str")
|
|
formatter_mod = ModuleType("swarms.utils.formatter")
|
|
token_mod = ModuleType("swarms.utils.litellm_tokenizer")
|
|
any_to_str_mod.any_to_str = lambda x: str(x)
|
|
formatter_mod.formatter = type(
|
|
"Formatter", (), {"print_panel": lambda *a, **k: None}
|
|
)()
|
|
token_mod.count_tokens = lambda s: len(str(s).split())
|
|
sys.modules["swarms.utils"] = utils_root
|
|
sys.modules["swarms.utils.any_to_str"] = any_to_str_mod
|
|
sys.modules["swarms.utils.formatter"] = formatter_mod
|
|
sys.modules["swarms.utils.litellm_tokenizer"] = token_mod
|
|
spec.loader.exec_module(conversation)
|
|
Conversation = conversation.Conversation
|
|
|
|
|
|
class OldConversation(Conversation):
|
|
def return_history_as_string(self):
|
|
formatted = [
|
|
f"{m['role']}: {m['content']}"
|
|
for m in self.conversation_history
|
|
]
|
|
return "\n\n".join(formatted)
|
|
|
|
def get_str(self):
|
|
return self.return_history_as_string()
|
|
|
|
|
|
def measure(conv_cls, messages=50, loops=1000):
|
|
conv = conv_cls(token_count=False)
|
|
for i in range(messages):
|
|
conv.add("user", f"msg{i}")
|
|
start = time.perf_counter()
|
|
for _ in range(loops):
|
|
conv.get_str()
|
|
return time.perf_counter() - start
|
|
|
|
|
|
def test_cache_perf_improvement():
|
|
old_time = measure(OldConversation)
|
|
new_time = measure(Conversation)
|
|
assert old_time / new_time >= 2
|