Skip to content

Commit

Permalink
lint + formating with black (#158)
Browse files Browse the repository at this point in the history 
* lint + formating with black

* add black as pre commit
  • Loading branch information
@cotran2
cotran2 authored Oct 9, 2024
1 parent 498e7f9 commit 5c4a6bc
Show file tree
Hide file tree
Showing 22 changed files with 582 additions and 296 deletions.
5 changes: 5 additions & 0 deletions .pre-commit-config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,3 +25,8 @@ repos:
# --lib is to only test the library, since when integration tests are made,
# they will be in a seperate tests directory
entry: bash -c "cd arch && cargo test -p intelligent-prompt-gateway --lib"
- repo: https://github.com/psf/black
rev: 23.1.0
hooks:
- id: black
language_version: python3
81 changes: 62 additions & 19 deletions arch/tools/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,26 +16,41 @@
/_/ \_\|_| \___||_| |_|
"""


@click.group(invoke_without_command=True)
@click.pass_context
def main(ctx):
if ctx.invoked_subcommand is None:
click.echo( """Arch (The Intelligent Prompt Gateway) CLI""")
click.echo("""Arch (The Intelligent Prompt Gateway) CLI""")
click.echo(logo)
click.echo(ctx.get_help())


# Command to build archgw and model_server Docker images
ARCHGW_DOCKERFILE = "./arch/Dockerfile"
MODEL_SERVER_BUILD_FILE = "./model_server/pyproject.toml"


@click.command()
def build():
"""Build Arch from source. Must be in root of cloned repo."""
# Check if /arch/Dockerfile exists
if os.path.exists(ARCHGW_DOCKERFILE):
click.echo("Building archgw image...")
try:
subprocess.run(["docker", "build", "-f", ARCHGW_DOCKERFILE, "-t", "archgw:latest", "."], check=True)
subprocess.run(
[
"docker",
"build",
"-f",
ARCHGW_DOCKERFILE,
"-t",
"archgw:latest",
".",
],
check=True,
)
click.echo("archgw image built successfully.")
except subprocess.CalledProcessError as e:
click.echo(f"Error building archgw image: {e}")
Expand All @@ -51,7 +66,11 @@ def build():
if os.path.exists(MODEL_SERVER_BUILD_FILE):
click.echo("Installing model server dependencies with Poetry...")
try:
subprocess.run(["poetry", "install", "--no-cache"], cwd=os.path.dirname(MODEL_SERVER_BUILD_FILE), check=True)
subprocess.run(
["poetry", "install", "--no-cache"],
cwd=os.path.dirname(MODEL_SERVER_BUILD_FILE),
check=True,
)
click.echo("Model server dependencies installed successfully.")
except subprocess.CalledProcessError as e:
click.echo(f"Error installing model server dependencies: {e}")
Expand All @@ -60,9 +79,12 @@ def build():
click.echo(f"Error: pyproject.toml not found in {MODEL_SERVER_BUILD_FILE}")
sys.exit(1)


@click.command()
@click.argument('file', required=False) # Optional file argument
@click.option('-path', default='.', help='Path to the directory containing arch_config.yml')
@click.argument("file", required=False) # Optional file argument
@click.option(
"-path", default=".", help="Path to the directory containing arch_config.yml"
)
def up(file, path):
"""Starts Arch."""
if file:
Expand All @@ -78,10 +100,15 @@ def up(file, path):
return

print(f"Validating {arch_config_file}")
arch_schema_config = pkg_resources.resource_filename(__name__, "config/arch_config_schema.yaml")
arch_schema_config = pkg_resources.resource_filename(
__name__, "config/arch_config_schema.yaml"
)

try:
config_generator.validate_prompt_config(arch_config_file=arch_config_file, arch_config_schema_file=arch_schema_config)
config_generator.validate_prompt_config(
arch_config_file=arch_config_file,
arch_config_schema_file=arch_schema_config,
)
except Exception as e:
print("Exiting archgw up")
sys.exit(1)
Expand All @@ -91,52 +118,67 @@ def up(file, path):
# Set the ARCH_CONFIG_FILE environment variable
env_stage = {}
env = os.environ.copy()
#check if access_keys are preesnt in the config file
# check if access_keys are preesnt in the config file
access_keys = get_llm_provider_access_keys(arch_config_file=arch_config_file)
if access_keys:
if file:
app_env_file = os.path.join(os.path.dirname(os.path.abspath(file)), ".env") #check the .env file in the path
app_env_file = os.path.join(
os.path.dirname(os.path.abspath(file)), ".env"
) # check the .env file in the path
else:
app_env_file = os.path.abspath(os.path.join(path, ".env"))

if not os.path.exists(app_env_file): #check to see if the environment variables in the current environment or not
if not os.path.exists(
app_env_file
): # check to see if the environment variables in the current environment or not
for access_key in access_keys:
if env.get(access_key) is None:
print (f"Access Key: {access_key} not found. Exiting Start")
print(f"Access Key: {access_key} not found. Exiting Start")
sys.exit(1)
else:
env_stage[access_key] = env.get(access_key)
else: #.env file exists, use that to send parameters to Arch
else: # .env file exists, use that to send parameters to Arch
env_file_dict = load_env_file_to_dict(app_env_file)
for access_key in access_keys:
if env_file_dict.get(access_key) is None:
print (f"Access Key: {access_key} not found. Exiting Start")
print(f"Access Key: {access_key} not found. Exiting Start")
sys.exit(1)
else:
env_stage[access_key] = env_file_dict[access_key]

with open(pkg_resources.resource_filename(__name__, "config/stage.env"), 'w') as file:
with open(
pkg_resources.resource_filename(__name__, "config/stage.env"), "w"
) as file:
for key, value in env_stage.items():
file.write(f"{key}={value}\n")

env.update(env_stage)
env['ARCH_CONFIG_FILE'] = arch_config_file
env["ARCH_CONFIG_FILE"] = arch_config_file

start_arch_modelserver()
start_arch(arch_config_file, env)


@click.command()
def down():
"""Stops Arch."""
stop_arch_modelserver()
stop_arch()


@click.command()
@click.option('-f', '--file', type=click.Path(exists=True), required=True, help="Path to the Python file")
@click.option(
"-f",
"--file",
type=click.Path(exists=True),
required=True,
help="Path to the Python file",
)
def generate_prompt_targets(file):
"""Generats prompt_targets from python methods.
Note: This works for simple data types like ['int', 'float', 'bool', 'str', 'list', 'tuple', 'set', 'dict']:
If you have a complex pydantic data type, you will have to flatten those manually until we add support for it."""
Note: This works for simple data types like ['int', 'float', 'bool', 'str', 'list', 'tuple', 'set', 'dict']:
If you have a complex pydantic data type, you will have to flatten those manually until we add support for it.
"""

print(f"Processing file: {file}")
if not file.endswith(".py"):
Expand All @@ -145,10 +187,11 @@ def generate_prompt_targets(file):

targets.generate_prompt_targets(file)


main.add_command(up)
main.add_command(down)
main.add_command(build)
main.add_command(generate_prompt_targets)

if __name__ == '__main__':
if __name__ == "__main__":
main()
67 changes: 39 additions & 28 deletions arch/tools/config_generator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,36 +3,44 @@
import yaml
from jsonschema import validate

ENVOY_CONFIG_TEMPLATE_FILE = os.getenv('ENVOY_CONFIG_TEMPLATE_FILE', 'envoy.template.yaml')
ARCH_CONFIG_FILE = os.getenv('ARCH_CONFIG_FILE', '/config/arch_config.yaml')
ENVOY_CONFIG_FILE_RENDERED = os.getenv('ENVOY_CONFIG_FILE_RENDERED', '/etc/envoy/envoy.yaml')
ARCH_CONFIG_SCHEMA_FILE = os.getenv('ARCH_CONFIG_SCHEMA_FILE', 'arch_config_schema.yaml')

def add_secret_key_to_llm_providers(config_yaml) :
ENVOY_CONFIG_TEMPLATE_FILE = os.getenv(
"ENVOY_CONFIG_TEMPLATE_FILE", "envoy.template.yaml"
)
ARCH_CONFIG_FILE = os.getenv("ARCH_CONFIG_FILE", "/config/arch_config.yaml")
ENVOY_CONFIG_FILE_RENDERED = os.getenv(
"ENVOY_CONFIG_FILE_RENDERED", "/etc/envoy/envoy.yaml"
)
ARCH_CONFIG_SCHEMA_FILE = os.getenv(
"ARCH_CONFIG_SCHEMA_FILE", "arch_config_schema.yaml"
)


def add_secret_key_to_llm_providers(config_yaml):
llm_providers = []
for llm_provider in config_yaml.get("llm_providers", []):
access_key_env_var = llm_provider.get('access_key', False)
access_key_env_var = llm_provider.get("access_key", False)
access_key_value = os.getenv(access_key_env_var, False)
if access_key_env_var and access_key_value:
llm_provider['access_key'] = access_key_value
llm_provider["access_key"] = access_key_value
llm_providers.append(llm_provider)
config_yaml["llm_providers"] = llm_providers
return config_yaml


def validate_and_render_schema():
env = Environment(loader=FileSystemLoader('./'))
template = env.get_template('envoy.template.yaml')
env = Environment(loader=FileSystemLoader("./"))
template = env.get_template("envoy.template.yaml")

try:
validate_prompt_config(ARCH_CONFIG_FILE, ARCH_CONFIG_SCHEMA_FILE)
except Exception as e:
print(e)
exit(1) # validate_prompt_config failed. Exit
exit(1) # validate_prompt_config failed. Exit

with open(ARCH_CONFIG_FILE, 'r') as file:
with open(ARCH_CONFIG_FILE, "r") as file:
arch_config = file.read()

with open(ARCH_CONFIG_SCHEMA_FILE, 'r') as file:
with open(ARCH_CONFIG_SCHEMA_FILE, "r") as file:
arch_config_schema = file.read()

config_yaml = yaml.safe_load(arch_config)
Expand All @@ -44,7 +52,7 @@ def validate_and_render_schema():
if name not in inferred_clusters:
inferred_clusters[name] = {
"name": name,
"port": 80, # default port
"port": 80, # default port
}

print(inferred_clusters)
Expand All @@ -55,14 +63,13 @@ def validate_and_render_schema():
if name in inferred_clusters:
print("updating cluster", endpoint_details)
inferred_clusters[name].update(endpoint_details)
endpoint = inferred_clusters[name]['endpoint']
if len(endpoint.split(':')) > 1:
inferred_clusters[name]['endpoint'] = endpoint.split(':')[0]
inferred_clusters[name]['port'] = int(endpoint.split(':')[1])
endpoint = inferred_clusters[name]["endpoint"]
if len(endpoint.split(":")) > 1:
inferred_clusters[name]["endpoint"] = endpoint.split(":")[0]
inferred_clusters[name]["port"] = int(endpoint.split(":")[1])
else:
inferred_clusters[name] = endpoint_details


print("updated clusters", inferred_clusters)

config_yaml = add_secret_key_to_llm_providers(config_yaml)
Expand All @@ -71,23 +78,24 @@ def validate_and_render_schema():
arch_config_string = yaml.dump(config_yaml)

data = {
'arch_config': arch_config_string,
'arch_clusters': inferred_clusters,
'arch_llm_providers': arch_llm_providers,
'arch_tracing': arch_tracing
"arch_config": arch_config_string,
"arch_clusters": inferred_clusters,
"arch_llm_providers": arch_llm_providers,
"arch_tracing": arch_tracing,
}

rendered = template.render(data)
print(rendered)
print(ENVOY_CONFIG_FILE_RENDERED)
with open(ENVOY_CONFIG_FILE_RENDERED, 'w') as file:
with open(ENVOY_CONFIG_FILE_RENDERED, "w") as file:
file.write(rendered)


def validate_prompt_config(arch_config_file, arch_config_schema_file):
with open(arch_config_file, 'r') as file:
with open(arch_config_file, "r") as file:
arch_config = file.read()

with open(arch_config_schema_file, 'r') as file:
with open(arch_config_schema_file, "r") as file:
arch_config_schema = file.read()

config_yaml = yaml.safe_load(arch_config)
Expand All @@ -96,8 +104,11 @@ def validate_prompt_config(arch_config_file, arch_config_schema_file):
try:
validate(config_yaml, config_schema_yaml)
except Exception as e:
print(f"Error validating arch_config file: {arch_config_file}, error: {e.message}")
print(
f"Error validating arch_config file: {arch_config_file}, error: {e.message}"
)
raise e

if __name__ == '__main__':

if __name__ == "__main__":
validate_and_render_schema()
Loading

0 comments on commit 5c4a6bc

Please sign in to comment.