diff --git a/.gitignore b/.gitignore index fa8a4622..7fc5f776 100644 --- a/.gitignore +++ b/.gitignore @@ -111,7 +111,7 @@ celerybeat.pid *.sage.py # Environments -.env +.env* .venv ^env/ venv/ @@ -172,3 +172,4 @@ mlruns/ *.out *.sh .aider* +rdagent/app/benchmark/factor/example.json diff --git a/rdagent/app/benchmark/model/eval.py b/rdagent/app/benchmark/model/eval.py index 0182e9a8..6b4c5d77 100644 --- a/rdagent/app/benchmark/model/eval.py +++ b/rdagent/app/benchmark/model/eval.py @@ -13,7 +13,7 @@ if __name__ == "__main__": from rdagent.components.coder.model_coder.benchmark.eval import ModelImpValEval from rdagent.components.coder.model_coder.one_shot import ModelCodeWriter - bench_folder = DIRNAME.parent.parent / "components" / "coder" / "model_coder" / "benchmark" + bench_folder = DIRNAME.parent.parent.parent / "components" / "coder" / "model_coder" / "benchmark" mtl = ModelTaskLoaderJson(str(bench_folder / "model_dict.json")) task_l = mtl.load() diff --git a/rdagent/components/loader/task_loader.py b/rdagent/components/loader/task_loader.py index 3ef80785..9fbb0ec6 100644 --- a/rdagent/components/loader/task_loader.py +++ b/rdagent/components/loader/task_loader.py @@ -75,6 +75,8 @@ class ModelTaskLoaderJson(ModelTaskLoader): formulation=model_data["formulation"], variables=model_data["variables"], model_type=model_data["model_type"], + architecture="", + hyperparameters="", ) model_impl_task_list.append(model_impl_task) return model_impl_task_list diff --git a/rdagent/oai/backend/__init__.py b/rdagent/oai/backend/__init__.py index e69de29b..5be02cc0 100644 --- a/rdagent/oai/backend/__init__.py +++ b/rdagent/oai/backend/__init__.py @@ -0,0 +1,2 @@ +from .deprec import DeprecBackend +from .litellm import LiteLLMAPIBackend diff --git a/rdagent/oai/backend/base.py b/rdagent/oai/backend/base.py index d9b57c31..8bb20a19 100644 --- a/rdagent/oai/backend/base.py +++ b/rdagent/oai/backend/base.py @@ -1,2 +1,53 @@ -class APIBackend: - """abstract""" +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Tuple, Union + + +class APIBackend(ABC): + """Abstract base class for LLM API backends""" + + @abstractmethod + def build_chat_session( + self, conversation_id: Optional[str] = None, session_system_prompt: Optional[str] = None + ) -> Any: + """Create a new chat session""" + pass + + @abstractmethod + def build_messages_and_create_chat_completion( + self, + user_prompt: str, + system_prompt: Optional[str] = None, + former_messages: Optional[List[Any]] = None, + chat_cache_prefix: str = "", + shrink_multiple_break: bool = False, + *args: Any, + **kwargs: Any, + ) -> str: + """Build messages and get chat completion""" + pass + + @abstractmethod + def create_embedding( + self, input_content: Union[str, List[str]], *args: Any, **kwargs: Any + ) -> Union[List[Any], Any]: + """Create embeddings for input text""" + pass + + @abstractmethod + def build_messages_and_calculate_token( + self, + user_prompt: str, + system_prompt: Optional[str], + former_messages: Optional[List[Dict[str, Any]]] = None, + *, + shrink_multiple_break: bool = False, + ) -> int: + """Build messages and calculate their token count""" + pass + + +# TODO: seperate cache layer. try to be tranparent. +class CachedAPIBackend(APIBackend): + ... + # @abstractmethod + # def none_cache_function ... diff --git a/rdagent/oai/backend/deprec.py b/rdagent/oai/backend/deprec.py new file mode 100644 index 00000000..9d3e6825 --- /dev/null +++ b/rdagent/oai/backend/deprec.py @@ -0,0 +1,803 @@ +from __future__ import annotations + +import inspect +import json +import os +import random +import re +import sqlite3 +import ssl +import time +import urllib.request +import uuid +from copy import deepcopy +from pathlib import Path +from typing import Any, Optional + +import numpy as np +import openai +import tiktoken + +from rdagent.core.utils import LLM_CACHE_SEED_GEN, SingletonBaseClass, import_class +from rdagent.log import LogColors +from rdagent.log import rdagent_logger as logger +from rdagent.oai.llm_conf import LLM_SETTINGS +from rdagent.utils import md5_hash + +DEFAULT_QLIB_DOT_PATH = Path("./") + +from rdagent.oai.backend.base import APIBackend + +try: + from azure.identity import DefaultAzureCredential, get_bearer_token_provider +except ImportError: + logger.warning("azure.identity is not installed.") + +try: + import openai +except ImportError: + logger.warning("openai is not installed.") + +try: + from llama import Llama +except ImportError: + if LLM_SETTINGS.use_llama2: + logger.warning("llama is not installed.") + + +class ConvManager: + """ + This is a conversation manager of LLM + It is for convenience of exporting conversation for debugging. + """ + + def __init__( + self, + path: Path | str = DEFAULT_QLIB_DOT_PATH / "llm_conv", + recent_n: int = 10, + ) -> None: + self.path = Path(path) + self.path.mkdir(parents=True, exist_ok=True) + self.recent_n = recent_n + + def _rotate_files(self) -> None: + pairs = [] + for f in self.path.glob("*.json"): + m = re.match(r"(\d+).json", f.name) + if m is not None: + n = int(m.group(1)) + pairs.append((n, f)) + pairs.sort(key=lambda x: x[0]) + for n, f in pairs[: self.recent_n][::-1]: + if (self.path / f"{n+1}.json").exists(): + (self.path / f"{n+1}.json").unlink() + f.rename(self.path / f"{n+1}.json") + + def append(self, conv: tuple[list, str]) -> None: + self._rotate_files() + with (self.path / "0.json").open("w") as file: + json.dump(conv, file) + # TODO: reseve line breaks to make it more convient to edit file directly. + + +class SQliteLazyCache(SingletonBaseClass): + def __init__(self, cache_location: str) -> None: + super().__init__() + self.cache_location = cache_location + db_file_exist = Path(cache_location).exists() + # TODO: sqlite3 does not support multiprocessing. + self.conn = sqlite3.connect(cache_location, timeout=20) + self.c = self.conn.cursor() + if not db_file_exist: + self.c.execute( + """ + CREATE TABLE chat_cache ( + md5_key TEXT PRIMARY KEY, + chat TEXT + ) + """, + ) + self.c.execute( + """ + CREATE TABLE embedding_cache ( + md5_key TEXT PRIMARY KEY, + embedding TEXT + ) + """, + ) + self.c.execute( + """ + CREATE TABLE message_cache ( + conversation_id TEXT PRIMARY KEY, + message TEXT + ) + """, + ) + self.conn.commit() + + def chat_get(self, key: str) -> str | None: + md5_key = md5_hash(key) + self.c.execute("SELECT chat FROM chat_cache WHERE md5_key=?", (md5_key,)) + result = self.c.fetchone() + return None if result is None else result[0] + + def embedding_get(self, key: str) -> list | dict | str | None: + md5_key = md5_hash(key) + self.c.execute("SELECT embedding FROM embedding_cache WHERE md5_key=?", (md5_key,)) + result = self.c.fetchone() + return None if result is None else json.loads(result[0]) + + def chat_set(self, key: str, value: str) -> None: + md5_key = md5_hash(key) + self.c.execute( + "INSERT OR REPLACE INTO chat_cache (md5_key, chat) VALUES (?, ?)", + (md5_key, value), + ) + self.conn.commit() + return None + + def embedding_set(self, content_to_embedding_dict: dict) -> None: + for key, value in content_to_embedding_dict.items(): + md5_key = md5_hash(key) + self.c.execute( + "INSERT OR REPLACE INTO embedding_cache (md5_key, embedding) VALUES (?, ?)", + (md5_key, json.dumps(value)), + ) + self.conn.commit() + + def message_get(self, conversation_id: str) -> list[dict[str, Any]]: + self.c.execute("SELECT message FROM message_cache WHERE conversation_id=?", (conversation_id,)) + result = self.c.fetchone() + return [] if result is None else json.loads(result[0]) + + def message_set(self, conversation_id: str, message_value: list[dict[str, Any]]) -> None: + self.c.execute( + "INSERT OR REPLACE INTO message_cache (conversation_id, message) VALUES (?, ?)", + (conversation_id, json.dumps(message_value)), + ) + self.conn.commit() + return None + + +class SessionChatHistoryCache(SingletonBaseClass): + def __init__(self) -> None: + """load all history conversation json file from self.session_cache_location""" + self.cache = SQliteLazyCache(cache_location=LLM_SETTINGS.prompt_cache_path) + + def message_get(self, conversation_id: str) -> list[dict[str, Any]]: + return self.cache.message_get(conversation_id) + + def message_set(self, conversation_id: str, message_value: list[dict[str, Any]]) -> None: + self.cache.message_set(conversation_id, message_value) + + +class ChatSession: + def __init__(self, api_backend: Any, conversation_id: str | None = None, system_prompt: str | None = None) -> None: + self.conversation_id = str(uuid.uuid4()) if conversation_id is None else conversation_id + self.system_prompt = system_prompt if system_prompt is not None else LLM_SETTINGS.default_system_prompt + self.api_backend = api_backend + + def build_chat_completion_message(self, user_prompt: str) -> list[dict[str, Any]]: + history_message = SessionChatHistoryCache().message_get(self.conversation_id) + messages = history_message + if not messages: + messages.append({"role": "system", "content": self.system_prompt}) + messages.append( + { + "role": "user", + "content": user_prompt, + }, + ) + return messages + + def build_chat_completion_message_and_calculate_token(self, user_prompt: str) -> Any: + messages = self.build_chat_completion_message(user_prompt) + return self.api_backend._calculate_token_from_messages(messages) + + def build_chat_completion(self, user_prompt: str, *args, **kwargs) -> str: # type: ignore[no-untyped-def] + """ + this function is to build the session messages + user prompt should always be provided + """ + messages = self.build_chat_completion_message(user_prompt) + + with logger.tag(f"session_{self.conversation_id}"): + response: str = self.api_backend._try_create_chat_completion_or_embedding( # noqa: SLF001 + *args, + messages=messages, + chat_completion=True, + **kwargs, + ) + logger.log_object({"user": user_prompt, "resp": response}, tag="debug_llm") + + messages.append( + { + "role": "assistant", + "content": response, + }, + ) + SessionChatHistoryCache().message_set(self.conversation_id, messages) + return response + + def get_conversation_id(self) -> str: + return self.conversation_id + + def display_history(self) -> None: + # TODO: Realize a beautiful presentation format for history messages + pass + + +class DeprecBackend(APIBackend): + """ + This is a unified interface for different backends. + + (xiao) thinks integrate all kinds of API in a single class is not a good design. + So we should split them into different classes in `oai/backends/` in the future. + """ + + # FIXME: (xiao) We should avoid using self.xxxx. + # Instead, we can use LLM_SETTINGS directly. If it's difficult to support different backend settings, we can split them into multiple BaseSettings. + def __init__( # noqa: C901, PLR0912, PLR0915 + self, + *, + chat_api_key: str | None = None, + chat_model: str | None = None, + chat_api_base: str | None = None, + chat_api_version: str | None = None, + embedding_api_key: str | None = None, + embedding_model: str | None = None, + embedding_api_base: str | None = None, + embedding_api_version: str | None = None, + use_chat_cache: bool | None = None, + dump_chat_cache: bool | None = None, + use_embedding_cache: bool | None = None, + dump_embedding_cache: bool | None = None, + ) -> None: + if LLM_SETTINGS.use_llama2: + self.generator = Llama.build( + ckpt_dir=LLM_SETTINGS.llama2_ckpt_dir, + tokenizer_path=LLM_SETTINGS.llama2_tokenizer_path, + max_seq_len=LLM_SETTINGS.chat_max_tokens, + max_batch_size=LLM_SETTINGS.llams2_max_batch_size, + ) + self.encoder = None + elif LLM_SETTINGS.use_gcr_endpoint: + gcr_endpoint_type = LLM_SETTINGS.gcr_endpoint_type + if gcr_endpoint_type == "llama2_70b": + self.gcr_endpoint_key = LLM_SETTINGS.llama2_70b_endpoint_key + self.gcr_endpoint_deployment = LLM_SETTINGS.llama2_70b_endpoint_deployment + self.gcr_endpoint = LLM_SETTINGS.llama2_70b_endpoint + elif gcr_endpoint_type == "llama3_70b": + self.gcr_endpoint_key = LLM_SETTINGS.llama3_70b_endpoint_key + self.gcr_endpoint_deployment = LLM_SETTINGS.llama3_70b_endpoint_deployment + self.gcr_endpoint = LLM_SETTINGS.llama3_70b_endpoint + elif gcr_endpoint_type == "phi2": + self.gcr_endpoint_key = LLM_SETTINGS.phi2_endpoint_key + self.gcr_endpoint_deployment = LLM_SETTINGS.phi2_endpoint_deployment + self.gcr_endpoint = LLM_SETTINGS.phi2_endpoint + elif gcr_endpoint_type == "phi3_4k": + self.gcr_endpoint_key = LLM_SETTINGS.phi3_4k_endpoint_key + self.gcr_endpoint_deployment = LLM_SETTINGS.phi3_4k_endpoint_deployment + self.gcr_endpoint = LLM_SETTINGS.phi3_4k_endpoint + elif gcr_endpoint_type == "phi3_128k": + self.gcr_endpoint_key = LLM_SETTINGS.phi3_128k_endpoint_key + self.gcr_endpoint_deployment = LLM_SETTINGS.phi3_128k_endpoint_deployment + self.gcr_endpoint = LLM_SETTINGS.phi3_128k_endpoint + else: + error_message = f"Invalid gcr_endpoint_type: {gcr_endpoint_type}" + raise ValueError(error_message) + self.headers = { + "Content-Type": "application/json", + "Authorization": ("Bearer " + self.gcr_endpoint_key), + } + self.gcr_endpoint_temperature = LLM_SETTINGS.gcr_endpoint_temperature + self.gcr_endpoint_top_p = LLM_SETTINGS.gcr_endpoint_top_p + self.gcr_endpoint_do_sample = LLM_SETTINGS.gcr_endpoint_do_sample + self.gcr_endpoint_max_token = LLM_SETTINGS.gcr_endpoint_max_token + if not os.environ.get("PYTHONHTTPSVERIFY", "") and hasattr(ssl, "_create_unverified_context"): + ssl._create_default_https_context = ssl._create_unverified_context # noqa: SLF001 + self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map) + self.chat_model = LLM_SETTINGS.chat_model if chat_model is None else chat_model + self.encoder = None + else: + self.chat_use_azure = LLM_SETTINGS.chat_use_azure or LLM_SETTINGS.use_azure + self.embedding_use_azure = LLM_SETTINGS.embedding_use_azure or LLM_SETTINGS.use_azure + self.chat_use_azure_token_provider = LLM_SETTINGS.chat_use_azure_token_provider + self.embedding_use_azure_token_provider = LLM_SETTINGS.embedding_use_azure_token_provider + self.managed_identity_client_id = LLM_SETTINGS.managed_identity_client_id + + # Priority: chat_api_key/embedding_api_key > openai_api_key > os.environ.get("OPENAI_API_KEY") + # TODO: Simplify the key design. Consider Pandatic's field alias & priority. + self.chat_api_key = ( + chat_api_key + or LLM_SETTINGS.chat_openai_api_key + or LLM_SETTINGS.openai_api_key + or os.environ.get("OPENAI_API_KEY") + ) + self.embedding_api_key = ( + embedding_api_key + or LLM_SETTINGS.embedding_openai_api_key + or LLM_SETTINGS.openai_api_key + or os.environ.get("OPENAI_API_KEY") + ) + + self.chat_model = LLM_SETTINGS.chat_model if chat_model is None else chat_model + self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map) + self.encoder = self._get_encoder() + self.chat_openai_base_url = LLM_SETTINGS.chat_openai_base_url + self.embedding_openai_base_url = LLM_SETTINGS.embedding_openai_base_url + self.chat_api_base = LLM_SETTINGS.chat_azure_api_base if chat_api_base is None else chat_api_base + self.chat_api_version = ( + LLM_SETTINGS.chat_azure_api_version if chat_api_version is None else chat_api_version + ) + self.chat_stream = LLM_SETTINGS.chat_stream + self.chat_seed = LLM_SETTINGS.chat_seed + + self.embedding_model = LLM_SETTINGS.embedding_model if embedding_model is None else embedding_model + self.embedding_api_base = ( + LLM_SETTINGS.embedding_azure_api_base if embedding_api_base is None else embedding_api_base + ) + self.embedding_api_version = ( + LLM_SETTINGS.embedding_azure_api_version if embedding_api_version is None else embedding_api_version + ) + + if (self.chat_use_azure or self.embedding_use_azure) and ( + self.chat_use_azure_token_provider or self.embedding_use_azure_token_provider + ): + dac_kwargs = {} + if self.managed_identity_client_id is not None: + dac_kwargs["managed_identity_client_id"] = self.managed_identity_client_id + credential = DefaultAzureCredential(**dac_kwargs) + token_provider = get_bearer_token_provider( + credential, + "https://cognitiveservices.azure.com/.default", + ) + self.chat_client: openai.OpenAI = ( + openai.AzureOpenAI( + azure_ad_token_provider=token_provider if self.chat_use_azure_token_provider else None, + api_key=self.chat_api_key if not self.chat_use_azure_token_provider else None, + api_version=self.chat_api_version, + azure_endpoint=self.chat_api_base, + ) + if self.chat_use_azure + else openai.OpenAI(api_key=self.chat_api_key, base_url=self.chat_openai_base_url) + ) + + self.embedding_client: openai.OpenAI = ( + openai.AzureOpenAI( + azure_ad_token_provider=token_provider if self.embedding_use_azure_token_provider else None, + api_key=self.embedding_api_key if not self.embedding_use_azure_token_provider else None, + api_version=self.embedding_api_version, + azure_endpoint=self.embedding_api_base, + ) + if self.embedding_use_azure + else openai.OpenAI(api_key=self.embedding_api_key, base_url=self.embedding_openai_base_url) + ) + + self.dump_chat_cache = LLM_SETTINGS.dump_chat_cache if dump_chat_cache is None else dump_chat_cache + self.use_chat_cache = LLM_SETTINGS.use_chat_cache if use_chat_cache is None else use_chat_cache + self.dump_embedding_cache = ( + LLM_SETTINGS.dump_embedding_cache if dump_embedding_cache is None else dump_embedding_cache + ) + self.use_embedding_cache = ( + LLM_SETTINGS.use_embedding_cache if use_embedding_cache is None else use_embedding_cache + ) + if self.dump_chat_cache or self.use_chat_cache or self.dump_embedding_cache or self.use_embedding_cache: + self.cache_file_location = LLM_SETTINGS.prompt_cache_path + self.cache = SQliteLazyCache(cache_location=self.cache_file_location) + + # transfer the config to the class if the config is not supposed to change during the runtime + self.use_llama2 = LLM_SETTINGS.use_llama2 + self.use_gcr_endpoint = LLM_SETTINGS.use_gcr_endpoint + self.retry_wait_seconds = LLM_SETTINGS.retry_wait_seconds + + def _get_encoder(self) -> tiktoken.Encoding: + """ + tiktoken.encoding_for_model(self.chat_model) does not cover all cases it should consider. + + This function attempts to handle several edge cases. + """ + + # 1) cases + def _azure_patch(model: str) -> str: + """ + When using Azure API, self.chat_model is the deployment name that can be any string. + For example, it may be `gpt-4o_2024-08-06`. But tiktoken.encoding_for_model can't handle this. + """ + return model.replace("_", "-") + + model = self.chat_model + try: + encoding = tiktoken.encoding_for_model(model) + except KeyError: + logger.warning(f"Failed to get encoder. Trying to patch the model name") + for patch_func in [_azure_patch]: + try: + encoding = tiktoken.encoding_for_model(patch_func(model)) + except KeyError: + logger.error(f"Failed to get encoder even after patching with {patch_func.__name__}") + raise + return encoding + + def build_chat_session( + self, + conversation_id: str | None = None, + session_system_prompt: str | None = None, + ) -> ChatSession: + """ + conversation_id is a 256-bit string created by uuid.uuid4() and is also + the file name under session_cache_folder/ for each conversation + """ + return ChatSession(self, conversation_id, session_system_prompt) + + def _build_messages( + self, + user_prompt: str, + system_prompt: str | None = None, + former_messages: list[dict[str, Any]] | None = None, + *, + shrink_multiple_break: bool = False, + ) -> list[dict[str, Any]]: + """ + build the messages to avoid implementing several redundant lines of code + + """ + if former_messages is None: + former_messages = [] + # shrink multiple break will recursively remove multiple breaks(more than 2) + if shrink_multiple_break: + while "\n\n\n" in user_prompt: + user_prompt = user_prompt.replace("\n\n\n", "\n\n") + if system_prompt is not None: + while "\n\n\n" in system_prompt: + system_prompt = system_prompt.replace("\n\n\n", "\n\n") + system_prompt = LLM_SETTINGS.default_system_prompt if system_prompt is None else system_prompt + messages = [ + { + "role": "system", + "content": system_prompt, + }, + ] + messages.extend(former_messages[-1 * LLM_SETTINGS.max_past_message_include :]) + messages.append( + { + "role": "user", + "content": user_prompt, + }, + ) + return messages + + def build_messages_and_create_chat_completion( # type: ignore[no-untyped-def] + self, + user_prompt: str, + system_prompt: str | None = None, + former_messages: list | None = None, + chat_cache_prefix: str = "", + shrink_multiple_break: bool = False, + *args, + **kwargs, + ) -> str: + if former_messages is None: + former_messages = [] + messages = self._build_messages( + user_prompt, + system_prompt, + former_messages, + shrink_multiple_break=shrink_multiple_break, + ) + + resp = self._try_create_chat_completion_or_embedding( # type: ignore[misc] + *args, + messages=messages, + chat_completion=True, + chat_cache_prefix=chat_cache_prefix, + **kwargs, + ) + if isinstance(resp, list): + raise ValueError("The response of _try_create_chat_completion_or_embedding should be a string.") + logger.log_object({"system": system_prompt, "user": user_prompt, "resp": resp}, tag="debug_llm") + return resp + + def create_embedding(self, input_content: str | list[str], *args, **kwargs) -> list[Any] | Any: # type: ignore[no-untyped-def] + input_content_list = [input_content] if isinstance(input_content, str) else input_content + resp = self._try_create_chat_completion_or_embedding( # type: ignore[misc] + input_content_list=input_content_list, + embedding=True, + *args, + **kwargs, + ) + if isinstance(input_content, str): + return resp[0] + return resp + + def _create_chat_completion_auto_continue(self, messages: list[dict[str, Any]], *args, **kwargs) -> str: # type: ignore[no-untyped-def] + """ + Call the chat completion function and automatically continue the conversation if the finish_reason is length. + TODO: This function only continues once, maybe need to continue more than once in the future. + """ + response, finish_reason = self._create_chat_completion_inner_function(messages, *args, **kwargs) + + if finish_reason == "length": + new_message = deepcopy(messages) + new_message.append({"role": "assistant", "content": response}) + new_message.append( + { + "role": "user", + "content": "continue the former output with no overlap", + }, + ) + new_response, finish_reason = self._create_chat_completion_inner_function(new_message, *args, **kwargs) + return response + new_response + return response + + def _try_create_chat_completion_or_embedding( # type: ignore[no-untyped-def] + self, + max_retry: int = 10, + chat_completion: bool = False, + embedding: bool = False, + *args, + **kwargs, + ) -> str | list[float]: + assert not (chat_completion and embedding), "chat_completion and embedding cannot be True at the same time" + max_retry = LLM_SETTINGS.max_retry if LLM_SETTINGS.max_retry is not None else max_retry + for i in range(max_retry): + try: + if embedding: + return self._create_embedding_inner_function(*args, **kwargs) + if chat_completion: + return self._create_chat_completion_auto_continue(*args, **kwargs) + except openai.BadRequestError as e: # noqa: PERF203 + logger.warning(str(e)) + logger.warning(f"Retrying {i+1}th time...") + if ( + "'messages' must contain the word 'json' in some form" in e.message + or "\\'messages\\' must contain the word \\'json\\' in some form" in e.message + ): + kwargs["add_json_in_prompt"] = True + elif embedding and "maximum context length" in e.message: + kwargs["input_content_list"] = [ + content[: len(content) // 2] for content in kwargs.get("input_content_list", []) + ] + except Exception as e: # noqa: BLE001 + logger.warning(str(e)) + logger.warning(f"Retrying {i+1}th time...") + time.sleep(self.retry_wait_seconds) + error_message = f"Failed to create chat completion after {max_retry} retries." + raise RuntimeError(error_message) + + def _create_embedding_inner_function( # type: ignore[no-untyped-def] + self, input_content_list: list[str], *args, **kwargs + ) -> list[Any]: # noqa: ARG002 + content_to_embedding_dict = {} + filtered_input_content_list = [] + if self.use_embedding_cache: + for content in input_content_list: + cache_result = self.cache.embedding_get(content) + if cache_result is not None: + content_to_embedding_dict[content] = cache_result + else: + filtered_input_content_list.append(content) + else: + filtered_input_content_list = input_content_list + + if len(filtered_input_content_list) > 0: + for sliced_filtered_input_content_list in [ + filtered_input_content_list[i : i + LLM_SETTINGS.embedding_max_str_num] + for i in range(0, len(filtered_input_content_list), LLM_SETTINGS.embedding_max_str_num) + ]: + if self.embedding_use_azure: + response = self.embedding_client.embeddings.create( + model=self.embedding_model, + input=sliced_filtered_input_content_list, + ) + else: + response = self.embedding_client.embeddings.create( + model=self.embedding_model, + input=sliced_filtered_input_content_list, + ) + for index, data in enumerate(response.data): + content_to_embedding_dict[sliced_filtered_input_content_list[index]] = data.embedding + + if self.dump_embedding_cache: + self.cache.embedding_set(content_to_embedding_dict) + return [content_to_embedding_dict[content] for content in input_content_list] + + def _build_log_messages(self, messages: list[dict[str, Any]]) -> str: + log_messages = "" + for m in messages: + log_messages += ( + f"\n{LogColors.MAGENTA}{LogColors.BOLD}Role:{LogColors.END}" + f"{LogColors.CYAN}{m['role']}{LogColors.END}\n" + f"{LogColors.MAGENTA}{LogColors.BOLD}Content:{LogColors.END} " + f"{LogColors.CYAN}{m['content']}{LogColors.END}\n" + ) + return log_messages + + def _create_chat_completion_inner_function( # type: ignore[no-untyped-def] # noqa: C901, PLR0912, PLR0915 + self, + messages: list[dict[str, Any]], + temperature: float | None = None, + max_tokens: int | None = None, + chat_cache_prefix: str = "", + frequency_penalty: float | None = None, + presence_penalty: float | None = None, + json_mode: bool = False, + add_json_in_prompt: bool = False, + seed: Optional[int] = None, + *args, + **kwargs, + ) -> tuple[str, str | None]: + """ + seed : Optional[int] + When retrying with cache enabled, it will keep returning the same results. + To make retries useful, we need to enable a seed. + This seed is different from `self.chat_seed` for GPT. It is for the local cache mechanism enabled by RD-Agent locally. + """ + if seed is None and LLM_SETTINGS.use_auto_chat_cache_seed_gen: + seed = LLM_CACHE_SEED_GEN.get_next_seed() + + # TODO: we can add this function back to avoid so much `self.cfg.log_llm_chat_content` + if LLM_SETTINGS.log_llm_chat_content: + logger.info(self._build_log_messages(messages), tag="llm_messages") + # TODO: fail to use loguru adaptor due to stream response + input_content_json = json.dumps(messages) + input_content_json = ( + chat_cache_prefix + input_content_json + f"" + ) # FIXME this is a hack to make sure the cache represents the round index + if self.use_chat_cache: + cache_result = self.cache.chat_get(input_content_json) + if cache_result is not None: + if LLM_SETTINGS.log_llm_chat_content: + logger.info(f"{LogColors.CYAN}Response:{cache_result}{LogColors.END}", tag="llm_messages") + return cache_result, None + + if temperature is None: + temperature = LLM_SETTINGS.chat_temperature + if max_tokens is None: + max_tokens = LLM_SETTINGS.chat_max_tokens + if frequency_penalty is None: + frequency_penalty = LLM_SETTINGS.chat_frequency_penalty + if presence_penalty is None: + presence_penalty = LLM_SETTINGS.chat_presence_penalty + + # Use index 4 to skip the current function and intermediate calls, + # and get the locals of the caller's frame. + caller_locals = inspect.stack()[4].frame.f_locals + if "self" in caller_locals: + tag = caller_locals["self"].__class__.__name__ + else: + tag = inspect.stack()[4].function + model = self.chat_model_map.get(tag, self.chat_model) + + finish_reason = None + if self.use_llama2: + response = self.generator.chat_completion( + messages, + max_gen_len=max_tokens, + temperature=temperature, + ) + resp = response[0]["generation"]["content"] + if LLM_SETTINGS.log_llm_chat_content: + logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages") + elif self.use_gcr_endpoint: + body = str.encode( + json.dumps( + { + "input_data": { + "input_string": messages, + "parameters": { + "temperature": self.gcr_endpoint_temperature, + "top_p": self.gcr_endpoint_top_p, + "max_new_tokens": self.gcr_endpoint_max_token, + }, + }, + }, + ), + ) + + req = urllib.request.Request(self.gcr_endpoint, body, self.headers) # noqa: S310 + response = urllib.request.urlopen(req) # noqa: S310 + resp = json.loads(response.read().decode())["output"] + if LLM_SETTINGS.log_llm_chat_content: + logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages") + else: + call_kwargs = dict( + model=model, + messages=messages, + max_tokens=max_tokens, + temperature=temperature, + stream=self.chat_stream, + seed=self.chat_seed, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ) + if json_mode: + if add_json_in_prompt: + for message in messages[::-1]: + message["content"] = message["content"] + "\nPlease respond in json format." + if message["role"] == "system": + break + call_kwargs["response_format"] = {"type": "json_object"} + response = self.chat_client.chat.completions.create(**call_kwargs) + + if self.chat_stream: + resp = "" + # TODO: with logger.config(stream=self.chat_stream): and add a `stream_start` flag to add timestamp for first message. + if LLM_SETTINGS.log_llm_chat_content: + logger.info(f"{LogColors.CYAN}Response:{LogColors.END}", tag="llm_messages") + + for chunk in response: + content = ( + chunk.choices[0].delta.content + if len(chunk.choices) > 0 and chunk.choices[0].delta.content is not None + else "" + ) + if LLM_SETTINGS.log_llm_chat_content: + logger.info(LogColors.CYAN + content + LogColors.END, raw=True, tag="llm_messages") + resp += content + if len(chunk.choices) > 0 and chunk.choices[0].finish_reason is not None: + finish_reason = chunk.choices[0].finish_reason + + if LLM_SETTINGS.log_llm_chat_content: + logger.info("\n", raw=True, tag="llm_messages") + + else: + resp = response.choices[0].message.content + finish_reason = response.choices[0].finish_reason + if LLM_SETTINGS.log_llm_chat_content: + logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages") + logger.info( + json.dumps( + { + "tag": tag, + "total_tokens": response.usage.total_tokens, + "prompt_tokens": response.usage.prompt_tokens, + "completion_tokens": response.usage.completion_tokens, + "model": model, + } + ), + tag="llm_messages", + ) + if json_mode: + json.loads(resp) + if self.dump_chat_cache: + self.cache.chat_set(input_content_json, resp) + return resp, finish_reason + + def _calculate_token_from_messages(self, messages: list[dict[str, Any]]) -> int: + if self.encoder is None: + raise ValueError("Encoder is not initialized.") + if self.use_llama2 or self.use_gcr_endpoint: + logger.warning("num_tokens_from_messages() is not implemented for model llama2.") + return 0 # TODO implement this function for llama2 + + if "gpt4" in self.chat_model or "gpt-4" in self.chat_model: + tokens_per_message = 3 + tokens_per_name = 1 + else: + tokens_per_message = 4 # every message follows {role/name}\n{content}\n + tokens_per_name = -1 # if there's a name, the role is omitted + num_tokens = 0 + for message in messages: + num_tokens += tokens_per_message + for key, value in message.items(): + num_tokens += len(self.encoder.encode(value)) + if key == "name": + num_tokens += tokens_per_name + num_tokens += 3 # every reply is primed with assistant + return num_tokens + + def build_messages_and_calculate_token( + self, + user_prompt: str, + system_prompt: str | None, + former_messages: list[dict[str, Any]] | None = None, + *, + shrink_multiple_break: bool = False, + ) -> int: + if former_messages is None: + former_messages = [] + messages = self._build_messages( + user_prompt, system_prompt, former_messages, shrink_multiple_break=shrink_multiple_break + ) + return self._calculate_token_from_messages(messages) diff --git a/rdagent/oai/backend/deprec/__init__.py b/rdagent/oai/backend/deprec/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/rdagent/oai/backend/deprec/conf.py b/rdagent/oai/backend/deprec/conf.py deleted file mode 100644 index e69de29b..00000000 diff --git a/rdagent/oai/backend/deprec/deprecated.py b/rdagent/oai/backend/deprec/deprecated.py deleted file mode 100644 index e69de29b..00000000 diff --git a/rdagent/oai/backend/litellm.py b/rdagent/oai/backend/litellm.py index e69de29b..998273c6 100644 --- a/rdagent/oai/backend/litellm.py +++ b/rdagent/oai/backend/litellm.py @@ -0,0 +1,149 @@ +import os +import uuid +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from litellm import acompletion, completion +from litellm import encode as encode_litellm +from litellm import token_counter + +from rdagent.core.conf import ExtendedBaseSettings +from rdagent.core.utils import LLM_CACHE_SEED_GEN, SingletonBaseClass, import_class +from rdagent.log import LogColors +from rdagent.log import rdagent_logger as logger +from rdagent.oai.backend.base import APIBackend +from rdagent.oai.llm_conf import LLM_SETTINGS + + +class LiteLLMSettings(ExtendedBaseSettings): + + class Config: + env_prefix = "LITELLM_" + """Use `LITELLM_` as prefix for environment variables""" + + # LiteLLM backend related config + chat_model: str = "openai/gpt-4o" + # LiteLLM embedding related config + embedding_model: str = "openai/text-embedding-3-small" + + +LITELLM_SETTINGS = LiteLLMSettings() + + +class LiteLLMAPIBackend(APIBackend): + """LiteLLM implementation of APIBackend interface""" + + def __init__(self, litellm_model_name: str = "", litellm_api_key: str = "", *args: Any, **kwargs: Any) -> None: + super().__init__() + if len(args) > 0 or len(kwargs) > 0: + logger.warning("LiteLLM backend does not support any additional arguments") + + def build_chat_session( + self, conversation_id: Optional[str] = None, session_system_prompt: Optional[str] = None + ) -> Any: + """Create a new chat session using LiteLLM""" + # return { + # "conversation_id": conversation_id or str(uuid.uuid4()), + # "system_prompt": session_system_prompt, + # "messages": [] + # } + raise NotImplementedError("LiteLLM backend does not support chat session creation") + # TODO: Implement the chat session creation logic , with ChatSession class + + def build_messages_and_create_chat_completion( + self, + user_prompt: str, + system_prompt: Optional[str] = None, + former_messages: Optional[List[Any]] = None, + chat_cache_prefix: str = "", + shrink_multiple_break: bool = False, + *args: Any, + **kwargs: Any, + ) -> str: + """Build messages and get LiteLLM chat completion""" + messages = [] + + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + + if former_messages: + messages.extend(former_messages) + + messages.append({"role": "user", "content": user_prompt}) + model_name = LITELLM_SETTINGS.chat_model + # Call LiteLLM completion + response = completion( + model=model_name, + messages=messages, + stream=kwargs.get("stream", False), + temperature=kwargs.get("temperature", 0.7), + max_tokens=kwargs.get("max_tokens", 1000), + **kwargs, + ) + logger.info( + f"{LogColors.GREEN}Using chat model{LogColors.END} {model_name}", + tag="debug_llm", + ) + + if system_prompt: + logger.info(f"{LogColors.RED}system:{LogColors.END} {system_prompt}", tag="debug_llm") + if former_messages: + for message in former_messages: + logger.info(f"{LogColors.CYAN}{message['role']}:{LogColors.END} {message['content']}", tag="debug_llm") + else: + logger.info( + f"{LogColors.RED}user:{LogColors.END} {user_prompt}\n{LogColors.BLUE}resp(next row):\n{LogColors.END} {response.choices[0].message.content}", + tag="debug_llm", + ) + + return str(response.choices[0].message.content) + + def create_embedding(self, input_content: str | list[str], *args: Any, **kwargs: Any) -> list[Any] | Any: + """Create embeddings using LiteLLM""" + from litellm import embedding + + single_input = False + if isinstance(input_content, str): + input_content = [input_content] + single_input = True + response_list = [] + for input_content_iter in input_content: + model_name = LITELLM_SETTINGS.embedding_model or "azure/text-embedding-3-small" + logger.info(f"{LogColors.GREEN}Using emb model{LogColors.END} {model_name}", tag="debug_litellm_emb") + logger.info(f"Creating embedding for: {input_content_iter}", tag="debug_litellm_emb") + if not isinstance(input_content_iter, str): + raise ValueError("Input content must be a string") + response = embedding( + model=model_name, + input=input_content_iter, + **kwargs, + ) + response_list.append(response.data[0]["embedding"]) + if single_input: + return response_list[0] + return response_list + + def build_messages_and_calculate_token( + self, + user_prompt: str, + system_prompt: Optional[str], + former_messages: Optional[List[Dict[str, Any]]] = None, + shrink_multiple_break: bool = False, + ) -> int: + """Build messages and calculate their token count using LiteLLM""" + messages = [] + + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + + if former_messages: + messages.extend(former_messages) + + messages.append({"role": "user", "content": user_prompt}) + + num_tokens = token_counter( + model=LITELLM_SETTINGS.chat_model, + messages=messages, + ) + logger.info(f"{LogColors.CYAN}Token count: {LogColors.END} {num_tokens}", tag="debug_litellm_token") + return num_tokens diff --git a/rdagent/oai/llm_utils.py b/rdagent/oai/llm_utils.py index e9c26ab8..ccbe73e1 100644 --- a/rdagent/oai/llm_utils.py +++ b/rdagent/oai/llm_utils.py @@ -1,811 +1,13 @@ from __future__ import annotations -import hashlib -import inspect -import json -import os -import random -import re -import sqlite3 -import ssl -import time -import urllib.request -import uuid -from copy import deepcopy -from pathlib import Path -from typing import Any, Optional, cast +from typing import Any, Type import numpy as np -import tiktoken -from rdagent.core.utils import LLM_CACHE_SEED_GEN, SingletonBaseClass -from rdagent.log import LogColors -from rdagent.log import rdagent_logger as logger +from rdagent.core.utils import import_class +from rdagent.oai.backend.base import APIBackend as BaseAPIBackend from rdagent.oai.llm_conf import LLM_SETTINGS - -DEFAULT_QLIB_DOT_PATH = Path("./") - - -def md5_hash(input_string: str) -> str: - hash_md5 = hashlib.md5(usedforsecurity=False) - input_bytes = input_string.encode("utf-8") - hash_md5.update(input_bytes) - return hash_md5.hexdigest() - - -try: - from azure.identity import DefaultAzureCredential, get_bearer_token_provider -except ImportError: - logger.warning("azure.identity is not installed.") - -try: - import openai -except ImportError: - logger.warning("openai is not installed.") - -try: - from llama import Llama -except ImportError: - if LLM_SETTINGS.use_llama2: - logger.warning("llama is not installed.") - - -class ConvManager: - """ - This is a conversation manager of LLM - It is for convenience of exporting conversation for debugging. - """ - - def __init__( - self, - path: Path | str = DEFAULT_QLIB_DOT_PATH / "llm_conv", - recent_n: int = 10, - ) -> None: - self.path = Path(path) - self.path.mkdir(parents=True, exist_ok=True) - self.recent_n = recent_n - - def _rotate_files(self) -> None: - pairs = [] - for f in self.path.glob("*.json"): - m = re.match(r"(\d+).json", f.name) - if m is not None: - n = int(m.group(1)) - pairs.append((n, f)) - pairs.sort(key=lambda x: x[0]) - for n, f in pairs[: self.recent_n][::-1]: - if (self.path / f"{n+1}.json").exists(): - (self.path / f"{n+1}.json").unlink() - f.rename(self.path / f"{n+1}.json") - - def append(self, conv: tuple[list, str]) -> None: - self._rotate_files() - with (self.path / "0.json").open("w") as file: - json.dump(conv, file) - # TODO: reseve line breaks to make it more convient to edit file directly. - - -class SQliteLazyCache(SingletonBaseClass): - def __init__(self, cache_location: str) -> None: - super().__init__() - self.cache_location = cache_location - db_file_exist = Path(cache_location).exists() - # TODO: sqlite3 does not support multiprocessing. - self.conn = sqlite3.connect(cache_location, timeout=20) - self.c = self.conn.cursor() - if not db_file_exist: - self.c.execute( - """ - CREATE TABLE chat_cache ( - md5_key TEXT PRIMARY KEY, - chat TEXT - ) - """, - ) - self.c.execute( - """ - CREATE TABLE embedding_cache ( - md5_key TEXT PRIMARY KEY, - embedding TEXT - ) - """, - ) - self.c.execute( - """ - CREATE TABLE message_cache ( - conversation_id TEXT PRIMARY KEY, - message TEXT - ) - """, - ) - self.conn.commit() - - def chat_get(self, key: str) -> str | None: - md5_key = md5_hash(key) - self.c.execute("SELECT chat FROM chat_cache WHERE md5_key=?", (md5_key,)) - result = self.c.fetchone() - return None if result is None else result[0] - - def embedding_get(self, key: str) -> list | dict | str | None: - md5_key = md5_hash(key) - self.c.execute("SELECT embedding FROM embedding_cache WHERE md5_key=?", (md5_key,)) - result = self.c.fetchone() - return None if result is None else json.loads(result[0]) - - def chat_set(self, key: str, value: str) -> None: - md5_key = md5_hash(key) - self.c.execute( - "INSERT OR REPLACE INTO chat_cache (md5_key, chat) VALUES (?, ?)", - (md5_key, value), - ) - self.conn.commit() - return None - - def embedding_set(self, content_to_embedding_dict: dict) -> None: - for key, value in content_to_embedding_dict.items(): - md5_key = md5_hash(key) - self.c.execute( - "INSERT OR REPLACE INTO embedding_cache (md5_key, embedding) VALUES (?, ?)", - (md5_key, json.dumps(value)), - ) - self.conn.commit() - - def message_get(self, conversation_id: str) -> list[dict[str, Any]]: - self.c.execute("SELECT message FROM message_cache WHERE conversation_id=?", (conversation_id,)) - result = self.c.fetchone() - return [] if result is None else cast(list[dict[str, Any]], json.loads(result[0])) - - def message_set(self, conversation_id: str, message_value: list[dict[str, Any]]) -> None: - self.c.execute( - "INSERT OR REPLACE INTO message_cache (conversation_id, message) VALUES (?, ?)", - (conversation_id, json.dumps(message_value)), - ) - self.conn.commit() - return None - - -class SessionChatHistoryCache(SingletonBaseClass): - def __init__(self) -> None: - """load all history conversation json file from self.session_cache_location""" - self.cache = SQliteLazyCache(cache_location=LLM_SETTINGS.prompt_cache_path) - - def message_get(self, conversation_id: str) -> list[dict[str, Any]]: - return self.cache.message_get(conversation_id) - - def message_set(self, conversation_id: str, message_value: list[dict[str, Any]]) -> None: - self.cache.message_set(conversation_id, message_value) - - -class ChatSession: - def __init__(self, api_backend: Any, conversation_id: str | None = None, system_prompt: str | None = None) -> None: - self.conversation_id = str(uuid.uuid4()) if conversation_id is None else conversation_id - self.system_prompt = system_prompt if system_prompt is not None else LLM_SETTINGS.default_system_prompt - self.api_backend = api_backend - - def build_chat_completion_message(self, user_prompt: str) -> list[dict[str, Any]]: - history_message = SessionChatHistoryCache().message_get(self.conversation_id) - messages = history_message - if not messages: - messages.append({"role": "system", "content": self.system_prompt}) - messages.append( - { - "role": "user", - "content": user_prompt, - }, - ) - return messages - - def build_chat_completion_message_and_calculate_token(self, user_prompt: str) -> Any: - messages = self.build_chat_completion_message(user_prompt) - return self.api_backend.calculate_token_from_messages(messages) - - def build_chat_completion(self, user_prompt: str, *args, **kwargs) -> str: # type: ignore[no-untyped-def] - """ - this function is to build the session messages - user prompt should always be provided - """ - messages = self.build_chat_completion_message(user_prompt) - - with logger.tag(f"session_{self.conversation_id}"): - response: str = self.api_backend._try_create_chat_completion_or_embedding( # noqa: SLF001 - *args, - messages=messages, - chat_completion=True, - **kwargs, - ) - logger.log_object({"user": user_prompt, "resp": response}, tag="debug_llm") - - messages.append( - { - "role": "assistant", - "content": response, - }, - ) - SessionChatHistoryCache().message_set(self.conversation_id, messages) - return response - - def get_conversation_id(self) -> str: - return self.conversation_id - - def display_history(self) -> None: - # TODO: Realize a beautiful presentation format for history messages - pass - - -class APIBackend: - """ - This is a unified interface for different backends. - - (xiao) thinks integrate all kinds of API in a single class is not a good design. - So we should split them into different classes in `oai/backends/` in the future. - """ - - # FIXME: (xiao) We should avoid using self.xxxx. - # Instead, we can use LLM_SETTINGS directly. If it's difficult to support different backend settings, we can split them into multiple BaseSettings. - def __init__( # noqa: C901, PLR0912, PLR0915 - self, - *, - chat_api_key: str | None = None, - chat_model: str | None = None, - chat_api_base: str | None = None, - chat_api_version: str | None = None, - embedding_api_key: str | None = None, - embedding_model: str | None = None, - embedding_api_base: str | None = None, - embedding_api_version: str | None = None, - use_chat_cache: bool | None = None, - dump_chat_cache: bool | None = None, - use_embedding_cache: bool | None = None, - dump_embedding_cache: bool | None = None, - ) -> None: - if LLM_SETTINGS.use_llama2: - self.generator = Llama.build( - ckpt_dir=LLM_SETTINGS.llama2_ckpt_dir, - tokenizer_path=LLM_SETTINGS.llama2_tokenizer_path, - max_seq_len=LLM_SETTINGS.chat_max_tokens, - max_batch_size=LLM_SETTINGS.llams2_max_batch_size, - ) - self.encoder = None - elif LLM_SETTINGS.use_gcr_endpoint: - gcr_endpoint_type = LLM_SETTINGS.gcr_endpoint_type - if gcr_endpoint_type == "llama2_70b": - self.gcr_endpoint_key = LLM_SETTINGS.llama2_70b_endpoint_key - self.gcr_endpoint_deployment = LLM_SETTINGS.llama2_70b_endpoint_deployment - self.gcr_endpoint = LLM_SETTINGS.llama2_70b_endpoint - elif gcr_endpoint_type == "llama3_70b": - self.gcr_endpoint_key = LLM_SETTINGS.llama3_70b_endpoint_key - self.gcr_endpoint_deployment = LLM_SETTINGS.llama3_70b_endpoint_deployment - self.gcr_endpoint = LLM_SETTINGS.llama3_70b_endpoint - elif gcr_endpoint_type == "phi2": - self.gcr_endpoint_key = LLM_SETTINGS.phi2_endpoint_key - self.gcr_endpoint_deployment = LLM_SETTINGS.phi2_endpoint_deployment - self.gcr_endpoint = LLM_SETTINGS.phi2_endpoint - elif gcr_endpoint_type == "phi3_4k": - self.gcr_endpoint_key = LLM_SETTINGS.phi3_4k_endpoint_key - self.gcr_endpoint_deployment = LLM_SETTINGS.phi3_4k_endpoint_deployment - self.gcr_endpoint = LLM_SETTINGS.phi3_4k_endpoint - elif gcr_endpoint_type == "phi3_128k": - self.gcr_endpoint_key = LLM_SETTINGS.phi3_128k_endpoint_key - self.gcr_endpoint_deployment = LLM_SETTINGS.phi3_128k_endpoint_deployment - self.gcr_endpoint = LLM_SETTINGS.phi3_128k_endpoint - else: - error_message = f"Invalid gcr_endpoint_type: {gcr_endpoint_type}" - raise ValueError(error_message) - self.headers = { - "Content-Type": "application/json", - "Authorization": ("Bearer " + self.gcr_endpoint_key), - } - self.gcr_endpoint_temperature = LLM_SETTINGS.gcr_endpoint_temperature - self.gcr_endpoint_top_p = LLM_SETTINGS.gcr_endpoint_top_p - self.gcr_endpoint_do_sample = LLM_SETTINGS.gcr_endpoint_do_sample - self.gcr_endpoint_max_token = LLM_SETTINGS.gcr_endpoint_max_token - if not os.environ.get("PYTHONHTTPSVERIFY", "") and hasattr(ssl, "_create_unverified_context"): - ssl._create_default_https_context = ssl._create_unverified_context # noqa: SLF001 - self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map) - self.chat_model = LLM_SETTINGS.chat_model if chat_model is None else chat_model - self.encoder = None - else: - self.chat_use_azure = LLM_SETTINGS.chat_use_azure or LLM_SETTINGS.use_azure - self.embedding_use_azure = LLM_SETTINGS.embedding_use_azure or LLM_SETTINGS.use_azure - self.chat_use_azure_token_provider = LLM_SETTINGS.chat_use_azure_token_provider - self.embedding_use_azure_token_provider = LLM_SETTINGS.embedding_use_azure_token_provider - self.managed_identity_client_id = LLM_SETTINGS.managed_identity_client_id - - # Priority: chat_api_key/embedding_api_key > openai_api_key > os.environ.get("OPENAI_API_KEY") - # TODO: Simplify the key design. Consider Pandatic's field alias & priority. - self.chat_api_key = ( - chat_api_key - or LLM_SETTINGS.chat_openai_api_key - or LLM_SETTINGS.openai_api_key - or os.environ.get("OPENAI_API_KEY") - ) - self.embedding_api_key = ( - embedding_api_key - or LLM_SETTINGS.embedding_openai_api_key - or LLM_SETTINGS.openai_api_key - or os.environ.get("OPENAI_API_KEY") - ) - - self.chat_model = LLM_SETTINGS.chat_model if chat_model is None else chat_model - self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map) - self.encoder = self._get_encoder() - self.chat_openai_base_url = LLM_SETTINGS.chat_openai_base_url - self.embedding_openai_base_url = LLM_SETTINGS.embedding_openai_base_url - self.chat_api_base = LLM_SETTINGS.chat_azure_api_base if chat_api_base is None else chat_api_base - self.chat_api_version = ( - LLM_SETTINGS.chat_azure_api_version if chat_api_version is None else chat_api_version - ) - self.chat_stream = LLM_SETTINGS.chat_stream - self.chat_seed = LLM_SETTINGS.chat_seed - - self.embedding_model = LLM_SETTINGS.embedding_model if embedding_model is None else embedding_model - self.embedding_api_base = ( - LLM_SETTINGS.embedding_azure_api_base if embedding_api_base is None else embedding_api_base - ) - self.embedding_api_version = ( - LLM_SETTINGS.embedding_azure_api_version if embedding_api_version is None else embedding_api_version - ) - - if (self.chat_use_azure or self.embedding_use_azure) and ( - self.chat_use_azure_token_provider or self.embedding_use_azure_token_provider - ): - dac_kwargs = {} - if self.managed_identity_client_id is not None: - dac_kwargs["managed_identity_client_id"] = self.managed_identity_client_id - credential = DefaultAzureCredential(**dac_kwargs) - token_provider = get_bearer_token_provider( - credential, - "https://cognitiveservices.azure.com/.default", - ) - self.chat_client: openai.OpenAI = ( - openai.AzureOpenAI( - azure_ad_token_provider=token_provider if self.chat_use_azure_token_provider else None, - api_key=self.chat_api_key if not self.chat_use_azure_token_provider else None, - api_version=self.chat_api_version, - azure_endpoint=self.chat_api_base, - ) - if self.chat_use_azure - else openai.OpenAI(api_key=self.chat_api_key, base_url=self.chat_openai_base_url) - ) - - self.embedding_client: openai.OpenAI = ( - openai.AzureOpenAI( - azure_ad_token_provider=token_provider if self.embedding_use_azure_token_provider else None, - api_key=self.embedding_api_key if not self.embedding_use_azure_token_provider else None, - api_version=self.embedding_api_version, - azure_endpoint=self.embedding_api_base, - ) - if self.embedding_use_azure - else openai.OpenAI(api_key=self.embedding_api_key, base_url=self.embedding_openai_base_url) - ) - - self.dump_chat_cache = LLM_SETTINGS.dump_chat_cache if dump_chat_cache is None else dump_chat_cache - self.use_chat_cache = LLM_SETTINGS.use_chat_cache if use_chat_cache is None else use_chat_cache - self.dump_embedding_cache = ( - LLM_SETTINGS.dump_embedding_cache if dump_embedding_cache is None else dump_embedding_cache - ) - self.use_embedding_cache = ( - LLM_SETTINGS.use_embedding_cache if use_embedding_cache is None else use_embedding_cache - ) - if self.dump_chat_cache or self.use_chat_cache or self.dump_embedding_cache or self.use_embedding_cache: - self.cache_file_location = LLM_SETTINGS.prompt_cache_path - self.cache = SQliteLazyCache(cache_location=self.cache_file_location) - - # transfer the config to the class if the config is not supposed to change during the runtime - self.use_llama2 = LLM_SETTINGS.use_llama2 - self.use_gcr_endpoint = LLM_SETTINGS.use_gcr_endpoint - self.retry_wait_seconds = LLM_SETTINGS.retry_wait_seconds - - def _get_encoder(self) -> tiktoken.Encoding: - """ - tiktoken.encoding_for_model(self.chat_model) does not cover all cases it should consider. - - This function attempts to handle several edge cases. - """ - - # 1) cases - def _azure_patch(model: str) -> str: - """ - When using Azure API, self.chat_model is the deployment name that can be any string. - For example, it may be `gpt-4o_2024-08-06`. But tiktoken.encoding_for_model can't handle this. - """ - return model.replace("_", "-") - - model = self.chat_model - try: - encoding = tiktoken.encoding_for_model(model) - except KeyError: - logger.warning(f"Failed to get encoder. Trying to patch the model name") - for patch_func in [_azure_patch]: - try: - encoding = tiktoken.encoding_for_model(patch_func(model)) - except KeyError: - logger.error(f"Failed to get encoder even after patching with {patch_func.__name__}") - raise - return encoding - - def build_chat_session( - self, - conversation_id: str | None = None, - session_system_prompt: str | None = None, - ) -> ChatSession: - """ - conversation_id is a 256-bit string created by uuid.uuid4() and is also - the file name under session_cache_folder/ for each conversation - """ - return ChatSession(self, conversation_id, session_system_prompt) - - def build_messages( - self, - user_prompt: str, - system_prompt: str | None = None, - former_messages: list[dict[str, Any]] | None = None, - *, - shrink_multiple_break: bool = False, - ) -> list[dict[str, Any]]: - """ - build the messages to avoid implementing several redundant lines of code - - """ - if former_messages is None: - former_messages = [] - # shrink multiple break will recursively remove multiple breaks(more than 2) - if shrink_multiple_break: - while "\n\n\n" in user_prompt: - user_prompt = user_prompt.replace("\n\n\n", "\n\n") - if system_prompt is not None: - while "\n\n\n" in system_prompt: - system_prompt = system_prompt.replace("\n\n\n", "\n\n") - system_prompt = LLM_SETTINGS.default_system_prompt if system_prompt is None else system_prompt - messages = [ - { - "role": "system", - "content": system_prompt, - }, - ] - messages.extend(former_messages[-1 * LLM_SETTINGS.max_past_message_include :]) - messages.append( - { - "role": "user", - "content": user_prompt, - }, - ) - return messages - - def build_messages_and_create_chat_completion( # type: ignore[no-untyped-def] - self, - user_prompt: str, - system_prompt: str | None = None, - former_messages: list | None = None, - chat_cache_prefix: str = "", - shrink_multiple_break: bool = False, - *args, - **kwargs, - ) -> str: - if former_messages is None: - former_messages = [] - messages = self.build_messages( - user_prompt, - system_prompt, - former_messages, - shrink_multiple_break=shrink_multiple_break, - ) - - resp = self._try_create_chat_completion_or_embedding( # type: ignore[misc] - *args, - messages=messages, - chat_completion=True, - chat_cache_prefix=chat_cache_prefix, - **kwargs, - ) - if isinstance(resp, list): - raise ValueError("The response of _try_create_chat_completion_or_embedding should be a string.") - logger.log_object({"system": system_prompt, "user": user_prompt, "resp": resp}, tag="debug_llm") - return resp - - def create_embedding(self, input_content: str | list[str], *args, **kwargs) -> list[Any] | Any: # type: ignore[no-untyped-def] - input_content_list = [input_content] if isinstance(input_content, str) else input_content - resp = self._try_create_chat_completion_or_embedding( # type: ignore[misc] - input_content_list=input_content_list, - embedding=True, - *args, - **kwargs, - ) - if isinstance(input_content, str): - return resp[0] - return resp - - def _create_chat_completion_auto_continue(self, messages: list[dict[str, Any]], *args, **kwargs) -> str: # type: ignore[no-untyped-def] - """ - Call the chat completion function and automatically continue the conversation if the finish_reason is length. - TODO: This function only continues once, maybe need to continue more than once in the future. - """ - response, finish_reason = self._create_chat_completion_inner_function(messages, *args, **kwargs) - - if finish_reason == "length": - new_message = deepcopy(messages) - new_message.append({"role": "assistant", "content": response}) - new_message.append( - { - "role": "user", - "content": "continue the former output with no overlap", - }, - ) - new_response, finish_reason = self._create_chat_completion_inner_function(new_message, *args, **kwargs) - return response + new_response - return response - - def _try_create_chat_completion_or_embedding( # type: ignore[no-untyped-def] - self, - max_retry: int = 10, - chat_completion: bool = False, - embedding: bool = False, - *args, - **kwargs, - ) -> str | list[float]: - assert not (chat_completion and embedding), "chat_completion and embedding cannot be True at the same time" - max_retry = LLM_SETTINGS.max_retry if LLM_SETTINGS.max_retry is not None else max_retry - for i in range(max_retry): - try: - if embedding: - return self._create_embedding_inner_function(*args, **kwargs) - if chat_completion: - return self._create_chat_completion_auto_continue(*args, **kwargs) - except openai.BadRequestError as e: # noqa: PERF203 - logger.warning(str(e)) - logger.warning(f"Retrying {i+1}th time...") - if ( - "'messages' must contain the word 'json' in some form" in e.message - or "\\'messages\\' must contain the word \\'json\\' in some form" in e.message - ): - kwargs["add_json_in_prompt"] = True - elif embedding and "maximum context length" in e.message: - kwargs["input_content_list"] = [ - content[: len(content) // 2] for content in kwargs.get("input_content_list", []) - ] - except Exception as e: # noqa: BLE001 - logger.warning(str(e)) - logger.warning(f"Retrying {i+1}th time...") - time.sleep(self.retry_wait_seconds) - error_message = f"Failed to create chat completion after {max_retry} retries." - raise RuntimeError(error_message) - - def _create_embedding_inner_function( # type: ignore[no-untyped-def] - self, input_content_list: list[str], *args, **kwargs - ) -> list[Any]: # noqa: ARG002 - content_to_embedding_dict = {} - filtered_input_content_list = [] - if self.use_embedding_cache: - for content in input_content_list: - cache_result = self.cache.embedding_get(content) - if cache_result is not None: - content_to_embedding_dict[content] = cache_result - else: - filtered_input_content_list.append(content) - else: - filtered_input_content_list = input_content_list - - if len(filtered_input_content_list) > 0: - for sliced_filtered_input_content_list in [ - filtered_input_content_list[i : i + LLM_SETTINGS.embedding_max_str_num] - for i in range(0, len(filtered_input_content_list), LLM_SETTINGS.embedding_max_str_num) - ]: - if self.embedding_use_azure: - response = self.embedding_client.embeddings.create( - model=self.embedding_model, - input=sliced_filtered_input_content_list, - ) - else: - response = self.embedding_client.embeddings.create( - model=self.embedding_model, - input=sliced_filtered_input_content_list, - ) - for index, data in enumerate(response.data): - content_to_embedding_dict[sliced_filtered_input_content_list[index]] = data.embedding - - if self.dump_embedding_cache: - self.cache.embedding_set(content_to_embedding_dict) - return [content_to_embedding_dict[content] for content in input_content_list] - - def _build_log_messages(self, messages: list[dict[str, Any]]) -> str: - log_messages = "" - for m in messages: - log_messages += ( - f"\n{LogColors.MAGENTA}{LogColors.BOLD}Role:{LogColors.END}" - f"{LogColors.CYAN}{m['role']}{LogColors.END}\n" - f"{LogColors.MAGENTA}{LogColors.BOLD}Content:{LogColors.END} " - f"{LogColors.CYAN}{m['content']}{LogColors.END}\n" - ) - return log_messages - - def _create_chat_completion_inner_function( # type: ignore[no-untyped-def] # noqa: C901, PLR0912, PLR0915 - self, - messages: list[dict[str, Any]], - temperature: float | None = None, - max_tokens: int | None = None, - chat_cache_prefix: str = "", - frequency_penalty: float | None = None, - presence_penalty: float | None = None, - json_mode: bool = False, - add_json_in_prompt: bool = False, - seed: Optional[int] = None, - *args, - **kwargs, - ) -> tuple[str, str | None]: - """ - seed : Optional[int] - When retrying with cache enabled, it will keep returning the same results. - To make retries useful, we need to enable a seed. - This seed is different from `self.chat_seed` for GPT. It is for the local cache mechanism enabled by RD-Agent locally. - """ - if seed is None and LLM_SETTINGS.use_auto_chat_cache_seed_gen: - seed = LLM_CACHE_SEED_GEN.get_next_seed() - - # TODO: we can add this function back to avoid so much `self.cfg.log_llm_chat_content` - if LLM_SETTINGS.log_llm_chat_content: - logger.info(self._build_log_messages(messages), tag="llm_messages") - # TODO: fail to use loguru adaptor due to stream response - input_content_json = json.dumps(messages) - input_content_json = ( - chat_cache_prefix + input_content_json + f"" - ) # FIXME this is a hack to make sure the cache represents the round index - if self.use_chat_cache: - cache_result = self.cache.chat_get(input_content_json) - if cache_result is not None: - if LLM_SETTINGS.log_llm_chat_content: - logger.info(f"{LogColors.CYAN}Response:{cache_result}{LogColors.END}", tag="llm_messages") - return cache_result, None - - if temperature is None: - temperature = LLM_SETTINGS.chat_temperature - if max_tokens is None: - max_tokens = LLM_SETTINGS.chat_max_tokens - if frequency_penalty is None: - frequency_penalty = LLM_SETTINGS.chat_frequency_penalty - if presence_penalty is None: - presence_penalty = LLM_SETTINGS.chat_presence_penalty - - # Use index 4 to skip the current function and intermediate calls, - # and get the locals of the caller's frame. - caller_locals = inspect.stack()[4].frame.f_locals - if "self" in caller_locals: - tag = caller_locals["self"].__class__.__name__ - else: - tag = inspect.stack()[4].function - model = self.chat_model_map.get(tag, self.chat_model) - - finish_reason = None - if self.use_llama2: - response = self.generator.chat_completion( - messages, - max_gen_len=max_tokens, - temperature=temperature, - ) - resp = response[0]["generation"]["content"] - if LLM_SETTINGS.log_llm_chat_content: - logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages") - elif self.use_gcr_endpoint: - body = str.encode( - json.dumps( - { - "input_data": { - "input_string": messages, - "parameters": { - "temperature": self.gcr_endpoint_temperature, - "top_p": self.gcr_endpoint_top_p, - "max_new_tokens": self.gcr_endpoint_max_token, - }, - }, - }, - ), - ) - - req = urllib.request.Request(self.gcr_endpoint, body, self.headers) # noqa: S310 - response = urllib.request.urlopen(req) # noqa: S310 - resp = json.loads(response.read().decode())["output"] - if LLM_SETTINGS.log_llm_chat_content: - logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages") - else: - call_kwargs = dict( - model=model, - messages=messages, - max_tokens=max_tokens, - temperature=temperature, - stream=self.chat_stream, - seed=self.chat_seed, - frequency_penalty=frequency_penalty, - presence_penalty=presence_penalty, - ) - if json_mode: - if add_json_in_prompt: - for message in messages[::-1]: - message["content"] = message["content"] + "\nPlease respond in json format." - if message["role"] == "system": - break - call_kwargs["response_format"] = {"type": "json_object"} - response = self.chat_client.chat.completions.create(**call_kwargs) - - if self.chat_stream: - resp = "" - # TODO: with logger.config(stream=self.chat_stream): and add a `stream_start` flag to add timestamp for first message. - if LLM_SETTINGS.log_llm_chat_content: - logger.info(f"{LogColors.CYAN}Response:{LogColors.END}", tag="llm_messages") - - for chunk in response: - content = ( - chunk.choices[0].delta.content - if len(chunk.choices) > 0 and chunk.choices[0].delta.content is not None - else "" - ) - if LLM_SETTINGS.log_llm_chat_content: - logger.info(LogColors.CYAN + content + LogColors.END, raw=True, tag="llm_messages") - resp += content - if len(chunk.choices) > 0 and chunk.choices[0].finish_reason is not None: - finish_reason = chunk.choices[0].finish_reason - - if LLM_SETTINGS.log_llm_chat_content: - logger.info("\n", raw=True, tag="llm_messages") - - else: - resp = response.choices[0].message.content - finish_reason = response.choices[0].finish_reason - if LLM_SETTINGS.log_llm_chat_content: - logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages") - logger.info( - json.dumps( - { - "tag": tag, - "total_tokens": response.usage.total_tokens, - "prompt_tokens": response.usage.prompt_tokens, - "completion_tokens": response.usage.completion_tokens, - "model": model, - } - ), - tag="llm_messages", - ) - if json_mode: - json.loads(resp) - if self.dump_chat_cache: - self.cache.chat_set(input_content_json, resp) - return resp, finish_reason - - def calculate_token_from_messages(self, messages: list[dict[str, Any]]) -> int: - if self.encoder is None: - raise ValueError("Encoder is not initialized.") - if self.use_llama2 or self.use_gcr_endpoint: - logger.warning("num_tokens_from_messages() is not implemented for model llama2.") - return 0 # TODO implement this function for llama2 - - if "gpt4" in self.chat_model or "gpt-4" in self.chat_model: - tokens_per_message = 3 - tokens_per_name = 1 - else: - tokens_per_message = 4 # every message follows {role/name}\n{content}\n - tokens_per_name = -1 # if there's a name, the role is omitted - num_tokens = 0 - for message in messages: - num_tokens += tokens_per_message - for key, value in message.items(): - num_tokens += len(self.encoder.encode(value)) - if key == "name": - num_tokens += tokens_per_name - num_tokens += 3 # every reply is primed with assistant - return num_tokens - - def build_messages_and_calculate_token( - self, - user_prompt: str, - system_prompt: str | None, - former_messages: list[dict[str, Any]] | None = None, - *, - shrink_multiple_break: bool = False, - ) -> int: - if former_messages is None: - former_messages = [] - messages = self.build_messages( - user_prompt, system_prompt, former_messages, shrink_multiple_break=shrink_multiple_break - ) - return self.calculate_token_from_messages(messages) +from rdagent.utils import md5_hash # for compatible with previous import def calculate_embedding_distance_between_str_list( @@ -828,3 +30,15 @@ def calculate_embedding_distance_between_str_list( similarity_matrix = np.dot(source_embeddings_np, target_embeddings_np.T) return similarity_matrix.tolist() # type: ignore[no-any-return] + + +def get_api_backend(*args: Any, **kwargs: Any) -> BaseAPIBackend: # TODO: import it from base.py + """ + get llm api backend based on settings dynamically. + """ + api_backend_cls: Type[BaseAPIBackend] = import_class(LLM_SETTINGS.backend) + return api_backend_cls(*args, **kwargs) + + +# Alias +APIBackend = get_api_backend diff --git a/rdagent/scenarios/data_science/dev/runner/__init__.py b/rdagent/scenarios/data_science/dev/runner/__init__.py index 3bfad756..2641ee0e 100644 --- a/rdagent/scenarios/data_science/dev/runner/__init__.py +++ b/rdagent/scenarios/data_science/dev/runner/__init__.py @@ -16,9 +16,8 @@ from rdagent.components.coder.CoSTEER.task import CoSTEERTask from rdagent.core.exception import RunnerError from rdagent.core.scenario import Scenario from rdagent.log import rdagent_logger as logger -from rdagent.oai.llm_utils import md5_hash +from rdagent.oai.llm_utils import APIBackend, md5_hash from rdagent.scenarios.data_science.dev.runner.eval import DSCoSTEERCoSTEEREvaluator -from rdagent.utils import APIBackend from rdagent.utils.agent.ret import BatchEditOut from rdagent.utils.agent.tpl import T from rdagent.utils.env import DockerEnv, MLEBDockerConf diff --git a/rdagent/utils/__init__.py b/rdagent/utils/__init__.py index 14f3bea7..6b5acaee 100644 --- a/rdagent/utils/__init__.py +++ b/rdagent/utils/__init__.py @@ -6,6 +6,7 @@ it is not binding to the scenarios or framework (So it is not placed in rdagent. # TODO: merge the common utils in `rdagent.core.utils` into this folder # TODO: split the utils in this module into different modules in the future. +import hashlib import importlib import json import re @@ -15,7 +16,6 @@ from types import ModuleType from typing import Union from rdagent.oai.llm_conf import LLM_SETTINGS -from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T @@ -78,6 +78,8 @@ def filter_progress_bar(stdout: str) -> str: """ Filter out progress bars from stdout using regex. """ + from rdagent.oai.llm_utils import APIBackend # avoid circular import + # Initial progress bar regex pattern progress_bar_re = ( r"(\d+/\d+\s+[━]+\s+\d+s?\s+\d+ms/step.*?\u0008+|" @@ -147,3 +149,10 @@ def remove_path_info_from_str(base_path: Path, target_string: str) -> str: target_string = re.sub(str(base_path), "...", target_string) target_string = re.sub(str(base_path.absolute()), "...", target_string) return target_string + + +def md5_hash(input_string: str) -> str: + hash_md5 = hashlib.md5(usedforsecurity=False) + input_bytes = input_string.encode("utf-8") + hash_md5.update(input_bytes) + return hash_md5.hexdigest() diff --git a/rdagent/utils/agent/workflow.py b/rdagent/utils/agent/workflow.py index 33a4362d..864fa9c1 100644 --- a/rdagent/utils/agent/workflow.py +++ b/rdagent/utils/agent/workflow.py @@ -3,7 +3,6 @@ from typing import Any, Callable, Type, TypeVar, Union, cast from rdagent.core.exception import FormatError from rdagent.log import rdagent_logger as logger -from rdagent.oai.llm_utils import APIBackend T = TypeVar("T") @@ -39,6 +38,8 @@ def build_cls_from_json_with_retry( T An instance of the specified class type created from the response data. """ + from rdagent.oai.llm_utils import APIBackend # avoid circular import + for i in range(retry_n): # currently, it only handle exception caused by initial class resp = APIBackend().build_messages_and_create_chat_completion( diff --git a/requirements.txt b/requirements.txt index 3320c1d0..d5808ae5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,7 @@ loguru fire fuzzywuzzy openai +litellm azure.identity numpy # we use numpy as default data format. So we have to install numpy @@ -46,5 +47,5 @@ kaggle nbformat # tool +setuptools-scm seaborn -setuptools-scm \ No newline at end of file diff --git a/test/oai/test_advanced.py b/test/oai/test_advanced.py new file mode 100644 index 00000000..03330a01 --- /dev/null +++ b/test/oai/test_advanced.py @@ -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() diff --git a/test/oai/test_completion.py b/test/oai/test_completion.py index 3b60f0f9..758f55aa 100644 --- a/test/oai/test_completion.py +++ b/test/oai/test_completion.py @@ -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__": diff --git a/test/oai/test_embedding_and_similarity.py b/test/oai/test_embedding_and_similarity.py index 8e06997c..08cc51b3 100644 --- a/test/oai/test_embedding_and_similarity.py +++ b/test/oai/test_embedding_and_similarity.py @@ -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