ozai-webui/flake.nix

{
  inputs.nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable";

  outputs = { self, nixpkgs}:
    let
      supportedSystems = [ "x86_64-linux"];
      forAllSystems = nixpkgs.lib.genAttrs supportedSystems;
      pkgs = forAllSystems (system: nixpkgs.legacyPackages.${system});
    in
    {
      packages = forAllSystems (system: let
        python = pkgs.${system}.python3.withPackages (ps: [ps.flask ps.flask-socketio]);
        ozaiWebui = import ./default.nix { inherit (pkgs.${system}) pkgs; };
      in {
        default = ozaiWebui;
      });

      devShells = forAllSystems (system: let
        python = pkgs.${system}.python3.withPackages (ps: [ps.flask ps.flask-socketio]);
      in {
        default = pkgs.${system}.mkShellNoCC {
          packages = with pkgs.${system}; [
            python
          ];
        };
      });

      nixosModules = forAllSystems (system: let
              ozaiWebui = import ./default.nix { inherit (pkgs.${system}) pkgs; };
            in {
              ozai-webui = { config, pkgs, ... }: {
                options.services.ozai-webui = {
                  enable = mkEnableOption "Ozai WebUI";
                  host = mkOption {
                    type = types.str;
                    default = "127.0.0.1";
                    description = "The host to run the server on";
                  };
                  port = mkOption {
                    type = types.int;
                    default = 8080;
                    description = "The port to run the server on";
                  };
                  ozaiUrl = mkOption {
                    type = types.str;
                    default = "http://localhost:8000/api/";
                    description = "The URL of the Ozai server";
                  };
                  secretKey = mkOption {
                    type = types.str;
                    default = "secret_key";
                    description = "The secret key for the Flask app";
                  };
                };
                config = mkIf config.services.ozai-webui.enable {
                  systemd.services.ozai-webui = {
                    description = "Ozai WebUI";
                    after = [ "network.target" ];
                    wantedBy = [ "multi-user.target" ];
                    serviceConfig = {
                      ExecStart = "${ozaiWebui}/bin/ozai-webui --port ${toString config.services.ozai-webui.port} --host '${config.services.ozai-webui.host}' --ozai_url '${config.services.ozai-webui.ozaiUrl}' --secret_key '${config.services.ozai-webui.secretKey}'";
                      Restart = "always";
                    };
                  };
                };
              };
            });



    };
}