ozai-webui/flake.nix

{
  description = "Ozai WebUI";

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

  outputs = { self, nixpkgs, ... }@inputs:
  let
    pkgs = import nixpkgs { system = "x86_64-linux"; };
    lib = pkgs.lib;
    types = lib.types;
    dependencies = [
      pkgs.python3
      pkgs.python3Packages.flask
      pkgs.python3Packages.flask-socketio
      pkgs.python3Packages.requests
      pkgs.python3Packages.gunicorn
      pkgs.python3Packages.eventlet
    ];
  in
  {
    packages.x86_64-linux.default = pkgs.python3Packages.buildPythonPackage rec {
      pname = "ozai-webui";
      version = "0.1.1";
      propagatedBuildInputs = dependencies;
      src = ./src;
      doCheck = false;

      postInstall = ''
        mkdir -p $out/share/ozai-webui/static
        mkdir -p $out/share/ozai-webui/templates
        install -Dm444 ${src}/static/* $out/share/ozai-webui/static/
        install -Dm444 ${src}/templates/* $out/share/ozai-webui/templates/
      '';
    };

    packages.x86_64-linux.ozai-webui-run = pkgs.stdenv.mkDerivation rec {
      pname = "ozai-webui-run";
      version = "0.1.1";
      src = ./src;
      buildInputs = [ self.packages.x86_64-linux.default pkgs.python3Packages.gunicorn ];
      installPhase = ''
        mkdir -p $out/bin
        echo 'gunicorn --worker-class eventlet -w 1 --chdir ${self.packages.x86_64-linux.default}/bin ozai_webui:app' > $out/bin/run
        chmod +x $out/bin/run
      '';
    };

    devShells.x86_64-linux.default  = pkgs.mkShell {
      buildInputs = dependencies;
      shellHook = ''
        OZAI_WEBUI_STATIC_FOLDER=result/share/ozai-webui/static
        OZAI_WEBUI_TEMPLATE_FOLDER=result/share/ozai-webui/templates
      '';
    };

    nixosModules.ozai-webui = { config, pkgs, ... }: {
      options.services.ozai-webui = {
        enable = lib.mkEnableOption "Ozai WebUI";
        host = lib.mkOption {
          type = types.str;
          default = "127.0.0.1";
          description = "The host to run the server on";
        };
        port = lib.mkOption {
          type = types.int;
          default = 8080;
          description = "The port to run the server on";
        };
        ozaiUrl = lib.mkOption {
          type = types.str;
          default = "http://localhost:8000/api/";
          description = "The URL of the Ozai server";
        };
      };

      config = lib.mkIf config.services.ozai-webui.enable {
        systemd.services.ozai-webui = {
          description = "Ozai WebUI server";
          after = [ "network.target" ];
          wantedBy = [ "multi-user.target" ];
	  environment = {
              OZAI_URL=  "${config.services.ozai-webui.ozaiUrl}";
              OZAI_WEBUI_HOST= "${config.services.ozai-webui.host}";
              OZAI_WEBUI_PORT= toString(config.services.ozai-webui.port);
              OZAI_WEBUI_STATIC_FOLDER= "${self.packages.x86_64-linux.default}/share/ozai-webui/static";
              OZAI_WEBUI_TEMPLATE_FOLDER="${self.packages.x86_64-linux.default}/share/ozai-webui/templates";
            };

          serviceConfig = {
            ExecStart = "${self.packages.x86_64-linux.ozai-webui-run}/bin/run";
            Restart = "always";
          };
        };
      };
    };
  };
}