From ac9015465c9c85211d506bd61928d837c31d9902 Mon Sep 17 00:00:00 2001 From: Jonathan Lessinger Date: Thu, 11 Jan 2024 18:31:54 -0500 Subject: [PATCH] [AIC-py][editor] server: persistent local cfg + telemetry flag, endpoints Test: - Start the server with and without flag and with no file, see that it gets created with correct contents - start the server with and without flag with existing file, with both true and false - hit both endpoints, check responses and that the file contents are changed https://github.com/lastmile-ai/aiconfig/assets/148090348/b36c3c00-c14b-4cde-8449-60386f2be1eb --- python/requirements.txt | 3 +- python/src/aiconfig/editor/server/server.py | 36 +++++++- .../aiconfig/editor/server/server_utils.py | 37 ++++++++- python/src/aiconfig/scripts/aiconfig_cli.py | 82 +++++++++++++++---- 4 files changed, 136 insertions(+), 22 deletions(-) diff --git a/python/requirements.txt b/python/requirements.txt index 0e3b7ffde..4294c1612 100644 --- a/python/requirements.txt +++ b/python/requirements.txt @@ -20,4 +20,5 @@ pytest-asyncio python-dotenv pyyaml requests -result \ No newline at end of file +result +ruamel.yaml \ No newline at end of file diff --git a/python/src/aiconfig/editor/server/server.py b/python/src/aiconfig/editor/server/server.py index fc091cd51..d4e5cfa86 100644 --- a/python/src/aiconfig/editor/server/server.py +++ b/python/src/aiconfig/editor/server/server.py @@ -15,6 +15,7 @@ from aiconfig.Config import AIConfigRuntime from aiconfig.editor.server.queue_iterator import STOP_STREAMING_SIGNAL, QueueIterator from aiconfig.editor.server.server_utils import ( + AIConfigRC, EditServerConfig, FlaskResponse, HttpResponseWithAIConfig, @@ -57,7 +58,7 @@ CORS(app, resources={r"/api/*": {"origins": "*"}}) -def run_backend_server(edit_config: EditServerConfig) -> Result[str, str]: +def run_backend_server(edit_config: EditServerConfig, aiconfigrc_path: str) -> Result[str, str]: LOGGER.setLevel(edit_config.log_level) LOGGER.info("Edit config: %s", edit_config.model_dump_json()) LOGGER.info(f"Starting server on http://localhost:{edit_config.server_port}") @@ -68,7 +69,7 @@ def run_backend_server(edit_config: EditServerConfig) -> Result[str, str]: LOGGER.warning(f"Failed to open browser: {e}. Please open http://localhost:{port} manually.") app.server_state = ServerState() # type: ignore - res_server_state_init = init_server_state(app, edit_config) + res_server_state_init = init_server_state(app, edit_config, aiconfigrc_path) match res_server_state_init: case Ok(_): LOGGER.info("Initialized server state") @@ -566,3 +567,34 @@ def _op(aiconfig_runtime: AIConfigRuntime, _op_args: OpArgs) -> Result[None, str signature: dict[str, Type[Any]] = {} return run_aiconfig_operation_with_request_json(aiconfig, request_json, f"method_", _op, signature) + + +@app.route("/api/get_aiconfigrc", methods=["GET"]) +def get_aiconfigrc() -> FlaskResponse: + state = get_server_state(app) + + yaml_mapping: Result[AIConfigRC, str] = core_utils.read_text_file(state.aiconfigrc_path).and_then(AIConfigRC.from_yaml) + match yaml_mapping: + case Ok(yaml_mapping_ok): + return FlaskResponse((yaml_mapping_ok.model_dump(), 200)) + case Err(e): + return FlaskResponse(({"message": f"Failed to load aiconfigrc: {e}"}, 400)) + + +@app.route("/api/set_aiconfigrc", methods=["POST"]) +def set_aiconfigrc() -> FlaskResponse: + state = get_server_state(app) + request_json = request.get_json() + # TODO: + # We might not need to implement this at all. + # + # If so: + # Assuming request_json["aiconfigrc"] is a yaml-formatted string + # (possibly with comments) + # Note that the file might already exist and have contents. + # + # here's how to write it to a file: + # from ruamel.yaml import YAML + # yaml = YAML() + # with open(state.aiconfigrc_path, "w") as f: + # yaml.dump(request_json["aiconfigrc"], f) diff --git a/python/src/aiconfig/editor/server/server_utils.py b/python/src/aiconfig/editor/server/server_utils.py index 072e03e9d..2a1e6490d 100644 --- a/python/src/aiconfig/editor/server/server_utils.py +++ b/python/src/aiconfig/editor/server/server_utils.py @@ -6,9 +6,10 @@ import typing from dataclasses import dataclass, field from enum import Enum +from textwrap import dedent +from threading import Event from types import ModuleType from typing import Any, Callable, NewType, Type, TypeVar, cast -from threading import Event import lastmile_utils.lib.core.api as core_utils import result @@ -17,6 +18,7 @@ from flask import Flask from pydantic import field_validator from result import Err, Ok, Result +from ruamel.yaml import YAML from aiconfig.schema import Prompt, PromptMetadata @@ -75,10 +77,40 @@ def convert_to_mode(cls, value: Any) -> ServerMode: # pylint: disable=no-self-a @dataclass class ServerState: + aiconfigrc_path: str = os.path.join(os.path.expanduser("~"), ".aiconfigrc") aiconfig: AIConfigRuntime | None = None events: dict[str, Event] = field(default_factory=dict) +class AIConfigRC(core_utils.Record): + allow_usage_data_sharing: bool + + class Config: + extra = "forbid" + + @classmethod + def from_yaml(cls: Type["AIConfigRC"], yaml: str) -> Result["AIConfigRC", str]: + try: + loaded = YAML().load(yaml) + loaded_dict = dict(loaded) + validated_model = cls.model_validate(loaded_dict) + return Ok(validated_model) + except Exception as e: + return core_utils.ErrWithTraceback(e) + + +DEFAULT_AICONFIGRC = YAML().load( + dedent( + """ + # Tip: make sure this file is called .aiconfigrc and is in your home directory. + + # Flag allowing or denying telemetry for product development purposes. + allow_usage_data_sharing: true + """ + ), +) + + FlaskResponse = NewType("FlaskResponse", tuple[core_utils.JSONObject, int]) @@ -200,10 +232,11 @@ def safe_load_from_disk(aiconfig_path: ValidatedPath) -> Result[AIConfigRuntime, return core_utils.ErrWithTraceback(e) -def init_server_state(app: Flask, edit_config: EditServerConfig) -> Result[None, str]: +def init_server_state(app: Flask, edit_config: EditServerConfig, aiconfigrc_path: str) -> Result[None, str]: LOGGER.info("Initializing server state") _load_user_parser_module_if_exists(edit_config.parsers_module_path) state = get_server_state(app) + state.aiconfigrc_path = aiconfigrc_path assert state.aiconfig is None if os.path.exists(edit_config.aiconfig_path): diff --git a/python/src/aiconfig/scripts/aiconfig_cli.py b/python/src/aiconfig/scripts/aiconfig_cli.py index f57de74dc..d74bf63ea 100644 --- a/python/src/aiconfig/scripts/aiconfig_cli.py +++ b/python/src/aiconfig/scripts/aiconfig_cli.py @@ -1,19 +1,22 @@ import asyncio import logging +import os import signal import socket import subprocess import sys import lastmile_utils.lib.core.api as core_utils -import result +from ruamel.yaml import YAML + from aiconfig.editor.server.server import run_backend_server -from aiconfig.editor.server.server_utils import EditServerConfig, ServerMode +from aiconfig.editor.server.server_utils import DEFAULT_AICONFIGRC, EditServerConfig, ServerMode from result import Err, Ok, Result class AIConfigCLIConfig(core_utils.Record): log_level: str | int = "WARNING" + aiconfigrc_path: str = os.path.join(os.path.expanduser("~"), ".aiconfigrc") logging.basicConfig(format=core_utils.LOGGER_FMT) @@ -36,28 +39,39 @@ def run_subcommand(argv: list[str]) -> Result[str, str]: subparser_record_types = {"edit": EditServerConfig} main_parser = core_utils.argparsify(AIConfigCLIConfig, subparser_record_types=subparser_record_types) - res_cli_config = core_utils.parse_args(main_parser, argv[1:], AIConfigCLIConfig) - res_cli_config.and_then(_process_cli_config) + # Try to parse the CLI args into a config. + cli_config: Result[AIConfigCLIConfig, str] = core_utils.parse_args(main_parser, argv[1:], AIConfigCLIConfig) + + # If cli_config is Ok(), pass its contents to _get_cli_process_result_from_config(). + # Otherwise, short circuit and assign process_result to the Err. + # Nothing gets mutated except for log level (see inside _get_cli_process_result_from_config() + process_result = cli_config.and_then(_set_log_level_and_create_default_yaml) + LOGGER.info(f"{process_result=}") subparser_name = core_utils.get_subparser_name(main_parser, argv[1:]) LOGGER.info(f"Running subcommand: {subparser_name}") if subparser_name == "edit": LOGGER.debug("Running edit subcommand") - res_edit_config = core_utils.parse_args(main_parser, argv[1:], EditServerConfig) - LOGGER.debug(f"{res_edit_config.is_ok()=}") - res_servers = res_edit_config.and_then(_run_editor_servers) - out: Result[str, str] = result.do( - # - Ok(",".join(res_servers_ok)) - # - for res_servers_ok in res_servers - ) + edit_config = core_utils.parse_args(main_parser, argv[1:], EditServerConfig) + LOGGER.debug(f"{edit_config.is_ok()=}") + out = _run_editor_servers_with_configs(edit_config, cli_config) return out else: return Err(f"Unknown subparser: {subparser_name}") +def _run_editor_servers_with_configs(edit_config: Result[EditServerConfig, str], cli_config: Result[AIConfigCLIConfig, str]) -> Result[str, str]: + if not (edit_config.is_ok() and cli_config.is_ok()): + return Err(f"Something went wrong with configs: {edit_config=}, {cli_config=}") + + server_outcomes = _run_editor_servers(edit_config.unwrap(), cli_config.unwrap().aiconfigrc_path) + if server_outcomes.is_err(): + return Err(f"Something went wrong with servers: {server_outcomes=}") + + return Ok(",".join(server_outcomes.unwrap())) + + def _sigint(procs: list[subprocess.Popen[bytes]]) -> Result[str, str]: LOGGER.info("sigint") for p in procs: @@ -76,7 +90,7 @@ def is_port_in_use(port: int) -> bool: return s.connect_ex(("localhost", port)) == 0 -def _run_editor_servers(edit_config: EditServerConfig) -> Result[list[str], str]: +def _run_editor_servers(edit_config: EditServerConfig, aiconfigrc_path: str) -> Result[list[str], str]: port = edit_config.server_port while is_port_in_use(port): @@ -100,7 +114,7 @@ def _run_editor_servers(edit_config: EditServerConfig) -> Result[list[str], str] return Err(e) results: list[Result[str, str]] = [] - backend_res = run_backend_server(edit_config) + backend_res = run_backend_server(edit_config, aiconfigrc_path) match backend_res: case Ok(_): pass @@ -114,8 +128,43 @@ def _run_editor_servers(edit_config: EditServerConfig) -> Result[list[str], str] return core_utils.result_reduce_list_all_ok(results) -def _process_cli_config(cli_config: AIConfigCLIConfig) -> Result[bool, str]: +def _set_log_level_and_create_default_yaml(cli_config: AIConfigCLIConfig) -> Result[bool, str]: + """ + This function has 2 jobs (currently): + 1. Set the log level + 2. Write the default aiconfigrc if it doesn't exist. + + It returns Ok(True) if everything went well. Currently, it never returns Ok(False). + As usual, we return an error with a message if something went wrong. + """ + aiconfigrc_path = cli_config.aiconfigrc_path + LOGGER.setLevel(cli_config.log_level) + try: + with open(aiconfigrc_path, "x") as f: + YAML().dump(DEFAULT_AICONFIGRC, f) + except FileExistsError: + try: + + def _read() -> str: + with open(aiconfigrc_path, "r") as f: + return f.read() + + contents = YAML().load(_read()) + with open(aiconfigrc_path, "w") as f: + if contents is None: + contents = {} + + for k, v in DEFAULT_AICONFIGRC.items(): + if k not in contents: + contents[k] = v + + YAML().dump(contents, f) + except Exception as e: + return core_utils.ErrWithTraceback(e) + except Exception as e: + return core_utils.ErrWithTraceback(e) + return Ok(True) @@ -142,7 +191,6 @@ def _run_frontend_server_background() -> Result[list[subprocess.Popen[bytes]], s def main() -> int: - print("Running main") argv = sys.argv return asyncio.run(main_with_args(argv))