mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 23:47:46 +00:00
feat(backend): integrate LiteLLM API Backend (#564)
* 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:
炼金术师华华
+2
-1
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
from .deprec import DeprecBackend
|
||||
from .litellm import LiteLLMAPIBackend
|
||||
|
||||
@@ -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 ...
|
||||
|
||||
@@ -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"<seed={seed}/>"
|
||||
) # 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 <start>{role/name}\n{content}<end>\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 <start>assistant<message>
|
||||
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)
|
||||
@@ -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
|
||||
|
||||
+16
-802
@@ -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"<seed={seed}/>"
|
||||
) # 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 <start>{role/name}\n{content}<end>\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 <start>assistant<message>
|
||||
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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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(
|
||||
|
||||
+2
-1
@@ -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
|
||||
@@ -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()
|
||||
+5
-143
@@ -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__":
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user