nixos/modules/services/web-apps/chatgpt-retrieval-plugin.nix

   1 {
   2   config,
   3   pkgs,
   4   lib,
   5   ...
   6 }:
   7
   8 with lib;
   9
  10 let
  11   cfg = config.services.chatgpt-retrieval-plugin;
  12 in
  13 {
  14   options.services.chatgpt-retrieval-plugin = {
  15     enable = mkEnableOption "chatgpt-retrieval-plugin service";
  16
  17     port = mkOption {
  18       type = types.port;
  19       default = 8080;
  20       description = "Port the chatgpt-retrieval-plugin service listens on.";
  21     };
  22
  23     host = mkOption {
  24       type = types.str;
  25       default = "127.0.0.1";
  26       example = "0.0.0.0";
  27       description = "The hostname or IP address for chatgpt-retrieval-plugin to bind to.";
  28     };
  29
  30     bearerTokenPath = mkOption {
  31       type = types.path;
  32       description = ''
  33         Path to the secret bearer token used for the http api authentication.
  34       '';
  35       default = "";
  36       example = "config.age.secrets.CHATGPT_RETRIEVAL_PLUGIN_BEARER_TOKEN.path";
  37     };
  38
  39     openaiApiKeyPath = mkOption {
  40       type = types.path;
  41       description = ''
  42         Path to the secret openai api key used for embeddings.
  43       '';
  44       default = "";
  45       example = "config.age.secrets.CHATGPT_RETRIEVAL_PLUGIN_OPENAI_API_KEY.path";
  46     };
  47
  48     datastore = mkOption {
  49       type = types.enum [
  50         "pinecone"
  51         "weaviate"
  52         "zilliz"
  53         "milvus"
  54         "qdrant"
  55         "redis"
  56       ];
  57       default = "qdrant";
  58       description = "This specifies the vector database provider you want to use to store and query embeddings.";
  59     };
  60
  61     qdrantCollection = mkOption {
  62       type = types.str;
  63       description = ''
  64         name of the qdrant collection used to store documents.
  65       '';
  66       default = "document_chunks";
  67     };
  68   };
  69
  70   config = mkIf cfg.enable {
  71
  72     assertions = [
  73       {
  74         assertion = cfg.bearerTokenPath != "";
  75         message = "services.chatgpt-retrieval-plugin.bearerTokenPath should not be an empty string.";
  76       }
  77       {
  78         assertion = cfg.openaiApiKeyPath != "";
  79         message = "services.chatgpt-retrieval-plugin.openaiApiKeyPath should not be an empty string.";
  80       }
  81     ];
  82
  83     systemd.services.chatgpt-retrieval-plugin = {
  84       description = "ChatGPT Retrieval Plugin";
  85       after = [ "network.target" ];
  86       wantedBy = [ "multi-user.target" ];
  87
  88       serviceConfig = {
  89         DynamicUser = true;
  90         Restart = "always";
  91         LoadCredential = [
  92           "BEARER_TOKEN:${cfg.bearerTokenPath}"
  93           "OPENAI_API_KEY:${cfg.openaiApiKeyPath}"
  94         ];
  95         StateDirectory = "chatgpt-retrieval-plugin";
  96         StateDirectoryMode = "0755";
  97       };
  98
  99       # it doesn't make sense to pass secrets as env vars, this is a hack until
 100       # upstream has proper secret management.
 101       script = ''
 102         export BEARER_TOKEN=$(${pkgs.systemd}/bin/systemd-creds cat BEARER_TOKEN)
 103         export OPENAI_API_KEY=$(${pkgs.systemd}/bin/systemd-creds cat OPENAI_API_KEY)
 104         exec ${pkgs.chatgpt-retrieval-plugin}/bin/start --host ${cfg.host} --port ${toString cfg.port}
 105       '';
 106
 107       environment = {
 108         DATASTORE = cfg.datastore;
 109         QDRANT_COLLECTION = mkIf (cfg.datastore == "qdrant") cfg.qdrantCollection;
 110       };
 111     };
 112
 113     systemd.tmpfiles.rules = [
 114       # create the directory for static files for fastapi
 115       "C /var/lib/chatgpt-retrieval-plugin/.well-known - - - - ${pkgs.chatgpt-retrieval-plugin}/${pkgs.python3Packages.python.sitePackages}/.well-known"
 116     ];
 117   };
 118 }