litellmlope / main.py
ka1kuk's picture
Update main.py
20a7d21 verified
raw
history blame
13.4 kB
import click
import subprocess, traceback, json
import os, sys
import random
import importlib
def run_ollama_serve():
try:
command = ["ollama", "serve"]
with open(os.devnull, "w") as devnull:
process = subprocess.Popen(command, stdout=devnull, stderr=devnull)
except Exception as e:
print(
f"""
LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve`
"""
) # noqa
def is_port_in_use(port):
import socket
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
return s.connect_ex(("localhost", port)) == 0
def run_server(
host = "0.0.0.0",
port = 8000,
api_base = None,
api_version = "2023-07-01-preview",
model = None,
alias = None,
add_key = None,
headers = None,
save = False,
debug = False,
detailed_debug = False,
temperature = 0.0,
max_tokens = 1000,
request_timeout = 10,
drop_params = True,
add_function_to_prompt = True,
config = None,
max_budget = 100,
telemetry = False,
test = False,
local = False,
num_workers = 1,
test_async = False,
num_requests = 1,
use_queue = False,
health = False,
version = False,
):
global feature_telemetry
args = locals()
if local:
from proxy_server import app, save_worker_config, usage_telemetry
else:
try:
from .litellm.proxy.proxy_server import app, save_worker_config, usage_telemetry
except ImportError as e:
if "litellm[proxy]" in str(e):
# user is missing a proxy dependency, ask them to pip install litellm[proxy]
raise e
else:
# this is just a local/relative import error, user git cloned litellm
from proxy_server import app, save_worker_config, usage_telemetry
feature_telemetry = usage_telemetry
if version == True:
pkg_version = importlib.metadata.version("litellm")
click.echo(f"\nLiteLLM: Current Version = {pkg_version}\n")
return
if model and "ollama" in model and api_base is None:
run_ollama_serve()
if test_async is True:
import requests, concurrent, time
api_base = f"http://{host}:{port}"
def _make_openai_completion():
data = {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "Write a short poem about the moon"}
],
}
response = requests.post("http://0.0.0.0:8000/queue/request", json=data)
response = response.json()
while True:
try:
url = response["url"]
polling_url = f"{api_base}{url}"
polling_response = requests.get(polling_url)
polling_response = polling_response.json()
print("\n RESPONSE FROM POLLING JOB", polling_response)
status = polling_response["status"]
if status == "finished":
llm_response = polling_response["result"]
break
print(
f"POLLING JOB{polling_url}\nSTATUS: {status}, \n Response {polling_response}"
) # noqa
time.sleep(0.5)
except Exception as e:
print("got exception in polling", e)
break
# Number of concurrent calls (you can adjust this)
concurrent_calls = num_requests
# List to store the futures of concurrent calls
futures = []
start_time = time.time()
# Make concurrent calls
with concurrent.futures.ThreadPoolExecutor(
max_workers=concurrent_calls
) as executor:
for _ in range(concurrent_calls):
futures.append(executor.submit(_make_openai_completion))
# Wait for all futures to complete
concurrent.futures.wait(futures)
# Summarize the results
successful_calls = 0
failed_calls = 0
for future in futures:
if future.done():
if future.result() is not None:
successful_calls += 1
else:
failed_calls += 1
end_time = time.time()
print(f"Elapsed Time: {end_time-start_time}")
print(f"Load test Summary:")
print(f"Total Requests: {concurrent_calls}")
print(f"Successful Calls: {successful_calls}")
print(f"Failed Calls: {failed_calls}")
return
if health != False:
import requests
print("\nLiteLLM: Health Testing models in config")
response = requests.get(url=f"http://{host}:{port}/health")
print(json.dumps(response.json(), indent=4))
return
if test != False:
request_model = model or "gpt-3.5-turbo"
click.echo(
f"\nLiteLLM: Making a test ChatCompletions request to your proxy. Model={request_model}"
)
import openai
if test == True: # flag value set
api_base = f"http://{host}:{port}"
else:
api_base = test
client = openai.OpenAI(api_key="My API Key", base_url=api_base)
response = client.chat.completions.create(
model=request_model,
messages=[
{
"role": "user",
"content": "this is a test request, write a short poem",
}
],
max_tokens=256,
)
click.echo(f"\nLiteLLM: response from proxy {response}")
print(
f"\n LiteLLM: Making a test ChatCompletions + streaming request to proxy. Model={request_model}"
)
response = client.chat.completions.create(
model=request_model,
messages=[
{
"role": "user",
"content": "this is a test request, write a short poem",
}
],
stream=True,
)
for chunk in response:
click.echo(f"LiteLLM: streaming response from proxy {chunk}")
print("\n making completion request to proxy")
response = client.completions.create(
model=request_model, prompt="this is a test request, write a short poem"
)
print(response)
return
else:
if headers:
headers = json.loads(headers)
save_worker_config(
model=model,
alias=alias,
api_base=api_base,
api_version=api_version,
debug=debug,
detailed_debug=detailed_debug,
temperature=temperature,
max_tokens=max_tokens,
request_timeout=request_timeout,
max_budget=max_budget,
telemetry=telemetry,
drop_params=drop_params,
add_function_to_prompt=add_function_to_prompt,
headers=headers,
save=save,
config=config,
use_queue=use_queue,
)
try:
import uvicorn
if os.name == "nt":
pass
else:
import gunicorn.app.base
except:
raise ImportError(
"Uvicorn, gunicorn needs to be imported. Run - `pip 'litellm[proxy]'`"
)
if config is not None:
"""
Allow user to pass in db url via config
read from there and save it to os.env['DATABASE_URL']
"""
try:
import yaml
except:
raise ImportError(
"yaml needs to be imported. Run - `pip install 'litellm[proxy]'`"
)
if os.path.exists(config):
with open(config, "r") as config_file:
config = yaml.safe_load(config_file)
general_settings = config.get("general_settings", {})
database_url = general_settings.get("database_url", None)
if database_url and database_url.startswith("os.environ/"):
original_dir = os.getcwd()
# set the working directory to where this script is
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path - for litellm local dev
import litellm
database_url = litellm.get_secret(database_url)
os.chdir(original_dir)
if database_url is not None and isinstance(database_url, str):
os.environ["DATABASE_URL"] = database_url
if os.getenv("DATABASE_URL", None) is not None:
try:
subprocess.run(["prisma"], capture_output=True)
is_prisma_runnable = True
except FileNotFoundError:
is_prisma_runnable = False
if is_prisma_runnable:
# run prisma db push, before starting server
# Save the current working directory
original_dir = os.getcwd()
# set the working directory to where this script is
abspath = os.path.abspath(__file__)
dname = os.path.dirname(abspath)
os.chdir(dname)
try:
subprocess.run(
["prisma", "db", "push", "--accept-data-loss"]
) # this looks like a weird edge case when prisma just wont start on render. we need to have the --accept-data-loss
finally:
os.chdir(original_dir)
else:
print(
f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found."
)
if port == 8000 and is_port_in_use(port):
port = random.randint(1024, 49152)
from litellm.proxy.proxy_server import app
uvicorn.run(app, host=host, port=port) # run uvicorn
# if os.name == "nt":
# else:
# import gunicorn.app.base
# # Gunicorn Application Class
# class StandaloneApplication(gunicorn.app.base.BaseApplication):
# def __init__(self, app, options=None):
# self.options = options or {} # gunicorn options
# self.application = app # FastAPI app
# super().__init__()
# _endpoint_str = (
# f"curl --location 'http://0.0.0.0:{port}/chat/completions' \\"
# )
# curl_command = (
# _endpoint_str
# + """
# --header 'Content-Type: application/json' \\
# --data ' {
# "model": "gpt-3.5-turbo",
# "messages": [
# {
# "role": "user",
# "content": "what llm are you"
# }
# ]
# }'
# \n
# """
# )
# print() # noqa
# print( # noqa
# f'\033[1;34mLiteLLM: Test your local proxy with: "litellm --test" This runs an openai.ChatCompletion request to your proxy [In a new terminal tab]\033[0m\n'
# )
# print( # noqa
# f"\033[1;34mLiteLLM: Curl Command Test for your local proxy\n {curl_command} \033[0m\n"
# )
# print(
# "\033[1;34mDocs: https://docs.litellm.ai/docs/simple_proxy\033[0m\n"
# ) # noqa
# print( # noqa
# f"\033[1;34mSee all Router/Swagger docs on http://0.0.0.0:{port} \033[0m\n"
# ) # noqa
# def load_config(self):
# # note: This Loads the gunicorn config - has nothing to do with LiteLLM Proxy config
# config = {
# key: value
# for key, value in self.options.items()
# if key in self.cfg.settings and value is not None
# }
# for key, value in config.items():
# self.cfg.set(key.lower(), value)
# def load(self):
# # gunicorn app function
# return self.application
# gunicorn_options = {
# "bind": f"{host}:{port}",
# "workers": num_workers, # default is 1
# "worker_class": "uvicorn.workers.UvicornWorker",
# "preload": True, # Add the preload flag,
# "accesslog": "-", # Log to stdout
# "access_log_format": '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s',
# }
# StandaloneApplication(
# app=app, options=gunicorn_options
# ).run() # Run gunicorn
if __name__ == "__main__":
run_server()