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;
    deps = [
      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 = deps;
      dependencies = deps;
      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 = [
		pkgs.bash
		self.packages.x86_64-linux.default
		pkgs.python3
      		pkgs.python3Packages.flask
                pkgs.python3Packages.flask-socketio
      		pkgs.python3Packages.requests
      		pkgs.python3Packages.gunicorn
      		pkgs.python3Packages.eventlet
	];
      #propagatedBuildInputs
      nativeBuildInputs = [ pkgs.makeWrapper ];
      installPhase = ''
        mkdir -p $out/bin
	echo '#!/usr/bin/env bash' > $out/bin/run
    	echo 'host="${1:-0.0.0.0}"' >> $out/bin/run
    	echo 'port="${2:-8000}"' >> $out/bin/run
    	echo 'gunicorn -w 1 --chdir ${self.packages.x86_64-linux.default}/lib/python3.11/site-packages/ozai_webui --bind "$host:$port" __init__:main' >> $out/bin/run
    	chmod +x $out/bin/run
	wrapProgram $out/bin/run \
        --prefix PATH : ${lib.makeBinPath [ pkgs.bash pkgs.python3 pkgs.python3Packages.flask pkgs.python3Packages.flask-socketio pkgs.python3Packages.requests pkgs.python3Packages.gunicorn pkgs.python3Packages.eventlet  ]}
      '';
    };

    devShells.x86_64-linux.default  = pkgs.mkShell {
      buildInputs = deps;
      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";
          };
        };
      };
    };
  };
}