aiben / openai_server /openai_client.py
abugaber's picture
Upload folder using huggingface_hub
3943768 verified
import ast
import contextlib
import gc
import os
import shutil
import tempfile
import threading
import traceback
import time
import base64
import mimetypes
import uuid
from enum import Enum
from pathlib import Path
from collections import defaultdict
import numpy as np
from pydantic import BaseModel
from .chat_history_render import chat_to_pretty_markdown
# control convert_to_pdf as expensive use of cores
num_convert_threads = max(min(10, os.cpu_count() or 1), 1)
convert_sem = threading.Semaphore(num_convert_threads)
class MyReturnType(BaseModel):
class Config:
extra = "allow"
# Local copy of minimal version from h2oGPT server
class LangChainAction(Enum):
"""LangChain action"""
QUERY = "Query"
SUMMARIZE_MAP = "Summarize"
EXTRACT = "Extract"
def get_files_from_ids(usage=None, client=None, file_ids=None, work_dir=None):
if usage is None and file_ids:
pass
elif hasattr(usage, "file_ids"):
file_ids = usage.file_ids
else:
return []
response_dict = {
file_id: dict(client.files.retrieve(file_id)) for file_id in file_ids
}
# sort file_ids by server ctime, so first is newest
file_ids = list(
reversed(sorted(file_ids, key=lambda x: response_dict[x]["created_at"]))
)
if work_dir is None:
temp_dir = tempfile.mkdtemp()
if os.path.exists(temp_dir):
shutil.rmtree(temp_dir)
os.makedirs(temp_dir, exist_ok=True)
work_dir = temp_dir
files = []
for file_id in file_ids:
new_filename = os.path.join(
work_dir, os.path.basename(response_dict[file_id]["filename"])
)
if os.path.exists(new_filename):
# FIXME: small chance different with same name
pass
else:
content = client.files.content(file_id).content
with open(new_filename, "wb") as f:
f.write(content)
files.append(new_filename)
return files
def file_to_base64(file_path, file_path_to_use=None):
# Detect the file's MIME type
mime_type, _ = mimetypes.guess_type(file_path)
if not mime_type:
mime_type = "unknown"
# Read the file and encode it in base64
with open(file_path, "rb") as file:
encoded_file = base64.b64encode(file.read()).decode("utf-8")
# Construct the data URL
data_url = f"data:{mime_type};base64,{encoded_file}"
if file_path_to_use is None:
file_path_to_use = file_path
return {file_path_to_use: data_url}
def clean_text_string(input_string):
lines = input_string.split("\n")
cleaned_lines = [
line for line in lines if line and line.strip() and line.strip() != "-"
]
return "\n".join(cleaned_lines)
def local_convert_to_pdf(convert_to_pdf, x, files_already_pdf, *args, **kwargs):
if x in files_already_pdf:
return x
try:
with convert_sem:
return convert_to_pdf(x, *args, **kwargs)
except Exception as e1:
print(f"Error converting {x} to PDF: {e1}")
return None
def group_files_by_base_name(file_names):
grouped_files = defaultdict(list)
for file in file_names:
base_name = Path(file).stem
grouped_files[base_name].append(file)
return grouped_files
def group_and_prioritize_files(file_names):
grouped_files = group_files_by_base_name(file_names)
prioritized_files = []
for base_name, files in grouped_files.items():
preferred_file = select_preferred_file(files)
# Put the preferred file first, then add all other files
prioritized_group = [preferred_file] + [f for f in files if f != preferred_file]
prioritized_files.extend(prioritized_group)
return prioritized_files
def select_preferred_file(files):
# Preference order: PDF, PNG, SVG, others
for ext in [".pdf", ".png", ".svg"]:
for file in files:
if file.lower().endswith(ext):
return file
# If no preferred format found, return the first file
return files[0]
def get_pdf_files(file_names, convert_to_pdf):
# Group files by base name
prioritized_files = group_and_prioritize_files(file_names)
# Filter out binary files with text-like extensions
# e.g. .txt but giant binary, then libreoffice will take too long to convert
selected_files = [
file
for file in prioritized_files
if not (is_binary(file) and Path(file).suffix.lower() in TEXT_EXTENSIONS)
]
# Filter out audio files
audio_exts = [
".mp3",
".wav",
".flac",
".ogg",
".m4a",
".aac",
".wma",
".aiff",
".mp4",
".mpeg",
".mpg",
".mpga",
".webm",
]
exclude_exts = audio_exts + [".zip", ".tar", ".gz", ".bz2", ".xz", ".7z", ".rar"]
selected_files = [
file
for file in selected_files
if not any(file.lower().endswith(ext) for ext in exclude_exts)
]
# 5MB limit to avoid long conversions
selected_files = [
f for f in selected_files if os.path.getsize(f) <= 5 * 1024 * 1024
]
# Convert files to PDF
pdf_file_names = []
pdf_base_names = set()
errors = []
def process_file(file, pdf_base_names, convert_to_pdf):
file_path = Path(file)
base_name = file_path.stem
ext_name = file_path.suffix.lower()
if file_path.suffix.lower() == ".pdf":
pdf_base_names.add(base_name)
return str(file_path), base_name, None
if base_name in pdf_base_names:
new_pdf_name = f"{base_name}{ext_name}.pdf"
else:
new_pdf_name = f"{base_name}.pdf"
pdf_base_names.add(base_name)
new_pdf_path = file_path.with_name(new_pdf_name)
new_dir = os.path.dirname(new_pdf_path)
temp_file = file_path.with_suffix(f".{uuid.uuid4()}{file_path.suffix}")
try:
if not os.path.exists(new_dir):
os.makedirs(new_dir, exist_ok=True)
shutil.copy(file_path, temp_file)
converted_pdf = local_convert_to_pdf(
convert_to_pdf,
temp_file,
set(),
correct_image=False,
)
if converted_pdf:
shutil.move(converted_pdf, str(new_pdf_path))
return str(new_pdf_path), base_name, None
except Exception as e:
return None, None, f"Error converting {file} to PDF: {e}"
finally:
if os.path.isfile(temp_file):
try:
os.remove(temp_file)
except Exception as e:
print(f"Error removing temp file {temp_file}: {e}")
return None, None, f"Failed to process {file}"
from concurrent.futures import ThreadPoolExecutor, as_completed, TimeoutError
# Set timeouts
timeout_seconds = 3 * 60
timeout_seconds_per_file = 30
t0 = time.time()
with ThreadPoolExecutor() as executor:
future_to_file = {
executor.submit(process_file, file, pdf_base_names, convert_to_pdf): file
for file in selected_files
}
while future_to_file:
# Re-check remaining time for the overall timeout
remaining_time = timeout_seconds - (time.time() - t0)
if remaining_time <= 0:
errors.append(f"Overall timeout of {timeout_seconds} seconds reached.")
break
# Check the futures as they complete or timeout
try:
for future in as_completed(future_to_file, timeout=remaining_time):
file = future_to_file[future] # Get the corresponding file
try:
# Wait for the result of each future with a per-file timeout
result, base_name, error = future.result(
timeout=timeout_seconds_per_file
)
# Only pop the future after successful completion
future_to_file.pop(future)
if error:
errors.append(f"Error processing {file}: {error}")
elif result:
pdf_file_names.append(result)
pdf_base_names.add(base_name)
except TimeoutError:
errors.append(
f"Timeout error processing {file}: operation took longer than {timeout_seconds_per_file} seconds"
)
except Exception as exc:
errors.append(f"Unexpected error processing {file}: {exc}")
# We still want to pop the future on failure
future_to_file.pop(future)
except TimeoutError:
errors.append(
f"Timeout error processing {file}: operation took longer than {timeout_seconds_per_file} seconds"
)
except Exception as exc:
errors.append(f"Unexpected error processing {file}: {exc}")
# If all futures are processed or timeout reached, break
if time.time() - t0 > timeout_seconds:
errors.append(
f"Overall timeout of {timeout_seconds} seconds reached. {len(future_to_file)} files remaining."
)
break
if errors:
print(errors)
return pdf_file_names
def completion_with_backoff(
get_client,
model,
messages,
stream_output,
hyper_kwargs,
extra_body,
timeout,
time_to_first_token_max,
ReturnType=None,
use_agent=False,
add_extra_endofturn=False,
max_chars_per_turn=1024 * 4,
):
t0_outer = time.time()
ntrials = 3
trial = 0
while True:
t0 = time.time()
responses = None
client = None
time_to_first_token = None
response = ""
usage = None
file_names = []
try:
client = get_client()
responses = client.chat.completions.create(
model=model,
messages=messages,
stream=stream_output,
**hyper_kwargs,
extra_body=extra_body,
timeout=timeout,
)
if not stream_output:
usage = responses.usage
if responses.choices:
response = responses.choices[-1].message.content
else:
response = ""
yield ReturnType(reply=response)
time_to_first_token = time.time() - t0
else:
response = ""
usages = []
for chunk in responses:
if chunk.usage is not None:
usages.append(chunk.usage)
if chunk.choices:
delta = chunk.choices[0].delta.content
if delta:
response += delta
# ensure if h2oGPTe wants full or delta, looks like delta from gradio code, except at very end?
yield ReturnType(reply=delta)
if time_to_first_token is None:
time_to_first_token = time.time() - t0
if use_agent and add_extra_endofturn:
splits = response.split("ENDOFTURN")
if splits and len(splits[-1]) > max_chars_per_turn:
# force end of turn for UI purposes
delta = "\n\nENDOFTURN\n\n"
response += delta
yield ReturnType(reply=delta)
time.sleep(0.005)
if (
time_to_first_token is None
and time.time() - t0 > time_to_first_token_max
):
raise TimeoutError(
f"LLM {model} timed out without any response after {time_to_first_token_max} seconds, for total {time.time() - t0_outer} seconds.."
)
if time.time() - t0 > timeout:
print("Timed out, but had response: %s" % response, flush=True)
raise TimeoutError(
f"LLM {model} timed out after {time.time() - t0} seconds, for total {time.time() - t0_outer} seconds."
)
assert len(usages) == 1, 'Missing usage"'
usage = usages[0]
# Get files
file_names = (
get_files_from_ids(usage=usage, client=client) if use_agent else []
)
return (
response,
usage,
file_names,
time_to_first_token or (time.time() - t0),
None,
None,
)
except (GeneratorExit, StopIteration):
# caller is trying to cancel
print(f"Caller initiated GeneratorExit in completion_with_backoff.")
raise
except Exception as e:
error_ex = traceback.format_exc()
error_e = str(e)
if trial == ntrials - 1 or "Output contains sensitive information" in str(
e
):
print(
f"{model} hit final error in completion_with_backoff: {e}. Retrying trial {trial}."
)
if os.getenv("HARD_ASSERTS"):
raise
# Note: response can be partial
return (
response,
usage,
file_names,
time_to_first_token or (time.time() - t0),
error_e,
error_ex,
)
else:
if trial == 0:
time.sleep(1)
elif trial == 1:
time.sleep(5)
else:
time.sleep(20)
trial += 1
print(
f"{model} hit error in completion_with_backoff: {e}. Retrying trial {trial}."
)
finally:
if responses is not None:
try:
responses.close()
del responses
gc.collect()
except Exception as e:
print("Failed to close OpenAI response: %s" % str(e), flush=True)
if client is not None:
try:
client.close()
del client
gc.collect()
except Exception as e:
print("Failed to close OpenAI client: %s" % str(e), flush=True)
def run_openai_client(
get_client=None,
ReturnType=None,
convert_to_pdf=None,
use_agent=False,
agent_accuracy="standard",
autogen_max_turns=80,
agent_chat_history=[],
agent_files=[],
agent_venv_dir=None,
agent_work_dir=None,
base64_encode_agent_files=True,
cute=False,
time_to_first_token_max=None,
**query_kwargs,
):
"""
Bsed upon test in h2oGPT OSS:
https://github.com/h2oai/h2ogpt/blob/ee3995865c85bf74f3644a4ebd007971c809de11/openai_server/test_openai_server.py#L189-L320
"""
if ReturnType is None:
ReturnType = MyReturnType
# pick correct prompt
# langchain_mode = query_kwargs.get("langchain_mode", "LLM")
langchain_action = query_kwargs.get("langchain_action", "Query")
# prompt will be "" for langchain_action = 'Summarize'
prompt = query_kwargs["instruction"]
model = query_kwargs["visible_models"]
stream_output = query_kwargs["stream_output"]
max_time = query_kwargs["max_time"]
time_to_first_token_max = time_to_first_token_max or max_time
text_context_list = query_kwargs["text_context_list"]
chat_conversation = query_kwargs["chat_conversation"]
image_files = query_kwargs["image_file"]
system_message = query_kwargs["system_prompt"]
from h2ogpte_core.backend_utils import structure_to_messages
if use_agent:
chat_conversation = None # don't include high-level history yet
file_ids = []
if agent_files:
client = get_client()
for file_path in agent_files:
with open(file_path, "rb") as file:
ret = client.files.create(
file=file,
purpose="assistants",
)
file_id = ret.id
file_ids.append(file_id)
assert ret.bytes > 0
extra_body = dict(
use_agent=use_agent,
agent_type="auto",
agent_accuracy=agent_accuracy,
autogen_stop_docker_executor=False,
autogen_run_code_in_docker=False,
autogen_max_consecutive_auto_reply=80,
autogen_max_turns=autogen_max_turns,
autogen_timeout=240,
autogen_cache_seed=None,
work_dir=agent_work_dir,
venv_dir=agent_venv_dir,
agent_verbose=True,
text_context_list=text_context_list,
agent_chat_history=agent_chat_history,
agent_files=file_ids,
client_metadata=query_kwargs.get("client_metadata", ""),
)
# agent needs room, else keep hitting continue
hyper_kwargs = dict(
temperature=query_kwargs["temperature"],
seed=query_kwargs["seed"],
max_tokens=8192 if "claude-3-5-sonnet" in model else 4096,
)
else:
extra_body = query_kwargs.copy()
from h2ogpte_core.src.evaluate_params import eval_func_param_names
extra_body = {k: v for k, v in extra_body.items() if k in eval_func_param_names}
hyper_kwargs = dict(
temperature=query_kwargs["temperature"],
top_p=query_kwargs["top_p"],
seed=query_kwargs["seed"],
max_tokens=query_kwargs["max_new_tokens"],
)
extra_body = {k: v for k, v in extra_body.items() if k not in hyper_kwargs}
# remove things that go through OpenAI API messages
keys_in_api = [
"visible_models",
"image_file",
"chat_conversation",
"system_prompt",
"instruction",
"stream_output",
]
for key in keys_in_api:
extra_body.pop(key, None)
# translate
if "response_format" in extra_body:
extra_body["response_format"] = dict(type=extra_body["response_format"])
time_to_first_token = None
t0 = time.time()
messages = structure_to_messages(
prompt, system_message, chat_conversation, image_files
)
timeout = 5 * max_time if use_agent else max_time
(
response,
usage,
file_names,
time_to_first_token,
error_e,
error_ex,
) = yield from completion_with_backoff(
get_client,
model,
messages,
stream_output,
hyper_kwargs,
extra_body,
timeout,
time_to_first_token_max,
ReturnType=ReturnType,
use_agent=use_agent,
)
# in case streaming had deletions not yet accounted for, recover at least final answer,
# e.g. for JSON {} then {}{"response": "yes"}
if hasattr(usage, "response"):
response = usage.response
tf = time.time()
# See if we can make text in case of no extension
for file_i, file in enumerate(file_names):
file_path = Path(file)
suffix = file_path.suffix.lower()
# If no suffix and not binary, rename to ".txt"
if not suffix and not is_binary(file):
new_file = file_path.with_suffix(".txt")
try:
file_path.rename(new_file) # Rename the file, overwriting if necessary
file_names[file_i] = str(new_file)
except OSError as e:
print(f"Error renaming {file} to {new_file}: {e}")
if base64_encode_agent_files:
files = [file_to_base64(x) for x in file_names]
files = update_file_names(files)
else:
files = file_names
# Process files and get PDF file names
pdf_file_names = get_pdf_files(files, convert_to_pdf)
if base64_encode_agent_files:
files_pdf = [file_to_base64(x, y) for x, y in zip(pdf_file_names, file_names)]
files_pdf = update_file_names(files_pdf)
# clean-up
[remove(x) for x in file_names if os.path.isfile(x)]
[remove(x) for x in pdf_file_names if os.path.isfile(x)]
else:
files_pdf = pdf_file_names
# Get usage
input_tokens = usage.prompt_tokens if usage else 0
output_tokens = usage.completion_tokens if usage else 0
if hasattr(usage, "cost") and usage.cost:
usage_no_caching = usage.cost["usage_excluding_cached_inference"]
assert model in usage_no_caching, "Missing model %s in %s" % (
model,
usage_no_caching,
)
input_tokens += usage_no_caching[model]["prompt_tokens"]
output_tokens += usage_no_caching[model]["completion_tokens"]
# Get internal chat history
chat_history = (
usage.chat_history
if hasattr(usage, "chat_history")
else [{"role": "assistant", "content": response}]
)
chat_history_md = (
chat_to_pretty_markdown(chat_history, cute=cute) if chat_history else ""
)
agent_work_dir = usage.agent_work_dir if hasattr(usage, "agent_work_dir") else None
agent_venv_dir = usage.agent_venv_dir if hasattr(usage, "agent_venv_dir") else None
# Get final answer
response_intermediate = response
if hasattr(usage, "summary"):
response = usage.summary
if not response:
split1 = response_intermediate.split(
"code_writer_agent(tocode_executor_agent):"
)
if split1 and split1[-1]:
split2 = split1[-1].split("code_executor_agent(tocode_writer_agent):")
if split2 and split1[0]:
response = split2[0]
response = clean_text_string(response)
if not response:
response = "The task is complete"
elif "ENDOFTURN" in response:
# show last turn as final response
split_responses = response.split("ENDOFTURN")
if len(split_responses) > 1:
response = split_responses[-1]
if not response:
response = "The task completed"
# estimate tokens per second
tokens_per_second = output_tokens / (tf - t0 + 1e-6)
t_taken_s = time.time() - t0
t_taken = "%.4f" % t_taken_s
if use_agent:
if not (response or response_intermediate or files or chat_history):
msg = f"No output from Agent with LLM {model} after {t_taken} seconds."
if error_e:
raise ValueError("Error: " + error_e + "\n" + msg)
else:
raise TimeoutError(msg)
else:
if not (response or response_intermediate):
msg = f"No response from LLM {model} after {t_taken} seconds."
if error_e:
raise ValueError("Error: " + error_e + "\n" + msg)
else:
raise TimeoutError(msg)
# extract other usages:
sources = usage.sources if hasattr(usage, "sources") else []
prompt_raw = usage.prompt_raw if hasattr(usage, "prompt_raw") else ""
save_dict = usage.save_dict if hasattr(usage, "save_dict") else {}
if not use_agent:
if not hasattr(usage, "sources"):
print("missing sources from usage: %s" % usage)
if not hasattr(usage, "prompt_raw"):
print("missing prompt_raw from usage: %s" % usage)
if not hasattr(usage, "save_dict"):
print("missing save_dict from usage: %s" % usage)
extra_dict = save_dict.get("extra_dict", {})
texts_out = [x["content"] for x in sources] if not use_agent else text_context_list
t_taken_s = time.time() - t0
t_taken = "%.4f" % t_taken_s
if langchain_action != LangChainAction.EXTRACT.value:
response = response.strip() if response else ""
response_intermediate = response_intermediate.strip()
else:
response = [r.strip() if r else "" for r in ast.literal_eval(response)]
response_intermediate = [
r.strip() if r else "" for r in ast.literal_eval(response_intermediate)
]
try:
actual_llm = save_dict["display_name"]
except Exception as e:
actual_llm = model
print(f"Unable to access save_dict to get actual_llm: {str(e)}")
reply = response_intermediate if use_agent else response
if not reply:
error_e = (
error_ex
) = f"No final response from LLM {actual_llm} after {t_taken} seconds\nError:{error_e}."
if "error" in save_dict and not prompt_raw:
msg = f"Error from LLM {actual_llm}: {save_dict['error']}"
if os.getenv("HARD_ASSERTS"):
if error_e:
raise ValueError("Error: " + error_e + "\n" + msg)
else:
raise ValueError(msg)
if not use_agent:
if not (prompt_raw or extra_dict):
msg = "LLM response failed to return final metadata."
if os.getenv("HARD_ASSERTS"):
if error_e:
raise ValueError("Error: " + error_e + "\n" + msg)
else:
raise ValueError(msg)
else:
prompt_raw = prompt
try:
input_tokens = extra_dict["num_prompt_tokens"]
output_tokens = extra_dict["ntokens"]
vision_visible_model = extra_dict.get("batch_vision_visible_model")
vision_batch_input_tokens = extra_dict.get("batch_num_prompt_tokens", 0)
vision_batch_output_tokens = extra_dict.get("batch_ntokens", 0)
tokens_per_second = np.round(extra_dict["tokens_persecond"], decimals=3)
vision_batch_tokens_per_second = extra_dict.get("batch_tokens_persecond", 0)
if vision_batch_tokens_per_second:
vision_batch_tokens_per_second = np.round(
vision_batch_tokens_per_second, decimals=3
)
except:
vision_visible_model = model
vision_batch_input_tokens = 0
vision_batch_output_tokens = 0
vision_batch_tokens_per_second = 0
if not use_agent and os.getenv("HARD_ASSERTS"):
raise
if use_agent and not response and reply:
# show streamed output then, to avoid confusion with whether had response
response = reply
if error_e or error_ex:
delta_error = f"\n\n**Partial Error:**\n\n {error_e}"
if use_agent:
yield ReturnType(reply="\nENDOFTURN\n" + delta_error)
response = delta_error
else:
yield ReturnType(reply=delta_error)
response += delta_error
# final yield
yield ReturnType(
reply=reply,
reply_final=response,
prompt_raw=prompt_raw,
actual_llm=actual_llm,
text_context_list=texts_out,
input_tokens=input_tokens,
output_tokens=output_tokens,
tokens_per_second=tokens_per_second,
time_to_first_token=time_to_first_token or (time.time() - t0),
trial=0, # Not required, OpenAI has retries
vision_visible_model=vision_visible_model,
vision_batch_input_tokens=vision_batch_input_tokens,
vision_batch_output_tokens=vision_batch_output_tokens,
vision_batch_tokens_per_second=vision_batch_tokens_per_second,
agent_work_dir=agent_work_dir,
agent_venv_dir=agent_venv_dir,
files=files,
files_pdf=files_pdf,
chat_history=chat_history,
chat_history_md=chat_history_md,
)
# List of common text file extensions
TEXT_EXTENSIONS = {
".txt",
".md",
".csv",
".html",
".xml",
".json",
".yaml",
".yml",
".log",
}
def is_binary(filename):
"""
Check if a file is binary or text using a quick check.
Args:
filename (str): The path to the file.
Returns:
bool: True if the file is binary, False otherwise.
"""
try:
with open(filename, "rb") as f:
chunk = f.read(1024) # Read the first 1KB of the file for a quick check
if b"\0" in chunk: # Null byte found, indicating binary content
return True
# Try decoding the chunk as UTF-8
try:
chunk.decode("utf-8")
except UnicodeDecodeError:
return True # Decoding failed, likely a binary file
except Exception as e:
print(f"Error reading file: {e}")
return True
return False # No null bytes and successful UTF-8 decoding, likely a text file
def update_file_names(file_list):
def process_item(item):
if isinstance(item, str):
return os.path.basename(item)
elif isinstance(item, dict):
old_key = list(item.keys())[0]
return {os.path.basename(old_key): item[old_key]}
else:
raise ValueError(f"Unsupported item type: {type(item)}")
return [process_item(item) for item in file_list]
def shutil_rmtree(*args, **kwargs):
path = args[0]
assert not os.path.samefile(
path, "/"
), "Should not be trying to remove entire root directory: %s" % str(path)
assert not os.path.samefile(
path, "./"
), "Should not be trying to remove entire local directory: %s" % str(path)
return shutil.rmtree(*args, **kwargs)
def remove(path: str):
try:
if path is not None and os.path.exists(path):
if os.path.isdir(path):
shutil_rmtree(path, ignore_errors=True)
else:
with contextlib.suppress(FileNotFoundError):
os.remove(path)
except BaseException as e:
print(f"Error removing {path}: {e}")
pass