apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-06-26 15:39:40 +02:00
Code Issues Projects Releases Packages Wiki Activity

lint + formating with black (#158)

Browse source 
* lint + formating with black

* add black as pre commit
This commit is contained in:
Co Tran 2024-10-09 11:25:07 -07:00 • committed by GitHub
parent 498e7f9724
commit 5c4a6bc8ff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 581 additions and 295 deletions
Download patch file Download diff file Expand all files Collapse all files

81
arch/tools/cli.py
View file

@ -16,18 +16,22 @@ logo = r"""
/_/ \_\|_| \___||_| |_|
"""
@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."""
@ -35,7 +39,18 @@ def build():
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}")
@ -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}")
@ -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:
@ -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)
@ -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"):
@ -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()

65
arch/tools/config_generator.py
View file

@ -3,36 +3,44 @@ from jinja2 import Environment, FileSystemLoader
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')
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) :
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)
@ -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)
@ -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)
@ -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)
@ -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()

77
arch/tools/core.py
View file

@ -5,6 +5,7 @@ import pkg_resources
import select
from utils import run_docker_compose_ps, print_service_status, check_services_state
def start_arch(arch_config_file, env, log_timeout=120):
"""
Start Docker Compose in detached mode and stream logs until services are healthy.
@ -14,22 +15,35 @@ def start_arch(arch_config_file, env, log_timeout=120):
log_timeout (int): Time in seconds to show logs before checking for healthy state.
"""
compose_file = pkg_resources.resource_filename(__name__, 'config/docker-compose.yaml')
compose_file = pkg_resources.resource_filename(
__name__, "config/docker-compose.yaml"
)
try:
# Run the Docker Compose command in detached mode (-d)
subprocess.run(
["docker", "compose", "-p", "arch", "up", "-d",],
cwd=os.path.dirname(compose_file), # Ensure the Docker command runs in the correct path
env=env, # Pass the modified environment
check=True # Raise an exception if the command fails
[
"docker",
"compose",
"-p",
"arch",
"up",
"-d",
],
cwd=os.path.dirname(
compose_file
), # Ensure the Docker command runs in the correct path
env=env, # Pass the modified environment
check=True, # Raise an exception if the command fails
)
print(f"Arch docker-compose started in detached.")
print("Monitoring `docker-compose ps` logs...")
start_time = time.time()
services_status = {}
services_running = False #assume that the services are not running at the moment
services_running = (
False # assume that the services are not running at the moment
)
while True:
current_time = time.time()
@ -40,16 +54,22 @@ def start_arch(arch_config_file, env, log_timeout=120):
print(f"Stopping log monitoring after {log_timeout} seconds.")
break
current_services_status = run_docker_compose_ps(compose_file=compose_file, env=env)
current_services_status = run_docker_compose_ps(
compose_file=compose_file, env=env
)
if not current_services_status:
print("Status for the services could not be detected. Something went wrong. Please run docker logs")
print(
"Status for the services could not be detected. Something went wrong. Please run docker logs"
)
break
if not services_status:
services_status = current_services_status #set the first time
print_service_status(services_status) #print the services status and proceed.
services_status = current_services_status # set the first time
print_service_status(
services_status
) # print the services status and proceed.
#check if anyone service is failed or exited state, if so print and break out
# check if anyone service is failed or exited state, if so print and break out
unhealthy_states = ["unhealthy", "exit", "exited", "dead", "bad"]
running_states = ["running", "up"]
@ -58,14 +78,23 @@ def start_arch(arch_config_file, env, log_timeout=120):
break
if check_services_state(current_services_status, unhealthy_states):
print("One or more Arch services are unhealthy. Please run `docker logs` for more information")
print_service_status(current_services_status) #print the services status and proceed.
print(
"One or more Arch services are unhealthy. Please run `docker logs` for more information"
)
print_service_status(
current_services_status
) # print the services status and proceed.
break
#check to see if the status of one of the services has changed from prior. Print and loop over until finish, or error
# check to see if the status of one of the services has changed from prior. Print and loop over until finish, or error
for service_name in services_status.keys():
if services_status[service_name]['State'] != current_services_status[service_name]['State']:
print("One or more Arch services have changed state. Printing current state")
if (
services_status[service_name]["State"]
!= current_services_status[service_name]["State"]
):
print(
"One or more Arch services have changed state. Printing current state"
)
print_service_status(current_services_status)
break
@ -82,7 +111,9 @@ def stop_arch():
Args:
path (str): The path where the docker-compose.yml file is located.
"""
compose_file = pkg_resources.resource_filename(__name__, 'config/docker-compose.yaml')
compose_file = pkg_resources.resource_filename(
__name__, "config/docker-compose.yaml"
)
try:
# Run `docker-compose down` to shut down all services
@ -96,6 +127,7 @@ def stop_arch():
except subprocess.CalledProcessError as e:
print(f"Failed to shut down services: {str(e)}")
def start_arch_modelserver():
"""
Start the model server. This assumes that the archgw_modelserver package is installed locally
@ -103,15 +135,14 @@ def start_arch_modelserver():
"""
try:
subprocess.run(
['archgw_modelserver', 'restart'],
check=True,
start_new_session=True
["archgw_modelserver", "restart"], check=True, start_new_session=True
)
print("Successfull run the archgw model_server")
except subprocess.CalledProcessError as e:
print (f"Failed to start model_server. Please check archgw_modelserver logs")
print(f"Failed to start model_server. Please check archgw_modelserver logs")
sys.exit(1)
def stop_arch_modelserver():
"""
Stop the model server. This assumes that the archgw_modelserver package is installed locally
@ -119,10 +150,10 @@ def stop_arch_modelserver():
"""
try:
subprocess.run(
['archgw_modelserver', 'stop'],
["archgw_modelserver", "stop"],
check=True,
)
print("Successfull stopped the archgw model_server")
except subprocess.CalledProcessError as e:
print (f"Failed to start model_server. Please check archgw_modelserver logs")
print(f"Failed to start model_server. Please check archgw_modelserver logs")
sys.exit(1)

22
arch/tools/setup.py
View file

@ -6,17 +6,29 @@ setup(
description="Python-based CLI tool to manage Arch and generate targets.",
author="Katanemo Labs, Inc.",
packages=find_packages(),
py_modules = ['cli', 'core', 'targets', 'utils', 'config_generator'],
py_modules=["cli", "core", "targets", "utils", "config_generator"],
include_package_data=True,
# Specify to include the docker-compose.yml file
package_data={
'': ['config/docker-compose.yaml', 'config/arch_config_schema.yaml', 'config/stage.env'] #Specify to include the docker-compose.yml file
"": [
"config/docker-compose.yaml",
"config/arch_config_schema.yaml",
"config/stage.env",
] # Specify to include the docker-compose.yml file
},
# Add dependencies here, e.g., 'PyYAML' for YAML processing
install_requires=['pyyaml', 'pydantic', 'click', 'jinja2','pyyaml','jsonschema', 'setuptools'],
install_requires=[
"pyyaml",
"pydantic",
"click",
"jinja2",
"pyyaml",
"jsonschema",
"setuptools",
],
entry_points={
'console_scripts': [
'archgw=cli:main',
"console_scripts": [
"archgw=cli:main",
],
},
)

175
arch/tools/targets.py
View file

@ -18,14 +18,20 @@ def detect_framework(tree: Any) -> str:
return "fastapi"
return "unknown"
def get_route_decorators(node: Any, framework: str) -> list:
"""Extract route decorators based on the framework."""
decorators = []
for decorator in node.decorator_list:
if isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Attribute):
if isinstance(decorator, ast.Call) and isinstance(
decorator.func, ast.Attribute
):
if framework == "flask" and decorator.func.attr in FLASK_ROUTE_DECORATORS:
decorators.append(decorator.func.attr)
elif framework == "fastapi" and decorator.func.attr in FASTAPI_ROUTE_DECORATORS:
elif (
framework == "fastapi"
and decorator.func.attr in FASTAPI_ROUTE_DECORATORS
):
decorators.append(decorator.func.attr)
return decorators
@ -36,6 +42,7 @@ def get_route_path(node: Any, framework: str) -> str:
if isinstance(decorator, ast.Call) and decorator.args:
return decorator.args[0].s # Assuming it's a string literal
def is_pydantic_model(annotation: ast.expr, tree: ast.AST) -> bool:
"""Check if a given type annotation is a Pydantic model."""
# We walk through the AST to find class definitions and check if they inherit from Pydantic's BaseModel
@ -47,6 +54,7 @@ def is_pydantic_model(annotation: ast.expr, tree: ast.AST) -> bool:
return True
return False
def get_pydantic_model_fields(model_name: str, tree: ast.AST) -> list:
"""Extract fields from a Pydantic model, handling list, tuple, set, dict types, and direct default values."""
fields = []
@ -62,15 +70,26 @@ def get_pydantic_model_fields(model_name: str, tree: ast.AST) -> list:
required = True # Assume the field is required initially
# Check if the field uses Field() with required status and description
if stmt.value and isinstance(stmt.value, ast.Call) and isinstance(stmt.value.func, ast.Name) and stmt.value.func.id == 'Field':
if (
stmt.value
and isinstance(stmt.value, ast.Call)
and isinstance(stmt.value.func, ast.Name)
and stmt.value.func.id == "Field"
):
# Extract the description argument inside the Field call
for keyword in stmt.value.keywords:
if keyword.arg == 'description' and isinstance(keyword.value, ast.Str):
if keyword.arg == "description" and isinstance(
keyword.value, ast.Str
):
description = keyword.value.s
if keyword.arg == 'default':
if keyword.arg == "default":
default_value = keyword.value
# If Ellipsis (...) is used, it means the field is required
if stmt.value.args and isinstance(stmt.value.args[0], ast.Constant) and stmt.value.args[0].value is Ellipsis:
if (
stmt.value.args
and isinstance(stmt.value.args[0], ast.Constant)
and stmt.value.args[0].value is Ellipsis
):
required = True
else:
required = False
@ -80,19 +99,28 @@ def get_pydantic_model_fields(model_name: str, tree: ast.AST) -> list:
if isinstance(stmt.value, ast.Constant):
# Set the default value from the assignment (e.g., name: str = "John Doe")
default_value = stmt.value.value
required = False # Not required since it has a default value
required = (
False # Not required since it has a default value
)
# Always extract the field type, even if there's a default value
if isinstance(stmt.annotation, ast.Subscript):
# Get the base type (list, tuple, set, dict)
base_type = stmt.annotation.value.id if isinstance(stmt.annotation.value, ast.Name) else "Unknown"
base_type = (
stmt.annotation.value.id
if isinstance(stmt.annotation.value, ast.Name)
else "Unknown"
)
# Handle only list, tuple, set, dict and ignore the inner types
if base_type.lower() in ['list', 'tuple', 'set', 'dict']:
if base_type.lower() in ["list", "tuple", "set", "dict"]:
field_type = base_type.lower()
# Handle the ellipsis '...' for required fields if no Field() call
elif isinstance(stmt.value, ast.Constant) and stmt.value.value is Ellipsis:
elif (
isinstance(stmt.value, ast.Constant)
and stmt.value.value is Ellipsis
):
required = True
# Handle simple types like str, int, etc.
@ -100,16 +128,17 @@ def get_pydantic_model_fields(model_name: str, tree: ast.AST) -> list:
field_type = stmt.annotation.id
field_info = {
"name": stmt.target.id,
"type": field_type, # Always set the field type
"description": description,
"default": default_value, # Handle direct default values
"required": required
"name": stmt.target.id,
"type": field_type, # Always set the field type
"description": description,
"default": default_value, # Handle direct default values
"required": required,
}
fields.append(field_info)
return fields
def get_function_parameters(node: ast.FunctionDef, tree: ast.AST) -> list:
"""Extract the parameters and their types from the function definition."""
parameters = []
@ -119,40 +148,68 @@ def get_function_parameters(node: ast.FunctionDef, tree: ast.AST) -> list:
arg_descriptions = extract_arg_descriptions_from_docstring(docstring)
# Extract default values
defaults = [None] * (len(node.args.args) - len(node.args.defaults)) + node.args.defaults # Align defaults with args
defaults = [None] * (
len(node.args.args) - len(node.args.defaults)
) + node.args.defaults # Align defaults with args
for arg, default in zip(node.args.args, defaults):
if arg.arg != "self": # Skip 'self' or 'cls' in class methods
param_info = {"name": arg.arg, "description": arg_descriptions.get(arg.arg, "[ADD DESCRIPTION]")}
param_info = {
"name": arg.arg,
"description": arg_descriptions.get(arg.arg, "[ADD DESCRIPTION]"),
}
# Handle Pydantic model types
if hasattr(arg, 'annotation') and is_pydantic_model(arg.annotation, tree):
if hasattr(arg, "annotation") and is_pydantic_model(arg.annotation, tree):
# Extract and flatten Pydantic model fields
pydantic_fields = get_pydantic_model_fields(arg.annotation.id, tree)
parameters.extend(pydantic_fields) # Flatten the model fields into the parameters list
parameters.extend(
pydantic_fields
) # Flatten the model fields into the parameters list
continue # Skip adding the current param_info for the model since we expand the fields
# Handle standard Python types (int, float, str, etc.)
elif hasattr(arg, 'annotation') and isinstance(arg.annotation, ast.Name):
if arg.annotation.id in ['int', 'float', 'bool', 'str', 'list', 'tuple', 'set', 'dict']:
elif hasattr(arg, "annotation") and isinstance(arg.annotation, ast.Name):
if arg.annotation.id in [
"int",
"float",
"bool",
"str",
"list",
"tuple",
"set",
"dict",
]:
param_info["type"] = arg.annotation.id
else:
param_info["type"] = "[UNKNOWN - PLEASE FIX]"
# Handle generic subscript types (e.g., Optional, List[Type], etc.)
elif hasattr(arg, 'annotation') and isinstance(arg.annotation, ast.Subscript):
if isinstance(arg.annotation.value, ast.Name) and arg.annotation.value.id in ['list', 'tuple', 'set', 'dict']:
param_info["type"] = f"{arg.annotation.value.id}" # e.g., "List", "Tuple", etc.
elif hasattr(arg, "annotation") and isinstance(
arg.annotation, ast.Subscript
):
if isinstance(
arg.annotation.value, ast.Name
) and arg.annotation.value.id in ["list", "tuple", "set", "dict"]:
param_info[
"type"
] = f"{arg.annotation.value.id}" # e.g., "List", "Tuple", etc.
else:
param_info["type"] = "[UNKNOWN - PLEASE FIX]"
# Default for unknown types
else:
param_info["type"] = "[UNKNOWN - PLEASE FIX]" # If unable to detect type
param_info[
"type"
] = "[UNKNOWN - PLEASE FIX]" # If unable to detect type
# Handle default values
if default is not None:
if isinstance(default, ast.Constant) or isinstance(default, ast.NameConstant):
param_info["default"] = default.value # Use the default value directly
if isinstance(default, ast.Constant) or isinstance(
default, ast.NameConstant
):
param_info[
"default"
] = default.value # Use the default value directly
else:
param_info["default"] = "[UNKNOWN DEFAULT]" # Unknown default type
param_info["required"] = False # Optional since it has a default value
@ -164,6 +221,7 @@ def get_function_parameters(node: ast.FunctionDef, tree: ast.AST) -> list:
return parameters
def get_function_docstring(node: Any) -> str:
"""Extract the function's docstring description if present."""
# Check if the first node is a docstring
@ -178,6 +236,7 @@ def get_function_docstring(node: Any) -> str:
return "No description provided."
def extract_arg_descriptions_from_docstring(docstring: str) -> dict:
"""Extract descriptions for function parameters from the 'Args' section of the docstring."""
descriptions = {}
@ -195,14 +254,14 @@ def extract_arg_descriptions_from_docstring(docstring: str) -> dict:
continue # Proceed to the next line after 'Args:'
# End of 'Args' section if no indentation and no colon
if in_args_section and not line.startswith(" ") and ':' not in line:
if in_args_section and not line.startswith(" ") and ":" not in line:
break # Stop processing if we reach a new section
# Process lines in the 'Args' section
if in_args_section:
if ':' in line:
if ":" in line:
# Extract parameter name and description
param_name, description = line.split(':', 1)
param_name, description = line.split(":", 1)
descriptions[param_name.strip()] = description.strip()
current_param = param_name.strip()
elif current_param and line.startswith(" "):
@ -230,43 +289,50 @@ def generate_prompt_targets(input_file_path: str) -> None:
route_decorators = get_route_decorators(node, framework)
if route_decorators:
route_path = get_route_path(node, framework)
function_params = get_function_parameters(node, tree) # Get parameters for the route
function_params = get_function_parameters(
node, tree
) # Get parameters for the route
function_docstring = get_function_docstring(node) # Extract docstring
routes.append({
'name': node.name,
'path': route_path,
'methods': route_decorators,
'parameters': function_params, # Add parameters to the route
'description': function_docstring # Add the docstring as the description
})
routes.append(
{
"name": node.name,
"path": route_path,
"methods": route_decorators,
"parameters": function_params, # Add parameters to the route
"description": function_docstring, # Add the docstring as the description
}
)
# Generate YAML structure
output_structure = {
"prompt_targets": []
}
output_structure = {"prompt_targets": []}
for route in routes:
target = {
"name": route['name'],
"name": route["name"],
"endpoint": [
{
"name": "app_server",
"path": route['path'],
"path": route["path"],
}
],
"description": route['description'], # Use extracted docstring
"description": route["description"], # Use extracted docstring
"parameters": [
{
"name": param['name'],
"type": param['type'],
"name": param["name"],
"type": param["type"],
"description": f"{param['description']}",
**({"default": param['default']} if "default" in param and param['default'] is not None else {}), # Only add default if it's set
"required": param['required']
} for param in route['parameters']
]
**(
{"default": param["default"]}
if "default" in param and param["default"] is not None
else {}
), # Only add default if it's set
"required": param["required"],
}
for param in route["parameters"]
],
}
if route['name'] == "default":
if route["name"] == "default":
# Special case for `information_extraction` based on your YAML format
target["type"] = "default"
target["auto-llm-dispatch-on-response"] = True
@ -274,9 +340,12 @@ def generate_prompt_targets(input_file_path: str) -> None:
output_structure["prompt_targets"].append(target)
# Output as YAML
print(yaml.dump(output_structure, sort_keys=False,default_flow_style=False, indent=3))
print(
yaml.dump(output_structure, sort_keys=False, default_flow_style=False, indent=3)
)
if __name__ == '__main__':
if __name__ == "__main__":
if len(sys.argv) != 2:
print("Usage: python targets.py <input_file>")
sys.exit(1)

20
arch/tools/test/fastapi_test.py
View file

@ -4,12 +4,23 @@ from typing import List, Dict, Set
app = FastAPI()
class User(BaseModel):
name: str = Field("John Doe", description="The name of the user.") # Default value and description for name
name: str = Field(
"John Doe", description="The name of the user."
) # Default value and description for name
location: int = None
age: int = Field(30, description="The age of the user.") # Default value and description for age
tags: Set[str] = Field(default_factory=set, description="A set of tags associated with the user.") # Default empty set and description for tags
metadata: Dict[str, int] = Field(default_factory=dict, description="A dictionary storing metadata about the user, with string keys and integer values.") # Default empty dict and description for metadata
age: int = Field(
30, description="The age of the user."
) # Default value and description for age
tags: Set[str] = Field(
default_factory=set, description="A set of tags associated with the user."
) # Default empty set and description for tags
metadata: Dict[str, int] = Field(
default_factory=dict,
description="A dictionary storing metadata about the user, with string keys and integer values.",
) # Default empty dict and description for metadata
@app.get("/agent/default")
async def default(request: User):
@ -19,6 +30,7 @@ async def default(request: User):
"""
return {"info": f"Query: {request.name}, Count: {request.age}"}
@app.post("/agent/action")
async def reboot_network_device(device_id: str, confirmation: str):
"""

52
arch/tools/utils.py
View file

@ -6,6 +6,7 @@ import shlex
import yaml
import json
def run_docker_compose_ps(compose_file, env):
"""
Check if all Docker Compose services are in a healthy state.
@ -16,20 +17,31 @@ def run_docker_compose_ps(compose_file, env):
try:
# Run `docker-compose ps` to get the health status of each service
ps_process = subprocess.Popen(
["docker", "compose", "-p", "arch", "ps", "--format", "table{{.Service}}\t{{.State}}\t{{.Ports}}"],
[
"docker",
"compose",
"-p",
"arch",
"ps",
"--format",
"table{{.Service}}\t{{.State}}\t{{.Ports}}",
],
cwd=os.path.dirname(compose_file),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
start_new_session=True,
env=env
env=env,
)
# Capture the output of `docker-compose ps`
services_status, error_output = ps_process.communicate()
# Check if there is any error output
if error_output:
print(f"Error while checking service status:\n{error_output}", file=os.sys.stderr)
print(
f"Error while checking service status:\n{error_output}",
file=os.sys.stderr,
)
return {}
services = parse_docker_compose_ps_output(services_status)
@ -39,26 +51,31 @@ def run_docker_compose_ps(compose_file, env):
print(f"Failed to check service status. Error:\n{e.stderr}")
return e
#Helper method to print service status
# Helper method to print service status
def print_service_status(services):
print(f"{'Service Name':<25} {'State':<20} {'Ports'}")
print("="*72)
print("=" * 72)
for service_name, info in services.items():
status = info['STATE']
ports = info['PORTS']
status = info["STATE"]
ports = info["PORTS"]
print(f"{service_name:<25} {status:<20} {ports}")
#check for states based on the states passed in
# check for states based on the states passed in
def check_services_state(services, states):
for service_name, service_info in services.items():
status = service_info['STATE'].lower() # Convert status to lowercase for easier comparison
status = service_info[
"STATE"
].lower() # Convert status to lowercase for easier comparison
if any(state in status for state in states):
return True
return False
def get_llm_provider_access_keys(arch_config_file):
with open(arch_config_file, 'r') as file:
with open(arch_config_file, "r") as file:
arch_config = file.read()
arch_config_yaml = yaml.safe_load(arch_config)
@ -70,22 +87,23 @@ def get_llm_provider_access_keys(arch_config_file):
return access_key_list
def load_env_file_to_dict(file_path):
env_dict = {}
# Open and read the .env file
with open(file_path, 'r') as file:
with open(file_path, "r") as file:
for line in file:
# Strip any leading/trailing whitespaces
line = line.strip()
# Skip empty lines and comments
if not line or line.startswith('#'):
if not line or line.startswith("#"):
continue
# Split the line into key and value at the first '=' sign
if '=' in line:
key, value = line.split('=', 1)
if "=" in line:
key, value = line.split("=", 1)
key = key.strip()
value = value.strip()
@ -94,6 +112,7 @@ def load_env_file_to_dict(file_path):
return env_dict
def parse_docker_compose_ps_output(output):
# Split the output into lines
lines = output.strip().splitlines()
@ -111,10 +130,7 @@ def parse_docker_compose_ps_output(output):
parts = line.split()
# Create a dictionary entry using the header names
service_info = {
headers[1]: parts[1], # State
headers[2]: parts[2] # Ports
}
service_info = {headers[1]: parts[1], headers[2]: parts[2]} # State # Ports
# Add to the result dictionary using the service name as the key
services[parts[0]] = service_info