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
        mkdir -p $out/lib/python3.11/site-packages/ozai_webui/templates
        mkdir -p $out/lib/python3.11/site-packages/ozai_webui/static
        mkdir -p $out/lib/python3.11/site-packages/templates
        mkdir -p $out/lib/python3.11/site-packages/static
        mkdir -p $out/lib/python3.11/templates
        mkdir -p $out/lib/python3.11/static

        install -Dm444 ${src}/static/* $out/share/ozai_webui/static/
        install -Dm444 ${src}/templates/* $out/share/ozai_webui/templates/
        cp -r $src/templates $out/lib/python3.11/site-packages/ozai_webui/templates
        cp -r $src/templates $out/lib/python3.11/site-packages/templates
        cp -r $src/templates $out/lib/python3.11/templates
        cp -r $src/static $out/lib/python3.11/site-packages/ozai_webui/static
        cp -r $src/static $out/lib/python3.11/site-packages/static
        cp -r $src/static $out/lib/python3.11/static
      '';
    };



packages.x86_64-linux.ozai-webui-run = pkgs.stdenv.mkDerivation rec {
  pname = "ozai-webui-run";
  version = "0.1.1";
  src = ./src;

  pythonEnv = pkgs.python3.withPackages (ps: with ps; [
    flask
    flask-socketio
    requests
    gunicorn
    eventlet
  ]);

  buildInputs = [
    pkgs.bash
    pythonEnv
  ];

  propagatedBuildInputs = buildInputs;

  nativeBuildInputs = [ 
    pkgs.makeWrapper
  ];

  installPhase = ''
    mkdir -p $out/bin
    mkdir -p $out/share/ozai_webui/static
    mkdir -p $out/share/ozai_webui/templates
    mkdir -p $out/bin/templates
    mkdir -p $out/bin/static
    cp -r $src/templates $out/bin/templates
    cp -r $src/static $out/bin/static
    cp $src/run.sh $out/bin/run
    chmod +x $out/bin/run
    wrapProgram $out/bin/run \
      --prefix PATH : ${lib.makeBinPath buildInputs}
  '';
};

    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 ${config.services.ozai-webui.host} ${toString(config.services.ozai-webui.port)} ${self.packages.x86_64-linux.default}/bin";
            ExecStart = "${self.packages.x86_64-linux.ozai-webui-run}/bin/run ${config.services.ozai-webui.host} ${toString(config.services.ozai-webui.port)} ${self.packages.x86_64-linux.default}/lib/python3.11/site-packages/";
            # ExecStart = "${self.packages.x86_64-linux.ozai-webui}/bin/ozai-webui";
            # ExecStart = "${self.packages.x86_64-linux.default}/bin/ozai_webui";
            Restart = "always";
          };
        };
      };
    };
  };
}