rags/flake.nix

{
  description = "LightRAG + Graphiti dev environments";

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

  outputs =
    { self, nixpkgs }:
    let
      system = "x86_64-linux";
      pkgs = nixpkgs.legacyPackages.${system};

      stdLibs = pkgs.lib.makeLibraryPath [
        pkgs.stdenv.cc.cc
        pkgs.zlib
      ];

      startNeo4j = pkgs.writeShellScript "start-neo4j" ''
        set -e
        : "''${RAGS_ROOT:=$PWD}"
        NEO4J_DATA="$RAGS_ROOT/data/neo4j"
        mkdir -p "$NEO4J_DATA"/{data,logs,run,plugins,import,conf}

        cat > "$NEO4J_DATA/conf/neo4j.conf" <<CONF
        server.directories.data=$NEO4J_DATA/data
        server.directories.logs=$NEO4J_DATA/logs
        server.directories.run=$NEO4J_DATA/run
        server.directories.plugins=$NEO4J_DATA/plugins
        server.directories.import=$NEO4J_DATA/import
        server.bolt.listen_address=:7687
        server.http.listen_address=:7474
        dbms.security.auth_enabled=false
        CONF

        export NEO4J_CONF="$NEO4J_DATA/conf"
        ${pkgs.neo4j}/bin/neo4j start
        echo ""
        echo "Neo4j running:"
        echo "  bolt: bolt://localhost:7687"
        echo "  http: http://localhost:7474"
        echo "  stop: RAGS_ROOT=$RAGS_ROOT nix run .#neo4j-stop"
        echo ""
      '';

      stopNeo4j = pkgs.writeShellScript "stop-neo4j" ''
        : "''${RAGS_ROOT:=$PWD}"
        export NEO4J_CONF="$RAGS_ROOT/data/neo4j/conf"
        ${pkgs.neo4j}/bin/neo4j stop
      '';

      startAll = pkgs.writeShellScript "start-all" ''
        set -e
        : "''${RAGS_ROOT:=$PWD}"

        if [ -f "$RAGS_ROOT/.env.lightrag" ]; then
          set -a; source "$RAGS_ROOT/.env.lightrag"; set +a
        fi

        LIGHTRAG_BIN="$RAGS_ROOT/lightrag/.venv/bin/lightrag-server"
        LOG_DIR="$RAGS_ROOT/logs"
        mkdir -p "$LOG_DIR"

        LIGHTRAG_PID=""
        cleanup() {
          echo ""
          echo "Shutting down..."
          [ -n "$LIGHTRAG_PID" ] && kill "$LIGHTRAG_PID" 2>/dev/null || true
          wait 2>/dev/null || true
        }
        trap cleanup EXIT INT TERM

        echo "Starting LightRAG server..."
        "$LIGHTRAG_BIN" >> "$LOG_DIR/lightrag.log" 2>&1 &
        LIGHTRAG_PID=$!

        wait_for() {
          local label=$1 url=$2 tries=0
          printf "  Waiting for %s" "$label"
          while ! ${pkgs.curl}/bin/curl -so /dev/null --max-time 2 "$url" 2>/dev/null; do
            tries=$((tries+1))
            [ $tries -ge 300 ] && { echo " TIMEOUT — check logs/$label.log"; exit 1; }
            printf "."
            sleep 1
          done
          echo " ready"
        }

        wait_for "lightrag" "http://localhost:9621/docs"

        echo ""
        echo "All services up:"
        echo "  LightRAG webui: http://localhost:9621/webui"
        echo "  LightRAG API:   http://localhost:9621/docs"
        echo "  Ollama LLM:     http://localhost:11434 (external)"
        echo "  Ollama embed:   http://localhost:11434/api/embed (external)"
        echo "  logs:           $LOG_DIR/"
        echo ""
        echo "Ctrl+C to stop everything."
        echo ""

        tail -f "$LOG_DIR/lightrag.log"
      '';

    in
    {
      devShells.${system} = {

        lightrag = pkgs.mkShell {
          packages = [
            pkgs.uv
            pkgs.python312
            pkgs.curl
            pkgs.bun
          ];

          shellHook = ''
            RAGS_ROOT="$PWD"
            export VIRTUAL_ENV="$RAGS_ROOT/lightrag/.venv"
            export UV_PROJECT_ENVIRONMENT="$VIRTUAL_ENV"
            export UV_PYTHON="${pkgs.python312}/bin/python3.12"
            export LD_LIBRARY_PATH="${stdLibs}:$LD_LIBRARY_PATH"

            echo "Syncing lightrag venv..."
            (cd "$RAGS_ROOT/lightrag" && uv sync --extra api --extra offline-llm --quiet)
            source "$RAGS_ROOT/lightrag/.venv/bin/activate"

            if [ ! -d "$RAGS_ROOT/lightrag/lightrag/api/webui" ]; then
              echo "Building frontend (first time, ~1min)..."
              (cd "$RAGS_ROOT/lightrag/lightrag_webui" && bun install --frozen-lockfile --silent && bun run build:bun)
            fi

            if [ -f "$RAGS_ROOT/.env.lightrag" ]; then
              set -a; source "$RAGS_ROOT/.env.lightrag"; set +a
            fi

            echo ""
            echo "LightRAG shell ready."
            echo "  start server:  lightrag-server"
            echo "  start all:     nix run .#start"
            echo "  config:        $RAGS_ROOT/.env.lightrag"
            echo ""
          '';
        };

        graphiti = pkgs.mkShell {
          packages = [
            pkgs.uv
            pkgs.python312
            pkgs.neo4j
            pkgs.curl
          ];

          shellHook = ''
            RAGS_ROOT="$PWD"
            export VIRTUAL_ENV="$RAGS_ROOT/graphiti/.venv"
            export UV_PROJECT_ENVIRONMENT="$VIRTUAL_ENV"
            export UV_PYTHON="${pkgs.python312}/bin/python3.12"
            export LD_LIBRARY_PATH="${stdLibs}:$LD_LIBRARY_PATH"
            cd "$RAGS_ROOT/graphiti"

            echo "Syncing graphiti venv..."
            uv sync --quiet
            source .venv/bin/activate

            if [ -f "$RAGS_ROOT/.env.graphiti" ]; then
              set -a; source "$RAGS_ROOT/.env.graphiti"; set +a
            fi

            echo ""
            echo "Graphiti shell ready."
            echo "  neo4j:  nix run .#neo4j-start   (in another terminal, run first)"
            echo "  config: $RAGS_ROOT/.env.graphiti"
            echo ""
          '';
        };

      };

      apps.${system} = {
        start = {
          type = "app";
          program = "${startAll}";
        };
        neo4j-start = {
          type = "app";
          program = "${startNeo4j}";
        };
        neo4j-stop = {
          type = "app";
          program = "${stopNeo4j}";
        };
      };
    };
}