diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000000000000000000000000000000000000..239b2099e5d7019085dd91dfc58e1a17815cc23b --- /dev/null +++ b/.dockerignore @@ -0,0 +1,56 @@ +*.pdf +*.url +*.jpg +*.png +*.ipynb +*.xls +*.xlsx +*.pyc +*.json.bak.* +examples/* +output/* +tools/__pycache__/* +build/* +dist/* +logs/* +usage/* +docs/* +guide/* +feedback/* +test_code/* +test/tmp/* +unsloth_compiled_cache/* +.vscode/* +llm_topic_modelling.egg-info/* +input/ +output/ +logs/ +usage/ +feedback/ +config/ +tmp/ +index.qmd +_quarto.yml +!config/pi_agent.env.example +!config/docker_app_config.env.example +!config/app_config.env.example +workspace/* +user_guide/* +_extensions/* +cdk/config/* +!cdk/config/app_config.env.example +!cdk/config/lambda/ +cdk/config/lambda/* +!cdk/config/lambda/lambda_function.py +!cdk/config/headless_s3_seed/ +cdk/config/headless_s3_seed/* +!cdk/config/headless_s3_seed/input/ +cdk/config/headless_s3_seed/input/* +!cdk/config/headless_s3_seed/input/config/ +cdk/config/headless_s3_seed/input/config/* +!cdk/config/headless_s3_seed/input/config/example_headless_env_file.env +cdk/cdk.out/* +cdk/archive/* +cdk/cdk.json +cdk/cdk.context.json +cdk/precheck.context.json \ No newline at end of file diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..ca21a586197f5e77387721c96acad48927264fde --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.xlsx filter=lfs diff=lfs merge=lfs -text diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000000000000000000000000000000000000..5169db8330382d1539672b859a4749b35dcdd30e --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,196 @@ +name: CI/CD Pipeline + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + #schedule: + # Run tests daily at 2 AM UTC + # - cron: '0 2 * * *' + +permissions: + contents: read + actions: read + pull-requests: write + issues: write + +env: + PYTHON_VERSION: "3.11" + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: ${{ env.PYTHON_VERSION }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install ruff black + + - name: Run Ruff linter + run: ruff check . + + - name: Run Black formatter check + run: black --check . + + test-unit: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: [3.11, 3.12, 3.13] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + + - name: Cache pip dependencies + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements*.txt', '**/pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-pip- + + - name: Install Python dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements_lightweight.txt + pip install pytest pytest-cov pytest-html pytest-xdist + + - name: Verify example data files + run: | + echo "Checking if example data directory exists:" + ls -la example_data/ || echo "example_data directory not found" + echo "Checking for specific CSV files:" + ls -la example_data/*.csv || echo "No CSV files found" + + - name: Run CLI and GUI tests + run: | + cd test + python run_tests.py + + - name: Run tests with pytest + run: | + pytest test/test.py test/test_gui_only.py -v --tb=short --junitxml=test-results.xml + + - name: Run tests with coverage + run: | + pytest test/test.py test/test_gui_only.py --cov=. --cov-report=xml --cov-report=html --cov-report=term + + - name: Upload test results + uses: actions/upload-artifact@v4 + if: always() + with: + name: test-results-python-${{ matrix.python-version }} + path: | + test-results.xml + htmlcov/ + coverage.xml + + test-integration: + runs-on: ubuntu-latest + needs: [lint, test-unit] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: ${{ env.PYTHON_VERSION }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements_lightweight.txt + pip install pytest pytest-cov + + - name: Verify example data files + run: | + echo "Checking if example data directory exists:" + ls -la example_data/ + echo "Checking for specific CSV files:" + ls -la example_data/*.csv || echo "No CSV files found" + + - name: Run integration tests (CLI and GUI) + run: | + cd test + python run_tests.py + + - name: Test CLI help + run: | + python cli_topics.py --help + + - name: Test CLI version + run: | + python -c "import sys; print(f'Python {sys.version}')" + + security: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: ${{ env.PYTHON_VERSION }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install bandit + + - name: Run bandit security check + run: | + bandit -r . -f json -o bandit-report.json || true + + - name: Upload security report + uses: actions/upload-artifact@v4 + if: always() + with: + name: security-report + path: bandit-report.json + + build: + runs-on: ubuntu-latest + needs: [lint, test-unit] + if: github.event_name == 'push' && github.ref == 'refs/heads/main' + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: ${{ env.PYTHON_VERSION }} + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build twine + + - name: Build package + run: | + python -m build + + - name: Check package + run: | + twine check dist/* + + - name: Upload build artifacts + uses: actions/upload-artifact@v4 + with: + name: dist + path: dist/ + diff --git a/.github/workflows/simple-test.yml b/.github/workflows/simple-test.yml new file mode 100644 index 0000000000000000000000000000000000000000..18a51fb255e3134adc71f00ddc5b5c47809b7389 --- /dev/null +++ b/.github/workflows/simple-test.yml @@ -0,0 +1,46 @@ +name: Simple Test Run + +on: + push: + branches: [ dev ] + pull_request: + branches: [ dev ] + +permissions: + contents: read + actions: read + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v4 + with: + python-version: "3.11" + + - name: Install Python dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements_lightweight.txt + pip install pytest pytest-cov + + - name: Verify example data files + run: | + echo "Checking if example data directory exists:" + ls -la example_data/ || echo "example_data directory not found" + echo "Checking for specific CSV files:" + ls -la example_data/*.csv || echo "No CSV files found" + + - name: Run CLI and GUI tests + run: | + cd test + python run_tests.py + + - name: Run tests with pytest + run: | + pytest test/test.py test/test_gui_only.py -v --tb=short + diff --git a/.github/workflows/sync_to_hf.yml b/.github/workflows/sync_to_hf.yml new file mode 100644 index 0000000000000000000000000000000000000000..848acf3929cacbd690cdbd0cdd28a0853559246e --- /dev/null +++ b/.github/workflows/sync_to_hf.yml @@ -0,0 +1,53 @@ +name: Sync to Hugging Face hub +on: + push: + branches: [dev] + +permissions: + contents: read + +jobs: + sync-to-hub: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 1 # Only get the latest state + lfs: true # Download actual LFS files so they can be pushed + + - name: Install Git LFS + run: git lfs install + + - name: Recreate repo history (single-commit force push) + run: | + # 1. Capture the message BEFORE we delete the .git folder + COMMIT_MSG=$(git log -1 --pretty=%B) + echo "Syncing commit message: $COMMIT_MSG" + + # 2. DELETE the .git folder. + # This turns the repo into a standard folder of files. + rm -rf .git + + # 3. Re-initialize a brand new git repo + git init -b main + git config --global user.name "$HF_USERNAME" + git config --global user.email "$HF_EMAIL" + + # 4. Re-install LFS (needs to be done after git init) + git lfs install + + # 5. Add the remote + git remote add hf https://$HF_USERNAME:$HF_TOKEN@huggingface.co/spaces/$HF_USERNAME/$HF_REPO_ID + + # 6. Add all files + # Since this is a fresh init, Git sees EVERY file as "New" + git add . + + # 7. Commit and Force Push + git commit -m "Sync: $COMMIT_MSG" + git push --force hf main + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + HF_USERNAME: ${{ secrets.HF_USERNAME }} + HF_EMAIL: ${{ secrets.HF_EMAIL }} + HF_REPO_ID: ${{ secrets.HF_REPO_ID }} \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..1508e0e0481f84b549aa2109fe775d00f35b43ac --- /dev/null +++ b/.gitignore @@ -0,0 +1,55 @@ +*.pdf +*.url +*.jpg +*.png +*.ipynb +*.xls +*.pyc +*.json.bak.* +examples/* +output/* +tools/__pycache__/* +build/* +dist/* +logs/* +docs/* +guide/* +usage/* +feedback/* +test_code/* +config/* +tmp/* +test/tmp/* +unsloth_compiled_cache/* +.vscode/* +llm_topic_modelling.egg-info/* +index.qmd +_quarto.yml +/.quarto/ +**/*.quarto_ipynb +config/* +!config/pi_agent.env.example +!config/docker_app_config.env.example +!config/app_config.env.example +workspace/* +user_guide/* +_extensions/* +cdk/config/* +!cdk/config/app_config.env.example +!cdk/config/lambda/ +cdk/config/lambda/* +!cdk/config/lambda/lambda_function.py +!cdk/config/headless_s3_seed/ +cdk/config/headless_s3_seed/* +!cdk/config/headless_s3_seed/input/ +cdk/config/headless_s3_seed/input/* +!cdk/config/headless_s3_seed/input/config/ +cdk/config/headless_s3_seed/input/config/* +!cdk/config/headless_s3_seed/input/config/example_headless_env_file.env +cdk/cdk.out/* +cdk/archive/* +cdk/cdk.out/* +cdk/cdk.json +cdk/cdk.context.json +cdk/precheck.context.json + diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..81ad29e62ab59bbe0cfa457ca3295d8092a1ab00 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,165 @@ +# This Dockerfile is optimised for AWS ECS using Python 3.12, and assumes CUDA 12.6 for local models. The Dockerfile will need to be modified to install all linux CUDA / GPU dependencies. + +# Stage 1: Build dependencies and download models +FROM public.ecr.aws/docker/library/python:3.12.12-slim-trixie AS builder + +# Install system dependencies. +RUN apt-get update && apt-get install -y \ + build-essential \ + gcc \ + g++ \ + cmake \ + #libopenblas-dev \ + pkg-config \ + python3-dev \ + libffi-dev \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /src + +COPY requirements_lightweight.txt . + +ARG INSTALL_TORCH=False +ENV INSTALL_TORCH=${INSTALL_TORCH} + +# Local torch install requires CUDA 12.6 +RUN if [ "$INSTALL_TORCH" = "True" ]; then \ + pip install --no-cache-dir --target=/install torch==2.9.1 --extra-index-url https://download.pytorch.org/whl/cu126; \ + fi + +ARG INSTALL_LLAMA_CPP_PYTHON=False +ENV INSTALL_LLAMA_CPP_PYTHON=${INSTALL_LLAMA_CPP_PYTHON} + +# Llama CPP Python install requires CUDA 12.4 +RUN if [ "$INSTALL_LLAMA_CPP_PYTHON" = "True" ]; then \ + pip install --no-cache-dir --target=/install https://github.com/abetlen/llama-cpp-python/releases/download/v0.3.16-cu124/llama_cpp_python-0.3.30-cp312-cp312-linux_x86_64.whl; \ + fi + +RUN pip install --no-cache-dir --target=/install -r requirements_lightweight.txt + +RUN rm requirements_lightweight.txt + +# =================================================================== +# Stage 2: A common 'base' for both Lambda and Gradio +# =================================================================== +FROM public.ecr.aws/docker/library/python:3.12.12-slim-trixie AS base + +# Set build-time and runtime environment variable for whether to run in Gradio mode or Lambda mode +ARG APP_MODE=gradio +ENV APP_MODE=${APP_MODE} + +# Install runtime system dependencies +RUN apt-get update && apt-get install -y --no-install-recommends \ + libopenblas0 \ + && apt-get clean && rm -rf /var/lib/apt/lists/* + +ENV APP_HOME=/home/user + +# Set env variables for Gradio & other apps +ENV GRADIO_TEMP_DIR=/tmp/gradio_tmp/ \ + MPLCONFIGDIR=/tmp/matplotlib_cache/ \ + GRADIO_OUTPUT_FOLDER=$APP_HOME/app/output/ \ + GRADIO_INPUT_FOLDER=$APP_HOME/app/input/ \ + FEEDBACK_LOGS_FOLDER=$APP_HOME/app/feedback/ \ + ACCESS_LOGS_FOLDER=$APP_HOME/app/logs/ \ + USAGE_LOGS_FOLDER=$APP_HOME/app/usage/ \ + CONFIG_FOLDER=$APP_HOME/app/config/ \ + GRADIO_SERVER_NAME=0.0.0.0 \ + GRADIO_SERVER_PORT=7860 \ + PATH=$APP_HOME/.local/bin:$PATH \ + PYTHONPATH=$APP_HOME/app \ + PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + GRADIO_ALLOW_FLAGGING=never \ + GRADIO_NUM_PORTS=1 \ + GRADIO_THEME=huggingface \ + SYSTEM=spaces + +# Copy Python packages from the builder stage +COPY --from=builder /install /usr/local/lib/python3.12/site-packages/ +COPY --from=builder /install/bin /usr/local/bin/ + +# Copy your application code and entrypoint +COPY . ${APP_HOME}/app +COPY entrypoint.sh ${APP_HOME}/app/entrypoint.sh +# Fix line endings and set execute permissions +RUN sed -i 's/\r$//' ${APP_HOME}/app/entrypoint.sh \ + && chmod +x ${APP_HOME}/app/entrypoint.sh + +WORKDIR ${APP_HOME}/app + +# =================================================================== +# FINAL Stage 3: The Lambda Image (runs as root for simplicity) +# =================================================================== +FROM base AS lambda +# Set runtime ENV for Lambda mode +ENV APP_MODE=lambda +ENTRYPOINT ["/home/user/app/entrypoint.sh"] +CMD ["lambda_entrypoint.lambda_handler"] + +# =================================================================== +# FINAL Stage 4: The Gradio Image (runs as a secure, non-root user) +# =================================================================== +FROM base AS gradio +# Set runtime ENV for Gradio mode +ENV APP_MODE=gradio + +# Create non-root user +RUN useradd -m -u 1000 user + +# Create the base application directory and set its ownership +RUN mkdir -p ${APP_HOME}/app && chown user:user ${APP_HOME}/app + +# Create required sub-folders within the app directory and set their permissions +RUN mkdir -p \ + ${APP_HOME}/app/output \ + ${APP_HOME}/app/input \ + ${APP_HOME}/app/logs \ + ${APP_HOME}/app/usage \ + ${APP_HOME}/app/feedback \ + ${APP_HOME}/app/config \ + && chown user:user \ + ${APP_HOME}/app/output \ + ${APP_HOME}/app/input \ + ${APP_HOME}/app/logs \ + ${APP_HOME}/app/usage \ + ${APP_HOME}/app/feedback \ + ${APP_HOME}/app/config \ + && chmod 755 \ + ${APP_HOME}/app/output \ + ${APP_HOME}/app/input \ + ${APP_HOME}/app/logs \ + ${APP_HOME}/app/usage \ + ${APP_HOME}/app/feedback \ + ${APP_HOME}/app/config + +# Now handle the /tmp directories +RUN mkdir -p /tmp/gradio_tmp /tmp/matplotlib_cache /tmp /var/tmp \ + && chown user:user /tmp /var/tmp /tmp/gradio_tmp /tmp/matplotlib_cache \ + && chmod 1777 /tmp /var/tmp /tmp/gradio_tmp /tmp/matplotlib_cache + +# Fix apply user ownership to all files in the home directory +RUN chown -R user:user /home/user + +# Set permissions for Python executable +RUN chmod 755 /usr/local/bin/python + +# Declare volumes +VOLUME ["/tmp/matplotlib_cache"] +VOLUME ["/tmp/gradio_tmp"] +VOLUME ["/home/user/app/output"] +VOLUME ["/home/user/app/input"] +VOLUME ["/home/user/app/logs"] +VOLUME ["/home/user/app/usage"] +VOLUME ["/home/user/app/feedback"] +VOLUME ["/home/user/app/config"] +VOLUME ["/tmp"] +VOLUME ["/var/tmp"] + +USER user + +EXPOSE $GRADIO_SERVER_PORT + +ENTRYPOINT ["/home/user/app/entrypoint.sh"] +CMD ["python", "app.py"] \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..be3f7b28e564e7dd05eaf59d64adba1a4065ac0e --- /dev/null +++ b/LICENSE @@ -0,0 +1,661 @@ + GNU AFFERO GENERAL PUBLIC LICENSE + Version 3, 19 November 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU Affero General Public License is a free, copyleft license for +software and other kinds of works, specifically designed to ensure +cooperation with the community in the case of network server software. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +our General Public Licenses are intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + Developers that use our General Public Licenses protect your rights +with two steps: (1) assert copyright on the software, and (2) offer +you this License which gives you legal permission to copy, distribute +and/or modify the software. + + A secondary benefit of defending all users' freedom is that +improvements made in alternate versions of the program, if they +receive widespread use, become available for other developers to +incorporate. Many developers of free software are heartened and +encouraged by the resulting cooperation. However, in the case of +software used on network servers, this result may fail to come about. +The GNU General Public License permits making a modified version and +letting the public access it on a server without ever releasing its +source code to the public. + + The GNU Affero General Public License is designed specifically to +ensure that, in such cases, the modified source code becomes available +to the community. It requires the operator of a network server to +provide the source code of the modified version running there to the +users of that server. Therefore, public use of a modified version, on +a publicly accessible server, gives the public access to the source +code of the modified version. + + An older license, called the Affero General Public License and +published by Affero, was designed to accomplish similar goals. This is +a different license, not a version of the Affero GPL, but Affero has +released a new version of the Affero GPL which permits relicensing under +this license. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU Affero General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Remote Network Interaction; Use with the GNU General Public License. + + Notwithstanding any other provision of this License, if you modify the +Program, your modified version must prominently offer all users +interacting with it remotely through a computer network (if your version +supports such interaction) an opportunity to receive the Corresponding +Source of your version by providing access to the Corresponding Source +from a network server at no charge, through some standard or customary +means of facilitating copying of software. This Corresponding Source +shall include the Corresponding Source for any work covered by version 3 +of the GNU General Public License that is incorporated pursuant to the +following paragraph. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the work with which it is combined will remain governed by version +3 of the GNU General Public License. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU Affero General Public License from time to time. Such new versions +will be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU Affero General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU Affero General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU Affero General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Affero General Public License for more details. + + You should have received a copy of the GNU Affero General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If your software can interact with users remotely through a computer +network, you should also make sure that it provides a way for users to +get its source. For example, if your program is a web application, its +interface could display a "Source" link that leads users to an archive +of the code. There are many ways you could offer source, and different +solutions will be better for different programs; see section 13 for the +specific requirements. + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU AGPL, see +. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c1c6aed6375a45197e7cf11b364a41858529d616 --- /dev/null +++ b/README.md @@ -0,0 +1,174 @@ +--- +title: Large language model topic modelling +emoji: 📚 +colorFrom: purple +colorTo: yellow +sdk: gradio +sdk_version: 6.19.0 +app_file: app.py +pinned: true +license: agpl-3.0 +short_description: Extract topics from open text data with LLMs +--- + +# Large language model topic modelling + +Version: 0.13.0 + +Extract topics and summarise outputs using Large Language Models (LLMs), either local, Gemini, Azure, or AWS Bedrock models (e.g. Claude, Nova models). The app will query the LLM with batches of responses to produce summary tables, which are then compared iteratively to output a table with the general topics, subtopics, topic sentiment, and a topic summary. Instructions on use can be found in the README.md file. You can try out examples by clicking on one of the example datasets on the main app page, which will show you example outputs from a local model run. API keys for AWS, Azure, and Gemini services can be entered on the settings page (note that Gemini has a free public API). + +NOTE: Large language models are not 100% accurate and may produce biased or harmful outputs. All outputs from this app **absolutely need to be checked by a human** to check for harmful outputs, hallucinations, and accuracy. + +Basic use: +1. On the front page, choose your model for inference. Gemma 3/GPT-OSS will use 'on-device' inference. Calls to Gemini or AWS will require an API key that can be input on the 'LLM and topic extraction' page. +1. Upload a csv/xlsx/parquet file containing at least one open text column. +2. Select the relevant open text column from the dropdown. +3. If you have your own suggested (zero shot) topics, upload this (see examples folder for an example file) +4. Write a one sentence description of the consultation/context of the open text. +5. Click 'Extract topics, deduplicate, and summarise'. This will run through the whole analysis process from topic extraction, to topic deduplication, to topic-level and overall summaries. +6. A summary xlsx file workbook will be created on the front page in the box 'Overall summary xlsx file'. This will combine all the results from the different processes into one workbook. + +# Installation guide + +Here is a step-by-step guide to clone the repository, create a virtual environment, and install dependencies from the relevant `requirements` file. This guide assumes you have **Git** and **Python 3.11** installed. + +----- + +### Step 1: Clone the Git Repository + +First, you need to copy the project files to your local machine. Navigate to the directory where you want to store the project using the `cd` (change directory) command. Then, use `git clone` with the repository's URL. + +1. **Clone the repo:** + + ```bash + git clone https://github.com/seanpedrick-case/llm_topic_modelling.git + ``` + +2. **Navigate into the new project folder:** + + ```bash + cd llm_topic_modelling + ``` +----- + +### Step 2: Create and Activate a Virtual Environment + +A virtual environment is a self-contained directory that holds a specific Python interpreter and its own set of installed packages. This is crucial for isolating your project's dependencies. + +NOTE: Alternatively you could also create and activate a Conda environment instead of using venv below. + +1. **Create the virtual environment:** We'll use Python's built-in `venv` module. It's common practice to name the environment folder `.venv`. + + ```bash + python -m venv .venv + ``` + + *This command tells Python to create a new virtual environment in a folder named `.venv`.* + +2. **Activate the environment:** You must "activate" the environment to start using it. The command differs based on your operating system and shell. + + * **On macOS / Linux (bash/zsh):** + + ```bash + source .venv/bin/activate + ``` + + * **On Windows (Command Prompt):** + + ```bash + .\.venv\Scripts\activate + ``` + + * **On Windows (PowerShell):** + + ```powershell + .\.venv\Scripts\Activate.ps1 + ``` + + You'll know it's active because your command prompt will be prefixed with `(.venv)`. + +----- + +### Step 3: Install Dependencies + +Now that your virtual environment is active, you can install all the required packages. Here you have two options, install from the pyproject.toml file (recommended), or install from requirements files. + +1. **Install from pyproject.toml (recommended)** + +You can install the 'lightweight' version of the app to access all available cloud provider or local inference (e.g. llama server, vLLM server) APIs. This version will not allow you to run local models such as Gemma 12b or GPT-OSS-20b 'in-app', i.e. accessible from the GUI interface directly. However, you will have access to AWS, Gemma, or Azure/OpenAI models with appropriate API keys. Use the following command in your environment to install the relevant packages: + +```bash +pip install . +``` + +#### Install torch (optional) + +If you want to run inference with transformers with full/quantised models, and the associated Unsloth package, you can run the following command for CPU inference. For GPU inference, please refer to the requirements_gpu.txt guide, and the 'Install from a requirements file' section below: + +```bash +pip install .[torch] +``` + +#### Install llama-cpp-python (optional) + +You can run quantised GGUF models in-app using llama-cpp-python. However, installation of this package is not always straightforward, particularly considering that wheels are not available for the latest version apart from for linux. This package is not being updated regularly, and so support may be removed for this package in future. Long term I would advise instead looking into running GGUF models using llama-server and calling the API from this app using the lightweight version (details here: https://github.com/ggml-org/llama.cpp). + +If you do want to install llama-cpp-python in app, first try the following command: + +```bash +pip install .[llamacpp] +``` + +This will install the CPU version of llama-cpp-python. If you want GPU support, first I would try using pip install with specific wheels for your system: See files in https://github.com/abetlen/llama-cpp-python/releases . If you are still struggling, see here for more details on installation here: https://llama-cpp-python.readthedocs.io/en/latest + +#### Install mcp version of gradio + +You can install an mcp-compatible version of gradio for this app with the following command: + +```bash +pip install .[mcp] +``` + +2. **Install from a requirements file (not recommended)** + +The repo provides several requirements files that are relevant for different situations. To start, I advise installing using the **requirements_lightweight.txt** file, which installs the app with access to all cloud provider or local inference (e.g. llama server, vLLM server) APIs. This approach is much simpler as a first step, and avoids issues with potentially complicated llama-cpp-python installation and GPU management described below. + +If you want to run models locally 'in app', then you have two further requirements files to choose from: + +- **requirements_cpu.txt**: Used for Python 3.11 CPU-only environments. Uncomment the requirements under 'Windows' for Windows compatibility. Make sure you have [Openblas](https://github.com/OpenMathLib/OpenBLAS) installed! +- **requirements_gpu.txt**: Used for Python 3.11 GPU-enabled environments. Uncomment the requirements under 'Windows' for Windows compatibility (CUDA 12.4). + +Example The below instructions will guide you in how to install the GPU-enabled version of the app for local inference. + +**Install packages for local model 'in-app' inference from the requirements file:** + ```bash + pip install -r requirements_gpu.txt + ``` + *This command reads every package name listed in the file and installs it into your `.venv` environment.* + +NOTE: If default llama-cpp-python installation does not work when installing from the above, go into the requirements_gpu.txt file and uncomment the lines to install a wheel for llama-cpp-python 0.3.16 relevant to your system. + +### Step 4: Verify CUDA compatibility (if using a GPU environment) + +Install the relevant toolkit for CUDA from here: https://developer.nvidia.com/ + +Restart your computer + +Ensure you have the latest drivers for your NVIDIA GPU. Check your current version and memory availability by running nvidia-smi + +In command line, CUDA compatibility can be checked by running nvcc --version + + +### Step 5: Ensure you have compatible NVIDIA drivers + +Make sure you have the latest NVIDIA drivers installed on your system for your GPU (be careful in particular if using WSL that you have drivers compatible with this). Official drivers can be found here: https://www.nvidia.com/en-us/drivers + +Current drivers can be found by running nvidia-smi in command line + +### Step 6: Run the app + +Go to the app project directory. Run python app.py + +### Step 7: (optional) change default configuration + +A number of configuration options can be seen the tools/config.py file. You can either pass in these variables as environment variables, or you can create a file in config/app_config.env to read this into the app on initialisation. diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..65171adb5fd0d8102e4619c40bd313aa0bce0f3f --- /dev/null +++ b/app.py @@ -0,0 +1,3224 @@ +import os +from datetime import datetime + +import gradio as gr +import pandas as pd + +from tools.auth import authenticate_user +from tools.aws_functions import ( + download_cost_codes_with_error_handling, + export_outputs_to_s3_with_prompt_response_logs, + upload_file_to_s3, +) +from tools.combine_sheets_into_xlsx import collect_output_csvs_and_create_excel_output +from tools.config import ( + ACCESS_LOG_DYNAMODB_TABLE_NAME, + ACCESS_LOGS_FOLDER, + API_URL, + AWS_ACCESS_KEY, + AWS_REGION, + AWS_SECRET_KEY, + AZURE_OPENAI_API_KEY, + AZURE_OPENAI_INFERENCE_ENDPOINT, + BATCH_SIZE_DEFAULT, + COGNITO_AUTH, + CONFIG_FOLDER, + COST_CODES_PATH, + CSV_ACCESS_LOG_HEADERS, + CSV_FEEDBACK_LOG_HEADERS, + CSV_USAGE_LOG_HEADERS, + DEFAULT_COST_CODE, + DIRECT_MODE_ADDITIONAL_SUMMARY_INSTRUCTIONS, + DIRECT_MODE_ADDITIONAL_VALIDATION_ISSUES, + DIRECT_MODE_BATCH_SIZE, + DIRECT_MODE_CANDIDATE_TOPICS, + DIRECT_MODE_CONTEXT, + DIRECT_MODE_CREATE_TOPICS_CSV, + DIRECT_MODE_CREATE_XLSX_OUTPUT, + DIRECT_MODE_DEDUP_METHOD, + DIRECT_MODE_DEFAULT_COST_CODE, + DIRECT_MODE_EXCEL_SHEETS, + DIRECT_MODE_FORCE_SINGLE_TOPIC, + DIRECT_MODE_FORCE_ZERO_SHOT, + DIRECT_MODE_GROUP_BY, + DIRECT_MODE_INFERENCE_SERVER_MODEL, + DIRECT_MODE_INPUT_FILE, + DIRECT_MODE_MAX_TIME_FOR_LOOP, + DIRECT_MODE_MAX_TOKENS, + DIRECT_MODE_MERGE_GENERAL_TOPICS, + DIRECT_MODE_MERGE_SENTIMENT, + DIRECT_MODE_MODEL_CHOICE, + DIRECT_MODE_NO_OF_SAMPLED_SUMMARIES, + DIRECT_MODE_OUTPUT_DIR, + DIRECT_MODE_PREVIOUS_OUTPUT_FILES, + DIRECT_MODE_PRODUCE_STRUCTURED_SUMMARY, + DIRECT_MODE_RANDOM_SEED, + DIRECT_MODE_SAMPLE_FRACTION, + DIRECT_MODE_SAMPLE_REFERENCE_TABLE, + DIRECT_MODE_SENTIMENT, + DIRECT_MODE_SHOW_PREVIOUS_TABLE, + DIRECT_MODE_SIMILARITY_THRESHOLD, + DIRECT_MODE_SUMMARY_FORMAT, + # Direct mode variables + DIRECT_MODE_TASK, + DIRECT_MODE_TEMPERATURE, + DIRECT_MODE_TEXT_COLUMN, + DIRECT_MODE_USERNAME, + DYNAMODB_ACCESS_LOG_HEADERS, + DYNAMODB_FEEDBACK_LOG_HEADERS, + DYNAMODB_USAGE_LOG_HEADERS, + ENFORCE_COST_CODES, + FEEDBACK_LOG_DYNAMODB_TABLE_NAME, + FEEDBACK_LOG_FILE_NAME, + FEEDBACK_LOGS_FOLDER, + FILE_INPUT_HEIGHT, + GEMINI_API_KEY, + GET_COST_CODES, + GRADIO_SERVER_PORT, + GRADIO_TEMP_DIR, + HF_TOKEN, + HOST_NAME, + INPUT_FOLDER, + INTRO_TEXT, + LLM_SEED, + LLM_TEMPERATURE, + LOG_FILE_NAME, + MAX_FILE_SIZE, + MAX_QUEUE_SIZE, + MAXIMUM_ALLOWED_TOPICS, + MPLCONFIGDIR, + OUTPUT_COST_CODES_PATH, + OUTPUT_DEBUG_FILES, + OUTPUT_FOLDER, + ROOT_PATH, + RUN_AWS_FUNCTIONS, + RUN_DIRECT_MODE, + RUN_INFERENCE_SERVER, + RUN_MCP_SERVER, + S3_ACCESS_LOGS_FOLDER, + S3_COST_CODES_PATH, + S3_FEEDBACK_LOGS_FOLDER, + S3_LOG_BUCKET, + S3_OUTPUTS_FOLDER, + S3_USAGE_LOGS_FOLDER, + SAVE_LOGS_TO_CSV, + SAVE_LOGS_TO_DYNAMODB, + SAVE_OUTPUTS_TO_S3, + SESSION_OUTPUT_FOLDER, + SHOW_ADDITIONAL_INSTRUCTION_TEXTBOXES, + SHOW_COSTS, + SHOW_EXAMPLES, + USAGE_LOG_DYNAMODB_TABLE_NAME, + USAGE_LOG_FILE_NAME, + USAGE_LOGS_FOLDER, + convert_string_to_boolean, + default_model_choice, + default_model_source, + default_source_models, + ensure_folder_exists, + model_name_map, + model_sources, +) +from tools.custom_csvlogger import CSVLogger_custom +from tools.dedup_summaries import ( + deduplicate_topics, + overall_summary, + wrapper_summarise_output_topics_per_group, +) +from tools.example_table_outputs import ( + case_notes_table, + case_notes_table_grouped, + case_notes_table_structured_summary, + dummy_consultation_table, + dummy_consultation_table_zero_shot, +) +from tools.helper_functions import ( + df_select_callback_cost, + empty_output_vars_extract_topics, + empty_output_vars_summarise, + enforce_cost_codes, + get_connection_params, + join_cols_onto_reference_df, + load_in_data_file, + load_in_default_cost_codes, + load_in_previous_data_files, + load_in_previous_reference_file, + move_overall_summary_output_files_to_front_page, + put_columns_in_df, + reset_base_dataframe, + update_cost_code_dataframe_from_dropdown_select, + update_model_choice, + view_table, +) +from tools.llm_api_call import ( + all_in_one_pipeline, + deduplicate_topics_llm_wrapper, + discover_topics_from_sample_wrapper, + modify_existing_output_tables, + validate_topics_wrapper, + wrapper_extract_topics_per_column_value, +) +from tools.prompts import ( + add_existing_topics_prompt, + add_existing_topics_system_prompt, + initial_table_prompt, + single_para_summary_format_prompt, + system_prompt, + two_para_summary_format_prompt, +) +from tools.view_logs import ( + load_log_file_handler, + update_log_display_on_filter, +) + +ensure_folder_exists(CONFIG_FOLDER) +ensure_folder_exists(OUTPUT_FOLDER) +ensure_folder_exists(INPUT_FOLDER) +ensure_folder_exists(GRADIO_TEMP_DIR) +ensure_folder_exists(MPLCONFIGDIR) +ensure_folder_exists(FEEDBACK_LOGS_FOLDER) +ensure_folder_exists(ACCESS_LOGS_FOLDER) +ensure_folder_exists(USAGE_LOGS_FOLDER) + +today_rev = datetime.now().strftime("%Y%m%d") + +# Placeholders for example variables +in_data_files = gr.File( + height=FILE_INPUT_HEIGHT, + label="Choose Excel or csv files", + file_count="multiple", + file_types=[".xlsx", ".xls", ".csv", ".parquet"], +) +in_colnames = gr.Dropdown( + choices=[""], + multiselect=False, + label="Select the open text column of interest. In an Excel file, this shows columns across all sheets.", + allow_custom_value=True, + interactive=True, +) +if SHOW_ADDITIONAL_INSTRUCTION_TEXTBOXES == "True": + context_textbox = gr.Textbox( + label="Write up to one sentence giving context to the large language model for your task (e.g. 'Consultation for the construction of flats on Main Street')", + visible=True, + ) +else: + context_textbox = gr.Textbox( + label="Write up to one sentence giving context to the large language model for your task (e.g. 'Consultation for the construction of flats on Main Street')", + visible=False, + ) +topic_extraction_output_files_xlsx = gr.File( + label="Overall summary xlsx file and suggested topics CSV (when enabled). CSV outputs are available on the 'Advanced' tab.", + scale=1, + interactive=False, + file_count="multiple", +) +display_topic_table_markdown = gr.Markdown(value="", buttons=["copy"]) +output_messages_textbox = gr.Textbox( + value="", label="Output messages", scale=1, interactive=False, lines=4 +) +candidate_topics = gr.File( + height=FILE_INPUT_HEIGHT, + label="Input topics from file (csv). File should have at least one column with a header, and all topic names below this. Using the headers 'General topic' and/or 'Subtopic' will allow for these columns to be suggested to the model. If a third column is present, it will be assumed to be a topic description.", + file_count="single", +) +create_topics_csv_radio = gr.Radio( + label="Create suggested topics CSV from analysis results (unique General topic and Subtopic pairs, for reuse in future analyses)", + value="Yes", + choices=["Yes", "No"], +) +topic_discovery_sample_fraction = gr.Number( + label="Sample percentage for topic discovery (1–100)", + value=20, + precision=0, + minimum=1, + maximum=100, +) +topic_discovery_random_seed = gr.Number( + label="Random seed for topic discovery sample", + value=LLM_SEED, + precision=0, +) +discover_topics_output_file = gr.File( + label="Suggested topics CSV from sample discovery", + scale=1, + interactive=False, + file_count="single", +) +produce_structured_summary_radio = gr.Radio( + label="Ask the model to produce structured summaries using the suggested topics as headers rather than extract topics", + value="No", + choices=["Yes", "No"], +) +in_group_col = gr.Dropdown( + multiselect=False, + label="Select the column to group results by", + allow_custom_value=True, + interactive=True, + visible=True, +) +batch_size_number = gr.Number( + label="Number of responses to submit in a single LLM query (batch size)", + value=BATCH_SIZE_DEFAULT, + precision=0, + minimum=1, + maximum=50, +) + +css = """ +/* Target tab navigation buttons only - not buttons inside tab content */ +/* Gradio renders tab buttons with role="tab" in the navigation area */ +button[role="tab"] { + font-size: 1.1em !important; + padding: 0.75em 1.2em !important; +} + +/* Alternative selectors for different Gradio versions */ +.tab-nav button, +nav button[role="tab"], +div[class*="tab-nav"] button { + font-size: 1.1em !important; + padding: 0.75em 1.2em !important; +} +""" + +# Create the gradio interface +app = gr.Blocks( + fill_width=False, + analytics_enabled=False, + title="LLM topic modelling", + delete_cache=(43200, 43200), +) + +with app: + + ### + # STATE VARIABLES + ### + + # Workaround for Gradio 6 issue where 'hidden' element are still sometimes visible as a thing line in the UI + with gr.Accordion(visible=False, elem_classes="hidden_component", open=False): + text_output_file_list_state = gr.Dropdown( + list(), + allow_custom_value=True, + visible=False, + label="text_output_file_list_state", + elem_classes="hidden_component", + ) + text_output_modify_file_list_state = gr.Dropdown( + list(), + allow_custom_value=True, + visible=False, + label="text_output_modify_file_list_state", + elem_classes="hidden_component", + ) + log_files_output_list_state = gr.Dropdown( + list(), + allow_custom_value=True, + visible=False, + label="log_files_output_list_state", + elem_classes="hidden_component", + ) + first_loop_state = gr.Checkbox( + True, visible=False, elem_classes="hidden_component" + ) + second_loop_state = gr.Checkbox( + False, visible=False, elem_classes="hidden_component" + ) + modified_unique_table_change_bool = gr.Checkbox( + True, visible=False, elem_classes="hidden_component" + ) # This boolean is used to flag whether a file upload should change just the modified unique table object on the second tab + + file_data_state = gr.Dataframe( + value=pd.DataFrame(), + label="file_data_state", + visible=False, + type="pandas", + interactive=True, + elem_classes="hidden_component", + ) + master_topic_df_state = gr.Dataframe( + value=pd.DataFrame(), + label="master_topic_df_state", + visible=False, + type="pandas", + interactive=True, + ) + master_unique_topics_df_state = gr.Dataframe( + value=pd.DataFrame(), + label="master_unique_topics_df_state", + visible=False, + type="pandas", + interactive=True, + elem_classes="hidden_component", + ) + master_reference_df_state = gr.Dataframe( + value=pd.DataFrame(), + label="master_reference_df_state", + visible=False, + type="pandas", + interactive=True, + elem_classes="hidden_component", + ) + missing_df_state = gr.Dataframe( + value=pd.DataFrame(), + label="missing_df_state", + visible=False, + type="pandas", + interactive=True, + elem_classes="hidden_component", + ) + + master_modify_unique_topics_df_state = gr.Dataframe( + value=pd.DataFrame(), + label="master_modify_unique_topics_df_state", + visible=False, + type="pandas", + interactive=True, + elem_classes="hidden_component", + ) + master_modify_reference_df_state = gr.Dataframe( + value=pd.DataFrame(), + label="master_modify_reference_df_state", + visible=False, + type="pandas", + interactive=True, + elem_classes="hidden_component", + ) + + # Blank placeholder for conversation metadata textbox, as logging file output can get too long for large amounts of calls + conversation_metadata_textbox_placeholder = gr.Textbox( + value="", + label="Query metadata - usage counts and other parameters", + lines=8, + visible=False, + elem_classes="hidden_component", + ) + + session_hash_state = gr.Textbox(visible=False, value=HOST_NAME) + output_folder_state = gr.Textbox( + visible=False, value=OUTPUT_FOLDER, elem_classes="hidden_component" + ) + input_folder_state = gr.Textbox( + visible=False, value=INPUT_FOLDER, elem_classes="hidden_component" + ) + + # s3 bucket name + s3_default_bucket = gr.Textbox( + label="Default S3 bucket", + value=S3_LOG_BUCKET, + visible=False, + elem_classes="hidden_component", + ) + s3_log_bucket_name = gr.Textbox( + visible=False, value=S3_LOG_BUCKET, elem_classes="hidden_component" + ) + + # S3 output settings + s3_output_folder_state = gr.Textbox( + label="s3_output_folder_state", + value=S3_OUTPUTS_FOLDER, + visible=False, + elem_classes="hidden_component", + ) + save_outputs_to_s3_checkbox = gr.Checkbox( + label="save_outputs_to_s3_checkbox", + value=convert_string_to_boolean(SAVE_OUTPUTS_TO_S3), + visible=False, + elem_classes="hidden_component", + ) + + # Logging variables + access_logs_state = gr.Textbox( + label="access_logs_state", + value=ACCESS_LOGS_FOLDER + LOG_FILE_NAME, + visible=False, + elem_classes="hidden_component", + ) + access_s3_logs_loc_state = gr.Textbox( + label="access_s3_logs_loc_state", + value=S3_ACCESS_LOGS_FOLDER, + visible=False, + elem_classes="hidden_component", + ) + feedback_logs_state = gr.Textbox( + label="feedback_logs_state", + value=FEEDBACK_LOGS_FOLDER + FEEDBACK_LOG_FILE_NAME, + visible=False, + elem_classes="hidden_component", + ) + feedback_s3_logs_loc_state = gr.Textbox( + label="feedback_s3_logs_loc_state", + value=S3_FEEDBACK_LOGS_FOLDER, + visible=False, + elem_classes="hidden_component", + ) + usage_logs_state = gr.Textbox( + label="usage_logs_state", + value=USAGE_LOGS_FOLDER + USAGE_LOG_FILE_NAME, + visible=False, + elem_classes="hidden_component", + ) + usage_s3_logs_loc_state = gr.Textbox( + label="usage_s3_logs_loc_state", + value=S3_USAGE_LOGS_FOLDER, + visible=False, + elem_classes="hidden_component", + ) + + # Logging for logged content + logged_content_df = gr.Dataframe( + label="logged_content_df", + value=pd.DataFrame(), + visible=False, + type="pandas", + elem_classes="hidden_component", + ) + + # Logging for input / output tokens + input_tokens_num = gr.Textbox( + "0", + visible=False, + label="Total input tokens", + elem_classes="hidden_component", + ) + output_tokens_num = gr.Textbox( + "0", + visible=False, + label="Total output tokens", + elem_classes="hidden_component", + ) + number_of_calls_num = gr.Textbox( + "0", + visible=False, + label="Total LLM calls", + elem_classes="hidden_component", + ) + + # Additional UI components for validation + max_tokens_num = gr.Number( + value=8192, + visible=False, + label="Max tokens", + elem_classes="hidden_component", + ) + reasoning_suffix_textbox = gr.Textbox( + value="", + visible=False, + label="Reasoning suffix", + elem_classes="hidden_component", + ) + output_debug_files_radio = gr.Radio( + value="False", + choices=["True", "False"], + visible=False, + label="Output debug files", + ) + max_time_for_loop_num = gr.Number( + value=99999, + visible=False, + label="Max time for loop", + elem_classes="hidden_component", + ) + + # Summary state objects + summary_reference_table_sample_state = gr.Dataframe( + value=pd.DataFrame(), + headers=None, + column_count=None, + label="summary_reference_table_sample_state", + visible=False, + type="pandas", + elem_classes="hidden_component", + ) + master_reference_df_revised_summaries_state = gr.Dataframe( + value=pd.DataFrame(), + headers=None, + column_count=None, + label="master_reference_df_revised_summaries_state", + visible=False, + type="pandas", + elem_classes="hidden_component", + ) + master_unique_topics_df_revised_summaries_state = gr.Dataframe( + value=pd.DataFrame(), + headers=None, + column_count=None, + label="master_unique_topics_df_revised_summaries_state", + visible=False, + type="pandas", + elem_classes="hidden_component", + ) + summarised_output_df = gr.Dataframe( + value=pd.DataFrame(), + headers=None, + column_count=None, + label="summarised_output_df", + visible=False, + type="pandas", + elem_classes="hidden_component", + ) + summarised_references_markdown = gr.Markdown( + "", visible=False, elem_classes="hidden_component" + ) + summarised_outputs_list = gr.Dropdown( + value=list(), + choices=list(), + visible=False, + label="List of summarised outputs", + allow_custom_value=True, + elem_classes="hidden_component", + ) + latest_summary_completed_num = gr.Number( + 0, visible=False, elem_classes="hidden_component" + ) + + summary_xlsx_output_files_list = gr.Dropdown( + value=list(), + choices=list(), + visible=False, + label="List of xlsx summary output files", + allow_custom_value=True, + elem_classes="hidden_component", + ) + + original_data_file_name_textbox = gr.Textbox( + label="Response ID data file name", + value="", + visible=False, + elem_classes="hidden_component", + ) + working_data_file_name_textbox = gr.Textbox( + label="Working data file name", + value="", + visible=False, + elem_classes="hidden_component", + ) + unique_topics_table_file_name_textbox = gr.Textbox( + label="Unique topics data file name textbox", + visible=False, + elem_classes="hidden_component", + ) + + dummy_consultation_table_textbox = gr.Textbox( + value=dummy_consultation_table, + visible=False, + label="Dummy consultation table", + elem_classes="hidden_component", + ) + case_notes_table_textbox = gr.Textbox( + value=case_notes_table, + visible=False, + label="Case notes table", + elem_classes="hidden_component", + ) + + model_name_map_state = gr.JSON( + model_name_map, + visible=False, + label="model_name_map_state", + elem_classes="hidden_component", + ) + + # Cost code elements + s3_default_cost_codes_file = gr.Textbox( + label="Default cost centre file", + value=S3_COST_CODES_PATH, + visible=False, + elem_classes="hidden_component", + ) + default_cost_codes_output_folder_location = gr.Textbox( + label="Output default cost centre location", + value=OUTPUT_COST_CODES_PATH, + visible=False, + elem_classes="hidden_component", + ) + enforce_cost_code_textbox = gr.Textbox( + label="Enforce cost code textbox", + value=ENFORCE_COST_CODES, + visible=False, + elem_classes="hidden_component", + ) + default_cost_code_textbox = gr.Textbox( + label="Default cost code textbox", + value=DEFAULT_COST_CODE, + visible=False, + elem_classes="hidden_component", + ) + + # Placeholders for elements that may be made visible later below depending on environment variables + cost_code_dataframe_base = gr.Dataframe( + value=pd.DataFrame(columns=["Cost code", "Description"]), + label="Cost codes", + type="pandas", + buttons=["fullscreen", "copy"], + show_search="filter", + wrap=True, + max_height=200, + visible=False, + interactive=True, + ) + cost_code_dataframe = gr.Dataframe( + value=pd.DataFrame(columns=["Cost code", "Description"]), + type="pandas", + visible=False, + wrap=True, + interactive=True, + elem_classes="hidden_component", + ) + cost_code_choice_drop = gr.Dropdown( + value=DEFAULT_COST_CODE, + label="Choose cost code for analysis. Please contact Finance if you can't find your cost code in the given list.", + choices=[DEFAULT_COST_CODE], + allow_custom_value=False, + visible=False, + elem_classes="hidden_component", + ) + + latest_batch_completed = gr.Number( + value=0, + label="Number of files prepared", + interactive=False, + visible=False, + elem_classes="hidden_component", + ) + # Duplicate version of the above variable for when you don't want to initiate the summarisation loop + latest_batch_completed_no_loop = gr.Number( + value=0, + label="Number of files prepared", + interactive=False, + visible=False, + elem_classes="hidden_component", + ) + + # Invisible text box to hold the session hash/username just for logging purposes + session_hash_textbox = gr.Textbox( + label="Session hash", + value="", + visible=False, + elem_classes="hidden_component", + ) + + estimated_time_taken_number = gr.Number( + label="Estimated time taken (seconds)", + value=0.0, + precision=1, + visible=False, + elem_classes="hidden_component", + ) # This keeps track of the time taken to redact files for logging purposes. + total_number_of_batches = gr.Number( + label="Current batch number", + value=1, + precision=0, + visible=False, + elem_classes="hidden_component", + ) + + text_output_logs = gr.Textbox( + label="Output summary logs", + visible=False, + elem_classes="hidden_component", + ) + + # State to store loaded log data + log_data_state = gr.State(value=[]) + + ### + # UI LAYOUT + ### + + gr.Markdown(INTRO_TEXT) + + if SHOW_EXAMPLES == "True": + + def show_info_box_on_click( + in_data_files, + in_colnames, + context_textbox, + original_data_file_name_textbox, + topic_extraction_output_files_xlsx, + display_topic_table_markdown, + output_messages_textbox, + candidate_topics, + produce_structured_summary_radio, + in_group_col, + batch_size_number, + ): + gr.Info( + "Example data loaded. Now click on the 'Extract topics...' button below to run the full suite of topic extraction, deduplication, and summarisation." + ) + + # Check if required example files exist before creating Examples + # This prevents errors in CI environments where example files may not be present + required_example_files = [ + "example_data/dummy_consultation_response.csv", + "example_data/combined_case_notes.csv", + ] + example_files_exist = all(os.path.exists(f) for f in required_example_files) + + # Only create Examples if files exist, otherwise create empty Examples to avoid errors + if example_files_exist: + try: + examples = gr.Examples( + examples=[ + [ + ["example_data/dummy_consultation_response.csv"], + "Response text", + "Consultation for the construction of flats on Main Street", + "dummy_consultation_response.csv", + [ + "example_data/dummy_consultation_r_col_Response_text_Gemma_3_4B_topic_analysis.xlsx" + ], + dummy_consultation_table, + "Example output from the dummy consultation dataset successfully loaded. Download the xlsx outputs to the right to see full outputs.", + None, + "No", + None, + 5, + ], + [ + ["example_data/combined_case_notes.csv"], + "Case Note", + "Social Care case notes for young people", + "combined_case_notes.csv", + [ + "example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_topic_analysis.xlsx" + ], + case_notes_table, + "Example output from the case notes dataset successfully loaded. Download the xlsx outputs to the right to see full outputs.", + None, + "No", + None, + 5, + ], + [ + ["example_data/dummy_consultation_response.csv"], + "Response text", + "Consultation for the construction of flats on Main Street", + "dummy_consultation_response.csv", + [ + "example_data/dummy_consultation_r_col_Response_text_Gemma_3_4B_topic_analysis_zero_shot.xlsx" + ], + dummy_consultation_table_zero_shot, + "Example output from the dummy consultation dataset with suggested topics successfully loaded. Download the xlsx outputs to the right to see full outputs.", + "example_data/dummy_consultation_response_themes.csv", + "No", + None, + 5, + ], + [ + ["example_data/combined_case_notes.csv"], + "Case Note", + "Social Care case notes for young people", + "combined_case_notes.csv", + [ + "example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_topic_analysis_grouped.xlsx" + ], + case_notes_table_grouped, + "Example data from the case notes dataset with groups successfully loaded. Download the xlsx outputs to the right to see full outputs.", + "example_data/case_note_headers_specific.csv", + "No", + "Client", + 5, + ], + [ + ["example_data/combined_case_notes.csv"], + "Case Note", + "Social Care case notes for young people", + "combined_case_notes.csv", + [ + "example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_structured_summaries.xlsx" + ], + case_notes_table_structured_summary, + "Example data from the case notes dataset for structured summaries successfully loaded. Download the xlsx outputs to the right to see full outputs.", + "example_data/case_note_headers_specific.csv", + "Yes", + "Client", + 50, + ], + ], + inputs=[ + in_data_files, + in_colnames, + context_textbox, + original_data_file_name_textbox, + topic_extraction_output_files_xlsx, + display_topic_table_markdown, + output_messages_textbox, + candidate_topics, + produce_structured_summary_radio, + in_group_col, + batch_size_number, + ], + example_labels=[ + "Main Street construction consultation", + "Case notes for young people", + "Main Street construction consultation with suggested topics", + "Case notes grouped by person with suggested topics", + "Case notes structured summary with suggested topics", + ], + label="Try topic extraction and summarisation with an example dataset. Example outputs are displayed. Click the 'Extract topics...' button below to rerun the analysis.", + fn=show_info_box_on_click, + run_on_click=True, + ) + except (FileNotFoundError, OSError) as e: + # If example files don't exist (e.g., in CI environment), create empty Examples + # This allows the app to load without errors + print( + f"Warning: Example files not found, skipping Examples creation: {e}" + ) + examples = gr.Examples( + examples=[], + inputs=[ + in_data_files, + in_colnames, + context_textbox, + original_data_file_name_textbox, + topic_extraction_output_files_xlsx, + display_topic_table_markdown, + output_messages_textbox, + candidate_topics, + produce_structured_summary_radio, + in_group_col, + batch_size_number, + ], + label="Examples not available (example files not found).", + ) + else: + # Example files don't exist, create empty Examples + print( + "Warning: Required example files not found, skipping Examples creation." + ) + examples = gr.Examples( + examples=[], + inputs=[ + in_data_files, + in_colnames, + context_textbox, + original_data_file_name_textbox, + topic_extraction_output_files_xlsx, + display_topic_table_markdown, + output_messages_textbox, + candidate_topics, + produce_structured_summary_radio, + in_group_col, + batch_size_number, + ], + label="Examples not available (example files not found).", + ) + + with gr.Tab(label="All in one topic extraction and summarisation"): + with gr.Row(): + model_source = gr.Dropdown( + value=default_model_source, + choices=model_sources, + label="Large language model family", + multiselect=False, + ) + model_choice = gr.Dropdown( + value=default_model_choice, + choices=default_source_models, + label="Large language model for topic extraction and summarisation", + multiselect=False, + allow_custom_value=True, + ) + + model_source.change( + fn=update_model_choice, + inputs=[model_source], + outputs=[model_choice], + api_visibility="undocumented", + ) + + with gr.Accordion("Upload xlsx, csv, or parquet file", open=True): + in_data_files.render() + + in_excel_sheets = gr.Dropdown( + multiselect=False, + label="Select the Excel sheet of interest", + visible=False, + allow_custom_value=True, + ) + in_colnames.render() + + with gr.Accordion("Group analysis by values in another column", open=False): + in_group_col.render() + + with gr.Accordion("Provide list of suggested topics", open=False): + candidate_topics.render() + create_topics_csv_radio.render() + with gr.Accordion("Discover topics from a random sample", open=False): + gr.Markdown( + "Run topic extraction on a random subsample to build a suggested topics CSV. " + "Upload that file above and run a full extraction on all responses with suggested topics." + ) + with gr.Row(equal_height=True): + topic_discovery_sample_fraction.render() + topic_discovery_random_seed.render() + discover_topics_btn = gr.Button( + "Discover topics from sample", variant="secondary" + ) + discover_topics_output_file.render() + + with gr.Row(equal_height=True): + force_zero_shot_radio = gr.Radio( + label="Force responses into suggested topics", + value="No", + choices=["Yes", "No"], + ) + force_single_topic_radio = gr.Radio( + label="Ask the model to assign responses to only a single topic", + value="No", + choices=["Yes", "No"], + ) + produce_structured_summary_radio.render() + + with gr.Accordion( + "Response sentiment analysis (default is Negative or Positive)", open=False + ): + sentiment_checkbox = gr.Radio( + label="Should the model assess the sentiment of responses?", + value="Negative or Positive", + choices=[ + "Negative or Positive", + "Negative, Neutral, or Positive", + "Do not assess sentiment", + ], + ) + + if GET_COST_CODES == "True" or ENFORCE_COST_CODES == "True": + with gr.Accordion("Assign task to cost code", open=True, visible=True): + gr.Markdown( + "Please ensure that you have approval from your budget holder before using this app for summarisation tasks that incur a cost." + ) + with gr.Row(equal_height=True): + with gr.Column(): + with gr.Accordion("Cost code table", open=False, visible=True): + cost_code_dataframe = gr.Dataframe( + value=pd.DataFrame(), + label="Existing cost codes", + type="pandas", + interactive=True, + buttons=["fullscreen", "copy"], + show_search="filter", + visible=True, + wrap=True, + max_height=200, + ) + reset_cost_code_dataframe_button = gr.Button( + value="Reset code code table filter" + ) + with gr.Column(): + cost_code_choice_drop = gr.Dropdown( + value=DEFAULT_COST_CODE, + label="Choose cost code for analysis", + choices=[DEFAULT_COST_CODE], + allow_custom_value=False, + visible=True, + ) + + all_in_one_btn = gr.Button( + "Extract topics, deduplicate, and summarise", variant="primary" + ) + + with gr.Row(equal_height=True): + output_messages_textbox.render() + + topic_extraction_output_files_xlsx.render() + + display_topic_table_markdown.render() + + data_feedback_title = gr.Markdown( + value="## Please give feedback", visible=False + ) + data_feedback_radio = gr.Radio( + label="Please give some feedback about the results of the topic extraction.", + choices=["The results were good", "The results were not good"], + visible=False, + ) + data_further_details_text = gr.Textbox( + label="Please give more detailed feedback about the results:", + visible=False, + ) + data_submit_feedback_btn = gr.Button(value="Submit feedback", visible=False) + + with gr.Row(): + s3_logs_output_textbox = gr.Textbox( + label="Feedback submission logs", visible=False + ) + + with gr.Tab(label="Advanced - Step by step topic extraction and summarisation"): + + with gr.Accordion( + "1. Extract topics - go to first tab for file upload, model choice, and other settings before clicking this button", + open=True, + ): + context_textbox.render() + if SHOW_ADDITIONAL_INSTRUCTION_TEXTBOXES == "True": + additional_summary_instructions_textbox = gr.Textbox( + value="", visible=True, label="Additional summary instructions" + ) + else: + additional_summary_instructions_textbox = gr.Textbox( + value="", visible=False, label="Additional summary instructions" + ) + + extract_topics_btn = gr.Button("1. Extract topics", variant="secondary") + topic_extraction_output_files = gr.File( + label="Extract topics output files", + scale=1, + interactive=True, + height=FILE_INPUT_HEIGHT, + file_count="multiple", + ) + + with gr.Accordion( + "1b. Validate topics - validate previous results with an LLM", open=False + ): + if SHOW_ADDITIONAL_INSTRUCTION_TEXTBOXES == "True": + with gr.Row(): + show_previous_table_radio = gr.Radio( + label="Provide response data to validation process", + value="Yes", + choices=["Yes", "No"], + visible=True, + scale=1, + ) + additional_validation_issues_textbox = gr.Textbox( + value="", + visible=True, + label="Additional validation issues for the model to consider (bullet-point list)", + scale=3, + ) + else: + with gr.Row(): + show_previous_table_radio = gr.Radio( + label="Provide response data to validation process", + value="Yes", + choices=["Yes", "No"], + visible=False, + scale=1, + ) + additional_validation_issues_textbox = gr.Textbox( + value="", + visible=False, + label="Additional validation issues for the model to consider (bullet-point list)", + scale=3, + ) + validate_topics_btn = gr.Button("1b. Validate topics", variant="secondary") + validation_output_files = gr.File( + label="Validation output files", + scale=1, + interactive=False, + height=FILE_INPUT_HEIGHT, + ) + + with gr.Accordion("2. Modify topics from topic extraction", open=False): + gr.Markdown( + """Load in previously completed Extract Topics output files ('reference_table', and 'unique_topics' files) to modify topics, deduplicate topics, or summarise the outputs. If you want pivot table outputs, please load in the original data file along with the selected open text column on the first tab before deduplicating or summarising.""" + ) + + modification_input_files = gr.File( + height=FILE_INPUT_HEIGHT, + label="Upload reference and unique topic files to modify topics", + file_count="multiple", + file_types=[".xlsx", ".xls", ".csv", ".parquet"], + ) + + modifiable_unique_topics_df_state = gr.Dataframe( + value=pd.DataFrame(), + headers=None, + column_count=(4, "fixed"), + row_count=(1, "fixed"), + visible=True, + type="pandas", + ) + + save_modified_files_button = gr.Button(value="Save modified topic names") + + with gr.Accordion( + "3. Deduplicate topics using fuzzy matching or LLMs", open=False + ): + ### DEDUPLICATION + deduplication_input_files = gr.File( + height=FILE_INPUT_HEIGHT, + label="Upload reference and unique topic files to deduplicate topics. Optionally upload suggested topics on the first tab to match to these where possible with LLM deduplication", + file_count="multiple", + file_types=[".xlsx", ".xls", ".csv", ".parquet"], + ) + deduplication_input_files_status = gr.Textbox( + value="", label="Previous file input", visible=False + ) + + with gr.Row(): + merge_general_topics_drop = gr.Dropdown( + label="Merge general topic values together for duplicate subtopics.", + value="Yes", + choices=["Yes", "No"], + ) + merge_sentiment_drop = gr.Dropdown( + label="Merge sentiment values together for duplicate subtopics.", + value="No", + choices=["Yes", "No"], + ) + deduplicate_score_threshold = gr.Number( + label="Similarity threshold with which to determine duplicates.", + value=90, + minimum=5, + maximum=100, + precision=0, + ) + + with gr.Row(): + deduplicate_previous_data_btn = gr.Button( + "3. Deduplicate topics (Fuzzy matching)", variant="primary" + ) + deduplicate_llm_previous_data_btn = gr.Button( + "3b. Deduplicate topics (LLM semantic)", variant="secondary" + ) + + with gr.Accordion("4. Summarise topics", open=False): + ### SUMMARISATION + summarisation_input_files = gr.File( + height=FILE_INPUT_HEIGHT, + label="Upload reference and unique topic files to summarise", + file_count="multiple", + file_types=[".xlsx", ".xls", ".csv", ".parquet"], + ) + + summarise_format_radio = gr.Radio( + label="Choose summary type (Note: this will also use the custom summary instructions from step 1 above if provided)", + value=two_para_summary_format_prompt, + choices=[ + two_para_summary_format_prompt, + single_para_summary_format_prompt, + ], + ) + + with gr.Row(): + sample_reference_table_checkbox = gr.Checkbox( + value=True, + label="Sample reference table (recommended for large datasets)", + ) + no_of_sampled_summaries_number = gr.Number( + value=150, + label="Number of summaries per group", + precision=0, + minimum=10, + maximum=500, + ) + random_seed_number = gr.Number( + value=42, label="Random seed", precision=0, minimum=1, maximum=9999 + ) + + summarise_previous_data_btn = gr.Button( + "4. Summarise topics", variant="primary" + ) + with gr.Row(): + summary_output_files = gr.File( + height=FILE_INPUT_HEIGHT, + label="Summarised output files", + interactive=False, + scale=3, + ) + summary_output_files_xlsx = gr.File( + height=FILE_INPUT_HEIGHT, + label="xlsx file summary", + interactive=False, + scale=1, + ) + + summarised_output_markdown = gr.Markdown( + value="### Summarised table will appear here", buttons=["copy"] + ) + + with gr.Accordion("5. Create overall summary", open=False): + gr.Markdown( + """### Create an overall summary from an existing topic summary table.""" + ) + + ### SUMMARISATION + overall_summarisation_input_files = gr.File( + height=FILE_INPUT_HEIGHT, + label="Upload a '...unique_topic' file to summarise", + file_count="multiple", + file_types=[".xlsx", ".xls", ".csv", ".parquet"], + ) + + overall_summarise_format_radio = gr.Radio( + label="Choose summary type", + value=two_para_summary_format_prompt, + choices=[ + two_para_summary_format_prompt, + single_para_summary_format_prompt, + ], + visible=False, + ) # This is currently an invisible placeholder in case in future I want to add in overall summarisation customisation + + overall_summarise_previous_data_btn = gr.Button( + "5. Create overall summary", variant="primary" + ) + + with gr.Row(): + overall_summary_output_files = gr.File( + height=FILE_INPUT_HEIGHT, + label="Summarised output files", + interactive=False, + scale=3, + ) + overall_summary_output_files_xlsx = gr.File( + height=FILE_INPUT_HEIGHT, + label="xlsx file summary", + interactive=False, + scale=1, + ) + + overall_summarised_output_markdown = gr.HTML( + value="### Overall summary will appear here" + ) + + with gr.Tab(label="Review outputs", visible=True): + with gr.Accordion( + "View LLM log files containing prompts and responses", open=True + ): + gr.Markdown( + """Upload a JSON log file to view prompts and responses by batch number and task type.""" + ) + + with gr.Row(): + in_view_log_file = gr.File( + height=FILE_INPUT_HEIGHT, + label="Choose JSON log file", + file_count="single", + file_types=[".json"], + ) + in_view_table = gr.File( + height=FILE_INPUT_HEIGHT, + label="Choose unique topic csv files (legacy)", + file_count="single", + file_types=[".csv", ".parquet"], + visible=False, + ) + with gr.Row(): + log_batch_dropdown = gr.Dropdown( + choices=[], + label="Select batch number", + multiselect=False, + interactive=True, + allow_custom_value=False, + ) + log_task_type_dropdown = gr.Dropdown( + choices=[], + label="Select task type", + multiselect=False, + interactive=True, + allow_custom_value=False, + ) + with gr.Row(): + log_group_dropdown = gr.Dropdown( + choices=[], + label="Select group", + multiselect=False, + interactive=True, + allow_custom_value=False, + ) + log_model_choice_dropdown = gr.Dropdown( + choices=[], + label="Select model choice", + multiselect=False, + interactive=True, + allow_custom_value=False, + ) + log_validated_dropdown = gr.Dropdown( + choices=[], + label="Select validated", + multiselect=False, + interactive=True, + allow_custom_value=False, + ) + + with gr.Row(): + log_prompt_markdown = gr.Markdown( + value="### Prompt\n\nUpload a JSON log file to view prompts.", + label="Prompt", + buttons=["copy"], + ) + log_response_markdown = gr.Markdown( + value="### Response\n\nUpload a JSON log file to view responses.", + label="Response", + buttons=["copy"], + ) + + view_table_markdown = gr.Markdown( + value="", label="View table (legacy)", buttons=["copy"], visible=False + ) + + with gr.Tab(label="Continue unfinished topic extraction", visible=False): + gr.Markdown( + """### Load in output files from a previous topic extraction process and continue topic extraction with new data.""" + ) + + with gr.Accordion( + "Upload reference data file and unique data files", open=True + ): + in_previous_data_files = gr.File( + height=FILE_INPUT_HEIGHT, + label="Choose output csv files", + file_count="multiple", + file_types=[".csv"], + ) + in_previous_data_files_status = gr.Textbox( + value="", label="Previous file input" + ) + continue_previous_data_files_btn = gr.Button( + value="Continue previous topic extraction", variant="primary" + ) + + with gr.Tab(label="LLM and topic extraction settings"): + gr.Markdown("""Define settings that affect large language model output.""") + with gr.Accordion("Settings for LLM generation", open=True): + with gr.Row(): + temperature_slide = gr.Slider( + minimum=0.0, + maximum=1.0, + value=LLM_TEMPERATURE, + label="Choose LLM temperature setting", + precision=1, + step=0.1, + ) + with gr.Row(equal_height=True): + batch_size_number.render() + max_topics_number = gr.Number( + value=MAXIMUM_ALLOWED_TOPICS, + label="Maximum number of topics allowed. If exceeded, the LLM will make efforts to deduplicate topics after every batch until the total number of topics is below this number (not foolproof).", + precision=0, + minimum=1, + maximum=1000, + ) + random_seed = gr.Number( + value=LLM_SEED, label="Random seed for LLM generation", visible=False + ) + + with gr.Accordion("AWS API keys", open=False): + gr.Markdown( + """Querying Bedrock models with API keys requires a role with IAM permissions for the bedrock:InvokeModel action.""" + ) + with gr.Row(): + aws_access_key_textbox = gr.Textbox( + value=AWS_ACCESS_KEY, + label="AWS access key", + lines=1, + type="password", + ) + aws_secret_key_textbox = gr.Textbox( + value=AWS_SECRET_KEY, + label="AWS secret key", + lines=1, + type="password", + ) + aws_region_textbox = gr.Textbox( + value=AWS_REGION, + label="AWS region", + lines=1, + ) + + with gr.Accordion("Gemini API keys", open=False): + google_api_key_textbox = gr.Textbox( + value=GEMINI_API_KEY, + label="Enter Gemini API key (only if using Google API models)", + lines=1, + type="password", + ) + + with gr.Accordion("Azure/OpenAI Inference", open=False): + with gr.Row(): + azure_api_key_textbox = gr.Textbox( + value=AZURE_OPENAI_API_KEY, + label="Enter Azure/OpenAI Inference API key (only if using Azure/OpenAI models)", + lines=1, + type="password", + ) + azure_endpoint_textbox = gr.Textbox( + value=AZURE_OPENAI_INFERENCE_ENDPOINT, + label="Enter Azure Inference endpoint URL (only if using Azure models)", + lines=1, + ) + + with gr.Accordion( + "Llama-server API", open=False, visible=RUN_INFERENCE_SERVER == "1" + ): + api_url_textbox = gr.Textbox( + value=API_URL, + label="Enter inference-server API URL (only if using inference-server models)", + lines=1, + ) + + with gr.Accordion( + "Hugging Face token for downloading gated models", open=False + ): + hf_api_key_textbox = gr.Textbox( + value=HF_TOKEN, + label="Enter Hugging Face API key (only for gated models that need a token to download)", + lines=1, + type="password", + ) + + with gr.Accordion("Log outputs", open=False): + log_files_output = gr.File( + height=FILE_INPUT_HEIGHT, label="Log file output", interactive=False + ) + conversation_metadata_textbox = gr.Textbox( + value="", + label="Query metadata - usage counts and other parameters", + lines=8, + ) + + with gr.Accordion("Prompt settings", open=False, visible=False): + number_of_prompts = gr.Number( + value=1, + label="Number of prompts to send to LLM in sequence", + minimum=1, + maximum=3, + visible=False, + ) + system_prompt_textbox = gr.Textbox( + label="Initial system prompt", lines=4, value=system_prompt + ) + initial_table_prompt_textbox = gr.Textbox( + label="Initial topics prompt", lines=8, value=initial_table_prompt + ) + add_to_existing_topics_system_prompt_textbox = gr.Textbox( + label="Additional topics system prompt", + lines=4, + value=add_existing_topics_system_prompt, + ) + add_to_existing_topics_prompt_textbox = gr.Textbox( + label="Additional topics prompt", + lines=8, + value=add_existing_topics_prompt, + ) + + with gr.Accordion( + "Join additional columns to reference file outputs", open=False + ): + join_colnames = gr.Dropdown( + choices=["Choose column with responses"], + multiselect=True, + label="Select the open text column of interest. In an Excel file, this shows columns across all sheets.", + allow_custom_value=True, + interactive=True, + ) + with gr.Row(): + in_join_files = gr.File( + height=FILE_INPUT_HEIGHT, + label="Response ID file should go here. Original data file should be loaded on the first tab.", + ) + join_cols_btn = gr.Button( + "Join columns to reference output", variant="primary" + ) + out_join_files = gr.File( + height=FILE_INPUT_HEIGHT, + label="Output joined reference files will go here.", + ) + + with gr.Accordion( + "Export output files to xlsx format", open=False, visible=False + ): + export_xlsx_btn = gr.Button( + "Export output files to xlsx format", variant="primary" + ) + out_xlsx_files = gr.File( + height=FILE_INPUT_HEIGHT, label="Output xlsx files will go here." + ) + + ### + # INTERACTIVE ELEMENT FUNCTIONS + ### + + ### + # INITIAL TOPIC EXTRACTION + ### + + # Tabular data upload + in_data_files.upload( + fn=put_columns_in_df, + inputs=[in_data_files], + outputs=[ + in_colnames, + in_excel_sheets, + original_data_file_name_textbox, + join_colnames, + in_group_col, + ], + api_visibility="undocumented", + ) + + # Click on cost code dataframe/dropdown fills in cost code textbox + # Allow user to select items from cost code dataframe for cost code + + if SHOW_COSTS == "True" and ( + GET_COST_CODES == "True" or ENFORCE_COST_CODES == "True" + ): + cost_code_dataframe.select( + df_select_callback_cost, + inputs=[cost_code_dataframe], + outputs=[cost_code_choice_drop], + api_visibility="undocumented", + ) + reset_cost_code_dataframe_button.click( + reset_base_dataframe, + inputs=[cost_code_dataframe_base], + outputs=[cost_code_dataframe], + api_visibility="undocumented", + ) + + cost_code_choice_drop.select( + update_cost_code_dataframe_from_dropdown_select, + inputs=[cost_code_choice_drop, cost_code_dataframe_base], + outputs=[cost_code_dataframe], + api_visibility="undocumented", + ) + + # Extract topics + extract_topics_btn.click( + fn=empty_output_vars_extract_topics, + inputs=None, + outputs=[ + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + topic_extraction_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + file_data_state, + working_data_file_name_textbox, + display_topic_table_markdown, + summary_output_files, + summarisation_input_files, + overall_summarisation_input_files, + overall_summary_output_files, + ], + api_visibility="undocumented", + ).success( + fn=enforce_cost_codes, + inputs=[ + enforce_cost_code_textbox, + cost_code_choice_drop, + cost_code_dataframe_base, + ], + api_visibility="undocumented", + ).success( + load_in_data_file, + inputs=[in_data_files, in_colnames, batch_size_number, in_excel_sheets], + outputs=[ + file_data_state, + working_data_file_name_textbox, + total_number_of_batches, + ], + api_name="load_data", + ).success( + fn=wrapper_extract_topics_per_column_value, + inputs=[ + in_group_col, + in_data_files, + file_data_state, + master_topic_df_state, + master_reference_df_state, + master_unique_topics_df_state, + display_topic_table_markdown, + original_data_file_name_textbox, + total_number_of_batches, + google_api_key_textbox, + temperature_slide, + in_colnames, + model_choice, + candidate_topics, + first_loop_state, + conversation_metadata_textbox, + latest_batch_completed, + estimated_time_taken_number, + initial_table_prompt_textbox, + system_prompt_textbox, + add_to_existing_topics_system_prompt_textbox, + add_to_existing_topics_prompt_textbox, + number_of_prompts, + batch_size_number, + context_textbox, + sentiment_checkbox, + force_zero_shot_radio, + in_excel_sheets, + force_single_topic_radio, + produce_structured_summary_radio, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + hf_api_key_textbox, + azure_api_key_textbox, + azure_endpoint_textbox, + output_folder_state, + logged_content_df, + additional_summary_instructions_textbox, + additional_validation_issues_textbox, + show_previous_table_radio, + api_url_textbox, + max_topics_number, + ], + outputs=[ + display_topic_table_markdown, + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + topic_extraction_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + deduplication_input_files, + summarisation_input_files, + modifiable_unique_topics_df_state, + modification_input_files, + in_join_files, + missing_df_state, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + output_messages_textbox, + logged_content_df, + ], + api_name="extract_topics", + show_progress_on=[output_messages_textbox, topic_extraction_output_files], + ).success( + lambda *args: usage_callback.flag( + list(args), + save_to_csv=SAVE_LOGS_TO_CSV, + save_to_dynamodb=SAVE_LOGS_TO_DYNAMODB, + dynamodb_table_name=USAGE_LOG_DYNAMODB_TABLE_NAME, + dynamodb_headers=DYNAMODB_USAGE_LOG_HEADERS, + replacement_headers=CSV_USAGE_LOG_HEADERS, + ), + [ + session_hash_textbox, + original_data_file_name_textbox, + in_colnames, + model_choice, + conversation_metadata_textbox_placeholder, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + cost_code_choice_drop, + ], + None, + preprocess=False, + api_name="usage_logs", + ).then( + collect_output_csvs_and_create_excel_output, + inputs=[ + in_data_files, + in_colnames, + original_data_file_name_textbox, + in_group_col, + model_choice, + master_reference_df_state, + master_unique_topics_df_state, + summarised_output_df, + missing_df_state, + in_excel_sheets, + usage_logs_state, + model_name_map_state, + output_folder_state, + produce_structured_summary_radio, + candidate_topics, + create_topics_csv_radio, + ], + outputs=[topic_extraction_output_files_xlsx, summary_xlsx_output_files_list], + api_visibility="undocumented", + ).success( + fn=export_outputs_to_s3_with_prompt_response_logs, + inputs=[ + summary_xlsx_output_files_list, + log_files_output_list_state, + s3_output_folder_state, + save_outputs_to_s3_checkbox, + in_data_files, + ], + outputs=None, + api_visibility="undocumented", + ) + + # Validate topics + validate_topics_btn.click( + fn=enforce_cost_codes, + inputs=[ + enforce_cost_code_textbox, + cost_code_choice_drop, + cost_code_dataframe_base, + ], + api_visibility="undocumented", + ).success( + load_in_data_file, + inputs=[in_data_files, in_colnames, batch_size_number, in_excel_sheets], + outputs=[ + file_data_state, + working_data_file_name_textbox, + total_number_of_batches, + ], + api_visibility="undocumented", + ).success( + load_in_previous_data_files, + inputs=[topic_extraction_output_files], + outputs=[ + master_reference_df_state, + master_unique_topics_df_state, + latest_batch_completed_no_loop, + deduplication_input_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ], + api_visibility="undocumented", + ).success( + fn=validate_topics_wrapper, + inputs=[ + file_data_state, + master_reference_df_state, + master_unique_topics_df_state, + working_data_file_name_textbox, + in_colnames, + batch_size_number, + model_choice, + google_api_key_textbox, + temperature_slide, + max_tokens_num, + azure_api_key_textbox, + azure_endpoint_textbox, + reasoning_suffix_textbox, + in_group_col, + produce_structured_summary_radio, + force_zero_shot_radio, + force_single_topic_radio, + context_textbox, + additional_summary_instructions_textbox, + output_folder_state, + output_debug_files_radio, + original_data_file_name_textbox, + additional_validation_issues_textbox, + max_time_for_loop_num, + in_data_files, + sentiment_checkbox, + logged_content_df, + show_previous_table_radio, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + api_url_textbox, + max_topics_number, + ], + outputs=[ + display_topic_table_markdown, + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + validation_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + deduplication_input_files, + summarisation_input_files, + modifiable_unique_topics_df_state, + modification_input_files, + in_join_files, + missing_df_state, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + output_messages_textbox, + logged_content_df, + ], + api_name="validate_topics", + show_progress_on=[output_messages_textbox, validation_output_files], + ).success( + lambda *args: usage_callback.flag( + list(args), + save_to_csv=SAVE_LOGS_TO_CSV, + save_to_dynamodb=SAVE_LOGS_TO_DYNAMODB, + dynamodb_table_name=USAGE_LOG_DYNAMODB_TABLE_NAME, + dynamodb_headers=DYNAMODB_USAGE_LOG_HEADERS, + replacement_headers=CSV_USAGE_LOG_HEADERS, + ), + [ + session_hash_textbox, + original_data_file_name_textbox, + in_colnames, + model_choice, + conversation_metadata_textbox_placeholder, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + cost_code_choice_drop, + ], + None, + preprocess=False, + api_name="undocumented", + ).then( + collect_output_csvs_and_create_excel_output, + inputs=[ + in_data_files, + in_colnames, + original_data_file_name_textbox, + in_group_col, + model_choice, + master_reference_df_state, + master_unique_topics_df_state, + summarised_output_df, + missing_df_state, + in_excel_sheets, + usage_logs_state, + model_name_map_state, + output_folder_state, + produce_structured_summary_radio, + candidate_topics, + create_topics_csv_radio, + ], + outputs=[topic_extraction_output_files_xlsx, summary_xlsx_output_files_list], + api_visibility="undocumented", + ).success( + fn=export_outputs_to_s3_with_prompt_response_logs, + inputs=[ + summary_xlsx_output_files_list, + log_files_output_list_state, + s3_output_folder_state, + save_outputs_to_s3_checkbox, + in_data_files, + ], + outputs=None, + api_visibility="undocumented", + ) + + ### + # DEDUPLICATION AND SUMMARISATION FUNCTIONS + ### + # If you upload data into the deduplication input box, the modifiable topic dataframe box is updated + modification_input_files.upload( + fn=load_in_previous_data_files, + inputs=[modification_input_files, modified_unique_table_change_bool], + outputs=[ + modifiable_unique_topics_df_state, + master_modify_reference_df_state, + master_modify_unique_topics_df_state, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + text_output_modify_file_list_state, + ], + api_visibility="undocumented", + ) + + # Modify output table with custom topic names + save_modified_files_button.click( + fn=modify_existing_output_tables, + inputs=[ + master_modify_unique_topics_df_state, + modifiable_unique_topics_df_state, + master_modify_reference_df_state, + text_output_modify_file_list_state, + output_folder_state, + ], + outputs=[ + master_unique_topics_df_state, + master_reference_df_state, + topic_extraction_output_files, + text_output_file_list_state, + deduplication_input_files, + summarisation_input_files, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + summarised_output_markdown, + ], + api_visibility="undocumented", + ) + + # When button pressed, deduplicate data + deduplicate_previous_data_btn.click( + load_in_previous_data_files, + inputs=[deduplication_input_files], + outputs=[ + master_reference_df_state, + master_unique_topics_df_state, + latest_batch_completed_no_loop, + deduplication_input_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ], + api_visibility="undocumented", + ).success( + deduplicate_topics, + inputs=[ + master_reference_df_state, + master_unique_topics_df_state, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + in_excel_sheets, + merge_sentiment_drop, + merge_general_topics_drop, + deduplicate_score_threshold, + in_data_files, + in_colnames, + output_folder_state, + ], + outputs=[ + master_reference_df_state, + master_unique_topics_df_state, + summarisation_input_files, + log_files_output, + summarised_output_markdown, + ], + scroll_to_output=True, + api_name="deduplicate_topics", + ) + + deduplicate_llm_previous_data_btn.click( + load_in_previous_data_files, + inputs=[deduplication_input_files], + outputs=[ + master_reference_df_state, + master_unique_topics_df_state, + latest_batch_completed_no_loop, + deduplication_input_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ], + api_visibility="undocumented", + ).success( + deduplicate_topics_llm_wrapper, + inputs=[ + master_reference_df_state, + master_unique_topics_df_state, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + model_choice, + google_api_key_textbox, + temperature_slide, + in_excel_sheets, + merge_sentiment_drop, + merge_general_topics_drop, + in_data_files, + in_colnames, + output_folder_state, + candidate_topics, + azure_endpoint_textbox, + api_url_textbox, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + azure_api_key_textbox, + sentiment_checkbox, + ], + outputs=[ + master_reference_df_state, + master_unique_topics_df_state, + summarisation_input_files, + log_files_output, + summarised_output_markdown, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + ], + scroll_to_output=True, + api_name="deduplicate_topics_llm", + ).success( + lambda *args: usage_callback.flag( + list(args), + save_to_csv=SAVE_LOGS_TO_CSV, + save_to_dynamodb=SAVE_LOGS_TO_DYNAMODB, + dynamodb_table_name=USAGE_LOG_DYNAMODB_TABLE_NAME, + dynamodb_headers=DYNAMODB_USAGE_LOG_HEADERS, + replacement_headers=CSV_USAGE_LOG_HEADERS, + ), + [ + session_hash_textbox, + original_data_file_name_textbox, + in_colnames, + model_choice, + conversation_metadata_textbox_placeholder, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + cost_code_choice_drop, + ], + None, + preprocess=False, + api_name="usage_logs_llm_dedup", + ) + + # When button pressed, summarise previous data + summarise_previous_data_btn.click( + empty_output_vars_summarise, + inputs=None, + outputs=[ + summary_reference_table_sample_state, + master_unique_topics_df_revised_summaries_state, + master_reference_df_revised_summaries_state, + summary_output_files, + summarised_outputs_list, + latest_summary_completed_num, + overall_summarisation_input_files, + ], + api_visibility="undocumented", + ).success( + fn=enforce_cost_codes, + inputs=[ + enforce_cost_code_textbox, + cost_code_choice_drop, + cost_code_dataframe_base, + ], + api_visibility="undocumented", + ).success( + load_in_previous_data_files, + inputs=[summarisation_input_files], + outputs=[ + master_reference_df_state, + master_unique_topics_df_state, + latest_batch_completed_no_loop, + deduplication_input_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ], + api_visibility="undocumented", + ).success( + wrapper_summarise_output_topics_per_group, + inputs=[ + in_group_col, + summary_reference_table_sample_state, + master_unique_topics_df_state, + master_reference_df_state, + model_choice, + google_api_key_textbox, + temperature_slide, + working_data_file_name_textbox, + summarised_outputs_list, + latest_summary_completed_num, + conversation_metadata_textbox, + in_data_files, + in_excel_sheets, + in_colnames, + log_files_output_list_state, + summarise_format_radio, + output_folder_state, + context_textbox, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + model_name_map_state, + hf_api_key_textbox, + azure_endpoint_textbox, + logged_content_df, + sample_reference_table_checkbox, + no_of_sampled_summaries_number, + random_seed_number, + api_url_textbox, + ], + outputs=[ + summary_reference_table_sample_state, + master_unique_topics_df_revised_summaries_state, + master_reference_df_revised_summaries_state, + summary_output_files, + summarised_outputs_list, + latest_summary_completed_num, + conversation_metadata_textbox, + summarised_output_markdown, + log_files_output, + overall_summarisation_input_files, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + output_messages_textbox, + logged_content_df, + ], + api_name="summarise_topics", + show_progress_on=[output_messages_textbox, summary_output_files], + ).success( + lambda *args: usage_callback.flag( + list(args), + save_to_csv=SAVE_LOGS_TO_CSV, + save_to_dynamodb=SAVE_LOGS_TO_DYNAMODB, + dynamodb_table_name=USAGE_LOG_DYNAMODB_TABLE_NAME, + dynamodb_headers=DYNAMODB_USAGE_LOG_HEADERS, + replacement_headers=CSV_USAGE_LOG_HEADERS, + ), + [ + session_hash_textbox, + original_data_file_name_textbox, + in_colnames, + model_choice, + conversation_metadata_textbox_placeholder, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + cost_code_choice_drop, + ], + None, + preprocess=False, + api_visibility="undocumented", + ).then( + collect_output_csvs_and_create_excel_output, + inputs=[ + in_data_files, + in_colnames, + original_data_file_name_textbox, + in_group_col, + model_choice, + master_reference_df_revised_summaries_state, + master_unique_topics_df_revised_summaries_state, + summarised_output_df, + missing_df_state, + in_excel_sheets, + usage_logs_state, + model_name_map_state, + output_folder_state, + produce_structured_summary_radio, + candidate_topics, + create_topics_csv_radio, + ], + outputs=[summary_output_files_xlsx, summary_xlsx_output_files_list], + api_visibility="undocumented", + ).success( + fn=export_outputs_to_s3_with_prompt_response_logs, + inputs=[ + summary_xlsx_output_files_list, + log_files_output, + s3_output_folder_state, + save_outputs_to_s3_checkbox, + in_data_files, + ], + outputs=None, + api_visibility="undocumented", + ) + # success(sample_reference_table_summaries, inputs=[master_reference_df_state, random_seed, sample_reference_table_checkbox], outputs=[summary_reference_table_sample_state, summarised_references_markdown], api_name="sample_summaries").\ + + # SUMMARISE WHOLE TABLE PAGE + overall_summarise_previous_data_btn.click( + fn=enforce_cost_codes, + inputs=[ + enforce_cost_code_textbox, + cost_code_choice_drop, + cost_code_dataframe_base, + ], + api_visibility="undocumented", + ).success( + load_in_previous_data_files, + inputs=[overall_summarisation_input_files], + outputs=[ + master_reference_df_state, + master_unique_topics_df_state, + latest_batch_completed_no_loop, + deduplication_input_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ], + api_visibility="undocumented", + ).success( + overall_summary, + inputs=[ + master_unique_topics_df_state, + model_choice, + google_api_key_textbox, + temperature_slide, + working_data_file_name_textbox, + output_folder_state, + in_colnames, + context_textbox, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + model_name_map_state, + hf_api_key_textbox, + azure_endpoint_textbox, + logged_content_df, + api_url_textbox, + ], + outputs=[ + overall_summary_output_files, + overall_summarised_output_markdown, + summarised_output_df, + conversation_metadata_textbox, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + output_messages_textbox, + logged_content_df, + ], + scroll_to_output=True, + api_name="overall_summary", + show_progress_on=[output_messages_textbox, overall_summary_output_files], + ).success( + lambda *args: usage_callback.flag( + list(args), + save_to_csv=SAVE_LOGS_TO_CSV, + save_to_dynamodb=SAVE_LOGS_TO_DYNAMODB, + dynamodb_table_name=USAGE_LOG_DYNAMODB_TABLE_NAME, + dynamodb_headers=DYNAMODB_USAGE_LOG_HEADERS, + replacement_headers=CSV_USAGE_LOG_HEADERS, + ), + [ + session_hash_textbox, + original_data_file_name_textbox, + in_colnames, + model_choice, + conversation_metadata_textbox_placeholder, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + cost_code_choice_drop, + ], + None, + preprocess=False, + api_visibility="undocumented", + ).then( + collect_output_csvs_and_create_excel_output, + inputs=[ + in_data_files, + in_colnames, + original_data_file_name_textbox, + in_group_col, + model_choice, + master_reference_df_state, + master_unique_topics_df_state, + summarised_output_df, + missing_df_state, + in_excel_sheets, + usage_logs_state, + model_name_map_state, + output_folder_state, + produce_structured_summary_radio, + candidate_topics, + create_topics_csv_radio, + ], + outputs=[overall_summary_output_files_xlsx, summary_xlsx_output_files_list], + api_visibility="undocumented", + ).success( + fn=export_outputs_to_s3_with_prompt_response_logs, + inputs=[ + summary_xlsx_output_files_list, + log_files_output_list_state, + s3_output_folder_state, + save_outputs_to_s3_checkbox, + in_data_files, + ], + outputs=None, + api_visibility="undocumented", + ) + + # Discover topics from random sample (extract-only on subsample) + discover_topics_btn.click( + fn=empty_output_vars_extract_topics, + inputs=None, + outputs=[ + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + topic_extraction_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + file_data_state, + working_data_file_name_textbox, + display_topic_table_markdown, + summary_output_files, + summarisation_input_files, + overall_summarisation_input_files, + overall_summary_output_files, + ], + api_visibility="undocumented", + ).success( + fn=enforce_cost_codes, + inputs=[ + enforce_cost_code_textbox, + cost_code_choice_drop, + cost_code_dataframe_base, + ], + api_visibility="undocumented", + ).success( + load_in_data_file, + inputs=[in_data_files, in_colnames, batch_size_number, in_excel_sheets], + outputs=[ + file_data_state, + working_data_file_name_textbox, + total_number_of_batches, + ], + api_visibility="undocumented", + ).success( + fn=discover_topics_from_sample_wrapper, + inputs=[ + in_group_col, + in_data_files, + file_data_state, + original_data_file_name_textbox, + total_number_of_batches, + google_api_key_textbox, + temperature_slide, + in_colnames, + model_choice, + candidate_topics, + topic_discovery_sample_fraction, + topic_discovery_random_seed, + context_textbox, + sentiment_checkbox, + force_zero_shot_radio, + in_excel_sheets, + force_single_topic_radio, + produce_structured_summary_radio, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + hf_api_key_textbox, + azure_api_key_textbox, + azure_endpoint_textbox, + output_folder_state, + initial_table_prompt_textbox, + system_prompt_textbox, + add_to_existing_topics_system_prompt_textbox, + add_to_existing_topics_prompt_textbox, + number_of_prompts, + batch_size_number, + additional_summary_instructions_textbox, + additional_validation_issues_textbox, + show_previous_table_radio, + api_url_textbox, + max_topics_number, + ], + outputs=[ + display_topic_table_markdown, + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + topic_extraction_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + deduplication_input_files, + summarisation_input_files, + modifiable_unique_topics_df_state, + modification_input_files, + in_join_files, + missing_df_state, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + output_messages_textbox, + logged_content_df, + discover_topics_output_file, + ], + api_name="discover_topics_from_sample", + show_progress_on=[output_messages_textbox, discover_topics_output_file], + ) + + # All in one button + # Extract topics - deduplicate and summarise using default settings + all_in_one_btn.click( + fn=empty_output_vars_extract_topics, + inputs=None, + outputs=[ + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + topic_extraction_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + file_data_state, + working_data_file_name_textbox, + display_topic_table_markdown, + summary_output_files, + summarisation_input_files, + overall_summarisation_input_files, + overall_summary_output_files, + ], + api_visibility="undocumented", + ).success( + fn=enforce_cost_codes, + inputs=[ + enforce_cost_code_textbox, + cost_code_choice_drop, + cost_code_dataframe_base, + ], + api_visibility="undocumented", + ).success( + load_in_data_file, + inputs=[in_data_files, in_colnames, batch_size_number, in_excel_sheets], + outputs=[ + file_data_state, + working_data_file_name_textbox, + total_number_of_batches, + ], + api_visibility="undocumented", + ).success( + fn=all_in_one_pipeline, + inputs=[ + in_group_col, + in_data_files, + file_data_state, + master_topic_df_state, + master_reference_df_state, + master_unique_topics_df_state, + display_topic_table_markdown, + original_data_file_name_textbox, + total_number_of_batches, + google_api_key_textbox, + temperature_slide, + in_colnames, + model_choice, + candidate_topics, + first_loop_state, + conversation_metadata_textbox, + latest_batch_completed, + estimated_time_taken_number, + initial_table_prompt_textbox, + system_prompt_textbox, + add_to_existing_topics_system_prompt_textbox, + add_to_existing_topics_prompt_textbox, + number_of_prompts, + batch_size_number, + context_textbox, + sentiment_checkbox, + force_zero_shot_radio, + in_excel_sheets, + force_single_topic_radio, + produce_structured_summary_radio, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + hf_api_key_textbox, + azure_api_key_textbox, + azure_endpoint_textbox, + output_folder_state, + merge_sentiment_drop, + merge_general_topics_drop, + deduplicate_score_threshold, + summarise_format_radio, + random_seed, + log_files_output_list_state, + model_name_map_state, + usage_logs_state, + logged_content_df, + additional_summary_instructions_textbox, + additional_validation_issues_textbox, + show_previous_table_radio, + sample_reference_table_checkbox, + api_url_textbox, + ], + outputs=[ + display_topic_table_markdown, + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + topic_extraction_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + deduplication_input_files, + summarisation_input_files, + modifiable_unique_topics_df_state, + modification_input_files, + in_join_files, + missing_df_state, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + output_messages_textbox, + summary_reference_table_sample_state, + summarised_references_markdown, + master_unique_topics_df_revised_summaries_state, + master_reference_df_revised_summaries_state, + summary_output_files, + summarised_outputs_list, + latest_summary_completed_num, + overall_summarisation_input_files, + overall_summary_output_files, + overall_summarised_output_markdown, + summarised_output_df, + logged_content_df, + ], + show_progress_on=[output_messages_textbox], + api_name="all_in_one_pipeline", + ).success( + lambda *args: usage_callback.flag( + list(args), + save_to_csv=SAVE_LOGS_TO_CSV, + save_to_dynamodb=SAVE_LOGS_TO_DYNAMODB, + dynamodb_table_name=USAGE_LOG_DYNAMODB_TABLE_NAME, + dynamodb_headers=DYNAMODB_USAGE_LOG_HEADERS, + replacement_headers=CSV_USAGE_LOG_HEADERS, + ), + [ + session_hash_textbox, + original_data_file_name_textbox, + in_colnames, + model_choice, + conversation_metadata_textbox_placeholder, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + cost_code_choice_drop, + ], + None, + preprocess=False, + api_visibility="undocumented", + ).then( + collect_output_csvs_and_create_excel_output, + inputs=[ + in_data_files, + in_colnames, + original_data_file_name_textbox, + in_group_col, + model_choice, + master_reference_df_revised_summaries_state, + master_unique_topics_df_revised_summaries_state, + summarised_output_df, + missing_df_state, + in_excel_sheets, + usage_logs_state, + model_name_map_state, + output_folder_state, + produce_structured_summary_radio, + candidate_topics, + create_topics_csv_radio, + ], + outputs=[overall_summary_output_files_xlsx, summary_xlsx_output_files_list], + api_visibility="undocumented", + ).success( + fn=export_outputs_to_s3_with_prompt_response_logs, + inputs=[ + summary_xlsx_output_files_list, + log_files_output_list_state, + s3_output_folder_state, + save_outputs_to_s3_checkbox, + in_data_files, + ], + outputs=None, + api_visibility="undocumented", + ).success( + move_overall_summary_output_files_to_front_page, + inputs=[summary_xlsx_output_files_list], + outputs=[topic_extraction_output_files_xlsx], + api_visibility="undocumented", + ) + + ### + # CONTINUE PREVIOUS TOPIC EXTRACTION PAGE + ### + + # If uploaded partially completed consultation files do this. This should then start up the 'latest_batch_completed' change action above to continue extracting topics. + continue_previous_data_files_btn.click( + load_in_data_file, + inputs=[in_data_files, in_colnames, batch_size_number, in_excel_sheets], + outputs=[ + file_data_state, + working_data_file_name_textbox, + total_number_of_batches, + ], + api_visibility="undocumented", + ).success( + load_in_previous_data_files, + inputs=[in_previous_data_files], + outputs=[ + master_reference_df_state, + master_unique_topics_df_state, + latest_batch_completed, + in_previous_data_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ], + api_visibility="undocumented", + ) + + ### + # VIEW TABLE PAGE + ### + + # Handle log file upload + in_view_log_file.upload( + fn=load_log_file_handler, + inputs=[in_view_log_file], + outputs=[ + log_data_state, + log_batch_dropdown, + log_task_type_dropdown, + log_group_dropdown, + log_model_choice_dropdown, + log_validated_dropdown, + log_prompt_markdown, + log_response_markdown, + ], + api_visibility="undocumented", + ) + + # Update display when any filter changes + def update_on_filter_change( + log_data, batch_str, task_type, group, model_choice, validated + ): + """Wrapper to update display when any filter changes.""" + return update_log_display_on_filter( + log_data, batch_str, task_type, group, model_choice, validated + ) + + log_batch_dropdown.change( + fn=update_on_filter_change, + inputs=[ + log_data_state, + log_batch_dropdown, + log_task_type_dropdown, + log_group_dropdown, + log_model_choice_dropdown, + log_validated_dropdown, + ], + outputs=[log_prompt_markdown, log_response_markdown], + api_visibility="undocumented", + ) + + log_task_type_dropdown.change( + fn=update_on_filter_change, + inputs=[ + log_data_state, + log_batch_dropdown, + log_task_type_dropdown, + log_group_dropdown, + log_model_choice_dropdown, + log_validated_dropdown, + ], + outputs=[log_prompt_markdown, log_response_markdown], + api_visibility="undocumented", + ) + + log_group_dropdown.change( + fn=update_on_filter_change, + inputs=[ + log_data_state, + log_batch_dropdown, + log_task_type_dropdown, + log_group_dropdown, + log_model_choice_dropdown, + log_validated_dropdown, + ], + outputs=[log_prompt_markdown, log_response_markdown], + api_visibility="undocumented", + ) + + log_model_choice_dropdown.change( + fn=update_on_filter_change, + inputs=[ + log_data_state, + log_batch_dropdown, + log_task_type_dropdown, + log_group_dropdown, + log_model_choice_dropdown, + log_validated_dropdown, + ], + outputs=[log_prompt_markdown, log_response_markdown], + api_visibility="undocumented", + ) + + log_validated_dropdown.change( + fn=update_on_filter_change, + inputs=[ + log_data_state, + log_batch_dropdown, + log_task_type_dropdown, + log_group_dropdown, + log_model_choice_dropdown, + log_validated_dropdown, + ], + outputs=[log_prompt_markdown, log_response_markdown], + api_visibility="undocumented", + ) + + # Legacy view_table handler (kept for backward compatibility) + in_view_table.upload( + view_table, + inputs=[in_view_table], + outputs=[view_table_markdown], + api_visibility="undocumented", + ) + + ### + # LLM SETTINGS PAGE + ### + + reference_df_data_file_name_textbox = gr.Textbox( + label="reference_df_data_file_name_textbox", visible=False + ) + master_reference_df_state_joined = gr.Dataframe(visible=False) + + join_cols_btn.click( + fn=load_in_previous_reference_file, + inputs=[in_join_files], + outputs=[master_reference_df_state, reference_df_data_file_name_textbox], + api_visibility="undocumented", + ).success( + load_in_data_file, + inputs=[in_data_files, in_colnames, batch_size_number, in_excel_sheets], + outputs=[ + file_data_state, + working_data_file_name_textbox, + total_number_of_batches, + ], + api_visibility="undocumented", + ).success( + fn=join_cols_onto_reference_df, + inputs=[ + master_reference_df_state, + file_data_state, + join_colnames, + reference_df_data_file_name_textbox, + ], + outputs=[master_reference_df_state_joined, out_join_files], + api_visibility="undocumented", + ) + + # Export to xlsx file + export_xlsx_btn.click( + collect_output_csvs_and_create_excel_output, + inputs=[ + in_data_files, + in_colnames, + original_data_file_name_textbox, + in_group_col, + model_choice, + master_reference_df_state, + master_unique_topics_df_state, + summarised_output_df, + missing_df_state, + in_excel_sheets, + usage_logs_state, + model_name_map_state, + output_folder_state, + produce_structured_summary_radio, + candidate_topics, + create_topics_csv_radio, + ], + outputs=[out_xlsx_files, summary_xlsx_output_files_list], + api_name="export_xlsx", + ).success( + fn=export_outputs_to_s3_with_prompt_response_logs, + inputs=[ + summary_xlsx_output_files_list, + log_files_output_list_state, + s3_output_folder_state, + save_outputs_to_s3_checkbox, + in_data_files, + ], + outputs=None, + api_visibility="undocumented", + ) + + # If relevant environment variable is set, load in the default cost code file from S3 or locally + if GET_COST_CODES == "True" and (COST_CODES_PATH or S3_COST_CODES_PATH): + if ( + not os.path.exists(COST_CODES_PATH) + and S3_COST_CODES_PATH + and RUN_AWS_FUNCTIONS == "1" + ): + print("Downloading cost codes from S3") + print( + f"Attempting to download from bucket: {S3_LOG_BUCKET}, key: {S3_COST_CODES_PATH}" + ) + + app.load( + download_cost_codes_with_error_handling, + inputs=[ + s3_default_bucket, + s3_default_cost_codes_file, + default_cost_codes_output_folder_location, + ], + ).success( + load_in_default_cost_codes, + inputs=[ + default_cost_codes_output_folder_location, + default_cost_code_textbox, + ], + outputs=[ + cost_code_dataframe, + cost_code_dataframe_base, + cost_code_choice_drop, + ], + api_visibility="undocumented", + ) + print("Successfully loaded cost codes from S3") + elif os.path.exists(COST_CODES_PATH): + print( + "Loading cost codes from default cost codes path location:", + COST_CODES_PATH, + ) + app.load( + load_in_default_cost_codes, + inputs=[ + default_cost_codes_output_folder_location, + default_cost_code_textbox, + ], + outputs=[ + cost_code_dataframe, + cost_code_dataframe_base, + cost_code_choice_drop, + ], + api_visibility="undocumented", + ) + else: + print("Could not load in cost code data") + + ### + # LOGGING AND ON APP LOAD FUNCTIONS + ### + + # Get connection parameters + app.load( + get_connection_params, + inputs=None, + outputs=[ + session_hash_state, + output_folder_state, + session_hash_textbox, + input_folder_state, + ], + api_visibility="undocumented", + ) + + # Log usernames and times of access to file (to know who is using the app when running on AWS) + access_callback = CSVLogger_custom(dataset_file_name=LOG_FILE_NAME) + access_callback.setup([session_hash_textbox], ACCESS_LOGS_FOLDER) + + session_hash_textbox.change( + lambda *args: access_callback.flag( + list(args), + save_to_csv=SAVE_LOGS_TO_CSV, + save_to_dynamodb=SAVE_LOGS_TO_DYNAMODB, + dynamodb_table_name=ACCESS_LOG_DYNAMODB_TABLE_NAME, + dynamodb_headers=DYNAMODB_ACCESS_LOG_HEADERS, + replacement_headers=CSV_ACCESS_LOG_HEADERS, + ), + [session_hash_textbox], + None, + preprocess=False, + api_visibility="undocumented", + ).success( + fn=upload_file_to_s3, + inputs=[ + access_logs_state, + access_s3_logs_loc_state, + s3_log_bucket_name, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + ], + outputs=[s3_logs_output_textbox], + api_visibility="undocumented", + ) + + # Log usage when making a query + usage_callback = CSVLogger_custom(dataset_file_name=USAGE_LOG_FILE_NAME) + usage_callback.setup( + [ + session_hash_textbox, + original_data_file_name_textbox, + in_colnames, + model_choice, + conversation_metadata_textbox_placeholder, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + cost_code_choice_drop, + ], + USAGE_LOGS_FOLDER, + ) + + # See extract topics and summarise calls to see the calls to usage logs + + # number_of_calls_num.change(lambda *args: usage_callback.flag(list(args), save_to_csv=SAVE_LOGS_TO_CSV, save_to_dynamodb=SAVE_LOGS_TO_DYNAMODB, dynamodb_table_name=USAGE_LOG_DYNAMODB_TABLE_NAME, dynamodb_headers=DYNAMODB_USAGE_LOG_HEADERS, replacement_headers=CSV_USAGE_LOG_HEADERS), [session_hash_textbox, original_data_file_name_textbox, in_colnames, model_choice, conversation_metadata_textbox, input_tokens_num, output_tokens_num, number_of_calls_num, estimated_time_taken_number, cost_code_choice_drop], None, preprocess=False, api_name="usage_logs").\ + # success(fn = upload_file_to_s3, inputs=[usage_logs_state, usage_s3_logs_loc_state, s3_log_bucket_name, aws_access_key_textbox, aws_secret_key_textbox], outputs=[s3_logs_output_textbox]) + + number_of_calls_num.change( + fn=upload_file_to_s3, + inputs=[ + usage_logs_state, + usage_s3_logs_loc_state, + s3_log_bucket_name, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + ], + outputs=[s3_logs_output_textbox], + api_visibility="undocumented", + ) + + # User submitted feedback + feedback_callback = CSVLogger_custom(dataset_file_name=FEEDBACK_LOG_FILE_NAME) + feedback_callback.setup( + [ + data_feedback_radio, + data_further_details_text, + original_data_file_name_textbox, + model_choice, + temperature_slide, + display_topic_table_markdown, + conversation_metadata_textbox_placeholder, + ], + FEEDBACK_LOGS_FOLDER, + ) + + data_submit_feedback_btn.click( + lambda *args: feedback_callback.flag( + list(args), + save_to_csv=SAVE_LOGS_TO_CSV, + save_to_dynamodb=SAVE_LOGS_TO_DYNAMODB, + dynamodb_table_name=FEEDBACK_LOG_DYNAMODB_TABLE_NAME, + dynamodb_headers=DYNAMODB_FEEDBACK_LOG_HEADERS, + replacement_headers=CSV_FEEDBACK_LOG_HEADERS, + ), + [ + data_feedback_radio, + data_further_details_text, + original_data_file_name_textbox, + model_choice, + temperature_slide, + display_topic_table_markdown, + conversation_metadata_textbox_placeholder, + ], + None, + preprocess=False, + api_visibility="undocumented", + ).success( + fn=upload_file_to_s3, + inputs=[ + feedback_logs_state, + feedback_s3_logs_loc_state, + s3_log_bucket_name, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + ], + outputs=[data_further_details_text], + api_visibility="undocumented", + ) + +### +# APP RUN +### + +if __name__ == "__main__": + if RUN_DIRECT_MODE == "1": + from cli_topics import main + + # Validate required direct mode configuration + if not DIRECT_MODE_INPUT_FILE and DIRECT_MODE_TASK in [ + "extract", + "validate", + "all_in_one", + ]: + print( + "Error: DIRECT_MODE_INPUT_FILE environment variable must be set for direct mode." + ) + print("Please set DIRECT_MODE_INPUT_FILE to the path of your input file.") + exit(1) + + if ( + DIRECT_MODE_TASK in ["extract", "validate", "all_in_one"] + and not DIRECT_MODE_TEXT_COLUMN + ): + print( + "Error: DIRECT_MODE_TEXT_COLUMN environment variable must be set for direct mode tasks: extract, validate, all_in_one." + ) + print( + "Please set DIRECT_MODE_TEXT_COLUMN to the name of the text column to process." + ) + exit(1) + + if ( + DIRECT_MODE_TASK + in ["validate", "deduplicate", "summarise", "overall_summary"] + and not DIRECT_MODE_PREVIOUS_OUTPUT_FILES + ): + print( + "Error: DIRECT_MODE_PREVIOUS_OUTPUT_FILES environment variable must be set for direct mode tasks: validate, deduplicate, summarise, overall_summary." + ) + print( + "Please set DIRECT_MODE_PREVIOUS_OUTPUT_FILES to a pipe-separated (|) list of previous output file paths." + ) + exit(1) + + # Parse previous_output_files if provided (pipe-separated string to handle paths with spaces) + previous_output_files_list = [] + if DIRECT_MODE_PREVIOUS_OUTPUT_FILES: + # Use pipe separator to handle file paths with spaces + previous_output_files_list = [ + f.strip() + for f in DIRECT_MODE_PREVIOUS_OUTPUT_FILES.split("|") + if f.strip() + ] + + # Parse excel_sheets if provided (comma-separated string) + excel_sheets_list = [] + if DIRECT_MODE_EXCEL_SHEETS: + excel_sheets_list = [ + s.strip() for s in DIRECT_MODE_EXCEL_SHEETS.split(",") if s.strip() + ] + + # Parse input_file if provided (pipe-separated string for multiple files to handle paths with spaces) + # Note: S3 downloads are now handled in cli_topics.py main() function + input_file_list = [] + if DIRECT_MODE_INPUT_FILE: + # Use pipe separator to handle file paths with spaces + # First check if it's a single file (no pipe), then split if multiple files + if "|" in DIRECT_MODE_INPUT_FILE: + input_file_list = [ + f.strip() for f in DIRECT_MODE_INPUT_FILE.split("|") if f.strip() + ] + else: + # Single file - use as-is to preserve paths with spaces + input_file_list = [DIRECT_MODE_INPUT_FILE.strip()] + + # Prepare direct mode arguments based on environment variables + direct_mode_args = { + # Task Selection + "task": DIRECT_MODE_TASK, + # General Arguments + "input_file": input_file_list if input_file_list else None, + "output_dir": DIRECT_MODE_OUTPUT_DIR, + "input_dir": INPUT_FOLDER, + "text_column": DIRECT_MODE_TEXT_COLUMN if DIRECT_MODE_TEXT_COLUMN else None, + "previous_output_files": ( + previous_output_files_list if previous_output_files_list else None + ), + "username": DIRECT_MODE_USERNAME, + "save_to_user_folders": SESSION_OUTPUT_FOLDER, + "excel_sheets": excel_sheets_list, + "group_by": DIRECT_MODE_GROUP_BY if DIRECT_MODE_GROUP_BY else None, + # Model Configuration + "model_choice": DIRECT_MODE_MODEL_CHOICE, + "model_source": default_model_source, + "temperature": float(DIRECT_MODE_TEMPERATURE), + "batch_size": int(DIRECT_MODE_BATCH_SIZE), + "max_tokens": int(DIRECT_MODE_MAX_TOKENS), + "google_api_key": GEMINI_API_KEY, + "aws_access_key": AWS_ACCESS_KEY, + "aws_secret_key": AWS_SECRET_KEY, + "aws_region": AWS_REGION, + "hf_token": HF_TOKEN, + "azure_api_key": AZURE_OPENAI_API_KEY, + "azure_endpoint": AZURE_OPENAI_INFERENCE_ENDPOINT, + "api_url": API_URL, + "inference_server_model": ( + DIRECT_MODE_INFERENCE_SERVER_MODEL + if DIRECT_MODE_INFERENCE_SERVER_MODEL + else None + ), + # Topic Extraction Arguments + "context": DIRECT_MODE_CONTEXT if DIRECT_MODE_CONTEXT else "", + "candidate_topics": ( + DIRECT_MODE_CANDIDATE_TOPICS if DIRECT_MODE_CANDIDATE_TOPICS else None + ), + "force_zero_shot": DIRECT_MODE_FORCE_ZERO_SHOT, + "force_single_topic": DIRECT_MODE_FORCE_SINGLE_TOPIC, + "produce_structured_summary": DIRECT_MODE_PRODUCE_STRUCTURED_SUMMARY, + "sentiment": DIRECT_MODE_SENTIMENT, + "additional_summary_instructions": ( + DIRECT_MODE_ADDITIONAL_SUMMARY_INSTRUCTIONS + if DIRECT_MODE_ADDITIONAL_SUMMARY_INSTRUCTIONS + else "" + ), + # Validation Arguments + "additional_validation_issues": ( + DIRECT_MODE_ADDITIONAL_VALIDATION_ISSUES + if DIRECT_MODE_ADDITIONAL_VALIDATION_ISSUES + else "" + ), + "show_previous_table": DIRECT_MODE_SHOW_PREVIOUS_TABLE, + "output_debug_files": OUTPUT_DEBUG_FILES, + "max_time_for_loop": int(DIRECT_MODE_MAX_TIME_FOR_LOOP), + # Deduplication Arguments + "method": DIRECT_MODE_DEDUP_METHOD, + "similarity_threshold": int(DIRECT_MODE_SIMILARITY_THRESHOLD), + "merge_sentiment": DIRECT_MODE_MERGE_SENTIMENT, + "merge_general_topics": DIRECT_MODE_MERGE_GENERAL_TOPICS, + # Summarisation Arguments + "summary_format": DIRECT_MODE_SUMMARY_FORMAT, + "sample_reference_table": DIRECT_MODE_SAMPLE_REFERENCE_TABLE, + "no_of_sampled_summaries": int(DIRECT_MODE_NO_OF_SAMPLED_SUMMARIES), + "random_seed": int(DIRECT_MODE_RANDOM_SEED), + # Output Format Arguments + "create_xlsx_output": DIRECT_MODE_CREATE_XLSX_OUTPUT == "True", + "create_topics_csv": DIRECT_MODE_CREATE_TOPICS_CSV == "True", + "sample_fraction": float(DIRECT_MODE_SAMPLE_FRACTION), + # Logging Arguments + "save_logs_to_csv": SAVE_LOGS_TO_CSV, + "save_logs_to_dynamodb": SAVE_LOGS_TO_DYNAMODB, + "usage_logs_folder": USAGE_LOGS_FOLDER, + "cost_code": ( + DIRECT_MODE_DEFAULT_COST_CODE + if DIRECT_MODE_DEFAULT_COST_CODE + else DEFAULT_COST_CODE + ), + } + + print(f"Running in direct mode with task: {DIRECT_MODE_TASK}") + if input_file_list: + print(f"Input file(s): {', '.join(input_file_list)}") + print(f"Output directory: {DIRECT_MODE_OUTPUT_DIR}") + if DIRECT_MODE_TEXT_COLUMN: + print(f"Text column: {DIRECT_MODE_TEXT_COLUMN}") + if previous_output_files_list: + print(f"Previous output files: {', '.join(previous_output_files_list)}") + if DIRECT_MODE_GROUP_BY: + print(f"Group by: {DIRECT_MODE_GROUP_BY}") + + # Run the CLI main function with direct mode arguments + main(direct_mode_args=direct_mode_args) + else: + if COGNITO_AUTH == "1": + app.queue(max_size=MAX_QUEUE_SIZE).launch( + show_error=True, + inbrowser=True, + auth=authenticate_user, + max_file_size=MAX_FILE_SIZE, + server_port=GRADIO_SERVER_PORT, + root_path=ROOT_PATH, + mcp_server=RUN_MCP_SERVER, + theme=gr.themes.Default(primary_hue="blue"), + css=css, + ) + else: + app.queue(max_size=MAX_QUEUE_SIZE).launch( + show_error=True, + inbrowser=True, + max_file_size=MAX_FILE_SIZE, + server_port=GRADIO_SERVER_PORT, + root_path=ROOT_PATH, + mcp_server=RUN_MCP_SERVER, + theme=gr.themes.Default(primary_hue="blue"), + css=css, + ) diff --git a/cdk/__init__.py b/cdk/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/cdk/app.py b/cdk/app.py new file mode 100644 index 0000000000000000000000000000000000000000..66dc50547c25ebc974d98be288930f523e02d241 --- /dev/null +++ b/cdk/app.py @@ -0,0 +1,123 @@ +import os + +from aws_cdk import App, Environment +from cdk_appregistry import register_llm_topic_application +from cdk_config import ( + ALB_NAME, + APPREGISTRY_APPLICATION_NAME, + APPREGISTRY_ATTRIBUTE_GROUP_NAME, + APPREGISTRY_DESCRIPTION, + APPREGISTRY_REPOSITORY_URL, + APPREGISTRY_STACK_NAME, + AWS_ACCOUNT_ID, + AWS_REGION, + CDK_CONTEXT_FILE, + CDK_PREFIX, + ENABLE_APPREGISTRY, + RUN_USEAST_STACK, + USE_CLOUDFRONT, +) +from cdk_functions import ( + create_basic_config_env, + is_resource_delete_protection_enabled, + load_context_from_file, + log_aws_credential_context, + purge_cdk_lookup_context, +) +from cdk_stack import CdkStack, CdkStackCloudfront # , CdkStackMain +from check_resources import CONTEXT_FILE, check_and_set_context + +# Initialize the CDK app +app = App() + +log_aws_credential_context( + expected_account_id=AWS_ACCOUNT_ID, + expected_region=AWS_REGION, +) + +# Drop stale CDK lookup cache entries (require bootstrap lookup role in target account). +purge_cdk_lookup_context(CDK_CONTEXT_FILE) + +# --- Pre-check context (boto3) — written to precheck.context.json, NOT cdk.context.json --- +print(f"Pre-check context file: {CONTEXT_FILE}") +print(f"CDK lookup cache file: {CDK_CONTEXT_FILE}") +if os.path.basename(CONTEXT_FILE.replace("\\", "/")) == os.path.basename( + CDK_CONTEXT_FILE.replace("\\", "/") +): + raise RuntimeError( + f"CONTEXT_FILE and CDK_CONTEXT_FILE must differ (got '{CONTEXT_FILE}' for both). " + "Set CONTEXT_FILE=precheck.context.json in config/cdk_config.env." + ) + +print("Running pre-check script to generate application context...") +try: + check_and_set_context() + if not os.path.exists(CONTEXT_FILE): + raise RuntimeError( + f"check_and_set_context() finished, but {CONTEXT_FILE} was not created." + ) + print(f"Context generated successfully at {CONTEXT_FILE}.") +except Exception as e: + raise RuntimeError(f"Failed to generate context via check_and_set_context(): {e}") + +# Pre-check must not repopulate CDK lookup keys; purge again if paths were ever shared. +purge_cdk_lookup_context(CDK_CONTEXT_FILE) + +if os.path.exists(CONTEXT_FILE): + load_context_from_file(app, CONTEXT_FILE) +else: + raise RuntimeError(f"Could not find {CONTEXT_FILE}.") + +create_basic_config_env("config") + +aws_env_regional = Environment(account=AWS_ACCOUNT_ID, region=AWS_REGION) + +_stack_delete_protection = is_resource_delete_protection_enabled() + +regional_stack = CdkStack( + app, "SummarisationStack", env=aws_env_regional, cross_region_references=True +) +regional_stack.termination_protection = _stack_delete_protection + +if ENABLE_APPREGISTRY == "True": + # Use pre-check context only — not regional_stack.params (avoids AppRegistry + # -> SummarisationStack dependency cycle during synth). + _alb_dns_context = app.node.try_get_context(f"dns:{ALB_NAME}") + _alb_dns_name = ( + _alb_dns_context.strip() + if isinstance(_alb_dns_context, str) and _alb_dns_context.strip() + else None + ) + appregistry_stack = register_llm_topic_application( + app, + aws_account_id=AWS_ACCOUNT_ID, + aws_region=AWS_REGION, + application_name=APPREGISTRY_APPLICATION_NAME, + application_description=APPREGISTRY_DESCRIPTION, + appregistry_stack_name=APPREGISTRY_STACK_NAME, + attribute_group_name=APPREGISTRY_ATTRIBUTE_GROUP_NAME, + repository_url=APPREGISTRY_REPOSITORY_URL, + cdk_prefix=CDK_PREFIX, + use_cloudfront=USE_CLOUDFRONT, + alb_dns_name=_alb_dns_name, + ) + appregistry_stack.termination_protection = _stack_delete_protection + +if USE_CLOUDFRONT == "True" and RUN_USEAST_STACK == "True": + aws_env_us_east_1 = Environment(account=AWS_ACCOUNT_ID, region="us-east-1") + + cloudfront_stack = CdkStackCloudfront( + app, + "SummarisationStackCloudfront", + env=aws_env_us_east_1, + alb_arn=regional_stack.params["alb_arn_output"], + alb_sec_group_id=regional_stack.params["alb_security_group_id"], + alb_dns_name=regional_stack.params["alb_dns_name"], + cross_region_references=True, + ) + cloudfront_stack.termination_protection = _stack_delete_protection + +# CDK CLI invokes this script and expects a cloud assembly in cdk.out. +# Without app.synth(), Python defines constructs but never writes manifest.json +# (ENOENT on deploy). See: https://github.com/aws/aws-cdk/issues/11023 +app.synth() diff --git a/cdk/cdk.json.example b/cdk/cdk.json.example new file mode 100644 index 0000000000000000000000000000000000000000..628a5a286ddd433808101af6f02314fdaa5c7e81 --- /dev/null +++ b/cdk/cdk.json.example @@ -0,0 +1,7 @@ +{ + "app": "python app.py", + "output": "cdk.out", + "context": { + "@aws-cdk/aws-ec2:restrictDefaultSecurityGroup": false + } +} diff --git a/cdk/cdk_appregistry.py b/cdk/cdk_appregistry.py new file mode 100644 index 0000000000000000000000000000000000000000..4ee9837103ef9e355c05b3ccef9abeaca2fe45df --- /dev/null +++ b/cdk/cdk_appregistry.py @@ -0,0 +1,74 @@ +"""AWS Console myApplications (Service Catalog AppRegistry) integration.""" + +from aws_cdk import App, Environment +from aws_cdk.aws_servicecatalogappregistry_alpha import ( + ApplicationAssociator, + TargetApplication, +) + + +def register_llm_topic_application( + app: App, + *, + aws_account_id: str, + aws_region: str, + application_name: str, + application_description: str, + appregistry_stack_name: str, + attribute_group_name: str, + repository_url: str, + cdk_prefix: str, + use_cloudfront: str, + alb_dns_name: str | None = None, +) -> ApplicationAssociator: + """ + Register regional CDK stacks with AWS Console myApplications. + + Only stacks in ``aws_region`` are associated (phase 1). Cross-region stacks + such as SummarisationStackCloudfront (us-east-1) are not included. + + ``alb_dns_name`` must be a plain string (e.g. from pre-check context). Do not + pass a CloudFormation token from SummarisationStack or synth will fail with a + dependency cycle against the associator stack. + """ + associator = ApplicationAssociator( + app, + "SummarisationAppRegistry", + applications=[ + TargetApplication.create_application_stack( + application_name=application_name, + application_description=application_description, + stack_name=appregistry_stack_name, + env=Environment(account=aws_account_id, region=aws_region), + ) + ], + ) + + attributes = { + "repository": repository_url, + "cdkPrefix": cdk_prefix, + "awsRegion": aws_region, + "useCloudFront": use_cloudfront, + "cloudFrontInAppRegistry": "false", + "cloudFrontNote": ( + "CloudFront/WAF (SummarisationStackCloudfront) is in us-east-1 and is " + "not linked to this myApplications entry in phase 1. View it in " + "CloudFormation (us-east-1) or the CloudFront console." + ), + } + if alb_dns_name: + attributes["albDnsName"] = alb_dns_name + + associator.app_registry_application.add_attribute_group( + "SummarisationAttributeGroup", + attribute_group_name=attribute_group_name, + description="llm_topic_modeller deployment metadata", + attributes=attributes, + ) + + return associator + + +def register_doc_summarisation_application(*args, **kwargs): + """Deprecated alias for register_llm_topic_application.""" + return register_llm_topic_application(*args, **kwargs) diff --git a/cdk/cdk_cloudfront_headers.py b/cdk/cdk_cloudfront_headers.py new file mode 100644 index 0000000000000000000000000000000000000000..ff3a5eabdb63af87c983ea51b2e1f007d554d330 --- /dev/null +++ b/cdk/cdk_cloudfront_headers.py @@ -0,0 +1,172 @@ +"""CloudFront response headers policy (CSP and related security headers).""" + +from __future__ import annotations + +from pathlib import Path +from urllib.parse import urlparse + +from aws_cdk import Duration +from aws_cdk import aws_cloudfront as cloudfront +from constructs import Construct + +# Template exported from AWS; placeholders {APP-URL} and {COGNITO-APP-CLIENT-LOGIN-URL}. +_CSP_TEMPLATE = ( + "default-src 'self'; script-src 'self' cdnjs.cloudflare.com 'unsafe-inline'; " + "style-src 'self' https://fonts.googleapis.com 'unsafe-inline'; " + "img-src 'self' data:; font-src 'self' https://fonts.gstatic.com data:; " + "connect-src 'self' wss://{app_hostname} https://cdnjs.cloudflare.com; " + "form-action 'self'; frame-ancestors 'none'; object-src 'none'; base-uri 'self'; " + "manifest-src 'self' {cognito_login_url}; upgrade-insecure-requests;" +) + +RESPONSE_HEADERS_POLICY_TEMPLATE_PATH = ( + Path(__file__).resolve().parent / "config" / "response-headers-policy-config.json" +) + + +def normalize_https_origin(url: str) -> str: + """Return a canonical https origin (scheme + host, no path).""" + value = (url or "").strip() + if not value: + return "" + if "://" not in value: + value = f"https://{value}" + parsed = urlparse(value) + if not parsed.hostname: + return value.rstrip("/") + scheme = parsed.scheme or "https" + netloc = parsed.netloc or parsed.hostname + return f"{scheme}://{netloc}".rstrip("/") + + +def hostname_from_origin(origin: str) -> str: + parsed = urlparse(origin if "://" in origin else f"https://{origin}") + return parsed.hostname or origin.strip() + + +def cognito_hosted_ui_base_url(domain_prefix: str, region: str) -> str: + prefix = (domain_prefix or "").strip() + if not prefix: + return "" + return f"https://{prefix}.auth.{region}.amazoncognito.com" + + +def build_content_security_policy( + *, + app_origin: str, + cognito_login_url: str, +) -> str: + origin = normalize_https_origin(app_origin) + hostname = hostname_from_origin(origin) + cognito_url = normalize_https_origin(cognito_login_url) + return _CSP_TEMPLATE.format(app_hostname=hostname, cognito_login_url=cognito_url) + + +def create_secure_cloudfront_response_headers_policy( + scope: Construct, + construct_id: str, + *, + policy_name: str, + app_origin: str, + cognito_login_url: str, + comment: str = "Secure response headers with CSP for llm_topic_modeller", +) -> cloudfront.ResponseHeadersPolicy: + """Response headers policy aligned with config/response-headers-policy-config.json.""" + cors_origin = normalize_https_origin(app_origin) + csp = build_content_security_policy( + app_origin=app_origin, + cognito_login_url=cognito_login_url, + ) + + return cloudfront.ResponseHeadersPolicy( + scope, + construct_id, + response_headers_policy_name=policy_name, + comment=comment, + cors_behavior=cloudfront.ResponseHeadersCorsBehavior( + access_control_allow_credentials=False, + access_control_allow_headers=["*"], + access_control_allow_methods=["ALL"], + access_control_allow_origins=[cors_origin], + access_control_max_age=Duration.seconds(600), + origin_override=True, + ), + security_headers_behavior=cloudfront.ResponseSecurityHeadersBehavior( + content_security_policy=cloudfront.ResponseHeadersContentSecurityPolicy( + content_security_policy=csp, + override=True, + ), + content_type_options=cloudfront.ResponseHeadersContentTypeOptions( + override=True + ), + frame_options=cloudfront.ResponseHeadersFrameOptions( + frame_option=cloudfront.HeadersFrameOption.SAMEORIGIN, + override=True, + ), + referrer_policy=cloudfront.ResponseHeadersReferrerPolicy( + referrer_policy=cloudfront.HeadersReferrerPolicy.STRICT_ORIGIN_WHEN_CROSS_ORIGIN, + override=True, + ), + strict_transport_security=cloudfront.ResponseHeadersStrictTransportSecurity( + access_control_max_age=Duration.seconds(31536000), + include_subdomains=True, + preload=False, + override=True, + ), + xss_protection=cloudfront.ResponseHeadersXSSProtection( + protection=True, + mode_block=True, + override=True, + ), + ), + custom_headers_behavior=cloudfront.ResponseCustomHeadersBehavior( + custom_headers=[ + cloudfront.ResponseCustomHeader( + header="Permissions-Policy", + value=( + "accelerometer=(), autoplay=(), camera=(), " + "cross-origin-isolated=(), display-capture=(), encrypted-media=(), " + "fullscreen=(), geolocation=(), gyroscope=(), keyboard-map=(), " + "magnetometer=(), microphone=(), midi=(), payment=(), " + "picture-in-picture=(), publickey-credentials-get=(), " + "screen-wake-lock=(), sync-xhr=(), usb=(), web-share=(), " + "xr-spatial-tracking=()" + ), + override=True, + ) + ] + ), + remove_headers=["Server"], + ) + + +def resolve_cloudfront_csp_urls( + *, + cognito_redirection_url: str, + cloudfront_domain: str, + cognito_user_pool_domain_prefix: str, + aws_region: str, + cognito_user_pool_login_url: str = "", + ssl_certificate_domain: str = "", +) -> tuple[str, str]: + """ + Return (app_origin, cognito_login_url) for CSP/CORS substitution. + + App origin prefers COGNITO_REDIRECTION_URL (canonical browser URL), then + https://SSL_CERTIFICATE_DOMAIN, then https://CLOUDFRONT_DOMAIN. + Cognito login URL uses COGNITO_USER_POOL_LOGIN_URL when set, else the + hosted UI base URL derived from COGNITO_USER_POOL_DOMAIN_PREFIX. + """ + app_origin = normalize_https_origin(cognito_redirection_url) + if not app_origin or "placeholder" in app_origin.lower(): + if ssl_certificate_domain.strip(): + app_origin = normalize_https_origin(ssl_certificate_domain) + elif cloudfront_domain.strip(): + app_origin = normalize_https_origin(cloudfront_domain) + + login_url = (cognito_user_pool_login_url or "").strip() + if not login_url: + login_url = cognito_hosted_ui_base_url( + cognito_user_pool_domain_prefix, aws_region + ) + return app_origin, normalize_https_origin(login_url) diff --git a/cdk/cdk_config.py b/cdk/cdk_config.py new file mode 100644 index 0000000000000000000000000000000000000000..e48d509cb5dee921e99c451d98ad99e938d1693f --- /dev/null +++ b/cdk/cdk_config.py @@ -0,0 +1,851 @@ +import os +import tempfile +from typing import List + +from dotenv import load_dotenv + +# Set or retrieve configuration variables for CDK llm_topic_modeller deployment + + +def convert_string_to_boolean(value: str) -> bool: + """Convert string to boolean, handling various formats.""" + if isinstance(value, bool): + return value + elif value in ["True", "1", "true", "TRUE"]: + return True + elif value in ["False", "0", "false", "FALSE"]: + return False + else: + raise ValueError(f"Invalid boolean value: {value}") + + +def parse_comma_separated_list(value: str) -> List[str]: + """Parse a comma-separated env value into a list of non-empty strings.""" + if not value or not str(value).strip(): + return [] + cleaned = str(value).strip().strip("[]") + return [ + part.strip().strip('"').strip("'") + for part in cleaned.split(",") + if part.strip() + ] + + +def get_or_create_env_var(var_name: str, default_value: str, print_val: bool = False): + """ + Get an environmental variable, and set it to a default value if it doesn't exist + """ + # Get the environment variable if it exists + value = os.environ.get(var_name) + + # If it doesn't exist, set the environment variable to the default value + if value is None: + os.environ[var_name] = default_value + value = default_value + + if print_val is True: + print(f"The value of {var_name} is {value}") + + return value + + +def ensure_folder_exists(output_folder: str): + """Checks if the specified folder exists, creates it if not.""" + + if not os.path.exists(output_folder): + # Create the folder if it doesn't exist + os.makedirs(output_folder, exist_ok=True) + print(f"Created the {output_folder} folder.") + else: + print(f"The {output_folder} folder already exists.") + + +def add_folder_to_path(folder_path: str): + """ + Check if a folder exists on your system. If so, get the absolute path and then add it to the system Path variable if it doesn't already exist. Function is only relevant for locally-created executable files based on this app (when using pyinstaller it creates a _internal folder that contains tesseract and poppler. These need to be added to the system path to enable the app to run) + """ + + if os.path.exists(folder_path) and os.path.isdir(folder_path): + print(folder_path, "folder exists.") + + # Resolve relative path to absolute path + absolute_path = os.path.abspath(folder_path) + + current_path = os.environ["PATH"] + if absolute_path not in current_path.split(os.pathsep): + full_path_extension = absolute_path + os.pathsep + current_path + os.environ["PATH"] = full_path_extension + # print(f"Updated PATH with: ", full_path_extension) + else: + print(f"Directory {folder_path} already exists in PATH.") + else: + print(f"Folder not found at {folder_path} - not added to PATH") + + +### +# LOAD CONFIG FROM ENV FILE +### +CONFIG_FOLDER = get_or_create_env_var("CONFIG_FOLDER", "config/") + +ensure_folder_exists(CONFIG_FOLDER) + +# If you have an aws_config env file in the config folder, you can load in app variables this way, e.g. 'config/cdk_config.env' +CDK_CONFIG_PATH = get_or_create_env_var( + "CDK_CONFIG_PATH", "config/cdk_config.env" +) # e.g. config/cdk_config.env + +if CDK_CONFIG_PATH: + if os.path.exists(CDK_CONFIG_PATH): + print(f"Loading CDK variables from config file {CDK_CONFIG_PATH}") + # override=True: stale empty defaults from an earlier cdk_config import in the + # same shell (e.g. cdk_install wizard calling cdk_functions) must not win. + load_dotenv(CDK_CONFIG_PATH, override=True) + else: + print("CDK config file not found at location:", CDK_CONFIG_PATH) + +### +# AWS OPTIONS +### +AWS_REGION = get_or_create_env_var("AWS_REGION", "") +AWS_ACCOUNT_ID = get_or_create_env_var("AWS_ACCOUNT_ID", "") + +### +# CDK OPTIONS +### +CDK_PREFIX = get_or_create_env_var("CDK_PREFIX", "") + +# When True (default): CloudFormation stack termination protection, AWS +# deletion_protection on ALB/DynamoDB/Cognito, RemovalPolicy.RETAIN, and no S3 +# auto-delete on stack-created resources. Set False for dev sandboxes where +# cdk destroy should remove resources cleanly. +ENABLE_RESOURCE_DELETE_PROTECTION = get_or_create_env_var( + "ENABLE_RESOURCE_DELETE_PROTECTION", "True" +) + +# AWS Console myApplications (Service Catalog AppRegistry) +ENABLE_APPREGISTRY = get_or_create_env_var("ENABLE_APPREGISTRY", "True") +APPREGISTRY_APPLICATION_NAME = get_or_create_env_var( + "APPREGISTRY_APPLICATION_NAME", f"{CDK_PREFIX}llm-topic-modeller" +) +APPREGISTRY_DESCRIPTION = get_or_create_env_var( + "APPREGISTRY_DESCRIPTION", + "LLM topic modelling app (ALB, ECS Fargate, Cognito, S3)", +) +APPREGISTRY_STACK_NAME = get_or_create_env_var( + "APPREGISTRY_STACK_NAME", f"{CDK_PREFIX}AppRegistryStack" +) +APPREGISTRY_ATTRIBUTE_GROUP_NAME = get_or_create_env_var( + "APPREGISTRY_ATTRIBUTE_GROUP_NAME", + f"{APPREGISTRY_APPLICATION_NAME}-metadata", +) +APPREGISTRY_REPOSITORY_URL = get_or_create_env_var( + "APPREGISTRY_REPOSITORY_URL", + "https://github.com/seanpedrick-case/llm_topic_modeller.git", +) + +_precheck_context_file = get_or_create_env_var("CONTEXT_FILE", "precheck.context.json") +# Never write boto3 pre-check output into CDK's lookup cache file (causes stale +# vpc-provider / load-balancer entries and wrong-account lookup validation errors). +if os.path.basename(_precheck_context_file.replace("\\", "/")) == "cdk.context.json": + print( + "WARNING: CONTEXT_FILE must not be 'cdk.context.json' (that file is CDK's " + "lookup cache). Using 'precheck.context.json' instead. Update " + "config/cdk_config.env and remove CONTEXT_FILE=cdk.context.json if set." + ) + _precheck_context_file = "precheck.context.json" +CONTEXT_FILE = _precheck_context_file +CDK_CONTEXT_FILE = get_or_create_env_var("CDK_CONTEXT_FILE", "cdk.context.json") +CDK_FOLDER = get_or_create_env_var( + "CDK_FOLDER", "" +) # FULL_PATH_TO_CDK_FOLDER_HERE (with forward slash) + +# App runtime config (uploaded to S3 for legacy Fargate; inlined for ECS Express Mode) +APP_CONFIG_ENV_BASENAME = "app_config.env" +_app_config_rel = os.path.join(CONFIG_FOLDER, APP_CONFIG_ENV_BASENAME).replace( + "\\", "/" +) +APP_CONFIG_ENV_FILE = get_or_create_env_var( + "APP_CONFIG_ENV_FILE", + ( + os.path.normpath(os.path.join(CDK_FOLDER, _app_config_rel)) + if CDK_FOLDER + else os.path.normpath(_app_config_rel) + ), +) +RUN_USEAST_STACK = get_or_create_env_var("RUN_USEAST_STACK", "False") + +### VPC and connections +VPC_NAME = get_or_create_env_var("VPC_NAME", "") +NEW_VPC_DEFAULT_NAME = get_or_create_env_var("NEW_VPC_DEFAULT_NAME", f"{CDK_PREFIX}vpc") +NEW_VPC_CIDR = get_or_create_env_var("NEW_VPC_CIDR", "") # "10.0.0.0/24" + + +# Internet Gateway for legacy VPC public subnets (attach + 0.0.0.0/0 routes via CDK when missing). +EXISTING_IGW_ID = get_or_create_env_var("EXISTING_IGW_ID", "") +SINGLE_NAT_GATEWAY_ID = get_or_create_env_var("SINGLE_NAT_GATEWAY_ID", "") + +### SUBNETS / ROUTE TABLES / NAT GATEWAY +PUBLIC_SUBNETS_TO_USE = get_or_create_env_var( + "PUBLIC_SUBNETS_TO_USE", "" +) # e.g. ['PublicSubnet1', 'PublicSubnet2'] +PUBLIC_SUBNET_CIDR_BLOCKS = get_or_create_env_var( + "PUBLIC_SUBNET_CIDR_BLOCKS", "" +) # e.g. ["10.0.1.0/24", "10.0.2.0/24"] +PUBLIC_SUBNET_AVAILABILITY_ZONES = get_or_create_env_var( + "PUBLIC_SUBNET_AVAILABILITY_ZONES", "" +) # e.g. ["eu-east-1b", "eu-east1b"] + +PRIVATE_SUBNETS_TO_USE = get_or_create_env_var( + "PRIVATE_SUBNETS_TO_USE", "" +) # e.g. ['PrivateSubnet1', 'PrivateSubnet2'] +PRIVATE_SUBNET_CIDR_BLOCKS = get_or_create_env_var( + "PRIVATE_SUBNET_CIDR_BLOCKS", "" +) # e.g. ["10.0.1.0/24", "10.0.2.0/24"] +PRIVATE_SUBNET_AVAILABILITY_ZONES = get_or_create_env_var( + "PRIVATE_SUBNET_AVAILABILITY_ZONES", "" +) # e.g. ["eu-east-1b", "eu-east1b"] + +ROUTE_TABLE_BASE_NAME = get_or_create_env_var( + "ROUTE_TABLE_BASE_NAME", f"{CDK_PREFIX}PrivateRouteTable" +) +NAT_GATEWAY_EIP_NAME = get_or_create_env_var( + "NAT_GATEWAY_EIP_NAME", f"{CDK_PREFIX}NatGatewayEip" +) +NAT_GATEWAY_NAME = get_or_create_env_var("NAT_GATEWAY_NAME", f"{CDK_PREFIX}NatGateway") + +# IAM roles — managed policy *names* (AWS managed) and JSON policy *files* (inline statements) +AWS_MANAGED_TASK_ROLES_LIST = get_or_create_env_var( + "AWS_MANAGED_TASK_ROLES_LIST", + '["AmazonCognitoReadOnly", "service-role/AmazonECSTaskExecutionRolePolicy", "AmazonDynamoDBFullAccess", "service-role/AWSAppSyncPushToCloudWatchLogs", "AmazonBedrockLimitedAccess"]', +) +ECS_EXECUTION_ROLE_MANAGED_POLICIES = get_or_create_env_var( + "ECS_EXECUTION_ROLE_MANAGED_POLICIES", + '["service-role/AmazonECSTaskExecutionRolePolicy"]', +) +# JSON IAM policy document paths (relative to CDK_FOLDER or absolute). Task role = app runtime. +POLICY_FILE_LOCATIONS = get_or_create_env_var( + "POLICY_FILE_LOCATIONS", + "[]", +) +# Optional extra JSON policies for the ECS task *execution* role (image pull / logs / secrets). +ECS_EXECUTION_ROLE_POLICY_FILES = get_or_create_env_var( + "ECS_EXECUTION_ROLE_POLICY_FILES", + "", +) +# Customer-managed policy ARNs (full ARN per entry), attached in addition to the lists above. +POLICY_FILE_ARNS = get_or_create_env_var("POLICY_FILE_ARNS", "") +ECS_EXECUTION_ROLE_POLICY_ARNS = get_or_create_env_var( + "ECS_EXECUTION_ROLE_POLICY_ARNS", "" +) + +AWS_MANAGED_TASK_ROLES_LIST = parse_comma_separated_list(AWS_MANAGED_TASK_ROLES_LIST) +ECS_EXECUTION_ROLE_MANAGED_POLICIES = parse_comma_separated_list( + ECS_EXECUTION_ROLE_MANAGED_POLICIES +) +POLICY_FILE_LOCATIONS = parse_comma_separated_list(POLICY_FILE_LOCATIONS) +ECS_EXECUTION_ROLE_POLICY_FILES = parse_comma_separated_list( + ECS_EXECUTION_ROLE_POLICY_FILES +) +POLICY_FILE_ARNS = parse_comma_separated_list(POLICY_FILE_ARNS) +ECS_EXECUTION_ROLE_POLICY_ARNS = parse_comma_separated_list( + ECS_EXECUTION_ROLE_POLICY_ARNS +) + +# GITHUB REPO +GITHUB_REPO_USERNAME = get_or_create_env_var("GITHUB_REPO_USERNAME", "seanpedrick-case") +GITHUB_REPO_NAME = get_or_create_env_var("GITHUB_REPO_NAME", "llm_topic_modeller") +GITHUB_REPO_BRANCH = get_or_create_env_var("GITHUB_REPO_BRANCH", "main") + +### CODEBUILD +CODEBUILD_ROLE_NAME = get_or_create_env_var( + "CODEBUILD_ROLE_NAME", f"{CDK_PREFIX}CodeBuildRole" +) +CODEBUILD_PROJECT_NAME = get_or_create_env_var( + "CODEBUILD_PROJECT_NAME", f"{CDK_PREFIX}CodeBuildProject" +) + +### ECR +ECR_REPO_NAME = get_or_create_env_var( + "ECR_REPO_NAME", "llm-topic-modeller" +) # Beware - cannot have underscores and must be lower case +ECR_CDK_REPO_NAME = get_or_create_env_var( + "ECR_CDK_REPO_NAME", f"{CDK_PREFIX}{ECR_REPO_NAME}".lower() +) + + +### S3 +def _resolve_s3_bucket_name(env_key: str, suffix: str) -> str: + """ + Bucket name default is ``{CDK_PREFIX}{suffix}`` (lowercase). + + If an earlier import cached a bare ``suffix`` in ``os.environ`` before + ``CDK_PREFIX`` was loaded from dotenv, upgrade it to the prefixed name. + """ + prefix = (os.environ.get("CDK_PREFIX") or "").lower() + default = f"{prefix}{suffix}" + value = get_or_create_env_var(env_key, default) + if prefix and value == suffix: + os.environ[env_key] = default + return default + return value + + +S3_LOG_CONFIG_BUCKET_NAME = _resolve_s3_bucket_name( + "S3_LOG_CONFIG_BUCKET_NAME", "s3-logs" +) # S3 bucket names need to be lower case +S3_OUTPUT_BUCKET_NAME = _resolve_s3_bucket_name("S3_OUTPUT_BUCKET_NAME", "s3-output") + +### VPC endpoints for ECS tasks in private subnets (ECR image pull, logs, secrets) +ENABLE_ECS_VPC_INTERFACE_ENDPOINTS = get_or_create_env_var( + "ENABLE_ECS_VPC_INTERFACE_ENDPOINTS", "True" +) + +### KMS KEYS FOR S3 AND SECRETS MANAGER +USE_CUSTOM_KMS_KEY = get_or_create_env_var("USE_CUSTOM_KMS_KEY", "1") +CUSTOM_KMS_KEY_NAME = get_or_create_env_var( + "CUSTOM_KMS_KEY_NAME", f"alias/{CDK_PREFIX}kms-key".lower() +) + +### ECS +FARGATE_TASK_DEFINITION_NAME = get_or_create_env_var( + "FARGATE_TASK_DEFINITION_NAME", f"{CDK_PREFIX}FargateTaskDefinition" +) +TASK_DEFINITION_FILE_LOCATION = get_or_create_env_var( + "TASK_DEFINITION_FILE_LOCATION", CDK_FOLDER + CONFIG_FOLDER + "task_definition.json" +) + +CLUSTER_NAME = get_or_create_env_var("CLUSTER_NAME", f"{CDK_PREFIX}Cluster") +ECS_SERVICE_NAME = get_or_create_env_var("ECS_SERVICE_NAME", f"{CDK_PREFIX}ECSService") +# Second Fargate service when ENABLE_PI_AGENT_ECS_SERVICE=True (legacy path only). +ECS_PI_SERVICE_NAME = get_or_create_env_var( + "ECS_PI_SERVICE_NAME", f"{CDK_PREFIX}PiAgentService" +) +ECS_TASK_ROLE_NAME = get_or_create_env_var( + "ECS_TASK_ROLE_NAME", f"{CDK_PREFIX}TaskRole" +) +ECS_TASK_EXECUTION_ROLE_NAME = get_or_create_env_var( + "ECS_TASK_EXECUTION_ROLE_NAME", f"{CDK_PREFIX}ExecutionRole" +) +ECS_SECURITY_GROUP_NAME = get_or_create_env_var( + "ECS_SECURITY_GROUP_NAME", f"{CDK_PREFIX}SecurityGroupECS" +) +ECS_LOG_GROUP_NAME = get_or_create_env_var( + "ECS_LOG_GROUP_NAME", f"/ecs/{ECS_SERVICE_NAME}-logs".lower() +) + +ECS_TASK_CPU_SIZE = get_or_create_env_var("ECS_TASK_CPU_SIZE", "1024") +ECS_TASK_MEMORY_SIZE = get_or_create_env_var("ECS_TASK_MEMORY_SIZE", "4096") +ECS_USE_FARGATE_SPOT = get_or_create_env_var("USE_FARGATE_SPOT", "False") +ECS_READ_ONLY_FILE_SYSTEM = get_or_create_env_var("ECS_READ_ONLY_FILE_SYSTEM", "True") +# ECS service AZ rebalancing (AWS defaults new services to ENABLED if omitted). +ECS_AVAILABILITY_ZONE_REBALANCING = get_or_create_env_var( + "ECS_AVAILABILITY_ZONE_REBALANCING", "DISABLED" +) +if ECS_AVAILABILITY_ZONE_REBALANCING not in ("ENABLED", "DISABLED"): + raise ValueError( + "ECS_AVAILABILITY_ZONE_REBALANCING must be ENABLED or DISABLED " + f"(got {ECS_AVAILABILITY_ZONE_REBALANCING!r})." + ) + +### Cognito +COGNITO_USER_POOL_NAME = get_or_create_env_var( + "COGNITO_USER_POOL_NAME", f"{CDK_PREFIX}UserPool" +) +COGNITO_USER_POOL_CLIENT_NAME = get_or_create_env_var( + "COGNITO_USER_POOL_CLIENT_NAME", f"{CDK_PREFIX}UserPoolClient" +) +COGNITO_USER_POOL_CLIENT_SECRET_NAME = get_or_create_env_var( + "COGNITO_USER_POOL_CLIENT_SECRET_NAME", f"{CDK_PREFIX}ParamCognitoSecret" +) +COGNITO_USER_POOL_DOMAIN_PREFIX = get_or_create_env_var( + "COGNITO_USER_POOL_DOMAIN_PREFIX", "llm-topic-app-domain" +) # Should change this to something unique or you'll probably hit an error + +COGNITO_REFRESH_TOKEN_VALIDITY = int( + get_or_create_env_var("COGNITO_REFRESH_TOKEN_VALIDITY", "480") +) # Minutes +COGNITO_ID_TOKEN_VALIDITY = int( + get_or_create_env_var("COGNITO_ID_TOKEN_VALIDITY", "60") +) # Minutes +COGNITO_ACCESS_TOKEN_VALIDITY = int( + get_or_create_env_var("COGNITO_ACCESS_TOKEN_VALIDITY", "60") +) # Minutes + +# Application load balancer +ALB_NAME = get_or_create_env_var( + "ALB_NAME", f"{CDK_PREFIX}Alb"[-32:] +) # Application load balancer name can be max 32 characters, so taking the last 32 characters of the suggested name +ALB_NAME_SECURITY_GROUP_NAME = get_or_create_env_var( + "ALB_SECURITY_GROUP_NAME", f"{CDK_PREFIX}SecurityGroupALB" +) +ALB_TARGET_GROUP_NAME = get_or_create_env_var( + "ALB_TARGET_GROUP_NAME", f"{CDK_PREFIX}-tg"[-32:] +) # Max 32 characters +EXISTING_LOAD_BALANCER_ARN = get_or_create_env_var("EXISTING_LOAD_BALANCER_ARN", "") +EXISTING_LOAD_BALANCER_DNS = get_or_create_env_var( + "EXISTING_LOAD_BALANCER_DNS", "placeholder_load_balancer_dns.net" +) + +## CLOUDFRONT +USE_CLOUDFRONT = get_or_create_env_var("USE_CLOUDFRONT", "True") +CLOUDFRONT_PREFIX_LIST_ID = get_or_create_env_var( + "CLOUDFRONT_PREFIX_LIST_ID", "pl-93a247fa" +) +CLOUDFRONT_GEO_RESTRICTION = get_or_create_env_var( + "CLOUDFRONT_GEO_RESTRICTION", "" +) # A country that Cloudfront restricts access to. See here: https://docs.aws.amazon.com/AmazonCloudFront/latest/DeveloperGuide/georestrictions.html +CLOUDFRONT_DISTRIBUTION_NAME = get_or_create_env_var( + "CLOUDFRONT_DISTRIBUTION_NAME", f"{CDK_PREFIX}CfDist" +) +CLOUDFRONT_DOMAIN = get_or_create_env_var( + "CLOUDFRONT_DOMAIN", "cloudfront_placeholder.net" +) +# Attach CSP / security response headers to the CDK CloudFront distribution (us-east-1 stack). +CLOUDFRONT_ENABLE_SECURE_RESPONSE_HEADERS = get_or_create_env_var( + "CLOUDFRONT_ENABLE_SECURE_RESPONSE_HEADERS", "True" +) +# Optional override for manifest-src (Cognito hosted UI). Default: https://{COGNITO_USER_POOL_DOMAIN_PREFIX}.auth.{AWS_REGION}.amazoncognito.com +COGNITO_USER_POOL_LOGIN_URL = get_or_create_env_var("COGNITO_USER_POOL_LOGIN_URL", "") + + +# Certificate for Application load balancer (optional, for HTTPS and logins through the ALB) +ACM_SSL_CERTIFICATE_ARN = get_or_create_env_var("ACM_SSL_CERTIFICATE_ARN", "") +SSL_CERTIFICATE_DOMAIN = get_or_create_env_var( + "SSL_CERTIFICATE_DOMAIN", "" +) # e.g. example.com or www.example.com + +# ECS Express Mode (opt-in HTTPS ingress without supplying ACM_SSL_CERTIFICATE_ARN). +# Pilot/dev: Express PrimaryContainer does not support S3 environmentFiles or Fargate mount points. +USE_ECS_EXPRESS_MODE = get_or_create_env_var("USE_ECS_EXPRESS_MODE", "False") +ECS_EXPRESS_SERVICE_NAME = get_or_create_env_var( + "ECS_EXPRESS_SERVICE_NAME", ECS_SERVICE_NAME +) +ECS_EXPRESS_HEALTH_CHECK_PATH = get_or_create_env_var( + "ECS_EXPRESS_HEALTH_CHECK_PATH", "/" +) +ECS_EXPRESS_INFRASTRUCTURE_ROLE_NAME = get_or_create_env_var( + "ECS_EXPRESS_INFRASTRUCTURE_ROLE_NAME", f"{CDK_PREFIX}ExpressInfraRole" +) +# After first deploy, set to ExpressServiceEndpoint output (https://...) if not using CloudFront. +# The installer updates Cognito callback URLs via API (no second CDK deploy). +ECS_EXPRESS_COGNITO_REDIRECT_BASE = get_or_create_env_var( + "ECS_EXPRESS_COGNITO_REDIRECT_BASE", "" +) +# Express networkConfiguration.subnets drives both tasks and the managed ALB. +# Public subnets (IGW route) → internet-facing ALB; private → internal ALB only. +ECS_EXPRESS_USE_PUBLIC_SUBNETS = get_or_create_env_var( + "ECS_EXPRESS_USE_PUBLIC_SUBNETS", "True" +) + +if USE_ECS_EXPRESS_MODE == "True" and ACM_SSL_CERTIFICATE_ARN: + raise ValueError( + "USE_ECS_EXPRESS_MODE=True cannot be used with ACM_SSL_CERTIFICATE_ARN set. " + "Clear ACM_SSL_CERTIFICATE_ARN or set USE_ECS_EXPRESS_MODE=False." + ) + +# ECS Service Connect (legacy Fargate only): VPC service-to-service HTTP to Gradio/FastAPI. +ENABLE_ECS_SERVICE_CONNECT = get_or_create_env_var( + "ENABLE_ECS_SERVICE_CONNECT", "False" +) +ECS_SERVICE_CONNECT_NAMESPACE = get_or_create_env_var( + "ECS_SERVICE_CONNECT_NAMESPACE", + (f"{CDK_PREFIX}local".lower().replace("_", "-").strip("-") or "llm-topic-local"), +) +ECS_SERVICE_CONNECT_DISCOVERY_NAME = get_or_create_env_var( + "ECS_SERVICE_CONNECT_DISCOVERY_NAME", "llm-topic" +) +# Optional friendly DNS label; defaults to discovery name when empty. +ECS_SERVICE_CONNECT_DNS_NAME = get_or_create_env_var("ECS_SERVICE_CONNECT_DNS_NAME", "") +# Client task security groups (at least one of IDs, names, or CDK prefixes required when SC on). +ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_IDS = get_or_create_env_var( + "ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_IDS", "" +) +ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_IDS_LIST = parse_comma_separated_list( + ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_IDS +) +ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES = get_or_create_env_var( + "ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES", "" +) +ECS_SERVICE_CONNECT_CLIENT_CDK_PREFIXES = get_or_create_env_var( + "ECS_SERVICE_CONNECT_CLIENT_CDK_PREFIXES", "" +) + + +def normalize_https_redirect_url(url: str) -> str: + """Ensure Cognito/OAuth redirect bases use an explicit https:// scheme.""" + raw = (url or "").strip() + if not raw: + return "" + if raw.startswith("https://"): + return raw.rstrip("/") + if raw.startswith("http://"): + return ("https://" + raw[len("http://") :]).rstrip("/") + return ("https://" + raw.lstrip("/")).rstrip("/") + + +# This should be the CloudFront domain, the domain linked to your ACM certificate, or the DNS of your application load balancer in console afterwards +if USE_CLOUDFRONT == "True": + COGNITO_REDIRECTION_URL = get_or_create_env_var( + "COGNITO_REDIRECTION_URL", "https://" + CLOUDFRONT_DOMAIN + ) +elif SSL_CERTIFICATE_DOMAIN: + COGNITO_REDIRECTION_URL = get_or_create_env_var( + "COGNITO_REDIRECTION_URL", "https://" + SSL_CERTIFICATE_DOMAIN + ) +elif USE_ECS_EXPRESS_MODE == "True": + _express_redirect_default = ECS_EXPRESS_COGNITO_REDIRECT_BASE or ( + "https://" + EXISTING_LOAD_BALANCER_DNS + ) + COGNITO_REDIRECTION_URL = get_or_create_env_var( + "COGNITO_REDIRECTION_URL", _express_redirect_default + ) +else: + COGNITO_REDIRECTION_URL = get_or_create_env_var( + "COGNITO_REDIRECTION_URL", "https://" + EXISTING_LOAD_BALANCER_DNS + ) + +COGNITO_REDIRECTION_URL = normalize_https_redirect_url(COGNITO_REDIRECTION_URL) +if ECS_EXPRESS_COGNITO_REDIRECT_BASE: + ECS_EXPRESS_COGNITO_REDIRECT_BASE = normalize_https_redirect_url( + ECS_EXPRESS_COGNITO_REDIRECT_BASE + ) + +# Custom headers e.g. if routing traffic through Cloudfront +CUSTOM_HEADER = get_or_create_env_var( + "CUSTOM_HEADER", "" +) # Retrieving or setting CUSTOM_HEADER +CUSTOM_HEADER_VALUE = get_or_create_env_var( + "CUSTOM_HEADER_VALUE", "" +) # Retrieving or setting CUSTOM_HEADER_VALUE + +# Firewall on top of load balancer +LOAD_BALANCER_WEB_ACL_NAME = get_or_create_env_var( + "LOAD_BALANCER_WEB_ACL_NAME", f"{CDK_PREFIX}alb-web-acl" +) + +# Firewall on top of CloudFront +WEB_ACL_NAME = get_or_create_env_var("WEB_ACL_NAME", f"{CDK_PREFIX}cloudfront-web-acl") + +### +# File I/O options +### + +OUTPUT_FOLDER = get_or_create_env_var("GRADIO_OUTPUT_FOLDER", "output/") # 'output/' +INPUT_FOLDER = get_or_create_env_var("GRADIO_INPUT_FOLDER", "input/") # 'input/' + +# Allow for files to be saved in a temporary folder for increased security in some instances +if OUTPUT_FOLDER == "TEMP" or INPUT_FOLDER == "TEMP": + # Create a temporary directory + with tempfile.TemporaryDirectory() as temp_dir: + print(f"Temporary directory created at: {temp_dir}") + + if OUTPUT_FOLDER == "TEMP": + OUTPUT_FOLDER = temp_dir + "/" + if INPUT_FOLDER == "TEMP": + INPUT_FOLDER = temp_dir + "/" + +### +# LOGGING OPTIONS +### + +SAVE_LOGS_TO_CSV = get_or_create_env_var("SAVE_LOGS_TO_CSV", "True") + +### DYNAMODB logs. Whether to save to DynamoDB, and the headers of the table +SAVE_LOGS_TO_DYNAMODB = get_or_create_env_var("SAVE_LOGS_TO_DYNAMODB", "True") +ACCESS_LOG_DYNAMODB_TABLE_NAME = get_or_create_env_var( + "ACCESS_LOG_DYNAMODB_TABLE_NAME", f"{CDK_PREFIX}dynamodb-access-logs".lower() +) +FEEDBACK_LOG_DYNAMODB_TABLE_NAME = get_or_create_env_var( + "FEEDBACK_LOG_DYNAMODB_TABLE_NAME", f"{CDK_PREFIX}dynamodb-feedback-logs".lower() +) +USAGE_LOG_DYNAMODB_TABLE_NAME = get_or_create_env_var( + "USAGE_LOG_DYNAMODB_TABLE_NAME", f"{CDK_PREFIX}dynamodb-usage-logs".lower() +) + +# Optional scheduled export of usage-log DynamoDB table to CSV in S3 (EventBridge -> Lambda). +ENABLE_DYNAMODB_USAGE_LOG_EXPORT = get_or_create_env_var( + "ENABLE_DYNAMODB_USAGE_LOG_EXPORT", "False" +) +DYNAMODB_USAGE_LOG_EXPORT_LAMBDA_NAME = get_or_create_env_var( + "DYNAMODB_USAGE_LOG_EXPORT_LAMBDA_NAME", + f"{CDK_PREFIX}DynamoUsageLogExport".lower().replace("_", "-"), +) +DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE = get_or_create_env_var( + "DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE", "cron(0 6 ? * * *)" +) +DYNAMODB_USAGE_LOG_EXPORT_S3_KEY = get_or_create_env_var( + "DYNAMODB_USAGE_LOG_EXPORT_S3_KEY", + "reports/dynamodb-usage/dynamodb_logs_export.csv", +) +DYNAMODB_USAGE_LOG_EXPORT_DATE_ATTRIBUTE = get_or_create_env_var( + "DYNAMODB_USAGE_LOG_EXPORT_DATE_ATTRIBUTE", "timestamp" +) +DYNAMODB_USAGE_LOG_EXPORT_OUTPUT_FILENAME = get_or_create_env_var( + "DYNAMODB_USAGE_LOG_EXPORT_OUTPUT_FILENAME", "dynamodb_logs_export.csv" +) + +if ENABLE_DYNAMODB_USAGE_LOG_EXPORT == "True": + if SAVE_LOGS_TO_DYNAMODB != "True": + raise ValueError( + "ENABLE_DYNAMODB_USAGE_LOG_EXPORT=True requires SAVE_LOGS_TO_DYNAMODB=True." + ) + schedule = (DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE or "").strip() + if not schedule.startswith("cron(") and not schedule.startswith("rate("): + raise ValueError( + "DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE must be an EventBridge cron(...) or " + f"rate(...) expression; got {schedule!r}." + ) + if not (DYNAMODB_USAGE_LOG_EXPORT_S3_KEY or "").strip(): + raise ValueError( + "DYNAMODB_USAGE_LOG_EXPORT_S3_KEY is required when " + "ENABLE_DYNAMODB_USAGE_LOG_EXPORT=True." + ) + +### +# APP OPTIONS +### + +# Get some environment variables and Launch the Gradio app +COGNITO_AUTH = get_or_create_env_var("COGNITO_AUTH", "False") + +GRADIO_SERVER_PORT = int(get_or_create_env_var("GRADIO_SERVER_PORT", "7860")) + +# Must match the named port mapping on the Fargate container (see cdk_stack.py). +ECS_SERVICE_CONNECT_PORT_MAPPING_NAME = get_or_create_env_var( + "ECS_SERVICE_CONNECT_PORT_MAPPING_NAME", f"port-{GRADIO_SERVER_PORT}" +) + +# Suffix used with ECS_SERVICE_CONNECT_CLIENT_CDK_PREFIXES (matches this stack's ECS SG name). +if ECS_SECURITY_GROUP_NAME.startswith(CDK_PREFIX): + _default_sc_client_sg_suffix = ECS_SECURITY_GROUP_NAME[len(CDK_PREFIX) :] +else: + _default_sc_client_sg_suffix = "SecurityGroupECS" +ECS_SERVICE_CONNECT_CLIENT_SG_NAME_SUFFIX = get_or_create_env_var( + "ECS_SERVICE_CONNECT_CLIENT_SG_NAME_SUFFIX", _default_sc_client_sg_suffix +) + +ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES_LIST = parse_comma_separated_list( + ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES +) +ECS_SERVICE_CONNECT_CLIENT_CDK_PREFIXES_LIST = parse_comma_separated_list( + ECS_SERVICE_CONNECT_CLIENT_CDK_PREFIXES +) + + +def build_service_connect_client_security_group_names() -> List[str]: + """Explicit SG names plus {prefix}{suffix} for each client CDK_PREFIX.""" + names: List[str] = list(ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES_LIST) + for prefix in ECS_SERVICE_CONNECT_CLIENT_CDK_PREFIXES_LIST: + names.append(f"{prefix}{ECS_SERVICE_CONNECT_CLIENT_SG_NAME_SUFFIX}") + deduped: List[str] = [] + seen = set() + for name in names: + if name and name not in seen: + seen.add(name) + deduped.append(name) + return deduped + + +ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES_TO_LOOKUP = ( + build_service_connect_client_security_group_names() +) + +if ENABLE_ECS_SERVICE_CONNECT == "True" and USE_ECS_EXPRESS_MODE == "True": + raise ValueError( + "ENABLE_ECS_SERVICE_CONNECT=True is only supported on the legacy Fargate " + "service path. Set USE_ECS_EXPRESS_MODE=False or disable Service Connect." + ) + +# Headless deployment: S3 job .env -> Lambda -> one-shot ECS Fargate (RUN_DIRECT_MODE). +# No ALB, CloudFront, or always-on ECS service. +ENABLE_HEADLESS_DEPLOYMENT = get_or_create_env_var( + "ENABLE_HEADLESS_DEPLOYMENT", "False" +) + +# S3-uploaded job .env files trigger one-shot ECS Fargate tasks (direct mode / cli_redact). +ENABLE_S3_BATCH_ECS_TRIGGER = get_or_create_env_var( + "ENABLE_S3_BATCH_ECS_TRIGGER", "False" +) +S3_BATCH_ENV_PREFIX = get_or_create_env_var("S3_BATCH_ENV_PREFIX", "input/config/") +S3_BATCH_ENV_SUFFIX = get_or_create_env_var("S3_BATCH_ENV_SUFFIX", ".env") +S3_BATCH_INPUT_PREFIX = get_or_create_env_var("S3_BATCH_INPUT_PREFIX", "input/") +S3_BATCH_CONFIG_PREFIX = get_or_create_env_var("S3_BATCH_CONFIG_PREFIX", "") +S3_BATCH_GENERAL_ENV_PREFIX = get_or_create_env_var( + "S3_BATCH_GENERAL_ENV_PREFIX", "general-config/" +) +S3_BATCH_DEFAULT_PARAMS_KEY = get_or_create_env_var( + "S3_BATCH_DEFAULT_PARAMS_KEY", "general-config/app_defaults.env" +) +S3_BATCH_LAMBDA_FUNCTION_NAME = get_or_create_env_var( + "S3_BATCH_LAMBDA_FUNCTION_NAME", f"{CDK_PREFIX}S3BatchEcsTrigger" +) + +if ENABLE_S3_BATCH_ECS_TRIGGER == "True" and USE_ECS_EXPRESS_MODE == "True": + raise ValueError( + "ENABLE_S3_BATCH_ECS_TRIGGER=True requires the legacy Fargate task definition " + "for ecs.run_task. Set USE_ECS_EXPRESS_MODE=False or disable the batch trigger." + ) + +if ENABLE_HEADLESS_DEPLOYMENT == "True": + if ENABLE_S3_BATCH_ECS_TRIGGER != "True": + raise ValueError( + "ENABLE_HEADLESS_DEPLOYMENT=True requires ENABLE_S3_BATCH_ECS_TRIGGER=True." + ) + if USE_ECS_EXPRESS_MODE == "True": + raise ValueError( + "ENABLE_HEADLESS_DEPLOYMENT=True requires USE_ECS_EXPRESS_MODE=False." + ) + if USE_CLOUDFRONT == "True": + raise ValueError( + "ENABLE_HEADLESS_DEPLOYMENT=True is incompatible with USE_CLOUDFRONT=True." + ) + +# Optional headless follow-on: S3 output PutRequests alarm -> SNS email + IAM user for downloads. +ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS = get_or_create_env_var( + "ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS", "False" +) +HEADLESS_OUTPUT_NOTIFY_EMAIL = get_or_create_env_var("HEADLESS_OUTPUT_NOTIFY_EMAIL", "") +HEADLESS_OUTPUT_IAM_USER_NAME = get_or_create_env_var( + "HEADLESS_OUTPUT_IAM_USER_NAME", f"{CDK_PREFIX}s3-output-reader" +) +HEADLESS_OUTPUT_S3_METRIC_FILTER_ID = get_or_create_env_var( + "HEADLESS_OUTPUT_S3_METRIC_FILTER_ID", + f"{CDK_PREFIX}s3-output-put".lower().replace("_", "-"), +) +HEADLESS_OUTPUT_SNS_TOPIC_NAME = get_or_create_env_var( + "HEADLESS_OUTPUT_SNS_TOPIC_NAME", + f"{CDK_PREFIX}llm-topic-s3-save-sns".lower().replace("_", "-"), +) +HEADLESS_OUTPUT_ALARM_NAME = get_or_create_env_var( + "HEADLESS_OUTPUT_ALARM_NAME", + f"{CDK_PREFIX}cloudwatch-alarm-new-output-s3".lower().replace("_", "-"), +) +HEADLESS_OUTPUT_S3_PREFIX = get_or_create_env_var( + "HEADLESS_OUTPUT_S3_PREFIX", "output/" +) +HEADLESS_OUTPUT_IAM_SECRET_NAME = get_or_create_env_var( + "HEADLESS_OUTPUT_IAM_SECRET_NAME", + f"{CDK_PREFIX}headless-output-reader-key", +) + +if ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS == "True": + if ENABLE_HEADLESS_DEPLOYMENT != "True": + raise ValueError( + "ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS=True requires " + "ENABLE_HEADLESS_DEPLOYMENT=True." + ) + if not (HEADLESS_OUTPUT_NOTIFY_EMAIL or "").strip(): + raise ValueError( + "HEADLESS_OUTPUT_NOTIFY_EMAIL is required when " + "ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS=True." + ) + +# Pi agent Gradio UI (second Fargate service; shared legacy ALB + Service Connect to main app). +ENABLE_PI_AGENT_ECS_SERVICE = get_or_create_env_var( + "ENABLE_PI_AGENT_ECS_SERVICE", "False" +) +ECR_PI_REPO_NAME = get_or_create_env_var( + "ECR_PI_REPO_NAME", f"{CDK_PREFIX}pi-agent".lower() +) +CODEBUILD_PI_PROJECT_NAME = get_or_create_env_var( + "CODEBUILD_PI_PROJECT_NAME", f"{CDK_PREFIX}CodeBuildPiAgent" +) +ECS_PI_TASK_DEFINITION_NAME = get_or_create_env_var( + "ECS_PI_TASK_DEFINITION_NAME", f"{CDK_PREFIX}PiAgentTaskDefinition" +) +ECS_PI_SECURITY_GROUP_NAME = get_or_create_env_var( + "ECS_PI_SECURITY_GROUP_NAME", f"{CDK_PREFIX}SecurityGroupPiAgent" +) +ECS_PI_LOG_GROUP_NAME = get_or_create_env_var( + "ECS_PI_LOG_GROUP_NAME", f"/ecs/{ECS_PI_SERVICE_NAME}-logs".lower() +) +ECS_PI_TASK_CPU_SIZE = get_or_create_env_var("ECS_PI_TASK_CPU_SIZE", "1024") +ECS_PI_TASK_MEMORY_SIZE = get_or_create_env_var("ECS_PI_TASK_MEMORY_SIZE", "2048") +PI_GRADIO_PORT = get_or_create_env_var("PI_GRADIO_PORT", "7862") +# Pi ALB routing: path (default /pi on shared host e.g. CloudFront), host, or both. +PI_ALB_ROUTING = get_or_create_env_var("PI_ALB_ROUTING", "path").strip().lower() +PI_ALB_PATH_PREFIX = get_or_create_env_var("PI_ALB_PATH_PREFIX", "/pi") +PI_ALB_HOST_HEADER = get_or_create_env_var("PI_ALB_HOST_HEADER", "") +PI_ALB_TARGET_GROUP_NAME = get_or_create_env_var( + "PI_ALB_TARGET_GROUP_NAME", f"{CDK_PREFIX}PiAgentTG"[-32:] +) +PI_ALB_LISTENER_RULE_PRIORITY = int( + get_or_create_env_var("PI_ALB_LISTENER_RULE_PRIORITY", "3") +) +PI_AGENT_ENV_S3_KEY = get_or_create_env_var("PI_AGENT_ENV_S3_KEY", "pi_agent.env") + + +def _normalize_pi_alb_path_prefix(raw: str) -> str: + segment = (raw or "pi").strip().strip("/") + return f"/{segment}" if segment else "/pi" + + +PI_ALB_PATH_PREFIX_NORMALIZED = _normalize_pi_alb_path_prefix(PI_ALB_PATH_PREFIX) +_PI_ALB_ROUTING_MODES = frozenset({"path", "host", "both"}) + + +def _validate_pi_alb_routing_for_enabled_pi() -> None: + if PI_ALB_ROUTING not in _PI_ALB_ROUTING_MODES: + raise ValueError( + f"PI_ALB_ROUTING must be one of {sorted(_PI_ALB_ROUTING_MODES)}; got '{PI_ALB_ROUTING}'." + ) + if PI_ALB_ROUTING in ("host", "both") and not PI_ALB_HOST_HEADER.strip(): + raise ValueError( + "PI_ALB_HOST_HEADER is required when PI_ALB_ROUTING is 'host' or 'both' " + "(dedicated hostname on the shared ALB)." + ) + if PI_ALB_ROUTING in ("path", "both") and not PI_ALB_PATH_PREFIX_NORMALIZED: + raise ValueError("PI_ALB_PATH_PREFIX must resolve to a non-empty path segment.") + + +# Pi on ECS Express Mode (second Express service on shared ALB; SC via ecs:UpdateService). +ENABLE_PI_AGENT_EXPRESS_SERVICE = get_or_create_env_var( + "ENABLE_PI_AGENT_EXPRESS_SERVICE", "False" +) +ECS_PI_EXPRESS_SERVICE_NAME = get_or_create_env_var( + "ECS_PI_EXPRESS_SERVICE_NAME", f"{CDK_PREFIX}PiExpressService" +) +ECS_PI_EXPRESS_HEALTH_CHECK_PATH = get_or_create_env_var( + "ECS_PI_EXPRESS_HEALTH_CHECK_PATH", "/health" +) +ECS_PI_EXPRESS_SECURITY_GROUP_NAME = get_or_create_env_var( + "ECS_PI_EXPRESS_SECURITY_GROUP_NAME", f"{CDK_PREFIX}SecurityGroupPiExpress" +) +# Service Connect port names for Express services (applied in post_cdk_build_quickstart.py). +ECS_EXPRESS_SC_PORT_NAME = get_or_create_env_var( + "ECS_EXPRESS_SC_PORT_NAME", ECS_SERVICE_CONNECT_PORT_MAPPING_NAME +) +ECS_PI_EXPRESS_SC_PORT_NAME = get_or_create_env_var( + "ECS_PI_EXPRESS_SC_PORT_NAME", f"port-{PI_GRADIO_PORT}" +) + +if ENABLE_PI_AGENT_ECS_SERVICE == "True" and ENABLE_PI_AGENT_EXPRESS_SERVICE == "True": + raise ValueError( + "Enable at most one Pi deployment mode: ENABLE_PI_AGENT_ECS_SERVICE (legacy Fargate) " + "or ENABLE_PI_AGENT_EXPRESS_SERVICE (Express), not both." + ) +if ENABLE_PI_AGENT_EXPRESS_SERVICE == "True" and USE_ECS_EXPRESS_MODE != "True": + raise ValueError( + "ENABLE_PI_AGENT_EXPRESS_SERVICE=True requires USE_ECS_EXPRESS_MODE=True " + "(no ACM_SSL_CERTIFICATE_ARN)." + ) +if ENABLE_PI_AGENT_ECS_SERVICE == "True" and USE_ECS_EXPRESS_MODE == "True": + raise ValueError( + "ENABLE_PI_AGENT_ECS_SERVICE=True requires legacy Fargate (USE_ECS_EXPRESS_MODE=False). " + "For Pi on Express, use ENABLE_PI_AGENT_EXPRESS_SERVICE=True instead." + ) +if ENABLE_PI_AGENT_ECS_SERVICE == "True" and ENABLE_ECS_SERVICE_CONNECT != "True": + raise ValueError( + "ENABLE_PI_AGENT_ECS_SERVICE=True requires ENABLE_ECS_SERVICE_CONNECT=True " + "so the Pi task can reach the main app at http://:7860." + ) +if ENABLE_PI_AGENT_ECS_SERVICE == "True": + _validate_pi_alb_routing_for_enabled_pi() + +### +# WHOLE DOCUMENT API OPTIONS +### + +DAYS_TO_DISPLAY_WHOLE_DOCUMENT_JOBS = get_or_create_env_var( + "DAYS_TO_DISPLAY_WHOLE_DOCUMENT_JOBS", "7" +) # How many days into the past should whole document Textract jobs be displayed? After that, the data is not deleted from the Textract jobs csv, but it is just filtered out. Included to align with S3 buckets where the file outputs will be automatically deleted after X days. diff --git a/cdk/cdk_functions.py b/cdk/cdk_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..66e42f87b9d19d8b69afa1648c3ed8f485371098 --- /dev/null +++ b/cdk/cdk_functions.py @@ -0,0 +1,4268 @@ +import ipaddress +import json +import os +import re +from typing import Any, Dict, FrozenSet, List, Literal, Optional, Tuple, Union + +S3BucketAvailability = Literal["owned", "available", "globally_taken"] +CognitoDomainAvailability = Literal["available", "taken"] + +import boto3 +import pandas as pd +from aws_cdk import ( + App, + CfnOutput, + CfnTag, + CustomResource, + Duration, + Fn, + RemovalPolicy, + Stack, + Tags, +) +from aws_cdk import aws_cloudwatch as cloudwatch +from aws_cdk import aws_cloudwatch_actions as cloudwatch_actions +from aws_cdk import aws_codebuild as codebuild +from aws_cdk import aws_cognito as cognito +from aws_cdk import aws_dynamodb as dynamodb +from aws_cdk import aws_ec2 as ec2 +from aws_cdk import aws_ecs as ecs +from aws_cdk import aws_elasticloadbalancingv2 as elb +from aws_cdk import aws_elasticloadbalancingv2_actions as elb_act +from aws_cdk import aws_events as events +from aws_cdk import aws_events_targets as targets +from aws_cdk import aws_iam as iam +from aws_cdk import aws_lambda as lambda_ +from aws_cdk import aws_logs as logs +from aws_cdk import aws_s3 as s3 +from aws_cdk import aws_s3_notifications as s3n +from aws_cdk import aws_secretsmanager as secretsmanager +from aws_cdk import aws_sns as sns +from aws_cdk import aws_sns_subscriptions as sns_subscriptions +from aws_cdk import aws_wafv2 as wafv2 +from aws_cdk import custom_resources as cr +from botocore.exceptions import ClientError, NoCredentialsError +from cdk_config import ( + ACCESS_LOG_DYNAMODB_TABLE_NAME, + APP_CONFIG_ENV_BASENAME, + AWS_REGION, + ECS_AVAILABILITY_ZONE_REBALANCING, + ENABLE_RESOURCE_DELETE_PROTECTION, + FEEDBACK_LOG_DYNAMODB_TABLE_NAME, + NAT_GATEWAY_EIP_NAME, + PRIVATE_SUBNET_AVAILABILITY_ZONES, + PRIVATE_SUBNET_CIDR_BLOCKS, + PRIVATE_SUBNETS_TO_USE, + PUBLIC_SUBNET_AVAILABILITY_ZONES, + PUBLIC_SUBNET_CIDR_BLOCKS, + PUBLIC_SUBNETS_TO_USE, + S3_LOG_CONFIG_BUCKET_NAME, + S3_OUTPUT_BUCKET_NAME, + USAGE_LOG_DYNAMODB_TABLE_NAME, +) +from constructs import Construct +from dotenv import dotenv_values, set_key + +# CDK CLI stores lookup-provider results under these key prefixes in cdk.context.json. +_CDK_LOOKUP_CONTEXT_PREFIXES = ( + "vpc-provider:", + "load-balancer:", + "availability-zones:", + "hosted-zone:", + "security-group:", + "key-provider:", + "ami:", +) + + +def _ensure_folder_exists(output_folder: str) -> None: + if not os.path.exists(output_folder): + os.makedirs(output_folder, exist_ok=True) + print(f"Created the {output_folder} folder.") + else: + print(f"The {output_folder} folder already exists.") + + +def is_resource_delete_protection_enabled() -> bool: + """Whether stack and resource delete protection is enabled (see ENABLE_RESOURCE_DELETE_PROTECTION).""" + return str(ENABLE_RESOURCE_DELETE_PROTECTION).strip().lower() in ( + "true", + "1", + "yes", + ) + + +def resource_deletion_protection_flag() -> bool: + """AWS deletion_protection attribute (ALB, DynamoDB tables, Cognito user pools).""" + return is_resource_delete_protection_enabled() + + +def ecs_availability_zone_rebalancing( + setting: str, +) -> ecs.AvailabilityZoneRebalancing: + """Map ``ECS_AVAILABILITY_ZONE_REBALANCING`` env value to the CDK enum.""" + if setting == "ENABLED": + return ecs.AvailabilityZoneRebalancing.ENABLED + return ecs.AvailabilityZoneRebalancing.DISABLED + + +def managed_resource_removal_policy() -> RemovalPolicy: + """Removal policy for CDK-managed resources without a native deletion_protection flag.""" + return ( + RemovalPolicy.RETAIN + if is_resource_delete_protection_enabled() + else RemovalPolicy.DESTROY + ) + + +def s3_auto_delete_objects_on_stack_destroy() -> bool: + """Empty S3 buckets automatically when the stack is destroyed (dev/teardown only).""" + return not is_resource_delete_protection_enabled() + + +def ecr_empty_on_delete() -> bool: + """Force-delete ECR images when the repository is removed on stack destroy.""" + return not is_resource_delete_protection_enabled() + + +def purge_cdk_lookup_context(file_path: str) -> int: + """Remove stale CDK lookup cache entries that require the bootstrap lookup role.""" + if not os.path.exists(file_path): + return 0 + with open(file_path, "r", encoding="utf-8") as f: + context_data = json.load(f) + cleaned = { + key: value + for key, value in context_data.items() + if not key.startswith(_CDK_LOOKUP_CONTEXT_PREFIXES) + } + removed = len(context_data) - len(cleaned) + if removed: + with open(file_path, "w", encoding="utf-8") as f: + json.dump(cleaned, f, indent=2) + print(f"Removed {removed} stale CDK lookup context key(s) from {file_path}.") + return removed + + +def log_aws_credential_context( + expected_account_id: Optional[str] = None, + expected_region: Optional[str] = None, +) -> Dict[str, Any]: + """ + Print the active AWS identity and non-secret credential hints for CDK debugging. + + Helps distinguish SSO/assumed-role sessions from long-lived access keys in + ~/.aws/credentials or environment variables. + """ + profile = os.environ.get("AWS_PROFILE") or "(not set — using default profile chain)" + default_region = ( + os.environ.get("AWS_REGION") + or os.environ.get("AWS_DEFAULT_REGION") + or "(not set in environment)" + ) + env_access_key_set = bool(os.environ.get("AWS_ACCESS_KEY_ID")) + env_secret_key_set = bool(os.environ.get("AWS_SECRET_ACCESS_KEY")) + env_session_token_set = bool(os.environ.get("AWS_SESSION_TOKEN")) + + print("\n--- AWS credential context (CDK / boto3) ---") + print(f"AWS_PROFILE: {profile}") + print(f"AWS_REGION / AWS_DEFAULT_REGION (env): {default_region}") + print( + "Environment credential variables: " + f"AWS_ACCESS_KEY_ID={'set' if env_access_key_set else 'not set'}, " + f"AWS_SECRET_ACCESS_KEY={'set' if env_secret_key_set else 'not set'}, " + f"AWS_SESSION_TOKEN={'set' if env_session_token_set else 'not set'}" + ) + if expected_account_id: + print(f"Configured CDK target account (AWS_ACCOUNT_ID): {expected_account_id}") + if expected_region: + print(f"Configured CDK target region (AWS_REGION): {expected_region}") + + session = boto3.Session() + active_profile = session.profile_name or "(default)" + print(f"boto3 session profile: {active_profile}") + print(f"boto3 session region: {session.region_name or '(not set)'}") + + credentials = session.get_credentials() + credential_summary: Dict[str, Any] = { + "profile": profile, + "session_profile": active_profile, + } + + if credentials is None: + print("WARNING: No AWS credentials found in the default provider chain.") + print("--- End AWS credential context ---\n") + credential_summary["error"] = "no_credentials" + return credential_summary + + frozen = credentials.get_frozen_credentials() + access_key = frozen.access_key or "" + access_key_prefix = (access_key[:4] + "...") if len(access_key) >= 4 else "(none)" + credential_summary["access_key_prefix"] = access_key_prefix + + if env_access_key_set: + credential_source = "environment variables (highest precedence)" + elif access_key.startswith("AKIA"): + credential_source = "long-lived access key (likely ~/.aws/credentials [default] or named profile)" + elif access_key.startswith("ASIA"): + credential_source = "temporary credentials (SSO, assumed role, or STS session)" + else: + credential_source = ( + "resolved credentials (source could not be classified from key prefix)" + ) + + print(f"Inferred credential type: {credential_source}") + credential_summary["inferred_credential_type"] = credential_source + + if env_access_key_set and profile != "(not set — using default profile chain)": + print( + "NOTE: AWS_ACCESS_KEY_ID is set in the environment, so it overrides " + f"profile '{profile}' and SSO." + ) + + try: + sts = session.client("sts", region_name=session.region_name or expected_region) + identity = sts.get_caller_identity() + except (ClientError, NoCredentialsError) as exc: + print(f"WARNING: sts:GetCallerIdentity failed: {exc}") + print("--- End AWS credential context ---\n") + credential_summary["error"] = str(exc) + return credential_summary + + account = identity.get("Account", "") + arn = identity.get("Arn", "") + user_id = identity.get("UserId", "") + + print(f"Caller account: {account}") + print(f"Caller ARN: {arn}") + print(f"Caller UserId: {user_id}") + + if ":assumed-role/" in arn: + principal_kind = "assumed IAM role (typical for SSO or role chaining)" + elif ":user/" in arn: + principal_kind = "IAM user (typical for static access keys in credentials file)" + elif ":federated-user/" in arn: + principal_kind = "federated user" + else: + principal_kind = "other IAM principal" + + print(f"Principal kind: {principal_kind}") + credential_summary.update( + { + "account": account, + "arn": arn, + "user_id": user_id, + "principal_kind": principal_kind, + } + ) + + if expected_account_id and account and account != str(expected_account_id): + print( + "WARNING: Caller account does not match configured AWS_ACCOUNT_ID. " + "CDK will target the configured account but act as this identity — " + "deployments and lookups may fail. Set AWS_PROFILE to your SSO profile " + "and unset AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY if needed." + ) + credential_summary["account_mismatch"] = True + elif expected_account_id and account == str(expected_account_id): + print("Caller account matches configured AWS_ACCOUNT_ID.") + + if profile == "(not set — using default profile chain)": + print( + "TIP: Set AWS_PROFILE to your SSO profile name so Python and the CDK CLI " + "(Node) use the same session. Example: " + '$env:AWS_PROFILE = "YourSsoProfileName"' + ) + + print("--- End AWS credential context ---\n") + return credential_summary + + +# --- Function to load context from file --- + + +def _context_value_for_cdk(value): + """CDK/JSII context cannot use JSON null; normalize for Windows synth.""" + if value is None: + return "" + if isinstance(value, dict): + return {k: _context_value_for_cdk(v) for k, v in value.items()} + if isinstance(value, list): + return [_context_value_for_cdk(v) for v in value] + return value + + +def load_context_from_file(app: App, file_path: str): + if os.path.exists(file_path): + with open(file_path, "r", encoding="utf-8") as f: + context_data = json.load(f) + for key, value in context_data.items(): + app.node.set_context(key, _context_value_for_cdk(value)) + print(f"Loaded context from {file_path}") + else: + print(f"Context file not found: {file_path}") + + +# --- Helper to parse environment variables into lists --- +def _get_env_list(env_var_name: str) -> List[str]: + """Parses a comma-separated environment variable into a list of strings.""" + value = env_var_name[1:-1].strip().replace('"', "").replace("'", "") + if not value: + return [] + # Split by comma and filter out any empty strings that might result from extra commas + return [s.strip() for s in value.split(",") if s.strip()] + + +# 1. Try to load CIDR/AZs from environment variables +if PUBLIC_SUBNETS_TO_USE: + PUBLIC_SUBNETS_TO_USE = _get_env_list(PUBLIC_SUBNETS_TO_USE) +if PRIVATE_SUBNETS_TO_USE: + PRIVATE_SUBNETS_TO_USE = _get_env_list(PRIVATE_SUBNETS_TO_USE) + +if PUBLIC_SUBNET_CIDR_BLOCKS: + PUBLIC_SUBNET_CIDR_BLOCKS = _get_env_list("PUBLIC_SUBNET_CIDR_BLOCKS") +if PUBLIC_SUBNET_AVAILABILITY_ZONES: + PUBLIC_SUBNET_AVAILABILITY_ZONES = _get_env_list("PUBLIC_SUBNET_AVAILABILITY_ZONES") +if PRIVATE_SUBNET_CIDR_BLOCKS: + PRIVATE_SUBNET_CIDR_BLOCKS = _get_env_list("PRIVATE_SUBNET_CIDR_BLOCKS") +if PRIVATE_SUBNET_AVAILABILITY_ZONES: + PRIVATE_SUBNET_AVAILABILITY_ZONES = _get_env_list( + "PRIVATE_SUBNET_AVAILABILITY_ZONES" + ) + + +def resolve_policy_file_paths( + paths: List[str], + *, + cdk_folder: str = "", +) -> List[str]: + """Resolve JSON policy paths relative to ``CDK_FOLDER`` when not absolute.""" + resolved: List[str] = [] + base = (cdk_folder or "").strip().rstrip("/\\") + for raw in paths: + path = (raw or "").strip() + if not path: + continue + if os.path.isabs(path): + resolved.append(os.path.normpath(path)) + elif base: + resolved.append(os.path.normpath(os.path.join(base, path))) + else: + resolved.append(os.path.normpath(path)) + return resolved + + +def attach_managed_policy_arns(role: iam.IRole, policy_arns: List[str]) -> None: + """Attach existing customer-managed policies by full ARN.""" + for arn in policy_arns: + arn = (arn or "").strip() + if not arn: + continue + role.add_managed_policy( + iam.ManagedPolicy.from_managed_policy_arn(role, arn, arn) + ) + print(f"Attached managed policy ARN to {role.node.id}: {arn}") + + +def check_for_existing_role(role_name: str): + try: + iam = boto3.client("iam") + # iam.get_role(RoleName=role_name) + + response = iam.get_role(RoleName=role_name) + role = response["Role"]["Arn"] + + print("Response Role:", role) + + return True, role, "" + except iam.exceptions.NoSuchEntityException: + return False, "", "" + except Exception as e: + raise Exception("Getting information on IAM role failed due to:", e) + + +from typing import List + + +def add_statement_to_policy(role: iam.IRole, policy_document: Dict[str, Any]): + """ + Adds individual policy statements from a parsed policy document to a CDK Role. + + Args: + role: The CDK Role construct to attach policies to. + policy_document: A Python dictionary representing an IAM policy document. + """ + # Ensure the loaded JSON is a valid policy document structure + if "Statement" not in policy_document or not isinstance( + policy_document["Statement"], list + ): + print("Warning: Policy document does not contain a 'Statement' list. Skipping.") + return # Do not return role, just log and exit + + for statement_dict in policy_document["Statement"]: + try: + # Create a CDK PolicyStatement from the dictionary + cdk_policy_statement = iam.PolicyStatement.from_json(statement_dict) + + # Add the policy statement to the role + role.add_to_policy(cdk_policy_statement) + print(f" - Added statement: {statement_dict.get('Sid', 'No Sid')}") + except Exception as e: + print( + f"Warning: Could not process policy statement: {statement_dict}. Error: {e}" + ) + + +def vpc_endpoint_aws_service_name(service_suffix: str, region: str) -> str: + """Full EC2 ``ServiceName`` for a VPC endpoint (matches describe_vpc_endpoints).""" + return f"com.amazonaws.{region}.{service_suffix}" + + +def list_existing_vpc_endpoint_service_names( + vpc_id: str, + *, + region_name: Optional[str] = None, +) -> FrozenSet[str]: + """Return AWS service names for non-deleted VPC endpoints in the given VPC.""" + ec2_client = boto3.client("ec2", region_name=region_name) + service_names: set[str] = set() + paginator = ec2_client.get_paginator("describe_vpc_endpoints") + for page in paginator.paginate( + Filters=[{"Name": "vpc-id", "Values": [vpc_id]}], + ): + for endpoint in page.get("VpcEndpoints", []): + state = endpoint.get("State") + if state in ("available", "pending", "pendingAcceptance"): + service_name = endpoint.get("ServiceName") + if service_name: + service_names.add(service_name) + return frozenset(service_names) + + +def resolve_ecs_vpc_endpoint_subnet_selection( + *, + use_express_ingress: bool, + express_use_public_subnets: bool, + public_subnets: List[ec2.ISubnet], + private_subnets: List[ec2.ISubnet], +) -> Optional[ec2.SubnetSelection]: + """ + Choose subnets for ECS-related **interface** VPC endpoints. + + Interface ENIs must sit in the same subnets ECS tasks use. Express Mode with + ``ECS_EXPRESS_USE_PUBLIC_SUBNETS=True`` runs tasks in public subnets; legacy + Fargate and Express-on-private use private. S3 gateway routes use + ``resolve_ecs_s3_gateway_subnet_selection`` instead. + """ + if use_express_ingress: + if express_use_public_subnets: + if not public_subnets: + return None + return ec2.SubnetSelection(subnets=public_subnets) + if not private_subnets: + return None + return ec2.SubnetSelection(subnets=private_subnets) + if private_subnets: + return ec2.SubnetSelection(subnets=private_subnets) + if public_subnets: + return ec2.SubnetSelection(subnets=public_subnets) + return None + + +def resolve_ecs_s3_gateway_subnet_selection( + *, + public_subnets: List[ec2.ISubnet], + private_subnets: List[ec2.ISubnet], +) -> Optional[ec2.SubnetSelection]: + """ + All stack-managed subnets for the S3 **gateway** endpoint. + + Gateway endpoints attach to route tables; every public and private subnet the + stack imports or creates should get the S3 prefix-list route, not only the ECS + task tier. + """ + all_subnets: List[ec2.ISubnet] = [] + seen_subnet_ids: set[str] = set() + for subnet in public_subnets + private_subnets: + subnet_id = subnet.subnet_id + if subnet_id in seen_subnet_ids: + continue + seen_subnet_ids.add(subnet_id) + all_subnets.append(subnet) + if not all_subnets: + return None + return ec2.SubnetSelection(subnets=all_subnets) + + +def list_vpc_associated_cidr_blocks(vpc: dict) -> List[str]: + """All associated IPv4 CIDR blocks for a VPC (primary and secondary).""" + cidrs: List[str] = [] + seen: set[str] = set() + for assoc in vpc.get("CidrBlockAssociationSet") or []: + state = (assoc.get("CidrBlockState") or {}).get("State") + if state and state != "associated": + continue + cidr = assoc.get("CidrBlock") + if cidr and cidr not in seen: + seen.add(cidr) + cidrs.append(cidr) + primary = vpc.get("CidrBlock") + if primary and primary not in seen: + cidrs.insert(0, primary) + return cidrs + + +def resolve_vpc_endpoint_ingress_cidr_blocks( + *, + vpc_cidr_block: Optional[str] = None, + vpc_cidr_blocks: Optional[List[str]] = None, + fallback_vpc_cidr: Optional[str] = None, +) -> List[str]: + """CIDR list for VPC endpoint SG ingress (deduplicated, stable order).""" + resolved: List[str] = [] + seen: set[str] = set() + for cidr in vpc_cidr_blocks or []: + if cidr and cidr not in seen: + seen.add(cidr) + resolved.append(cidr) + if not resolved: + single = vpc_cidr_block or fallback_vpc_cidr + if single: + resolved = [single] + return resolved + + +def add_vpc_endpoint_https_ingress_from_vpc_cidrs( + endpoint_sg: ec2.SecurityGroup, + *, + vpc_cidr_block: Optional[str] = None, + vpc_cidr_blocks: Optional[List[str]] = None, + fallback_vpc_cidr: Optional[str] = None, +) -> None: + """Allow HTTPS from every CIDR associated with the VPC.""" + cidrs = resolve_vpc_endpoint_ingress_cidr_blocks( + vpc_cidr_block=vpc_cidr_block, + vpc_cidr_blocks=vpc_cidr_blocks, + fallback_vpc_cidr=fallback_vpc_cidr, + ) + if not cidrs: + raise ValueError( + "VPC CIDR block(s) are required for the VPC endpoint security group. " + "Re-run check_resources.py so vpc_cidr_blocks is stored in precheck context." + ) + for cidr in cidrs: + endpoint_sg.add_ingress_rule( + peer=ec2.Peer.ipv4(cidr), + connection=ec2.Port.tcp(443), + description=( + "HTTPS from VPC workloads" + if len(cidrs) == 1 + else f"HTTPS from VPC workloads ({cidr})" + ), + ) + + +def create_ecs_vpc_endpoints_for_private_subnets( + scope: Construct, + *, + vpc: ec2.IVpc, + subnets: Optional[ec2.SubnetSelection], + logical_id_prefix: str = "Ecs", + include_secrets_and_kms: bool = True, + vpc_cidr_block: Optional[str] = None, + vpc_cidr_blocks: Optional[List[str]] = None, + skip_service_names: Optional[FrozenSet[str]] = None, + aws_region: str, + s3_gateway_subnets: Optional[ec2.SubnetSelection] = None, +) -> None: + """ + Interface and S3 gateway VPC endpoints for ECS workloads. + + Interface endpoints use ``subnets`` (ECS task tier). The S3 gateway uses + ``s3_gateway_subnets`` when provided, otherwise ``subnets``. Without + ``ecr.api`` / ``ecr.dkr`` endpoints (or a working NAT path), tasks fail with + ``GetAuthorizationToken`` timeouts. + + ``skip_service_names`` should list full AWS endpoint service names already present + in the VPC (from pre-check) so shared VPCs do not fail on duplicate private DNS. + """ + s3_subnets = s3_gateway_subnets or subnets + if not subnets and not s3_subnets: + return + skip = skip_service_names or frozenset() + + interface_services: List[Tuple[str, str, ec2.InterfaceVpcEndpointAwsService]] = [ + ("EcrApi", "ecr.api", ec2.InterfaceVpcEndpointAwsService.ECR), + ("EcrDkr", "ecr.dkr", ec2.InterfaceVpcEndpointAwsService.ECR_DOCKER), + ("CloudWatchLogs", "logs", ec2.InterfaceVpcEndpointAwsService.CLOUDWATCH_LOGS), + ] + if include_secrets_and_kms: + interface_services.extend( + [ + ( + "SecretsManager", + "secretsmanager", + ec2.InterfaceVpcEndpointAwsService.SECRETS_MANAGER, + ), + ("Kms", "kms", ec2.InterfaceVpcEndpointAwsService.KMS), + ] + ) + + interface_services_to_create = [] + for suffix, service_suffix, service in interface_services: + service_name = vpc_endpoint_aws_service_name(service_suffix, aws_region) + if service_name in skip: + print( + f"Skipping {logical_id_prefix}{suffix}Endpoint " + f"({service_name} already exists in this VPC)." + ) + continue + interface_services_to_create.append((suffix, service)) + + s3_service_name = vpc_endpoint_aws_service_name("s3", aws_region) + s3_service = ec2.GatewayVpcEndpointAwsService.S3 + create_s3_gateway = s3_service_name not in skip + if not create_s3_gateway: + print( + f"Skipping {logical_id_prefix}S3GatewayEndpoint " + f"({s3_service_name} already exists in this VPC)." + ) + + if not interface_services_to_create and not create_s3_gateway: + print( + "All ECS-related VPC endpoints already exist in this VPC; " + "nothing to create." + ) + return + + endpoint_sg = None + if interface_services_to_create and subnets: + endpoint_sg = ec2.SecurityGroup( + scope, + f"{logical_id_prefix}VpcEndpointSecurityGroup", + vpc=vpc, + description="HTTPS ingress for ECS-related VPC interface endpoints", + allow_all_outbound=True, + ) + add_vpc_endpoint_https_ingress_from_vpc_cidrs( + endpoint_sg, + vpc_cidr_block=vpc_cidr_block, + vpc_cidr_blocks=vpc_cidr_blocks, + fallback_vpc_cidr=vpc.vpc_cidr_block, + ) + + if subnets: + for suffix, service in interface_services_to_create: + vpc.add_interface_endpoint( + f"{logical_id_prefix}{suffix}Endpoint", + service=service, + subnets=subnets, + security_groups=[endpoint_sg], + private_dns_enabled=True, + ) + elif interface_services_to_create: + print( + "Skipping ECS interface VPC endpoints: no task-tier subnet selection " + "was provided." + ) + + if create_s3_gateway and s3_subnets: + try: + vpc.add_gateway_endpoint( + f"{logical_id_prefix}S3GatewayEndpoint", + service=s3_service, + subnets=[s3_subnets], + ) + except Exception as exc: + print( + "Note: could not add S3 gateway VPC endpoint (one may already exist on " + f"this VPC): {exc}" + ) + + +def add_s3_enforce_ssl_policy(bucket: s3.IBucket) -> None: + """Deny non-TLS S3 requests (Security Hub S3.5). Compatible with all CDK versions.""" + bucket.add_to_resource_policy( + iam.PolicyStatement( + effect=iam.Effect.DENY, + principals=[iam.AnyPrincipal()], + actions=["s3:*"], + resources=[bucket.bucket_arn, f"{bucket.bucket_arn}/*"], + conditions={"Bool": {"aws:SecureTransport": "false"}}, + ) + ) + + +def add_custom_policies( + scope: Construct, # Not strictly used here, but good practice if you expand to ManagedPolicies + role: iam.IRole, + policy_file_locations: Optional[List[str]] = None, + custom_policy_text: Optional[str] = None, +) -> iam.IRole: + """ + Loads custom policies from JSON files or a string and attaches them to a CDK Role. + + Args: + scope: The scope in which to define constructs (if needed, e.g., for iam.ManagedPolicy). + role: The CDK Role construct to attach policies to. + policy_file_locations: List of file paths to JSON policy documents. + custom_policy_text: A JSON string representing a policy document. + + Returns: + The modified CDK Role construct. + """ + if policy_file_locations is None: + policy_file_locations = [] + + current_source = "unknown source" # For error messages + + try: + if policy_file_locations: + print(f"Attempting to add policies from files to role {role.node.id}...") + for path in policy_file_locations: + current_source = f"file: {path}" + try: + with open(path, "r") as f: + policy_document = json.load(f) + print(f"Processing policy from {current_source}...") + add_statement_to_policy(role, policy_document) + except FileNotFoundError: + print(f"Warning: Policy file not found at {path}. Skipping.") + except json.JSONDecodeError as e: + print( + f"Warning: Invalid JSON in policy file {path}: {e}. Skipping." + ) + except Exception as e: + print( + f"An unexpected error occurred processing policy from {path}: {e}. Skipping." + ) + + if custom_policy_text: + current_source = "custom policy text string" + print( + f"Attempting to add policy from custom text to role {role.node.id}..." + ) + try: + # *** FIX: Parse the JSON string into a Python dictionary *** + policy_document = json.loads(custom_policy_text) + print(f"Processing policy from {current_source}...") + add_statement_to_policy(role, policy_document) + except json.JSONDecodeError as e: + print(f"Warning: Invalid JSON in custom_policy_text: {e}. Skipping.") + except Exception as e: + print( + f"An unexpected error occurred processing policy from custom_policy_text: {e}. Skipping." + ) + + # You might want a final success message, but individual processing messages are also good. + print(f"Finished processing custom policies for role {role.node.id}.") + + except Exception as e: + print( + f"An unhandled error occurred during policy addition for {current_source}: {e}" + ) + + return role + + +# Import the S3 Bucket class if you intend to return a CDK object later +# from aws_cdk import aws_s3 as s3 + + +def _s3_bucket_listed_in_account(s3_client: Any, bucket_name: str) -> bool: + token = None + while True: + kwargs: Dict[str, Any] = {} + if token: + kwargs["ContinuationToken"] = token + response = s3_client.list_buckets(**kwargs) + for entry in response.get("Buckets", []): + if entry.get("Name") == bucket_name: + return True + token = response.get("NextContinuationToken") + if not token: + return False + + +def resolve_s3_bucket_availability( + bucket_name: str, + *, + s3_client: Any = None, +) -> Tuple[S3BucketAvailability, str]: + """ + Resolve whether an S3 bucket name is owned in this account, globally free, + or taken by another AWS account. + + S3 names are globally unique. ``head_bucket`` returns 404 when the name is + free, 200 when this principal can access it, and 403 when the name is taken + elsewhere (or access is denied in-account). + """ + client = s3_client or boto3.client("s3") + try: + client.head_bucket(Bucket=bucket_name) + print(f"Bucket '{bucket_name}' exists and is accessible in this account.") + return "owned", bucket_name + except ClientError as e: + error_code = e.response["Error"]["Code"] + if error_code in ("404", "NoSuchBucket", "NotFound"): + print(f"Bucket '{bucket_name}' is available (name not taken globally).") + return "available", bucket_name + if error_code in ("403", "AccessDenied"): + if _s3_bucket_listed_in_account(client, bucket_name): + print( + f"Bucket '{bucket_name}' exists in this account " + "(head_bucket denied; confirmed via list_buckets)." + ) + return "owned", bucket_name + print( + f"Bucket '{bucket_name}' is taken globally by another AWS account " + f"(head_bucket returned {error_code}; not listed in this account)." + ) + return "globally_taken", bucket_name + print( + f"An unexpected AWS ClientError occurred checking bucket '{bucket_name}': {e}" + ) + raise + except Exception as e: + print( + f"An unexpected non-ClientError occurred checking bucket '{bucket_name}': {e}" + ) + raise + + +def check_s3_bucket_exists( + bucket_name: str, +) -> Tuple[bool, Optional[str]]: + """ + Return whether the bucket exists in **this** AWS account (for CDK import). + + For global availability (including cross-account collisions), use + ``resolve_s3_bucket_availability``. + """ + status, name = resolve_s3_bucket_availability(bucket_name) + if status == "owned": + return True, name + return False, None + + +def resolve_cognito_domain_prefix_availability( + domain_prefix: str, + *, + region_name: Optional[str] = None, + cognito_client: Any = None, +) -> CognitoDomainAvailability: + """Return whether a Cognito hosted UI domain prefix is free in the region. + + Cognito prefixes are unique per region across all AWS accounts. When a prefix + is owned elsewhere, ``describe_user_pool_domain`` raises + ``ResourceNotFoundException`` ("does not exist in this account") — the same + symptom CloudFormation reports as ``AlreadyExists`` on create. + """ + prefix = (domain_prefix or "").strip().lower() + if not prefix: + return "available" + region = region_name or os.environ.get("AWS_REGION") + client = cognito_client or boto3.client("cognito-idp", region_name=region) + try: + response = client.describe_user_pool_domain(Domain=prefix) + except ClientError as e: + if e.response["Error"]["Code"] == "ResourceNotFoundException": + print( + f"Cognito domain prefix {prefix!r} is not available in {region} " + f"(likely taken by another AWS account; " + f"{e.response['Error'].get('Message', 'ResourceNotFoundException')})." + ) + return "taken" + raise + description = response.get("DomainDescription") or {} + if description.get("UserPoolId"): + print( + f"Cognito domain prefix {prefix!r} is in use by user pool " + f"{description['UserPoolId']} in this account." + ) + return "taken" + print(f"Cognito domain prefix {prefix!r} is available in {region}.") + return "available" + + +def default_secrets_manager_kms_key_arn(region: str, account_id: str) -> str: + """AWS managed CMK alias used by Secrets Manager when no customer key is set.""" + return f"arn:aws:kms:{region}:{account_id}:key/aws/secretsmanager" + + +def get_secret_kms_key_arn( + secret_id: str, + *, + region_name: Optional[str] = None, +) -> Optional[str]: + """Return the KMS key ARN that encrypts a Secrets Manager secret.""" + client = boto3.client("secretsmanager", region_name=region_name) + try: + description = client.describe_secret(SecretId=secret_id) + except Exception as exc: + print(f"Warning: could not describe secret for KMS key: {exc}") + return None + kms_key_id = description.get("KmsKeyId") + if not kms_key_id: + return None + if str(kms_key_id).startswith("arn:"): + return str(kms_key_id) + region = region_name or AWS_REGION + account_id = ( + description.get("OwningAccount") + or boto3.client("sts", region_name=region).get_caller_identity()["Account"] + ) + return f"arn:aws:kms:{region}:{account_id}:key/{kms_key_id}" + + +ECS_TASK_ROLE_S3_ACTIONS = [ + "s3:GetObject*", + "s3:GetBucket*", + "s3:PutObject", + "s3:DeleteObject", + "s3:List*", +] + + +def build_ecs_task_role_s3_statement(*, sid: str, bucket_name: str) -> Dict[str, Any]: + """Scoped S3 access for one bucket (identity policy on the ECS task role).""" + return { + "Sid": sid, + "Effect": "Allow", + "Action": list(ECS_TASK_ROLE_S3_ACTIONS), + "Resource": [ + f"arn:aws:s3:::{bucket_name}", + f"arn:aws:s3:::{bucket_name}/*", + ], + } + + +def build_ecs_task_role_inline_policy( + *, + output_bucket_name: str, + log_config_bucket_name: Optional[str] = None, + shared_kms_key_arn: Optional[str] = None, +) -> Dict[str, Any]: + """Task role inline policy: STS, scoped S3 buckets, optional shared CMK for S3.""" + statements: List[Dict[str, Any]] = [ + { + "Sid": "STSCallerIdentity", + "Effect": "Allow", + "Action": ["sts:GetCallerIdentity"], + "Resource": "*", + }, + build_ecs_task_role_s3_statement( + sid="S3Output", + bucket_name=output_bucket_name, + ), + ] + if log_config_bucket_name: + statements.append( + build_ecs_task_role_s3_statement( + sid="S3LogConfig", + bucket_name=log_config_bucket_name, + ) + ) + if shared_kms_key_arn: + statements.append( + { + "Sid": "KMSS3Access", + "Effect": "Allow", + "Action": [ + "kms:Encrypt", + "kms:Decrypt", + "kms:GenerateDataKey", + "kms:DescribeKey", + ], + "Resource": shared_kms_key_arn, + } + ) + return {"Version": "2012-10-17", "Statement": statements} + + +def build_ecs_task_role_kms_policy( + *, + shared_kms_key_arn: Optional[str] = None, +) -> Dict[str, Any]: + """Task role KMS-only policy (tests / legacy callers without S3 bucket names).""" + statements: List[Dict[str, Any]] = [ + { + "Sid": "STSCallerIdentity", + "Effect": "Allow", + "Action": ["sts:GetCallerIdentity"], + "Resource": "*", + }, + ] + if shared_kms_key_arn: + statements.append( + { + "Sid": "KMSS3Access", + "Effect": "Allow", + "Action": [ + "kms:Encrypt", + "kms:Decrypt", + "kms:GenerateDataKey", + "kms:DescribeKey", + ], + "Resource": shared_kms_key_arn, + } + ) + return {"Version": "2012-10-17", "Statement": statements} + + +def build_ecs_execution_role_kms_policy( + *, + secret_kms_key_arn: str, +) -> Dict[str, Any]: + """Execution role: decrypt the CMK that encrypts the Cognito secret only.""" + return { + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "STSCallerIdentity", + "Effect": "Allow", + "Action": ["sts:GetCallerIdentity"], + "Resource": "*", + }, + { + "Sid": "KMSSecretDecrypt", + "Effect": "Allow", + "Action": ["kms:Decrypt"], + "Resource": secret_kms_key_arn, + }, + ], + } + + +# Example usage in your check_resources.py: +# exists, bucket_name_if_exists = check_s3_bucket_exists(log_bucket_name) +# context_data[f"exists:{log_bucket_name}"] = exists +# # You don't necessarily need to store the name in context if using from_bucket_name + + +# Delete an S3 bucket +def delete_s3_bucket(bucket_name: str): + s3 = boto3.client("s3") + + try: + # List and delete all objects + response = s3.list_object_versions(Bucket=bucket_name) + versions = response.get("Versions", []) + response.get("DeleteMarkers", []) + for version in versions: + s3.delete_object( + Bucket=bucket_name, Key=version["Key"], VersionId=version["VersionId"] + ) + + # Delete the bucket + s3.delete_bucket(Bucket=bucket_name) + return {"Status": "SUCCESS"} + except Exception as e: + return {"Status": "FAILED", "Reason": str(e)} + + +# Function to get subnet ID from subnet name +def get_subnet_id(vpc: str, ec2_client: str, subnet_name: str): + response = ec2_client.describe_subnets( + Filters=[{"Name": "vpc-id", "Values": [vpc.vpc_id]}] + ) + + for subnet in response["Subnets"]: + if subnet["Tags"] and any( + tag["Key"] == "Name" and tag["Value"] == subnet_name + for tag in subnet["Tags"] + ): + return subnet["SubnetId"] + + return None + + +def check_ecr_repo_exists(repo_name: str) -> tuple[bool, dict]: + """ + Checks if an ECR repository with the given name exists. + + Args: + repo_name: The name of the ECR repository to check. + + Returns: + True if the repository exists, False otherwise. + """ + ecr_client = boto3.client("ecr") + try: + print("ecr repo_name to check:", repo_name) + response = ecr_client.describe_repositories(repositoryNames=[repo_name]) + # If describe_repositories succeeds and returns a list of repositories, + # and the list is not empty, the repository exists. + return len(response["repositories"]) > 0, response["repositories"][0] + except ClientError as e: + # Check for the specific error code indicating the repository doesn't exist + if e.response["Error"]["Code"] == "RepositoryNotFoundException": + return False, {} + else: + # Re-raise other exceptions to handle unexpected errors + raise + except Exception as e: + print(f"An unexpected error occurred: {e}") + return False, {} + + +def check_codebuild_project_exists( + project_name: str, +): # Adjust return type hint as needed + """ + Checks if a CodeBuild project with the given name exists. + + Args: + project_name: The name of the CodeBuild project to check. + + Returns: + A tuple: + - The first element is True if the project exists, False otherwise. + - The second element is the project object (dictionary) if found, + None otherwise. + """ + codebuild_client = boto3.client("codebuild") + try: + # Use batch_get_projects with a list containing the single project name + response = codebuild_client.batch_get_projects(names=[project_name]) + + # The response for batch_get_projects includes 'projects' (found) + # and 'projectsNotFound' (not found). + if response["projects"]: + # If the project is found in the 'projects' list + print(f"CodeBuild project '{project_name}' found.") + project = response["projects"][0] + return ( + True, + project["arn"], + project.get("serviceRole"), + ) + elif ( + response["projectsNotFound"] + and project_name in response["projectsNotFound"] + ): + # If the project name is explicitly in the 'projectsNotFound' list + print(f"CodeBuild project '{project_name}' not found.") + return False, None, None + else: + # This case is less expected for a single name lookup, + # but could happen if there's an internal issue or the response + # structure is slightly different than expected for an error. + # It's safer to assume it wasn't found if not in 'projects'. + print( + f"CodeBuild project '{project_name}' not found (not in 'projects' list)." + ) + return False, None, None + + except ClientError as e: + # Catch specific ClientErrors. batch_get_projects might not throw + # 'InvalidInputException' for a non-existent project name if the + # name format is valid. It typically just lists it in projectsNotFound. + # However, other ClientErrors are possible (e.g., permissions). + print( + f"An AWS ClientError occurred checking CodeBuild project '{project_name}': {e}" + ) + # Decide how to handle other ClientErrors - raising might be safer + raise # Re-raise the original exception + except Exception as e: + print( + f"An unexpected non-ClientError occurred checking CodeBuild project '{project_name}': {e}" + ) + # Decide how to handle other errors + raise # Re-raise the original exception + + +def public_github_repository_url(owner: str, repo: str) -> str: + """HTTPS clone URL for a public GitHub repository.""" + owner_clean = owner.strip().strip("/") + repo_clean = repo.strip().strip("/") + return f"https://github.com/{owner_clean}/{repo_clean}.git" + + +def public_github_codebuild_source( + owner: str, + repo: str, + branch_or_ref: str, +) -> codebuild.ISource: + """CodeBuild source for a public GitHub repo (no CodeConnections/OAuth).""" + return codebuild.Source.git_hub( + owner=owner, + repo=repo, + branch_or_ref=branch_or_ref, + webhook=False, + report_build_status=False, + ) + + +def configure_public_github_codebuild_source( + project: codebuild.Project, + owner: str, + repo: str, + branch_or_ref: str = "main", +) -> None: + """Ensure CodeBuild clones a public repo without account-level GitHub credentials.""" + cfn = project.node.default_child + if not isinstance(cfn, codebuild.CfnProject): + return + location = public_github_repository_url(owner, repo) + cfn.add_property_deletion_override("Source.Auth") + cfn.add_property_override("Source.Type", "GITHUB") + cfn.add_property_override("Source.Location", location) + cfn.add_property_override("Source.ReportBuildStatus", False) + cfn.add_property_override("Source.GitCloneDepth", 1) + cfn.add_property_override("SourceVersion", branch_or_ref) + cfn.add_property_override("Triggers.Webhook", False) + + +def ensure_codebuild_public_github_source( + project_name: str, + owner: str, + repo: str, + branch_or_ref: str = "main", + *, + aws_region: Optional[str] = None, +) -> bool: + """ + Point an existing CodeBuild project at a public GitHub HTTPS URL. + + Use after deploy or when a pre-existing project was imported via + ``from_project_arn`` (CDK does not manage its source). Removes CodeConnections + / OAuth ``auth`` blocks so clones work without account-linked GitHub. + """ + from botocore.exceptions import ClientError + + region = aws_region or AWS_REGION + if not (project_name and owner and repo): + return False + + client = boto3.client("codebuild", region_name=region) + try: + response = client.batch_get_projects(names=[project_name]) + except ClientError as exc: + print( + f"Warning: could not read CodeBuild project '{project_name}' " + f"for public GitHub source fixup: {exc}" + ) + return False + + projects = response.get("projects") or [] + if not projects: + return False + + project = projects[0] + source = project.get("source") or {} + desired_location = public_github_repository_url(owner, repo) + auth = source.get("auth") + needs_update = ( + source.get("type") != "GITHUB" + or source.get("location") != desired_location + or bool(auth) + or source.get("reportBuildStatus") + or (project.get("sourceVersion") or "") != branch_or_ref + ) + if not needs_update: + print( + f"CodeBuild project '{project_name}' already uses public GitHub source " + f"({desired_location})." + ) + return False + + new_source: Dict[str, Any] = { + "type": "GITHUB", + "location": desired_location, + "gitCloneDepth": source.get("gitCloneDepth") or 1, + "reportBuildStatus": False, + "insecureSsl": False, + } + if source.get("buildspec"): + new_source["buildspec"] = source["buildspec"] + + update_kwargs: Dict[str, Any] = { + "name": project_name, + "source": new_source, + "sourceVersion": branch_or_ref, + } + triggers = project.get("triggers") or {} + if triggers.get("webhook"): + update_kwargs["triggers"] = {"webhook": False} + + try: + client.update_project(**update_kwargs) + except ClientError as exc: + print( + f"Warning: could not update CodeBuild project '{project_name}' " + f"to public GitHub source: {exc}" + ) + return False + + print( + f"Updated CodeBuild project '{project_name}' to public GitHub source " + f"({desired_location}, ref {branch_or_ref!r})." + ) + return True + + +def get_vpc_id_by_name(vpc_name: str): + """ + Finds a VPC ID by its 'Name' tag. + + Returns ``(vpc_id, nat_gateways, vpc_cidr_block, vpc_cidr_blocks)`` or ``None`` + if not found. ``vpc_cidr_block`` is the primary block; ``vpc_cidr_blocks`` lists + every associated IPv4 CIDR (primary + secondary). + """ + ec2_client = boto3.client("ec2") + try: + response = ec2_client.describe_vpcs( + Filters=[{"Name": "tag:Name", "Values": [vpc_name]}] + ) + if response and response["Vpcs"]: + vpc = response["Vpcs"][0] + vpc_id = vpc["VpcId"] + vpc_cidr_block = vpc.get("CidrBlock") + vpc_cidr_blocks = list_vpc_associated_cidr_blocks(vpc) + cidr_summary = ( + ", ".join(vpc_cidr_blocks) if vpc_cidr_blocks else vpc_cidr_block + ) + print( + f"VPC '{vpc_name}' found with ID: {vpc_id}" + + (f", CIDR(s): {cidr_summary}" if cidr_summary else "") + ) + + # In get_vpc_id_by_name, after finding VPC ID: + + # Look for NAT Gateways in this VPC + ec2_client = boto3.client("ec2") + nat_gateways = [] + try: + response = ec2_client.describe_nat_gateways( + Filters=[ + {"Name": "vpc-id", "Values": [vpc_id]}, + # Optional: Add a tag filter if you consistently tag your NATs + # {'Name': 'tag:Name', 'Values': [f"{prefix}-nat-gateway"]} + ] + ) + nat_gateways = response.get("NatGateways", []) + except Exception as e: + print( + f"Warning: Could not describe NAT Gateways in VPC '{vpc_id}': {e}" + ) + # Decide how to handle this error - proceed or raise? + + # Decide how to identify the specific NAT Gateway you want to check for. + + return vpc_id, nat_gateways, vpc_cidr_block, vpc_cidr_blocks + else: + print(f"VPC '{vpc_name}' not found.") + return None + except Exception as e: + print(f"An unexpected error occurred finding VPC '{vpc_name}': {e}") + raise + + +# --- Helper to fetch all existing subnets in a VPC once --- +def _get_existing_subnets_in_vpc(vpc_id: str) -> Dict[str, Any]: + """ + Fetches all subnets in a given VPC. + Returns a dictionary with 'by_name' (map of name to subnet data), + 'by_id' (map of id to subnet data), and 'cidr_networks' (list of ipaddress.IPv4Network). + """ + ec2_client = boto3.client("ec2") + existing_subnets_data = { + "by_name": {}, # {subnet_name: {'id': 'subnet-id', 'cidr': 'x.x.x.x/x'}} + "by_id": {}, # {subnet_id: {'name': 'subnet-name', 'cidr': 'x.x.x.x/x/x'}} + "cidr_networks": [], # List of ipaddress.IPv4Network objects + } + try: + subnet_to_route_table: Dict[str, str] = {} + rt_response = ec2_client.describe_route_tables( + Filters=[{"Name": "vpc-id", "Values": [vpc_id]}] + ) + for route_table in rt_response.get("RouteTables", []): + route_table_id = route_table["RouteTableId"] + for association in route_table.get("Associations", []): + associated_subnet_id = association.get("SubnetId") + if associated_subnet_id: + subnet_to_route_table[associated_subnet_id] = route_table_id + + response = ec2_client.describe_subnets( + Filters=[{"Name": "vpc-id", "Values": [vpc_id]}] + ) + for s in response.get("Subnets", []): + subnet_id = s["SubnetId"] + cidr_block = s.get("CidrBlock") + # Extract 'Name' tag, which is crucial for lookup by name + name_tag = next( + (tag["Value"] for tag in s.get("Tags", []) if tag["Key"] == "Name"), + None, + ) + + subnet_info = { + "id": subnet_id, + "cidr": cidr_block, + "name": name_tag, + "az": s.get("AvailabilityZone"), + "route_table_id": subnet_to_route_table.get(subnet_id), + } + + if name_tag: + existing_subnets_data["by_name"][name_tag] = subnet_info + existing_subnets_data["by_id"][subnet_id] = subnet_info + + if cidr_block: + try: + existing_subnets_data["cidr_networks"].append( + ipaddress.ip_network(cidr_block, strict=False) + ) + except ValueError: + print( + f"Warning: Existing subnet {subnet_id} has an invalid CIDR: {cidr_block}. Skipping for overlap check." + ) + + print( + f"Fetched {len(response.get('Subnets', []))} existing subnets from VPC '{vpc_id}'." + ) + except Exception as e: + print( + f"Error describing existing subnets in VPC '{vpc_id}': {e}. Cannot perform full validation." + ) + raise # Re-raise if this essential step fails + + return existing_subnets_data + + +def get_internet_gateways_attached_to_vpc(vpc_id: str) -> List[str]: + """Return Internet Gateway IDs currently attached to the VPC.""" + ec2_client = boto3.client("ec2") + response = ec2_client.describe_internet_gateways( + Filters=[{"Name": "attachment.vpc-id", "Values": [vpc_id]}] + ) + return [ + igw["InternetGatewayId"] + for igw in response.get("InternetGateways", []) + if igw.get("InternetGatewayId") + ] + + +def internet_gateway_exists(igw_id: str) -> bool: + ec2_client = boto3.client("ec2") + response = ec2_client.describe_internet_gateways(InternetGatewayIds=[igw_id]) + return bool(response.get("InternetGateways")) + + +def route_table_default_internet_gateway(route_table_id: str) -> Optional[str]: + """ + Return the Internet Gateway ID for 0.0.0.0/0 on this route table, if any. + """ + ec2_client = boto3.client("ec2") + response = ec2_client.describe_route_tables(RouteTableIds=[route_table_id]) + tables = response.get("RouteTables", []) + if not tables: + return None + for route in tables[0].get("Routes", []): + if route.get("DestinationCidrBlock") != "0.0.0.0/0": + continue + gateway_id = route.get("GatewayId") or "" + if gateway_id.startswith("igw-"): + return gateway_id + return None + + +def route_table_has_non_igw_default_route(route_table_id: str) -> bool: + """True if 0.0.0.0/0 exists but does not target an Internet Gateway.""" + ec2_client = boto3.client("ec2") + response = ec2_client.describe_route_tables(RouteTableIds=[route_table_id]) + tables = response.get("RouteTables", []) + if not tables: + return False + for route in tables[0].get("Routes", []): + if route.get("DestinationCidrBlock") != "0.0.0.0/0": + continue + gateway_id = route.get("GatewayId") or "" + if gateway_id.startswith("igw-"): + return False + # Active default route via NAT instance, TGW, etc. + if ( + route.get("NatGatewayId") + or route.get("TransitGatewayId") + or route.get("GatewayId") + ): + return True + return False + + +def audit_public_subnet_internet_connectivity( + vpc_id: str, + configured_igw_id: str, + public_subnet_entries: List[Dict[str, Any]], +) -> Dict[str, Any]: + """ + Validate / discover Internet Gateway usage for legacy public subnets. + + Returns context fields for CDK: + internet_gateway_id, internet_gateway_needs_vpc_attachment, + public_subnets_needing_igw_route (list of {name, subnet_id, route_table_id}). + """ + configured_igw_id = (configured_igw_id or "").strip() + attached_igws = get_internet_gateways_attached_to_vpc(vpc_id) + + resolved_igw_id = configured_igw_id + needs_attachment = False + + if configured_igw_id: + if not internet_gateway_exists(configured_igw_id): + raise ValueError( + f"EXISTING_IGW_ID '{configured_igw_id}' was not found in this account/region." + ) + if configured_igw_id not in attached_igws: + # Ensure it is not attached to another VPC + ec2_client = boto3.client("ec2") + detail = ec2_client.describe_internet_gateways( + InternetGatewayIds=[configured_igw_id] + ) + for attachment in detail.get("InternetGateways", [{}])[0].get( + "Attachments", [] + ): + other_vpc = attachment.get("VpcId") + if other_vpc and other_vpc != vpc_id: + raise ValueError( + f"EXISTING_IGW_ID '{configured_igw_id}' is attached to VPC " + f"'{other_vpc}', not target VPC '{vpc_id}'. Detach it or choose " + "the IGW attached to this VPC." + ) + needs_attachment = True + elif attached_igws: + if len(attached_igws) > 1: + raise ValueError( + f"VPC '{vpc_id}' has multiple attached Internet Gateways " + f"({', '.join(attached_igws)}). Set EXISTING_IGW_ID to the one to use." + ) + resolved_igw_id = attached_igws[0] + print( + f"EXISTING_IGW_ID not set; using Internet Gateway attached to VPC: " + f"{resolved_igw_id}" + ) + elif public_subnet_entries: + raise ValueError( + f"VPC '{vpc_id}' has no Internet Gateway attached and EXISTING_IGW_ID is " + "empty. Set EXISTING_IGW_ID to an existing IGW for this VPC (CDK will " + "attach it if detached)." + ) + + subnets_needing_route: List[Dict[str, str]] = [] + for entry in public_subnet_entries: + name = entry.get("name") or "unknown" + route_table_id = entry.get("route_table_id") + subnet_id = entry.get("subnet_id") or entry.get("id") or "" + if not route_table_id: + print( + f"Warning: public subnet '{name}' has no route table association in " + "pre-check; skipping IGW route audit (CDK may still add routes after create)." + ) + continue + existing_igw = route_table_default_internet_gateway(route_table_id) + if existing_igw: + if resolved_igw_id and existing_igw != resolved_igw_id: + raise ValueError( + f"Public subnet '{name}' route table '{route_table_id}' has " + f"0.0.0.0/0 -> {existing_igw}, but EXISTING_IGW_ID / resolved IGW " + f"is '{resolved_igw_id}'. Fix the route table manually or align " + "EXISTING_IGW_ID." + ) + continue + if route_table_has_non_igw_default_route(route_table_id): + raise ValueError( + f"Public subnet '{name}' route table '{route_table_id}' has a default " + "route that does not use an Internet Gateway (e.g. NAT/TGW). Remove " + "or change it before adding 0.0.0.0/0 -> IGW for an internet-facing ALB." + ) + subnets_needing_route.append( + { + "name": name, + "subnet_id": subnet_id, + "route_table_id": route_table_id, + } + ) + + return { + "internet_gateway_id": resolved_igw_id, + "internet_gateway_needs_vpc_attachment": needs_attachment, + "public_subnets_needing_igw_route": subnets_needing_route, + } + + +def wire_public_subnet_internet_access( + scope: Construct, + logical_id_prefix: str, + *, + vpc_id: str, + internet_gateway_id: str, + needs_igw_vpc_attachment: bool, + subnets_needing_route: List[Dict[str, str]], +) -> Optional[ec2.CfnVPCGatewayAttachment]: + """ + Attach the Internet Gateway to the VPC (if needed) and add 0.0.0.0/0 routes on + imported public subnet route tables that lack an IGW default route. + """ + if not internet_gateway_id: + return None + + attachment = None + if needs_igw_vpc_attachment: + attachment = ec2.CfnVPCGatewayAttachment( + scope, + f"{logical_id_prefix}IgwVpcAttachment", + vpc_id=vpc_id, + internet_gateway_id=internet_gateway_id, + ) + print( + f"CDK: will attach Internet Gateway '{internet_gateway_id}' to VPC '{vpc_id}'." + ) + + seen_route_tables: set[str] = set() + for i, entry in enumerate(subnets_needing_route): + route_table_id = entry.get("route_table_id") + if not route_table_id or route_table_id in seen_route_tables: + continue + seen_route_tables.add(route_table_id) + safe_name = (entry.get("name") or f"rt{i}").replace("-", "")[:40] + route = ec2.CfnRoute( + scope, + f"{logical_id_prefix}IgwRoute{safe_name}{i}", + route_table_id=route_table_id, + destination_cidr_block="0.0.0.0/0", + gateway_id=internet_gateway_id, + ) + if attachment is not None: + route.add_dependency(attachment) + print( + f"CDK: will add 0.0.0.0/0 -> {internet_gateway_id} on route table " + f"'{route_table_id}' (subnet '{entry.get('name', '')}')." + ) + + return attachment + + +# --- Modified validate_subnet_creation_parameters to take pre-fetched data --- +def validate_subnet_creation_parameters( + vpc_id: str, + proposed_subnets_data: List[ + Dict[str, str] + ], # e.g., [{'name': 'my-public-subnet', 'cidr': '10.0.0.0/24', 'az': 'us-east-1a'}] + existing_aws_subnets_data: Dict[ + str, Any + ], # Pre-fetched data from _get_existing_subnets_in_vpc +) -> None: + """ + Validates proposed subnet names and CIDR blocks against existing AWS subnets + in the specified VPC and against each other. + This function uses pre-fetched AWS subnet data. + + Args: + vpc_id: The ID of the VPC (for logging/error messages). + proposed_subnets_data: A list of dictionaries, where each dict represents + a proposed subnet with 'name', 'cidr', and 'az'. + existing_aws_subnets_data: Dictionary containing existing AWS subnet data + (e.g., from _get_existing_subnets_in_vpc). + + Raises: + ValueError: If any proposed subnet name or CIDR block + conflicts with existing AWS resources or other proposed resources. + """ + if not proposed_subnets_data: + print("No proposed subnet data provided for validation. Skipping.") + return + + print( + f"--- Starting pre-synth validation for VPC '{vpc_id}' with proposed subnets ---" + ) + + print("Existing subnet data:", pd.DataFrame(existing_aws_subnets_data["by_name"])) + + existing_aws_subnet_names = set(existing_aws_subnets_data["by_name"].keys()) + existing_aws_cidr_networks = existing_aws_subnets_data["cidr_networks"] + + # Sets to track names and list to track networks for internal batch consistency + proposed_names_seen: set[str] = set() + proposed_cidr_networks_seen: List[ipaddress.IPv4Network] = [] + + for i, proposed_subnet in enumerate(proposed_subnets_data): + subnet_name = proposed_subnet.get("name") + cidr_block_str = proposed_subnet.get("cidr") + availability_zone = proposed_subnet.get("az") + + if not all([subnet_name, cidr_block_str, availability_zone]): + raise ValueError( + f"Proposed subnet at index {i} is incomplete. Requires 'name', 'cidr', and 'az'." + ) + + # 1. Check for duplicate names within the proposed batch + if subnet_name in proposed_names_seen: + raise ValueError( + f"Proposed subnet name '{subnet_name}' is duplicated within the input list." + ) + proposed_names_seen.add(subnet_name) + + # 2. Check for duplicate names against existing AWS subnets + if subnet_name in existing_aws_subnet_names: + print( + f"Proposed subnet name '{subnet_name}' already exists in VPC '{vpc_id}'." + ) + + # Parse proposed CIDR + try: + proposed_net = ipaddress.ip_network(cidr_block_str, strict=False) + except ValueError as e: + raise ValueError( + f"Invalid CIDR format '{cidr_block_str}' for proposed subnet '{subnet_name}': {e}" + ) + + # 3. Check for overlapping CIDRs within the proposed batch + for existing_proposed_net in proposed_cidr_networks_seen: + if proposed_net.overlaps(existing_proposed_net): + raise ValueError( + f"Proposed CIDR '{cidr_block_str}' for subnet '{subnet_name}' " + f"overlaps with another proposed CIDR '{str(existing_proposed_net)}' " + f"within the same batch." + ) + + # 4. Check for overlapping CIDRs against existing AWS subnets + for existing_aws_net in existing_aws_cidr_networks: + if proposed_net.overlaps(existing_aws_net): + raise ValueError( + f"Proposed CIDR '{cidr_block_str}' for subnet '{subnet_name}' " + f"overlaps with an existing AWS subnet CIDR '{str(existing_aws_net)}' " + f"in VPC '{vpc_id}'." + ) + + # If all checks pass for this subnet, add its network to the list for subsequent checks + proposed_cidr_networks_seen.append(proposed_net) + print( + f"Validation successful for proposed subnet '{subnet_name}' with CIDR '{cidr_block_str}'." + ) + + print( + f"--- All proposed subnets passed pre-synth validation checks for VPC '{vpc_id}'. ---" + ) + + +# --- Modified check_subnet_exists_by_name (Uses pre-fetched data) --- +def check_subnet_exists_by_name( + subnet_name: str, existing_aws_subnets_data: Dict[str, Any] +) -> Tuple[bool, Optional[str]]: + """ + Checks if a subnet with the given name exists within the pre-fetched data. + + Args: + subnet_name: The 'Name' tag value of the subnet to check. + existing_aws_subnets_data: Dictionary containing existing AWS subnet data + (e.g., from _get_existing_subnets_in_vpc). + + Returns: + A tuple: + - The first element is True if the subnet exists, False otherwise. + - The second element is the Subnet ID if found, None otherwise. + """ + subnet_info = existing_aws_subnets_data["by_name"].get(subnet_name) + if subnet_info: + print(f"Subnet '{subnet_name}' found with ID: {subnet_info['id']}") + return True, subnet_info["id"] + else: + print(f"Subnet '{subnet_name}' not found.") + return False, None + + +def create_nat_gateway( + scope: Construct, + public_subnet_for_nat: ec2.ISubnet, # Expects a proper ISubnet + nat_gateway_name: str, + nat_gateway_id_context_key: str, +) -> str: + """ + Creates a single NAT Gateway in the specified public subnet. + It does not handle lookup from context; the calling stack should do that. + Returns the CloudFormation Ref of the NAT Gateway ID. + """ + print( + f"Defining a new NAT Gateway '{nat_gateway_name}' in subnet '{public_subnet_for_nat.subnet_id}'." + ) + + # Create an Elastic IP for the NAT Gateway + eip = ec2.CfnEIP( + scope, + NAT_GATEWAY_EIP_NAME, + tags=[CfnTag(key="Name", value=NAT_GATEWAY_EIP_NAME)], + ) + + # Create the NAT Gateway + nat_gateway_logical_id = nat_gateway_name.replace("-", "") + "NatGateway" + nat_gateway = ec2.CfnNatGateway( + scope, + nat_gateway_logical_id, + subnet_id=public_subnet_for_nat.subnet_id, # Associate with the public subnet + allocation_id=eip.attr_allocation_id, # Associate with the EIP + tags=[CfnTag(key="Name", value=nat_gateway_name)], + ) + # The NAT GW depends on the EIP. The dependency on the subnet is implicit via subnet_id. + nat_gateway.add_dependency(eip) + + # *** CRUCIAL: Use CfnOutput to export the ID after deployment *** + # This is how you will get the ID to put into cdk.context.json + CfnOutput( + scope, + "SingleNatGatewayIdOutput", + value=nat_gateway.ref, + description=f"Physical ID of the Single NAT Gateway. Add this to cdk.context.json under the key '{nat_gateway_id_context_key}'.", + export_name=f"{scope.stack_name}-NatGatewayId", # Make export name unique + ) + + print( + f"CDK: Defined new NAT Gateway '{nat_gateway.ref}'. Its physical ID will be available in the stack outputs after deployment." + ) + # Return the tokenised reference for use within this synthesis + return nat_gateway.ref + + +def create_subnets( + scope: Construct, + vpc: ec2.IVpc, + prefix: str, + subnet_names: List[str], + cidr_blocks: List[str], + availability_zones: List[str], + is_public: bool, + internet_gateway_id: Optional[str] = None, + single_nat_gateway_id: Optional[str] = None, + internet_gateway_attachment: Optional[ec2.CfnVPCGatewayAttachment] = None, +) -> Tuple[List[ec2.CfnSubnet], List[ec2.CfnRouteTable]]: + """ + Creates subnets using L2 constructs but returns the underlying L1 Cfn objects + for backward compatibility. + """ + # --- Validations remain the same --- + if not (len(subnet_names) == len(cidr_blocks) == len(availability_zones) > 0): + raise ValueError( + "Subnet names, CIDR blocks, and Availability Zones lists must be non-empty and match in length." + ) + if is_public and not internet_gateway_id: + raise ValueError("internet_gateway_id must be provided for public subnets.") + if not is_public and not single_nat_gateway_id: + raise ValueError( + "single_nat_gateway_id must be provided for private subnets when using a single NAT Gateway." + ) + + # --- We will populate these lists with the L1 objects to return --- + created_subnets: List[ec2.CfnSubnet] = [] + created_route_tables: List[ec2.CfnRouteTable] = [] + + subnet_type_tag = "public" if is_public else "private" + + for i, subnet_name in enumerate(subnet_names): + logical_id = f"{prefix}{subnet_type_tag.capitalize()}Subnet{i+1}" + + # 1. Create the L2 Subnet (this is the easy part) + subnet = ec2.Subnet( + scope, + logical_id, + vpc_id=vpc.vpc_id, + cidr_block=cidr_blocks[i], + availability_zone=availability_zones[i], + map_public_ip_on_launch=is_public, + ) + Tags.of(subnet).add("Name", subnet_name) + Tags.of(subnet).add("Type", subnet_type_tag) + + if is_public and internet_gateway_attachment is not None: + subnet.node.add_dependency(internet_gateway_attachment) + + if is_public: + try: + subnet.add_route( + "DefaultInternetRoute", + router_id=internet_gateway_id, + router_type=ec2.RouterType.GATEWAY, + ) + except Exception as e: + raise RuntimeError( + f"Could not create 0.0.0.0/0 -> Internet Gateway route for public " + f"subnet '{subnet_name}'. Ensure EXISTING_IGW_ID is attached to this " + f"VPC ({internet_gateway_id}): {e}" + ) from e + print(f"CDK: Defined public L2 subnet '{subnet_name}' and added IGW route.") + else: + try: + subnet.add_route( + "DefaultNatRoute", + router_id=single_nat_gateway_id, + router_type=ec2.RouterType.NAT_GATEWAY, + ) + except Exception as e: + raise RuntimeError( + f"Could not create 0.0.0.0/0 -> NAT Gateway route for private " + f"subnet '{subnet_name}': {e}" + ) from e + print( + f"CDK: Defined private L2 subnet '{subnet_name}' and added NAT GW route." + ) + + route_table = subnet.route_table + + created_subnets.append(subnet) + created_route_tables.append(route_table) + + return created_subnets, created_route_tables + + +def ingress_rule_exists(security_group: str, peer: str, port: str): + for rule in security_group.connections.security_groups: + if port: + if rule.peer == peer and rule.connection == port: + return True + else: + if rule.peer == peer: + return True + return False + + +def check_for_existing_user_pool(user_pool_name: str): + cognito_client = boto3.client("cognito-idp") + list_pools_response = cognito_client.list_user_pools( + MaxResults=60 + ) # MaxResults up to 60 + + # ListUserPools might require pagination if you have more than 60 pools + # This simple example doesn't handle pagination, which could miss your pool + + existing_user_pool_id = "" + + for pool in list_pools_response.get("UserPools", []): + if pool.get("Name") == user_pool_name: + existing_user_pool_id = pool["Id"] + print( + f"Found existing user pool by name '{user_pool_name}' with ID: {existing_user_pool_id}" + ) + break # Found the one we're looking for + + if existing_user_pool_id: + return True, existing_user_pool_id, pool + else: + return False, "", "" + + +def check_for_existing_user_pool_client(user_pool_id: str, user_pool_client_name: str): + """ + Checks if a Cognito User Pool Client with the given name exists in the specified User Pool. + + Args: + user_pool_id: The ID of the Cognito User Pool. + user_pool_client_name: The name of the User Pool Client to check for. + + Returns: + A tuple: + - True, client_id, client_details if the client exists. + - False, "", {} otherwise. + """ + cognito_client = boto3.client("cognito-idp") + next_token = None + + while True: + try: + kwargs = {"UserPoolId": user_pool_id, "MaxResults": 60} + if next_token: + kwargs["NextToken"] = next_token + response = cognito_client.list_user_pool_clients(**kwargs) + except cognito_client.exceptions.ResourceNotFoundException: + print(f"Error: User pool with ID '{user_pool_id}' not found.") + return False, "", {} + + except Exception as e: + print( + f"Could not list app clients for pool '{user_pool_id}' " + f"(client name '{user_pool_client_name}'): {e}" + ) + return False, "", {} + + for client in response.get("UserPoolClients", []): + if client.get("ClientName") == user_pool_client_name: + print( + f"Found existing user pool client '{user_pool_client_name}' with ID: {client['ClientId']}" + ) + return True, client["ClientId"], client + + next_token = response.get("NextToken") + if not next_token: + break + + print( + f"No app client named '{user_pool_client_name}' in user pool '{user_pool_id}'." + ) + return False, "", {} + + +def check_for_secret(secret_name: str, secret_value: dict = ""): + """ + Checks if a Secrets Manager secret with the given name exists. + If it doesn't exist, it creates the secret. + + Args: + secret_name: The name of the Secrets Manager secret. + secret_value: A dictionary containing the key-value pairs for the secret. + + Returns: + Tuple of (exists, response). When exists is True, response is the + ``get_secret_value`` API dict (includes ``ARN`` for IAM grants). + """ + secretsmanager_client = boto3.client("secretsmanager") + + try: + # Try to get the secret. If it doesn't exist, a ResourceNotFoundException will be raised. + secret_response = secretsmanager_client.get_secret_value(SecretId=secret_name) + print("Secret already exists.") + return True, secret_response + except secretsmanager_client.exceptions.ResourceNotFoundException: + print("Secret not found") + return False, {} + except Exception as e: + # Handle other potential exceptions during the get operation + print(f"Error checking for secret: {e}") + return False, {} + + +def get_security_group_id_by_name( + group_name: str, + vpc_id: str, + region_name: str = AWS_REGION, +) -> Tuple[bool, str]: + """Look up a security group ID by name within a VPC.""" + if not group_name or not vpc_id: + return False, "" + try: + ec2_client = boto3.client("ec2", region_name=region_name) + response = ec2_client.describe_security_groups( + Filters=[ + {"Name": "group-name", "Values": [group_name]}, + {"Name": "vpc-id", "Values": [vpc_id]}, + ] + ) + groups = response.get("SecurityGroups") or [] + if groups: + return True, groups[0]["GroupId"] + return False, "" + except ClientError as e: + print(f"Error looking up security group '{group_name}': {e}") + return False, "" + + +def resolve_service_connect_client_security_group_ids( + explicit_ids: List[str], + security_group_names: List[str], + get_context_str, +) -> List[str]: + """ + Merge explicit sg- IDs with IDs resolved from pre-check context (security_group_id:{name}). + """ + resolved: List[str] = [] + for sg_id in explicit_ids: + if not sg_id.startswith("sg-"): + raise ValueError( + f"ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_IDS entry '{sg_id}' " + "must be a security group ID (sg-...)." + ) + if sg_id not in resolved: + resolved.append(sg_id) + + missing_names: List[str] = [] + for sg_name in security_group_names: + sg_id = get_context_str(f"security_group_id:{sg_name}") + if sg_id: + if sg_id not in resolved: + resolved.append(sg_id) + else: + missing_names.append(sg_name) + + if missing_names: + raise ValueError( + "Could not resolve Service Connect client security group(s) in VPC " + f"{get_context_str('vpc_id') or '(unknown)'}: " + + ", ".join(missing_names) + + ". Set ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_IDS, fix " + "ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES / " + "ECS_SERVICE_CONNECT_CLIENT_CDK_PREFIXES, and re-run check_resources.py." + ) + + return resolved + + +def check_alb_exists( + load_balancer_name: str, region_name: str = None +) -> tuple[bool, dict]: + """ + Checks if an Application Load Balancer (ALB) with the given name exists. + + Args: + load_balancer_name: The name of the ALB to check. + region_name: The AWS region to check in. If None, uses the default + session region. + + Returns: + A tuple: + - The first element is True if the ALB exists, False otherwise. + - The second element is the ALB object (dictionary) if found, + None otherwise. Specifically, it returns the first element of + the LoadBalancers list from the describe_load_balancers response. + """ + if region_name: + elbv2_client = boto3.client("elbv2", region_name=region_name) + else: + elbv2_client = boto3.client("elbv2") + try: + response = elbv2_client.describe_load_balancers(Names=[load_balancer_name]) + if response["LoadBalancers"]: + return ( + True, + response["LoadBalancers"][0], + ) # Return True and the first ALB object + else: + return False, {} + except ClientError as e: + # If the error indicates the ALB doesn't exist, return False + if e.response["Error"]["Code"] == "LoadBalancerNotFound": + return False, {} + else: + # Re-raise other exceptions + raise + except Exception as e: + print(f"An unexpected error occurred: {e}") + return False, {} + + +def check_fargate_task_definition_exists( + task_definition_name: str, region_name: str = None +) -> tuple[bool, dict]: + """ + Checks if a Fargate task definition with the given name exists. + + Args: + task_definition_name: The name or ARN of the task definition to check. + region_name: The AWS region to check in. If None, uses the default + session region. + + Returns: + A tuple: + - The first element is True if the task definition exists, False otherwise. + - The second element is the task definition object (dictionary) if found, + None otherwise. Specifically, it returns the first element of the + taskDefinitions list from the describe_task_definition response. + """ + if region_name: + ecs_client = boto3.client("ecs", region_name=region_name) + else: + ecs_client = boto3.client("ecs") + try: + response = ecs_client.describe_task_definition( + taskDefinition=task_definition_name + ) + # If describe_task_definition succeeds, it returns the task definition. + # We can directly return True and the task definition. + return True, response["taskDefinition"] + except ClientError as e: + # Check for the error code indicating the task definition doesn't exist. + if ( + e.response["Error"]["Code"] == "ClientException" + and "Task definition" in e.response["Message"] + and "does not exist" in e.response["Message"] + ): + return False, {} + else: + # Re-raise other exceptions. + raise + except Exception as e: + print(f"An unexpected error occurred: {e}") + return False, {} + + +def check_ecs_service_exists( + cluster_name: str, service_name: str, region_name: str = None +) -> tuple[bool, dict]: + """ + Checks if an ECS service with the given name exists in the specified cluster. + + Args: + cluster_name: The name or ARN of the ECS cluster. + service_name: The name of the ECS service to check. + region_name: The AWS region to check in. If None, uses the default + session region. + + Returns: + A tuple: + - The first element is True if the service exists, False otherwise. + - The second element is the service object (dictionary) if found, + None otherwise. + """ + if region_name: + ecs_client = boto3.client("ecs", region_name=region_name) + else: + ecs_client = boto3.client("ecs") + try: + response = ecs_client.describe_services( + cluster=cluster_name, services=[service_name] + ) + if response["services"]: + return ( + True, + response["services"][0], + ) # Return True and the first service object + else: + return False, {} + except ClientError as e: + # Check for the error code indicating the service doesn't exist. + if e.response["Error"]["Code"] == "ClusterNotFoundException": + return False, {} + elif e.response["Error"]["Code"] == "ServiceNotFoundException": + return False, {} + else: + # Re-raise other exceptions. + raise + except Exception as e: + print(f"An unexpected error occurred: {e}") + return False, {} + + +def check_cloudfront_distribution_exists( + distribution_name: str, region_name: str = None +) -> tuple[bool, dict | None]: + """ + Checks if a CloudFront distribution with the given name exists. + + Args: + distribution_name: The name of the CloudFront distribution to check. + region_name: The AWS region to check in. If None, uses the default + session region. Note: CloudFront is a global service, + so the region is usually 'us-east-1', but this parameter + is included for completeness. + + Returns: + A tuple: + - The first element is True if the distribution exists, False otherwise. + - The second element is the distribution object (dictionary) if found, + None otherwise. Specifically, it returns the first element of the + DistributionList from the ListDistributions response. + """ + if region_name: + cf_client = boto3.client("cloudfront", region_name=region_name) + else: + cf_client = boto3.client("cloudfront") + try: + response = cf_client.list_distributions() + if "Items" in response["DistributionList"]: + for distribution in response["DistributionList"]["Items"]: + # CloudFront doesn't directly filter by name, so we have to iterate. + if ( + distribution["AliasSet"]["Items"] + and distribution["AliasSet"]["Items"][0] == distribution_name + ): + return True, distribution + return False, None + else: + return False, None + except ClientError as e: + # If the error indicates the Distribution doesn't exist, return False + if e.response["Error"]["Code"] == "NoSuchDistribution": + return False, None + else: + # Re-raise other exceptions + raise + except Exception as e: + print(f"An unexpected error occurred: {e}") + return False, None + + +def create_web_acl_with_common_rules( + scope: Construct, web_acl_name: str, waf_scope: str = "CLOUDFRONT" +): + """ + Use CDK to create a web ACL based on an AWS common rule set with overrides. + This function now expects a 'scope' argument, typically 'self' from your stack, + as CfnWebACL requires a construct scope. + """ + + # Create full list of rules + rules = [] + aws_ruleset_names = [ + "AWSManagedRulesCommonRuleSet", + "AWSManagedRulesKnownBadInputsRuleSet", + "AWSManagedRulesAmazonIpReputationList", + ] + + # Use a separate counter to assign unique priorities sequentially + priority_counter = 1 + + for aws_rule_name in aws_ruleset_names: + current_rule_action_overrides = None + + # All managed rule groups need an override_action. + # 'none' means use the managed rule group's default action. + current_override_action = wafv2.CfnWebACL.OverrideActionProperty(none={}) + + current_priority = priority_counter + priority_counter += 1 + + if aws_rule_name == "AWSManagedRulesCommonRuleSet": + current_rule_action_overrides = [ + wafv2.CfnWebACL.RuleActionOverrideProperty( + name="SizeRestrictions_BODY", + action_to_use=wafv2.CfnWebACL.RuleActionProperty(allow={}), + ) + ] + # No need to set current_override_action here, it's already set above. + # If you wanted this specific rule to have a *fixed* priority, you'd handle it differently + # For now, it will get priority 1 from the counter. + + rule_property = wafv2.CfnWebACL.RuleProperty( + name=aws_rule_name, + priority=current_priority, + statement=wafv2.CfnWebACL.StatementProperty( + managed_rule_group_statement=wafv2.CfnWebACL.ManagedRuleGroupStatementProperty( + vendor_name="AWS", + name=aws_rule_name, + rule_action_overrides=current_rule_action_overrides, + ) + ), + visibility_config=wafv2.CfnWebACL.VisibilityConfigProperty( + cloud_watch_metrics_enabled=True, + metric_name=aws_rule_name, + sampled_requests_enabled=True, + ), + override_action=current_override_action, # THIS IS THE CRUCIAL PART FOR ALL MANAGED RULES + ) + + rules.append(rule_property) + + # Add the rate limit rule + rate_limit_priority = priority_counter # Use the next available priority + rules.append( + wafv2.CfnWebACL.RuleProperty( + name="RateLimitRule", + priority=rate_limit_priority, + statement=wafv2.CfnWebACL.StatementProperty( + rate_based_statement=wafv2.CfnWebACL.RateBasedStatementProperty( + limit=1000, aggregate_key_type="IP" + ) + ), + visibility_config=wafv2.CfnWebACL.VisibilityConfigProperty( + cloud_watch_metrics_enabled=True, + metric_name="RateLimitRule", + sampled_requests_enabled=True, + ), + action=wafv2.CfnWebACL.RuleActionProperty(block={}), + ) + ) + + web_acl = wafv2.CfnWebACL( + scope, + "WebACL", + name=web_acl_name, + default_action=wafv2.CfnWebACL.DefaultActionProperty(allow={}), + scope=waf_scope, + visibility_config=wafv2.CfnWebACL.VisibilityConfigProperty( + cloud_watch_metrics_enabled=True, + metric_name="webACL", + sampled_requests_enabled=True, + ), + rules=rules, + ) + + CfnOutput(scope, "WebACLArn", value=web_acl.attr_arn) + web_acl.apply_removal_policy(managed_resource_removal_policy()) + + return web_acl + + +def check_web_acl_exists( + web_acl_name: str, scope: str, region_name: str = None +) -> tuple[bool, dict]: + """ + Checks if a Web ACL with the given name and scope exists. + + Args: + web_acl_name: The name of the Web ACL to check. + scope: The scope of the Web ACL ('CLOUDFRONT' or 'REGIONAL'). + region_name: The AWS region to check in. Required for REGIONAL scope. + If None, uses the default session region. For CLOUDFRONT, + the region should be 'us-east-1'. + + Returns: + A tuple: + - The first element is True if the Web ACL exists, False otherwise. + - The second element is the Web ACL object (dictionary) if found, + None otherwise. + """ + if scope not in ["CLOUDFRONT", "REGIONAL"]: + raise ValueError("Scope must be either 'CLOUDFRONT' or 'REGIONAL'") + + if scope == "REGIONAL" and not region_name: + raise ValueError("Region name is required for REGIONAL scope") + + if scope == "CLOUDFRONT": + region_name = "us-east-1" # CloudFront scope requires us-east-1 + + if region_name: + waf_client = boto3.client("wafv2", region_name=region_name) + else: + waf_client = boto3.client("wafv2") + try: + response = waf_client.list_web_acls(Scope=scope) + if "WebACLs" in response: + for web_acl in response["WebACLs"]: + if web_acl["Name"] == web_acl_name: + # Describe the Web ACL to get the full object. + describe_response = waf_client.describe_web_acl( + Name=web_acl_name, Scope=scope + ) + return True, describe_response["WebACL"] + return False, {} + else: + return False, {} + except ClientError as e: + # Check for the error code indicating the web ACL doesn't exist. + if e.response["Error"]["Code"] == "ResourceNotFoundException": + return False, {} + else: + # Re-raise other exceptions. + raise + except Exception as e: + print(f"An unexpected error occurred: {e}") + return False, {} + + +def add_alb_https_listener_with_cert( + scope: Construct, + logical_id: str, # A unique ID for this listener construct + alb: elb.ApplicationLoadBalancer, + acm_certificate_arn: Optional[ + str + ], # Optional: If None, no HTTPS listener will be created + default_target_group: elb.ITargetGroup, # Mandatory: The target group to forward traffic to + listener_port_https: int = 443, + listener_open_to_internet: bool = False, # Be cautious with True, ensure ALB security group restricts access + # --- Cognito Authentication Parameters --- + enable_cognito_auth: bool = False, + cognito_user_pool: Optional[cognito.IUserPool] = None, + cognito_user_pool_client: Optional[cognito.IUserPoolClient] = None, + cognito_user_pool_domain: Optional[ + str + ] = None, # E.g., "my-app-domain" for "my-app-domain.auth.region.amazoncognito.com" + cognito_auth_scope: Optional[ + str + ] = "openid profile email", # Default recommended scope + cognito_auth_on_unauthenticated_request: elb.UnauthenticatedAction = elb.UnauthenticatedAction.AUTHENTICATE, + stickiness_cookie_duration=None, + # --- End Cognito Parameters --- +) -> Optional[elb.ApplicationListener]: + """ + Conditionally adds an HTTPS listener to an ALB with an ACM certificate, + and optionally enables Cognito User Pool authentication. + + Args: + scope (Construct): The scope in which to define this construct (e.g., your CDK Stack). + logical_id (str): A unique logical ID for the listener construct within the stack. + alb (elb.ApplicationLoadBalancer): The Application Load Balancer to add the listener to. + acm_certificate_arn (Optional[str]): The ARN of the ACM certificate to attach. + If None, the HTTPS listener will NOT be created. + default_target_group (elb.ITargetGroup): The default target group for the listener to forward traffic to. + This is mandatory for a functional listener. + listener_port_https (int): The HTTPS port to listen on (default: 443). + listener_open_to_internet (bool): Whether the listener should allow connections from all sources. + If False (recommended), ensure your ALB's security group allows + inbound traffic on this port from desired sources. + enable_cognito_auth (bool): Set to True to enable Cognito User Pool authentication. + cognito_user_pool (Optional[cognito.IUserPool]): The Cognito User Pool object. Required if enable_cognito_auth is True. + cognito_user_pool_client (Optional[cognito.IUserPoolClient]): The Cognito User Pool App Client object. Required if enable_cognito_auth is True. + cognito_user_pool_domain (Optional[str]): The domain prefix for your Cognito User Pool. Required if enable_cognito_auth is True. + cognito_auth_scope (Optional[str]): The scope for the Cognito authentication. + cognito_auth_on_unauthenticated_request (elb.UnauthenticatedAction): Action for unauthenticated requests. + Defaults to AUTHENTICATE (redirect to login). + + Returns: + Optional[elb.ApplicationListener]: The created ApplicationListener if successful, + None if no ACM certificate ARN was provided. + """ + https_listener = None + if acm_certificate_arn: + certificates_list = [elb.ListenerCertificate.from_arn(acm_certificate_arn)] + print( + f"Attempting to add ALB HTTPS listener on port {listener_port_https} with ACM certificate: {acm_certificate_arn}" + ) + + # Determine the default action based on whether Cognito auth is enabled + default_action = None + if enable_cognito_auth is True: + if not all( + [cognito_user_pool, cognito_user_pool_client, cognito_user_pool_domain] + ): + raise ValueError( + "Cognito User Pool, Client, and Domain must be provided if enable_cognito_auth is True." + ) + print( + f"Enabling Cognito authentication with User Pool: {cognito_user_pool.user_pool_id}" + ) + + default_action = elb_act.AuthenticateCognitoAction( + next=elb.ListenerAction.forward( + [default_target_group] + ), # After successful auth, forward to TG + user_pool=cognito_user_pool, + user_pool_client=cognito_user_pool_client, + user_pool_domain=cognito_user_pool_domain, + scope=cognito_auth_scope, + on_unauthenticated_request=cognito_auth_on_unauthenticated_request, + session_timeout=stickiness_cookie_duration, + # Additional options you might want to configure: + # session_cookie_name="AWSELBCookies" + ) + else: + default_action = elb.ListenerAction.forward([default_target_group]) + print("Cognito authentication is NOT enabled for this listener.") + + # Add the HTTPS listener + https_listener = alb.add_listener( + logical_id, + port=listener_port_https, + open=listener_open_to_internet, + certificates=certificates_list, + default_action=default_action, # Use the determined default action + ) + print(f"ALB HTTPS listener on port {listener_port_https} defined.") + else: + print("ACM_CERTIFICATE_ARN is not provided. Skipping HTTPS listener creation.") + + return https_listener + + +def create_ecs_express_infrastructure_role( + scope: Construct, + logical_id: str, + role_name: str, +) -> iam.Role: + """IAM role for ECS Express Mode to provision ALB, ACM cert, and autoscaling.""" + role = iam.Role( + scope, + logical_id, + role_name=role_name, + assumed_by=iam.ServicePrincipal("ecs.amazonaws.com"), + ) + role.add_managed_policy( + iam.ManagedPolicy.from_aws_managed_policy_name( + "service-role/AmazonECSInfrastructureRoleforExpressGatewayServices" + ) + ) + return role + + +def _secret_value_from_arn(secret_arn: str, json_key: str) -> str: + return f"{secret_arn}:{json_key}::" + + +def express_ingress_listener_arn( + express_service: ecs.CfnExpressGatewayService, +) -> str: + return express_service.attr_ecs_managed_resource_arns_ingress_path_listener_arn + + +def express_ingress_load_balancer_arn( + express_service: ecs.CfnExpressGatewayService, +) -> str: + return express_service.attr_ecs_managed_resource_arns_ingress_path_load_balancer_arn + + +def express_ingress_first_target_group_arn( + express_service: ecs.CfnExpressGatewayService, +) -> str: + """First target group ARN; use typed list attr (get_att returns a scalar Reference).""" + return Fn.select( + 0, + express_service.attr_ecs_managed_resource_arns_ingress_path_target_group_arns, + ) + + +def express_ingress_first_load_balancer_security_group_arn( + express_service: ecs.CfnExpressGatewayService, +) -> str: + """First Express-managed ALB security group ARN.""" + return Fn.select( + 0, + express_service.attr_ecs_managed_resource_arns_ingress_path_load_balancer_security_groups, + ) + + +def _security_group_id_from_arn(security_group_arn: str) -> str: + """EC2 APIs expect sg- IDs; Express managed-resource attrs return full ARNs.""" + return Fn.select(1, Fn.split("security-group/", security_group_arn)) + + +def express_ingress_first_load_balancer_security_group( + express_service: ecs.CfnExpressGatewayService, +) -> str: + """First ALB security group ID (sg-...) for EC2/ELB imports.""" + return _security_group_id_from_arn( + express_ingress_first_load_balancer_security_group_arn(express_service) + ) + + +# Injected via Express `secrets`, not plain environment (avoid duplication/leakage). +_EXPRESS_SECRET_ENV_NAMES = frozenset( + {"AWS_USER_POOL_ID", "AWS_CLIENT_ID", "AWS_CLIENT_SECRET"} +) + + +def create_basic_config_env( + out_dir: str = "config", + s3_log_config_bucket_name: str = S3_LOG_CONFIG_BUCKET_NAME, + s3_output_bucket_name: str = S3_OUTPUT_BUCKET_NAME, + access_log_dynamodb_table_name: str = ACCESS_LOG_DYNAMODB_TABLE_NAME, + feedback_log_dynamodb_table_name: str = FEEDBACK_LOG_DYNAMODB_TABLE_NAME, + usage_log_dynamodb_table_name: str = USAGE_LOG_DYNAMODB_TABLE_NAME, + *, + headless: bool = False, + alb_cognito: bool = False, + pi_express_backend: bool = False, +): + """ + Create a basic app_config.env file for the deployed llm_topic_modeller app. + + ``alb_cognito=True`` disables in-app Gradio Cognito login when the ALB + ``authenticate-cognito`` action already protects the service (Express Mode). + + ``pi_express_backend=True`` disables in-app login on the main app when + Pi Express calls it over Service Connect (users authenticate on the Pi UI only). + """ + variables = { + "COGNITO_AUTH": ( + "False" if headless or alb_cognito or pi_express_backend else "True" + ), + "RUN_AWS_FUNCTIONS": "True", + "RUN_AWS_BEDROCK_MODELS": "True", + "RUN_LOCAL_MODEL": "False", + "RUN_GEMINI_MODELS": "False", + "RUN_AZURE_MODELS": "False", + "PRIORITISE_SSO_OVER_AWS_ENV_ACCESS_KEYS": "True", + "DISPLAY_FILE_NAMES_IN_LOGS": "False", + "SESSION_OUTPUT_FOLDER": "True", + "SAVE_LOGS_TO_CSV": "True", + "SAVE_LOGS_TO_DYNAMODB": "True", + "SHOW_COSTS": "True", + "S3_LOG_BUCKET": s3_log_config_bucket_name, + "S3_OUTPUTS_BUCKET": s3_output_bucket_name if headless else "", + "S3_OUTPUTS_FOLDER": "output/" if headless else "", + "SAVE_OUTPUTS_TO_S3": "True" if headless else "False", + "ACCESS_LOG_DYNAMODB_TABLE_NAME": access_log_dynamodb_table_name, + "FEEDBACK_LOG_DYNAMODB_TABLE_NAME": feedback_log_dynamodb_table_name, + "USAGE_LOG_DYNAMODB_TABLE_NAME": usage_log_dynamodb_table_name, + } + + _ensure_folder_exists(out_dir + "/") + env_file_path = os.path.abspath(os.path.join(out_dir, APP_CONFIG_ENV_BASENAME)) + + if not os.path.exists(env_file_path): + with open(env_file_path, "w", encoding="utf-8"): + pass + + for key, value in variables.items(): + set_key(env_file_path, key, str(value), quote_mode="never") + + return variables + + +def load_app_config_env_for_express( + config_env_path: str, + *, + exclude_names: Optional[FrozenSet[str]] = None, + overrides: Optional[Dict[str, str]] = None, +) -> List[ecs.CfnExpressGatewayService.KeyValuePairProperty]: + """ + Load KEY=VALUE pairs from config/app_config.env for Express PrimaryContainer.environment. + + Uses the same file written by create_basic_config_env() and uploaded to S3 on the + legacy Fargate path (environmentFiles). ``overrides`` replace keys after loading + (e.g. ``COGNITO_AUTH=False`` when ALB handles Cognito). + """ + exclude = exclude_names or _EXPRESS_SECRET_ENV_NAMES + path = os.path.abspath(config_env_path) + if not os.path.isfile(path): + print( + f"Warning: app config env file not found at {path}; " + "Express container will not receive app config environment variables." + ) + return [] + + raw = dict(dotenv_values(path)) + if overrides: + raw.update(overrides) + environment: List[ecs.CfnExpressGatewayService.KeyValuePairProperty] = [] + for name, value in sorted(raw.items()): + if not name or value is None or name in exclude: + continue + environment.append( + ecs.CfnExpressGatewayService.KeyValuePairProperty( + name=name, + value=str(value), + ) + ) + print( + f"Loaded {len(environment)} environment variables from {path} for ECS Express Mode." + ) + return environment + + +def build_express_gateway_primary_container( + *, + image_uri: str, + container_port: int, + log_group_name: str, + aws_region: str, + secret: secretsmanager.ISecret, + environment: Optional[ + List[ecs.CfnExpressGatewayService.KeyValuePairProperty] + ] = None, +) -> ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty: + secret_arn = secret.secret_arn + return ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image=image_uri, + container_port=container_port, + aws_logs_configuration=ecs.CfnExpressGatewayService.ExpressGatewayServiceAwsLogsConfigurationProperty( + log_group=log_group_name, + log_stream_prefix="ecs", + ), + environment=environment or None, + secrets=[ + ecs.CfnExpressGatewayService.SecretProperty( + name="AWS_USER_POOL_ID", + value_from=_secret_value_from_arn( + secret_arn, "SUMMARISATION_USER_POOL_ID" + ), + ), + ecs.CfnExpressGatewayService.SecretProperty( + name="AWS_CLIENT_ID", + value_from=_secret_value_from_arn( + secret_arn, "SUMMARISATION_CLIENT_ID" + ), + ), + ecs.CfnExpressGatewayService.SecretProperty( + name="AWS_CLIENT_SECRET", + value_from=_secret_value_from_arn( + secret_arn, "SUMMARISATION_CLIENT_SECRET" + ), + ), + ], + ) + + +def express_gateway_idle_scaling_target( + *, + max_task_count: int = 1, +) -> ecs.CfnExpressGatewayService.ExpressGatewayScalingTargetProperty: + """Defer running tasks until post-deploy image build (legacy Fargate uses desired_count=0).""" + return ecs.CfnExpressGatewayService.ExpressGatewayScalingTargetProperty( + min_task_count=0, + max_task_count=max_task_count, + auto_scaling_metric="AVERAGE_CPU", + auto_scaling_target_value=60, + ) + + +def create_express_gateway_service( + scope: Construct, + logical_id: str, + *, + service_name: str, + cluster_name: str, + execution_role_arn: str, + infrastructure_role_arn: str, + task_role_arn: str, + cpu: str, + memory: str, + health_check_path: str, + primary_container: ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty, + subnet_ids: List[str], + security_group_ids: List[str], + scaling_target: Optional[ + ecs.CfnExpressGatewayService.ExpressGatewayScalingTargetProperty + ] = None, +) -> ecs.CfnExpressGatewayService: + network = None + if subnet_ids or security_group_ids: + network = ecs.CfnExpressGatewayService.ExpressGatewayServiceNetworkConfigurationProperty( + subnets=subnet_ids or None, + security_groups=security_group_ids or None, + ) + express_service = ecs.CfnExpressGatewayService( + scope, + logical_id, + service_name=service_name, + cluster=cluster_name, + execution_role_arn=execution_role_arn, + infrastructure_role_arn=infrastructure_role_arn, + task_role_arn=task_role_arn, + cpu=cpu, + memory=memory, + health_check_path=health_check_path, + primary_container=primary_container, + network_configuration=network, + scaling_target=scaling_target or express_gateway_idle_scaling_target(), + ) + return express_service + + +def _forward_target_group_action( + target_group_arn: str, + stickiness_seconds: int, +) -> Dict[str, Any]: + action: Dict[str, Any] = { + "Type": "forward", + "Order": 2, + "ForwardConfig": { + "TargetGroups": [{"TargetGroupArn": target_group_arn}], + }, + } + if stickiness_seconds > 0: + action["ForwardConfig"]["TargetGroupStickinessConfig"] = { + "Enabled": True, + "DurationSeconds": stickiness_seconds, + } + return action + + +def elbv2_cognito_auth_custom_resource_policy() -> cr.AwsCustomResourcePolicy: + """ + IAM policy for AwsCustomResource calls that configure authenticate-cognito on ALB. + + ELB validates the user pool client during modifyListener/createRule; the caller + (the custom-resource Lambda) needs cognito-idp:DescribeUserPoolClient. + """ + return cr.AwsCustomResourcePolicy.from_statements( + [ + iam.PolicyStatement( + actions=[ + "elasticloadbalancing:DescribeRules", + "elasticloadbalancing:ModifyListener", + "elasticloadbalancing:CreateRule", + "elasticloadbalancing:ModifyRule", + "elasticloadbalancing:DeleteRule", + ], + resources=["*"], + ), + iam.PolicyStatement( + actions=["cognito-idp:DescribeUserPoolClient"], + resources=["*"], + ), + ] + ) + + +def build_cognito_default_listener_actions( + *, + user_pool_arn: str, + user_pool_client_id: str, + user_pool_domain_prefix: str, + target_group_arn: str, + stickiness_seconds: int = 28800, + scope: str = "openid email profile", +) -> List[Dict[str, Any]]: + """Default actions for ELBv2 ModifyListener (authenticate-cognito + forward).""" + return [ + { + "Type": "authenticate-cognito", + "Order": 1, + "AuthenticateCognitoConfig": { + "UserPoolArn": user_pool_arn, + "UserPoolClientId": user_pool_client_id, + "UserPoolDomain": user_pool_domain_prefix, + "Scope": scope, + "OnUnauthenticatedRequest": "authenticate", + "SessionTimeout": stickiness_seconds, + }, + }, + _forward_target_group_action(target_group_arn, stickiness_seconds), + ] + + +def configure_express_listener_cognito_and_cloudfront( + scope: Construct, + logical_id_prefix: str, + *, + express_service: ecs.CfnExpressGatewayService, + user_pool_arn: str, + user_pool_client_id: str, + user_pool_domain_prefix: str, + use_cloudfront: bool, + cloudfront_host_header: str, + stickiness_seconds: int = 28800, + allow_cloudfront_origin_without_cognito: bool = False, +) -> None: + """ + Attach Cognito auth to the Express-managed HTTPS listener. + + By default, **no** forward-only host-header rules are added. CloudFront (and all + other) traffic uses the listener default action: ``authenticate-cognito`` then + forward. Set ``allow_cloudfront_origin_without_cognito=True`` only when the ALB + must accept CloudFront origin requests without a Cognito session (legacy pattern). + """ + listener_arn = express_ingress_listener_arn(express_service) + target_group_arn = express_ingress_first_target_group_arn(express_service) + default_actions = build_cognito_default_listener_actions( + user_pool_arn=user_pool_arn, + user_pool_client_id=user_pool_client_id, + user_pool_domain_prefix=user_pool_domain_prefix, + target_group_arn=target_group_arn, + stickiness_seconds=stickiness_seconds, + ) + modify_listener = cr.AwsCustomResource( + scope, + f"{logical_id_prefix}ModifyExpressListener", + on_create=cr.AwsSdkCall( + service="ELBv2", + action="modifyListener", + parameters={ + "ListenerArn": listener_arn, + "DefaultActions": default_actions, + }, + physical_resource_id=cr.PhysicalResourceId.of( + f"express-listener-cognito-{logical_id_prefix}" + ), + ), + on_update=cr.AwsSdkCall( + service="ELBv2", + action="modifyListener", + parameters={ + "ListenerArn": listener_arn, + "DefaultActions": default_actions, + }, + physical_resource_id=cr.PhysicalResourceId.of( + f"express-listener-cognito-{logical_id_prefix}" + ), + ), + policy=elbv2_cognito_auth_custom_resource_policy(), + ) + modify_listener.node.add_dependency(express_service) + + if ( + use_cloudfront + and cloudfront_host_header + and allow_cloudfront_origin_without_cognito + ): + forward_only = [ + { + "Type": "forward", + "Order": 1, + "ForwardConfig": { + "TargetGroups": [{"TargetGroupArn": target_group_arn}], + "TargetGroupStickinessConfig": { + "Enabled": True, + "DurationSeconds": stickiness_seconds, + }, + }, + } + ] + _elbv2_listener_rule_custom_resource( + scope, + f"{logical_id_prefix}ExpressCloudFrontHostRule", + listener_arn=listener_arn, + priority=1, + conditions=[ + { + "Field": "host-header", + "HostHeaderConfig": {"Values": [cloudfront_host_header]}, + } + ], + rule_actions=forward_only, + dependencies=[modify_listener], + ) + + +def allow_express_load_balancer_to_ecs_security_group( + scope: Construct, + logical_id: str, + *, + express_service: ecs.CfnExpressGatewayService, + ecs_security_group: ec2.ISecurityGroup, + container_port: int, +) -> None: + """Allow traffic from the Express-managed ALB security group to the task SG.""" + lb_sg_id = express_ingress_first_load_balancer_security_group(express_service) + ec2.CfnSecurityGroupIngress( + scope, + logical_id, + group_id=ecs_security_group.security_group_id, + ip_protocol="tcp", + from_port=container_port, + to_port=container_port, + source_security_group_id=lb_sg_id, + description="Express Mode ALB to ECS tasks", + ) + + +def _dict_env_to_express_key_value_pairs( + environment: Dict[str, str], +) -> List[ecs.CfnExpressGatewayService.KeyValuePairProperty]: + return [ + ecs.CfnExpressGatewayService.KeyValuePairProperty(name=k, value=str(v)) + for k, v in environment.items() + if v is not None and str(v) != "" + ] + + +def normalize_pi_alb_path_prefix(raw: str, *, default: str = "pi") -> str: + """Return a leading-slash path prefix (no trailing slash), e.g. '/pi'.""" + segment = (raw or default).strip().strip("/") + return f"/{segment}" if segment else f"/{default}" + + +def normalize_pi_alb_routing_mode(raw: str) -> str: + mode = (raw or "path").strip().lower() + allowed = frozenset({"path", "host", "both"}) + if mode not in allowed: + raise ValueError( + f"PI_ALB_ROUTING must be one of {sorted(allowed)}; got '{raw}'." + ) + return mode + + +def pi_alb_path_patterns(path_prefix: str) -> List[str]: + """ALB path-pattern values for a Pi path prefix (exact + subtree).""" + prefix = normalize_pi_alb_path_prefix(path_prefix) + return [prefix, f"{prefix}/*"] + + +def pi_alb_health_check_path(path_prefix: str, routing_mode: str) -> str: + if normalize_pi_alb_routing_mode(routing_mode) in ("path", "both"): + return f"{normalize_pi_alb_path_prefix(path_prefix)}/" + return "/" + + +def pi_alb_root_path_for_container(path_prefix: str, routing_mode: str) -> str: + """Gradio/FastAPI ROOT_PATH to set on Pi tasks when path routing is enabled.""" + if normalize_pi_alb_routing_mode(routing_mode) in ("path", "both"): + return normalize_pi_alb_path_prefix(path_prefix) + return "" + + +def pi_listener_rule_count(routing_mode: str) -> int: + mode = normalize_pi_alb_routing_mode(routing_mode) + count = 0 + if mode in ("path", "both"): + count += 1 + if mode in ("host", "both"): + count += 1 + return count + + +def format_express_pi_public_url(express_endpoint: str) -> str: + """Public Pi UI URL for a dedicated ECS Express managed HTTPS endpoint.""" + base = (express_endpoint or "").strip().rstrip("/") + return f"{base}/" if base else "" + + +def format_pi_public_urls( + *, + routing_mode: str, + path_prefix: str, + host_header: str, + cloudfront_domain: str = "", + use_https: bool = True, +) -> List[str]: + """Human-facing Pi UI URLs for stack outputs.""" + scheme = "https" if use_https else "http" + urls: List[str] = [] + mode = normalize_pi_alb_routing_mode(routing_mode) + prefix = normalize_pi_alb_path_prefix(path_prefix) + if mode in ("path", "both"): + if cloudfront_domain.strip(): + urls.append(f"{scheme}://{cloudfront_domain.strip()}{prefix}/") + else: + urls.append(f"{scheme}://{prefix}/") + if mode in ("host", "both") and host_header.strip(): + urls.append(f"{scheme}://{host_header.strip()}/") + return urls + + +def _apply_pi_root_path_env(env: Dict[str, str], pi_root_path: str) -> None: + if pi_root_path: + env["PI_ROOT_PATH"] = pi_root_path + env["ROOT_PATH"] = pi_root_path + env["FASTAPI_ROOT_PATH"] = pi_root_path + + +def build_pi_express_container_environment( + *, + service_connect_discovery_name: str, + main_app_port: Union[str, int], + pi_gradio_port: Union[str, int], + cognito_auth: bool = True, +) -> Dict[str, str]: + """Inline env for Pi on Express (no volume mounts; workspace under /tmp).""" + port = int(main_app_port) + pi_port = int(pi_gradio_port) + env = { + "APP_TYPE": "pi", + "APP_CONFIG_PATH": "/workspace/doc_summarisation/config/pi_agent.env.example", + "PI_DEPLOYMENT_PROFILE": "aws-ecs", + "PI_DEFAULT_PROVIDER": "amazon-bedrock", + "DOC_SUMMARISATION_GRADIO_URL": f"http://{service_connect_discovery_name}:{port}", + "PI_GRADIO_PORT": str(pi_port), + "GRADIO_SERVER_PORT": str(pi_port), + "GRADIO_SERVER_NAME": "0.0.0.0", + "PI_WORKSPACE_DIR": "/tmp/pi-workspace", + "PI_WORKDIR": "/workspace/doc_summarisation", + "PI_UPLOAD_ROOT": "/tmp/gradio", + "PI_SESSION_DIR": "/tmp/pi-sessions", + "PI_CODING_AGENT_DIR": "/tmp/pi-agent", + "ACCESS_LOGS_FOLDER": "/tmp/pi-logs/", + "USAGE_LOGS_FOLDER": "/tmp/pi-usage/", + "FEEDBACK_LOGS_FOLDER": "/tmp/pi-feedback/", + "RUN_FASTAPI": "True", + "RUN_AWS_FUNCTIONS": "True", + "COGNITO_AUTH": "True" if cognito_auth else "False", + } + return env + + +def build_express_pi_primary_container( + *, + image_uri: str, + container_port: int, + log_group_name: str, + aws_region: str, + environment: Optional[Dict[str, str]] = None, + secret: Optional[secretsmanager.ISecret] = None, + cognito_auth: bool = True, +) -> ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty: + """Express PrimaryContainer for Pi (inline env; Cognito creds from Secrets Manager).""" + env_pairs = ( + _dict_env_to_express_key_value_pairs(environment) if environment else None + ) + secrets: Optional[List[ecs.CfnExpressGatewayService.SecretProperty]] = None + if secret is not None and cognito_auth: + secret_arn = secret.secret_arn + secrets = [ + ecs.CfnExpressGatewayService.SecretProperty( + name="AWS_USER_POOL_ID", + value_from=_secret_value_from_arn( + secret_arn, "SUMMARISATION_USER_POOL_ID" + ), + ), + ecs.CfnExpressGatewayService.SecretProperty( + name="AWS_CLIENT_ID", + value_from=_secret_value_from_arn( + secret_arn, "SUMMARISATION_CLIENT_ID" + ), + ), + ecs.CfnExpressGatewayService.SecretProperty( + name="AWS_CLIENT_SECRET", + value_from=_secret_value_from_arn( + secret_arn, "SUMMARISATION_CLIENT_SECRET" + ), + ), + ] + return ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image=image_uri, + container_port=container_port, + aws_logs_configuration=ecs.CfnExpressGatewayService.ExpressGatewayServiceAwsLogsConfigurationProperty( + log_group=log_group_name, + log_stream_prefix="ecs-pi", + ), + environment=env_pairs, + secrets=secrets, + ) + + +_ELBV2_RULE_DELETE_IGNORE = ( + ".*(not a valid listener rule ARN|RuleNotFound|ResourceNotFound|ValidationError).*" +) +_ELBV2_LISTENER_RULE_UPSERT_PROVIDER_ID = "Elbv2ListenerRuleUpsertProvider" + + +def _elbv2_listener_rule_upsert_provider(scope: Construct) -> cr.Provider: + """Shared Provider for upserting ALB listener rules (one Lambda per stack).""" + stack = Stack.of(scope) + existing = stack.node.try_find_child(_ELBV2_LISTENER_RULE_UPSERT_PROVIDER_ID) + if existing is not None: + return existing + + asset_dir = os.path.join( + os.path.dirname(__file__), "lambda_elbv2_listener_rule_upsert" + ) + fn = lambda_.Function( + stack, + "Elbv2ListenerRuleUpsertFn", + runtime=lambda_.Runtime.PYTHON_3_12, + handler="lambda_function.handler", + code=lambda_.Code.from_asset(asset_dir), + timeout=Duration.seconds(120), + description="Upsert ALB listener rules for Express Cognito routing", + ) + fn.add_to_role_policy( + iam.PolicyStatement( + actions=[ + "elasticloadbalancing:DescribeRules", + "elasticloadbalancing:CreateRule", + "elasticloadbalancing:ModifyRule", + "elasticloadbalancing:DeleteRule", + ], + resources=["*"], + ) + ) + fn.add_to_role_policy( + iam.PolicyStatement( + actions=["cognito-idp:DescribeUserPoolClient"], + resources=["*"], + ) + ) + return cr.Provider( + stack, + _ELBV2_LISTENER_RULE_UPSERT_PROVIDER_ID, + on_event_handler=fn, + ) + + +def _elbv2_listener_rule_custom_resource( + scope: Construct, + logical_id: str, + *, + listener_arn: str, + priority: int, + conditions: List[Dict[str, Any]], + rule_actions: List[Dict[str, Any]], + dependencies: Optional[List[Any]] = None, +) -> CustomResource: + """ + Create or update a numbered listener rule (upsert by priority + conditions). + + Reuses an existing rule at the same priority when conditions match, which + avoids PriorityInUse failures after partial deploy rollbacks. + """ + provider = _elbv2_listener_rule_upsert_provider(scope) + resource = CustomResource( + scope, + logical_id, + service_token=provider.service_token, + resource_type="Custom::Elbv2ListenerRuleUpsert", + properties={ + "ListenerArn": listener_arn, + "Priority": priority, + "Conditions": conditions, + "Actions": rule_actions, + }, + ) + for dep in dependencies or []: + resource.node.add_dependency(dep) + return resource + + +def _express_pi_listener_rule_custom_resource( + scope: Construct, + logical_id: str, + *, + listener_arn: str, + priority: int, + conditions: List[Dict[str, Any]], + rule_actions: List[Dict[str, Any]], + express_main_service: ecs.CfnExpressGatewayService, + express_pi_service: ecs.CfnExpressGatewayService, +) -> CustomResource: + return _elbv2_listener_rule_custom_resource( + scope, + logical_id, + listener_arn=listener_arn, + priority=priority, + conditions=conditions, + rule_actions=rule_actions, + dependencies=[express_pi_service, express_main_service], + ) + + +def configure_express_pi_listener_rules( + scope: Construct, + logical_id_prefix: str, + *, + express_main_service: ecs.CfnExpressGatewayService, + express_pi_service: ecs.CfnExpressGatewayService, + routing_mode: str, + path_prefix: str, + pi_host_header: str, + rule_priority: int, + user_pool_arn: str, + user_pool_client_id: str, + user_pool_domain_prefix: str, + stickiness_seconds: int = 28800, +) -> int: + """ + Path and/or host-header rules on the shared Express HTTPS listener → Pi TG. + Returns the next free listener rule priority after Pi rules. + """ + mode = normalize_pi_alb_routing_mode(routing_mode) + listener_arn = express_ingress_listener_arn(express_main_service) + pi_target_group_arn = express_ingress_first_target_group_arn(express_pi_service) + rule_actions = build_cognito_default_listener_actions( + user_pool_arn=user_pool_arn, + user_pool_client_id=user_pool_client_id, + user_pool_domain_prefix=user_pool_domain_prefix, + target_group_arn=pi_target_group_arn, + stickiness_seconds=stickiness_seconds, + ) + priority = rule_priority + + if mode in ("path", "both"): + path_patterns = pi_alb_path_patterns(path_prefix) + _express_pi_listener_rule_custom_resource( + scope, + f"{logical_id_prefix}ExpressPiPathRule", + listener_arn=listener_arn, + priority=priority, + conditions=[ + { + "Field": "path-pattern", + "PathPatternConfig": {"Values": path_patterns}, + } + ], + rule_actions=rule_actions, + express_main_service=express_main_service, + express_pi_service=express_pi_service, + ) + priority += 1 + + if mode in ("host", "both") and pi_host_header.strip(): + _express_pi_listener_rule_custom_resource( + scope, + f"{logical_id_prefix}ExpressPiHostRule", + listener_arn=listener_arn, + priority=priority, + conditions=[ + { + "Field": "host-header", + "HostHeaderConfig": {"Values": [pi_host_header.strip()]}, + } + ], + rule_actions=rule_actions, + express_main_service=express_main_service, + express_pi_service=express_pi_service, + ) + priority += 1 + + return priority + + +def _express_service_connect_configuration( + *, + namespace: str, + port_name: Optional[str] = None, + discovery_name: Optional[str] = None, + port: Optional[int] = None, +) -> Dict[str, Any]: + """ECS API serviceConnectConfiguration payload for updateService.""" + cfg: Dict[str, Any] = {"enabled": True, "namespace": namespace} + if port_name and discovery_name and port is not None: + cfg["services"] = [ + { + "portName": port_name, + "discoveryName": discovery_name, + "clientAliases": [ + {"port": int(port), "dnsName": discovery_name}, + ], + } + ] + return cfg + + +def apply_service_connect_to_express_service( + scope: Construct, + logical_id: str, + *, + cluster_name: str, + service_name: str, + namespace: str, + express_service: ecs.CfnExpressGatewayService, + port_name: Optional[str] = None, + discovery_name: Optional[str] = None, + port: Optional[int] = None, +) -> cr.AwsCustomResource: + """ + Enable Service Connect on an Express gateway service after create (AWS does not + support SC at Express create time). Server config when port_name/discovery_name/port + are set; client-only when they are omitted. + """ + sc_cfg = _express_service_connect_configuration( + namespace=namespace, + port_name=port_name, + discovery_name=discovery_name, + port=port, + ) + physical_id = f"{cluster_name}/{service_name}/service-connect" + custom = cr.AwsCustomResource( + scope, + logical_id, + on_create=cr.AwsSdkCall( + service="ECS", + action="updateService", + parameters={ + "cluster": cluster_name, + "service": service_name, + "serviceConnectConfiguration": sc_cfg, + "forceNewDeployment": True, + }, + physical_resource_id=cr.PhysicalResourceId.of(physical_id), + ), + on_update=cr.AwsSdkCall( + service="ECS", + action="updateService", + parameters={ + "cluster": cluster_name, + "service": service_name, + "serviceConnectConfiguration": sc_cfg, + "forceNewDeployment": True, + }, + physical_resource_id=cr.PhysicalResourceId.of(physical_id), + ), + policy=cr.AwsCustomResourcePolicy.from_sdk_calls( + resources=cr.AwsCustomResourcePolicy.ANY_RESOURCE + ), + ) + custom.node.add_dependency(express_service) + return custom + + +def create_s3_batch_ecs_trigger_lambda( + scope: Construct, + logical_id: str, + *, + function_name: Optional[str], + lambda_asset_path: str, + output_bucket: s3.IBucket, + config_bucket: s3.IBucket, + cluster_name: str, + task_definition_arn: str, + container_name: str, + subnet_ids: List[str], + security_group_id: str, + execution_role: iam.IRole, + task_role: iam.IRole, + env_prefix: str, + env_suffix: str, + input_prefix: str, + config_prefix: str, + default_params_key: str, + general_env_prefix: str = "general-config/", + default_task_type: str = "extract", + default_input_s3_uri: str = "", + assign_public_ip: bool = False, +) -> lambda_.Function: + """ + Lambda triggered by job .env uploads on the output bucket; runs one-shot Fargate tasks. + """ + lambda_role = iam.Role( + scope, + f"{logical_id}Role", + assumed_by=iam.ServicePrincipal("lambda.amazonaws.com"), + managed_policies=[ + iam.ManagedPolicy.from_aws_managed_policy_name( + "service-role/AWSLambdaBasicExecutionRole" + ) + ], + ) + + lambda_role.add_to_policy( + iam.PolicyStatement( + actions=["ecs:RunTask"], + resources=[task_definition_arn], + ) + ) + lambda_role.add_to_policy( + iam.PolicyStatement( + actions=["ecs:RunTask"], + resources=[ + f"arn:aws:ecs:*:*:cluster/{cluster_name}", + ], + ) + ) + lambda_role.add_to_policy( + iam.PolicyStatement( + actions=["iam:PassRole"], + resources=[execution_role.role_arn, task_role.role_arn], + conditions={ + "StringEquals": {"iam:PassedToService": "ecs-tasks.amazonaws.com"} + }, + ) + ) + output_bucket.grant_read(lambda_role, f"{env_prefix}*") + if general_env_prefix: + output_bucket.grant_read(lambda_role, f"{general_env_prefix}*") + config_bucket.grant_read(lambda_role) + if default_params_key: + output_bucket.grant_read(lambda_role, default_params_key) + + bucket_name = output_bucket.bucket_name + if not default_input_s3_uri: + default_input_s3_uri = f"s3://{bucket_name}/{input_prefix.rstrip('/')}/dummy_consultation_response.xlsx" + + fn_kwargs: Dict[str, Any] = { + "runtime": lambda_.Runtime.PYTHON_3_12, + "handler": "lambda_function.lambda_handler", + "code": lambda_.Code.from_asset(lambda_asset_path), + "role": lambda_role, + "timeout": Duration.seconds(60), + "memory_size": 256, + "environment": { + "BUCKET": bucket_name, + "INPUT_PREFIX": input_prefix, + "ENV_PREFIX": env_prefix, + "GENERAL_ENV_PREFIX": general_env_prefix, + "ENV_SUFFIX": env_suffix, + "DEFAULT_PARAMS_KEY": default_params_key, + "ECS_CLUSTER": cluster_name, + "ECS_TASK_DEF": task_definition_arn, + "SUBNETS": ",".join(subnet_ids), + "SECURITY_GROUPS": security_group_id, + "CONTAINER_NAME": container_name, + "DEFAULT_TASK_TYPE": default_task_type, + "DEFAULT_INPUT_S3_URI": default_input_s3_uri, + "ECS_ASSIGN_PUBLIC_IP": "ENABLED" if assign_public_ip else "DISABLED", + }, + } + if function_name: + fn_kwargs["function_name"] = function_name + + batch_fn = lambda_.Function(scope, logical_id, **fn_kwargs) + + output_bucket.add_event_notification( + s3.EventType.OBJECT_CREATED, + s3n.LambdaDestination(batch_fn), + s3.NotificationKeyFilter(prefix=env_prefix, suffix=env_suffix), + ) + + return batch_fn + + +def create_dynamo_usage_log_export_lambda( + scope: Construct, + logical_id: str, + *, + function_name: Optional[str], + lambda_asset_path: str, + dynamodb_table: dynamodb.ITable, + output_bucket: s3.IBucket, + s3_output_key: str, + schedule_expression: str, + dynamodb_table_name: str, + date_attribute: str = "timestamp", + output_filename: str = "dynamodb_logs_export.csv", + shared_kms_key_arn: Optional[str] = None, +) -> lambda_.Function: + """ + Scheduled Lambda: scan usage-log DynamoDB table, export CSV, upload to S3. + + Triggered by EventBridge on ``schedule_expression`` (cron or rate). + """ + lambda_role = iam.Role( + scope, + f"{logical_id}Role", + assumed_by=iam.ServicePrincipal("lambda.amazonaws.com"), + managed_policies=[ + iam.ManagedPolicy.from_aws_managed_policy_name( + "service-role/AWSLambdaBasicExecutionRole" + ) + ], + ) + dynamodb_table.grant_read_data(lambda_role) + output_bucket.grant_put(lambda_role, s3_output_key) + if shared_kms_key_arn: + lambda_role.add_to_policy( + iam.PolicyStatement( + actions=[ + "kms:Encrypt", + "kms:Decrypt", + "kms:GenerateDataKey", + "kms:DescribeKey", + ], + resources=[shared_kms_key_arn], + ) + ) + + fn_kwargs: Dict[str, Any] = { + "runtime": lambda_.Runtime.PYTHON_3_12, + "handler": "lambda_function.lambda_handler", + "code": lambda_.Code.from_asset(lambda_asset_path), + "role": lambda_role, + "timeout": Duration.minutes(15), + "memory_size": 512, + "environment": { + "DYNAMODB_TABLE_NAME": dynamodb_table_name, + "USAGE_LOG_DYNAMODB_TABLE_NAME": dynamodb_table_name, + "OUTPUT_FOLDER": "/tmp", + "OUTPUT_FILENAME": output_filename, + "DATE_ATTRIBUTE": date_attribute, + "S3_OUTPUT_BUCKET": output_bucket.bucket_name, + "S3_OUTPUT_KEY": s3_output_key, + }, + } + if function_name: + fn_kwargs["function_name"] = function_name + + export_fn = lambda_.Function(scope, logical_id, **fn_kwargs) + + events.Rule( + scope, + f"{logical_id}Schedule", + schedule=events.Schedule.expression(schedule_expression), + description=( + "Export DynamoDB usage logs to CSV in S3 " f"({schedule_expression})" + ), + targets=[targets.LambdaFunction(export_fn)], + ) + + CfnOutput( + scope, + f"{logical_id}LambdaArn", + value=export_fn.function_arn, + description="Lambda ARN for scheduled DynamoDB usage log export to S3", + ) + CfnOutput( + scope, + f"{logical_id}S3Uri", + value=f"s3://{output_bucket.bucket_name}/{s3_output_key}", + description="S3 URI for the scheduled DynamoDB usage log CSV export", + ) + + return export_fn + + +def sanitize_headless_metric_filter_id(raw: str) -> str: + """S3 request metrics configuration Id (alphanumeric, hyphen, underscore).""" + cleaned = re.sub(r"[^A-Za-z0-9_-]", "-", (raw or "").strip()) + cleaned = re.sub(r"-{2,}", "-", cleaned).strip("-") + return (cleaned or "s3-output-put")[:64] + + +def create_headless_output_notifications( + scope: Construct, + logical_id_prefix: str, + *, + output_bucket: s3.IBucket, + output_prefix: str, + notify_email: str, + iam_user_name: str, + metric_filter_id: str, + sns_topic_name: str, + alarm_name: str, + kms_key_arn: Optional[str] = None, +) -> Dict[str, str]: + """ + Headless follow-on: S3 PutRequests metric on ``output_prefix`` -> CloudWatch alarm + -> SNS email, plus an IAM user that can list/get/put/delete objects in the bucket. + + Mirrors the pattern documented under cdk/alarms_and_user/. + """ + output_bucket_name = output_bucket.bucket_name + metric_id = sanitize_headless_metric_filter_id(metric_filter_id) + prefix = (output_prefix or "output/").strip() + if prefix and not prefix.endswith("/"): + prefix += "/" + email = (notify_email or "").strip() + if not email: + raise ValueError("notify_email is required for headless output notifications") + + metrics_resource = cr.AwsCustomResource( + scope, + f"{logical_id_prefix}OutputBucketPutMetrics", + on_create=cr.AwsSdkCall( + service="S3", + action="putBucketMetricsConfiguration", + parameters={ + "Bucket": output_bucket_name, + "Id": metric_id, + "MetricsConfiguration": { + "Id": metric_id, + "Filter": {"Prefix": prefix}, + }, + }, + physical_resource_id=cr.PhysicalResourceId.of( + f"{output_bucket_name}-{metric_id}" + ), + ), + on_update=cr.AwsSdkCall( + service="S3", + action="putBucketMetricsConfiguration", + parameters={ + "Bucket": output_bucket_name, + "Id": metric_id, + "MetricsConfiguration": { + "Id": metric_id, + "Filter": {"Prefix": prefix}, + }, + }, + physical_resource_id=cr.PhysicalResourceId.of( + f"{output_bucket_name}-{metric_id}" + ), + ), + on_delete=cr.AwsSdkCall( + service="S3", + action="deleteBucketMetricsConfiguration", + parameters={"Bucket": output_bucket_name, "Id": metric_id}, + ), + policy=cr.AwsCustomResourcePolicy.from_statements( + [ + iam.PolicyStatement( + actions=[ + "s3:PutMetricsConfiguration", + "s3:GetMetricsConfiguration", + "s3:DeleteMetricsConfiguration", + "s3:ListBucket", + ], + resources=[ + f"arn:aws:s3:::{output_bucket_name}", + f"arn:aws:s3:::{output_bucket_name}/*", + ], + ) + ] + ), + ) + + topic = sns.Topic( + scope, + f"{logical_id_prefix}OutputNotifyTopic", + topic_name=sns_topic_name[:256], + display_name=sns_topic_name[:100], + ) + topic.add_to_resource_policy( + iam.PolicyStatement( + sid="AllowPublishAlarms", + effect=iam.Effect.ALLOW, + principals=[iam.ServicePrincipal("cloudwatch.amazonaws.com")], + actions=["sns:Publish"], + resources=[topic.topic_arn], + ) + ) + topic.add_subscription(sns_subscriptions.EmailSubscription(email)) + + alarm = cloudwatch.Alarm( + scope, + f"{logical_id_prefix}OutputPutAlarm", + alarm_name=alarm_name[:255], + metric=cloudwatch.Metric( + namespace="AWS/S3", + metric_name="PutRequests", + dimensions_map={ + "BucketName": output_bucket_name, + "FilterId": metric_id, + }, + statistic="Sum", + period=Duration.minutes(1), + ), + threshold=0, + evaluation_periods=1, + comparison_operator=cloudwatch.ComparisonOperator.GREATER_THAN_THRESHOLD, + treat_missing_data=cloudwatch.TreatMissingData.NOT_BREACHING, + alarm_description=( + f"Notify when new analysis outputs are written under s3://" + f"{output_bucket_name}/{prefix}" + ), + ) + alarm.add_alarm_action(cloudwatch_actions.SnsAction(topic)) + alarm.node.add_dependency(metrics_resource) + + reader_user = iam.User( + scope, + f"{logical_id_prefix}OutputReaderUser", + user_name=iam_user_name[:64], + ) + reader_user.add_to_policy( + iam.PolicyStatement( + effect=iam.Effect.ALLOW, + actions=[ + "s3:ListBucket", + "s3:GetObject", + "s3:PutObject", + "s3:DeleteObject", + ], + resources=[ + f"arn:aws:s3:::{output_bucket_name}", + f"arn:aws:s3:::{output_bucket_name}/*", + ], + ) + ) + if kms_key_arn: + reader_user.add_to_policy( + iam.PolicyStatement( + effect=iam.Effect.ALLOW, + actions=[ + "kms:Encrypt", + "kms:Decrypt", + "kms:GenerateDataKey", + "kms:DescribeKey", + ], + resources=[kms_key_arn], + ) + ) + + # Resource-based policy on the bucket (matches TaskRole grants; required for some + # KMS-encrypted buckets and org policies that expect explicit bucket principals). + output_bucket.add_to_resource_policy( + iam.PolicyStatement( + effect=iam.Effect.ALLOW, + principals=[reader_user], + actions=["s3:GetObject", "s3:PutObject", "s3:DeleteObject"], + resources=[f"{output_bucket.bucket_arn}/*"], + ) + ) + output_bucket.add_to_resource_policy( + iam.PolicyStatement( + effect=iam.Effect.ALLOW, + principals=[reader_user], + actions=["s3:ListBucket"], + resources=[output_bucket.bucket_arn], + ) + ) + + CfnOutput( + scope, + f"{logical_id_prefix}OutputNotifyTopicArn", + value=topic.topic_arn, + description="SNS topic for headless output-bucket PutRequests alarms", + ) + CfnOutput( + scope, + f"{logical_id_prefix}OutputReaderUserArn", + value=reader_user.user_arn, + description="IAM user for programmatic download of headless analysis outputs", + ) + CfnOutput( + scope, + f"{logical_id_prefix}OutputPutAlarmName", + value=alarm.alarm_name, + description="CloudWatch alarm on S3 PutRequests for new analysis outputs", + ) + + return { + "sns_topic_arn": topic.topic_arn, + "iam_user_name": reader_user.user_name, + "alarm_name": alarm.alarm_name, + "metric_filter_id": metric_id, + } + + +def build_headless_app_defaults_env_content( + seed_asset_directory: str, + *, + s3_outputs_bucket_name: str, +) -> str: + """ + Render general-config/app_defaults.env for headless batch jobs. + + Merges the static seed template with deployment-specific values (notably + ``S3_OUTPUTS_BUCKET`` from the stack output bucket). + """ + template_path = os.path.join( + seed_asset_directory, "general-config", "app_defaults.env" + ) + out_lines: List[str] = [] + wrote_outputs_bucket = False + + if os.path.isfile(template_path): + with open(template_path, encoding="utf-8") as handle: + for raw_line in handle: + line = raw_line.rstrip("\n") + stripped = line.strip() + if ( + stripped + and not stripped.startswith("#") + and "=" in stripped + and stripped.split("=", 1)[0].strip() == "S3_OUTPUTS_BUCKET" + ): + out_lines.append(f"S3_OUTPUTS_BUCKET={s3_outputs_bucket_name}") + wrote_outputs_bucket = True + else: + out_lines.append(line) + + if not wrote_outputs_bucket: + if out_lines and out_lines[-1].strip(): + out_lines.append("") + out_lines.append(f"S3_OUTPUTS_BUCKET={s3_outputs_bucket_name}") + + return "\n".join(out_lines).rstrip() + "\n" + + +def create_headless_s3_batch_seed( + scope: Construct, + logical_id: str, + *, + destination_bucket: s3.IBucket, + seed_asset_directory: str, + s3_outputs_bucket_name: str, +) -> None: + """Upload input/ and input/config/ markers plus example job .env to the output bucket.""" + from aws_cdk import aws_s3_deployment as s3deploy + + app_defaults_content = build_headless_app_defaults_env_content( + seed_asset_directory, + s3_outputs_bucket_name=s3_outputs_bucket_name, + ) + + s3deploy.BucketDeployment( + scope, + logical_id, + sources=[ + s3deploy.Source.asset(seed_asset_directory), + s3deploy.Source.data( + "general-config/app_defaults.env", + app_defaults_content, + ), + ], + destination_bucket=destination_bucket, + prune=False, + ) + + +def build_pi_agent_container_environment( + *, + service_connect_discovery_name: str, + main_app_port: Union[str, int], + pi_gradio_port: Union[str, int], + pi_root_path: str = "", +) -> Dict[str, str]: + """Inline env for Pi agent tasks (overrides image defaults; SC URL for main app).""" + port = int(main_app_port) + pi_port = int(pi_gradio_port) + env = { + "APP_TYPE": "pi", + "APP_CONFIG_PATH": "/workspace/doc_summarisation/config/pi_agent.env", + "PI_DEPLOYMENT_PROFILE": "aws-ecs", + "PI_DEFAULT_PROVIDER": "amazon-bedrock", + "DOC_SUMMARISATION_GRADIO_URL": f"http://{service_connect_discovery_name}:{port}", + "PI_GRADIO_PORT": str(pi_port), + "GRADIO_SERVER_PORT": str(pi_port), + "GRADIO_SERVER_NAME": "0.0.0.0", + "PI_WORKSPACE_DIR": "/home/user/app/workspace", + "PI_WORKDIR": "/workspace/doc_summarisation", + "PI_UPLOAD_ROOT": "/tmp/gradio", + "PI_SESSION_DIR": "/tmp/pi-sessions", + "PI_CODING_AGENT_DIR": "/tmp/pi-agent", + "ACCESS_LOGS_FOLDER": "/tmp/pi-logs/", + "USAGE_LOGS_FOLDER": "/tmp/pi-usage/", + "FEEDBACK_LOGS_FOLDER": "/tmp/pi-feedback/", + "RUN_FASTAPI": "True", + "RUN_AWS_FUNCTIONS": "True", + "SAVE_OUTPUTS_TO_S3": "True", + "S3_OUTPUTS_BUCKET": S3_OUTPUT_BUCKET_NAME, + "COGNITO_AUTH": "False", + } + _apply_pi_root_path_env(env, pi_root_path) + return env + + +# Gradio mounted on FastAPI (tools.gradio_platform.mount_or_launch); matches agent-redact/pi/start.sh. +PI_ECS_APP_START_CMD = ( + "python3 agent-redact/pi/pi_agent_config.py && " + "exec uvicorn gradio_app:app --app-dir agent-redact/pi " + "--host 0.0.0.0 --port ${PI_GRADIO_PORT:-7862} " + '--proxy-headers --forwarded-allow-ips "*"' +) + +# Fargate volume mounts are root-owned; chown as root, then run the app as user (see entrypoint-ecs.sh). +PI_ECS_CONTAINER_USER = "root" +PI_ECS_CONTAINER_COMMAND = [ + "/usr/local/bin/entrypoint-ecs.sh", + PI_ECS_APP_START_CMD, +] +# Inline fallback when the image predates entrypoint-ecs.sh (same behaviour via bash). +PI_ECS_CONTAINER_COMMAND_FALLBACK = [ + "bash", + "-c", + "mkdir -p /tmp/pi-agent /tmp/pi-logs /tmp/pi-usage /tmp/pi-feedback " + "/home/user/app/workspace /tmp/gradio /tmp/pi-sessions && " + "chown -R user:user /tmp/pi-agent /tmp/pi-logs /tmp/pi-usage /tmp/pi-feedback " + "/home/user/app/workspace /tmp/gradio /tmp/pi-sessions && " + "cd /workspace/doc_summarisation && " + f"exec su -s /bin/bash user -c '{PI_ECS_APP_START_CMD}'", +] + + +def create_pi_agent_ecs_resources( + scope: Construct, + logical_id_prefix: str, + *, + vpc: ec2.IVpc, + cluster: ecs.ICluster, + private_subnets: List[ec2.ISubnet], + pi_ecr_image_uri: str, + container_name: str, + task_role: iam.IRole, + execution_role: iam.IRole, + config_bucket: s3.IBucket, + pi_agent_env_s3_key: str, + service_name: str, + task_family: str, + security_group_name: str, + log_group_name: str, + cpu: int, + memory_mib: int, + pi_gradio_port: int, + service_connect_namespace: str, + service_connect_discovery_name: str, + main_app_port: int, + use_fargate_spot: str, + pi_root_path: str = "", +) -> Tuple[ecs.FargateService, ec2.SecurityGroup, ecs.FargateTaskDefinition]: + """Second Fargate service for the Pi agent (joins Service Connect namespace as a client).""" + pi_security_group = ec2.SecurityGroup( + scope, + f"{logical_id_prefix}SecurityGroup", + vpc=vpc, + security_group_name=security_group_name, + description="Pi agent ECS tasks", + ) + + pi_log_group = logs.LogGroup( + scope, + f"{logical_id_prefix}LogGroup", + log_group_name=log_group_name, + retention=logs.RetentionDays.ONE_MONTH, + removal_policy=managed_resource_removal_policy(), + ) + + pi_volume = ecs.Volume(name="piEphemeralVolume") + pi_task_definition = ecs.FargateTaskDefinition( + scope, + f"{logical_id_prefix}TaskDefinition", + family=task_family, + cpu=cpu, + memory_limit_mib=memory_mib, + task_role=task_role, + execution_role=execution_role, + runtime_platform=ecs.RuntimePlatform( + cpu_architecture=ecs.CpuArchitecture.X86_64, + operating_system_family=ecs.OperatingSystemFamily.LINUX, + ), + ephemeral_storage_gib=21, + volumes=[pi_volume], + ) + + env_files: List[ecs.EnvironmentFile] = [] + if pi_agent_env_s3_key: + env_files.append( + ecs.EnvironmentFile.from_bucket(config_bucket, pi_agent_env_s3_key) + ) + + pi_container = pi_task_definition.add_container( + container_name, + image=ecs.ContainerImage.from_registry(f"{pi_ecr_image_uri}:latest"), + logging=ecs.LogDriver.aws_logs( + stream_prefix="ecs-pi", + log_group=pi_log_group, + ), + environment_files=env_files if env_files else None, + environment=build_pi_agent_container_environment( + service_connect_discovery_name=service_connect_discovery_name, + main_app_port=main_app_port, + pi_gradio_port=pi_gradio_port, + pi_root_path=pi_root_path, + ), + command=PI_ECS_CONTAINER_COMMAND_FALLBACK, + user=PI_ECS_CONTAINER_USER, + essential=True, + ) + + pi_container.add_mount_points( + ecs.MountPoint( + source_volume=pi_volume.name, + container_path="/home/user/app/workspace", + read_only=False, + ), + ecs.MountPoint( + source_volume=pi_volume.name, + container_path="/tmp/gradio", + read_only=False, + ), + ecs.MountPoint( + source_volume=pi_volume.name, + container_path="/tmp/pi-sessions", + read_only=False, + ), + ) + + pi_container.add_port_mappings( + ecs.PortMapping( + container_port=pi_gradio_port, + host_port=pi_gradio_port, + name=f"port-{pi_gradio_port}", + protocol=ecs.Protocol.TCP, + app_protocol=ecs.AppProtocol.http, + ) + ) + + pi_service = ecs.FargateService( + scope, + f"{logical_id_prefix}Service", + service_name=service_name, + cluster=cluster, + task_definition=pi_task_definition, + security_groups=[pi_security_group], + vpc_subnets=ec2.SubnetSelection(subnets=private_subnets), + platform_version=ecs.FargatePlatformVersion.LATEST, + capacity_provider_strategies=[ + ecs.CapacityProviderStrategy( + capacity_provider=use_fargate_spot, + base=0, + weight=1, + ) + ], + min_healthy_percent=0, + max_healthy_percent=100, + desired_count=0, + availability_zone_rebalancing=ecs_availability_zone_rebalancing( + ECS_AVAILABILITY_ZONE_REBALANCING + ), + service_connect_configuration=ecs.ServiceConnectProps( + namespace=service_connect_namespace, + ), + ) + + return pi_service, pi_security_group, pi_task_definition + + +def attach_pi_agent_to_shared_alb( + scope: Construct, + logical_id_prefix: str, + *, + vpc: ec2.IVpc, + alb_security_group: ec2.ISecurityGroup, + pi_security_group: ec2.SecurityGroup, + pi_service: ecs.FargateService, + pi_port: int, + routing_mode: str, + path_prefix: str, + pi_host_header: str, + listener_rule_priority: int, + target_group_name: str, + stickiness_cookie_duration: Duration, + https_listener: Optional[elb.IApplicationListener], + http_listener: Optional[elb.IApplicationListener], + acm_certificate_arn: str, + enable_cognito_auth: bool, + cognito_user_pool: Optional[cognito.IUserPool], + cognito_user_pool_client: Optional[cognito.IUserPoolClient], + cognito_user_pool_domain: Optional[cognito.IUserPoolDomain], +) -> Tuple[elb.ApplicationTargetGroup, int]: + """Register Pi on the shared legacy ALB (path and/or host-header listener rules).""" + pi_security_group.add_ingress_rule( + peer=alb_security_group, + connection=ec2.Port.tcp(pi_port), + description="Shared ALB to Pi agent", + ) + + pi_target_group = elb.ApplicationTargetGroup( + scope, + f"{logical_id_prefix}TargetGroup", + target_group_name=target_group_name, + port=pi_port, + protocol=elb.ApplicationProtocol.HTTP, + targets=[pi_service], + stickiness_cookie_duration=stickiness_cookie_duration, + vpc=vpc, + health_check=elb.HealthCheck( + path=pi_alb_health_check_path(path_prefix, routing_mode), + healthy_http_codes="200-399", + ), + ) + + if ( + enable_cognito_auth + and acm_certificate_arn + and cognito_user_pool + and cognito_user_pool_client + and cognito_user_pool_domain + and https_listener + ): + forward_action = elb_act.AuthenticateCognitoAction( + next=elb.ListenerAction.forward( + [pi_target_group], + stickiness_duration=stickiness_cookie_duration, + ), + user_pool=cognito_user_pool, + user_pool_client=cognito_user_pool_client, + user_pool_domain=cognito_user_pool_domain, + scope="openid profile email", + on_unauthenticated_request=elb.UnauthenticatedAction.AUTHENTICATE, + session_timeout=stickiness_cookie_duration, + ) + else: + forward_action = elb.ListenerAction.forward( + [pi_target_group], + stickiness_duration=stickiness_cookie_duration, + ) + + mode = normalize_pi_alb_routing_mode(routing_mode) + priority = listener_rule_priority + + def _add_rules(listener: elb.IApplicationListener, id_prefix: str) -> None: + nonlocal priority + if mode in ("path", "both"): + listener.add_action( + f"{id_prefix}PathRule", + priority=priority, + conditions=[ + elb.ListenerCondition.path_patterns( + pi_alb_path_patterns(path_prefix) + ) + ], + action=forward_action, + ) + priority += 1 + if mode in ("host", "both") and pi_host_header.strip(): + listener.add_action( + f"{id_prefix}HostRule", + priority=priority, + conditions=[ + elb.ListenerCondition.host_headers([pi_host_header.strip()]) + ], + action=forward_action, + ) + priority += 1 + + if https_listener: + _add_rules(https_listener, f"{logical_id_prefix}Https") + elif http_listener: + _add_rules(http_listener, f"{logical_id_prefix}Http") + + if ( + http_listener + and acm_certificate_arn + and pi_host_header.strip() + and mode in ("host", "both") + ): + redirect_priority = listener_rule_priority + if mode in ("path", "both"): + redirect_priority += 1 + http_listener.add_action( + f"{logical_id_prefix}HttpRedirectRule", + priority=redirect_priority, + conditions=[elb.ListenerCondition.host_headers([pi_host_header.strip()])], + action=elb.ListenerAction.redirect( + protocol="HTTPS", + port="443", + host="#{host}", + path="/#{path}", + query="#{query}", + ), + ) + + return pi_target_group, priority + + +def ensure_folder_exists(output_folder: str): + """Checks if the specified folder exists, creates it if not.""" + + if not os.path.exists(output_folder): + # Create the folder if it doesn't exist + os.makedirs(output_folder, exist_ok=True) + print(f"Created the {output_folder} folder.") + else: + print(f"The {output_folder} folder already exists.") + + +# Re-export for app.py and other CDK entrypoints (implementation is boto3-only). diff --git a/cdk/cdk_install.py b/cdk/cdk_install.py new file mode 100644 index 0000000000000000000000000000000000000000..5e78ea7215540482937463a52f91cd53c439ac21 --- /dev/null +++ b/cdk/cdk_install.py @@ -0,0 +1,3975 @@ +#!/usr/bin/env python3 +""" +Interactive CDK installer for llm_topic_modeller. + +Walks through demo vs production deployment profiles, writes config/cdk_config.env +and cdk.json, optionally runs cdk deploy, post-deploy Cognito callback fixups (API, +no second deploy), and post_cdk_build_quickstart.py. + +Usage examples:: + + # Full interactive install + deploy + quickstart + python cdk_install.py + + # Demo sandbox, config only + python cdk_install.py --profile demo --config-only --yes + + # Production with flags (partial non-interactive) + python cdk_install.py --profile production --vpc-name my-vpc \\ + --cert-arn arn:aws:acm:eu-west-2:123:certificate/abc \\ + --domain llm-topic.example.com --yes + + # Redeploy using existing config, skip quickstart + python cdk_install.py --deploy-only --skip-quickstart + + # Remove existing stacks before a clean install (non-interactive) + python cdk_install.py --profile headless --vpc-name my-vpc \\ + --force-delete-stacks --yes + + # Headless batch (S3 → Lambda → one-shot ECS direct mode) + python cdk_install.py --profile headless --vpc-name my-vpc --yes + python cdk_install.py --profile production --headless --vpc-name my-vpc --yes + + # Optional scheduled DynamoDB usage log export to S3 + python cdk_install.py --profile demo --vpc-name my-vpc --yes \\ + --dynamo-usage-log-export --dynamo-export-schedule-time 06:00 \\ + --dynamo-export-schedule-days weekdays + + # Headless with S3 output email alerts + IAM reader user + python cdk_install.py --profile headless --vpc-name my-vpc --yes \\ + --headless-output-notifications \\ + --headless-notify-email analyst@example.com +""" + +from __future__ import annotations + +import argparse +import ipaddress +import json +import os +import re +import secrets +import shutil +import subprocess +import sys +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence, Tuple + +CDK_DIR = Path(__file__).resolve().parent +REPO_ROOT = CDK_DIR.parent +CONFIG_DIR = CDK_DIR / "config" +ENV_PATH = CONFIG_DIR / "cdk_config.env" +APP_CONFIG_ENV_PATH = CONFIG_DIR / "app_config.env" +APP_CONFIG_ENV_EXAMPLE = CONFIG_DIR / "app_config.env.example" +PI_AGENT_ENV_PATH = CONFIG_DIR / "pi_agent.env" +PI_AGENT_ENV_EXAMPLE = REPO_ROOT / "config" / "pi_agent.env.example" +PI_ALB_ROUTING_MODES = ("path", "host", "both") +CDK_JSON_PATH = CDK_DIR / "cdk.json" +CDK_JSON_EXAMPLE = CDK_DIR / "cdk.json.example" +QUICKSTART_SCRIPT = CDK_DIR / "post_cdk_build_quickstart.py" + +REGIONAL_STACK = "SummarisationStack" +CLOUDFRONT_STACK = "SummarisationStackCloudfront" +CLOUDFRONT_STACK_REGION = "us-east-1" +APPREGISTRY_STACK_SUFFIX = "AppRegistryStack" + +DEMO_PRESET: Dict[str, str] = { + "USE_ECS_EXPRESS_MODE": "True", + "ECS_EXPRESS_USE_PUBLIC_SUBNETS": "True", + "ENABLE_ECS_VPC_INTERFACE_ENDPOINTS": "True", + "USE_CLOUDFRONT": "False", + "RUN_USEAST_STACK": "False", + "ENABLE_RESOURCE_DELETE_PROTECTION": "False", + "ENABLE_APPREGISTRY": "True", + "ACM_SSL_CERTIFICATE_ARN": "", + "SSL_CERTIFICATE_DOMAIN": "", +} + +PRODUCTION_PRESET: Dict[str, str] = { + "USE_ECS_EXPRESS_MODE": "False", + "USE_CLOUDFRONT": "True", + "RUN_USEAST_STACK": "True", + "ENABLE_RESOURCE_DELETE_PROTECTION": "True", + "ENABLE_APPREGISTRY": "True", +} + +HEADLESS_PRESET: Dict[str, str] = { + "USE_ECS_EXPRESS_MODE": "False", + "ECS_EXPRESS_USE_PUBLIC_SUBNETS": "True", + "USE_CLOUDFRONT": "False", + "RUN_USEAST_STACK": "False", + "ENABLE_RESOURCE_DELETE_PROTECTION": "False", + "ENABLE_APPREGISTRY": "True", + "ENABLE_HEADLESS_DEPLOYMENT": "True", + "ENABLE_S3_BATCH_ECS_TRIGGER": "True", + "COGNITO_AUTH": "False", + "ACM_SSL_CERTIFICATE_ARN": "", + "SSL_CERTIFICATE_DOMAIN": "", +} + +DEFAULT_CDK_JSON_CONTEXT: Dict[str, Any] = { + "@aws-cdk/aws-ec2:restrictDefaultSecurityGroup": False, +} + + +# --------------------------------------------------------------------------- +# Wizard helpers +# --------------------------------------------------------------------------- + + +def ask(prompt: str, default: str = "") -> str: + suffix = f" [{default}]" if default else "" + value = input(f"{prompt}{suffix}: ").strip() + return value or default + + +def ask_yes_no(prompt: str, default: bool = True) -> bool: + default_label = "Y/n" if default else "y/N" + while True: + raw = input(f"{prompt} ({default_label}): ").strip().lower() + if not raw: + return default + if raw in ("y", "yes"): + return True + if raw in ("n", "no"): + return False + print("Please enter y or n.") + + +def ask_choice(prompt: str, options: Sequence[str], default_index: int = 0) -> int: + print(prompt) + for i, opt in enumerate(options): + marker = "*" if i == default_index else " " + print(f" {marker} {i + 1}) {opt}") + while True: + raw = input( + f"Choice [1-{len(options)}] (default {default_index + 1}): " + ).strip() + if not raw: + return default_index + try: + idx = int(raw) - 1 + if 0 <= idx < len(options): + return idx + except ValueError: + pass + print("Invalid choice.") + + +def timestamp() -> str: + return datetime.now().strftime("%Y%m%d-%H%M%S") + + +def backup_file(path: Path) -> Optional[Path]: + if not path.is_file(): + return None + backup = path.with_suffix(path.suffix + f".bak.{timestamp()}") + shutil.copy2(path, backup) + print(f"Backed up {path.name} -> {backup.name}") + return backup + + +# --------------------------------------------------------------------------- +# Python / cdk.json +# --------------------------------------------------------------------------- + + +def _venv_python_paths() -> List[Path]: + candidates: List[Path] = [] + virtual_env = os.environ.get("VIRTUAL_ENV", "").strip() + if virtual_env: + root = Path(virtual_env) + if os.name == "nt": + candidates.append(root / "Scripts" / "python.exe") + else: + candidates.append(root / "bin" / "python") + + rel_venvs = [ + CDK_DIR / ".venv", + REPO_ROOT / ".venv", + REPO_ROOT.parent / ".venv", + ] + for venv_root in rel_venvs: + if os.name == "nt": + candidates.append(venv_root / "Scripts" / "python.exe") + else: + candidates.append(venv_root / "bin" / "python") + + for name in ("python3", "python"): + found = shutil.which(name) + if found: + candidates.append(Path(found)) + + seen: set[str] = set() + unique: List[Path] = [] + for p in [Path(sys.executable)] + candidates: + key = str(p.resolve()) if p.exists() else str(p) + if key not in seen: + seen.add(key) + unique.append(p) + return unique + + +def _resolve_node_executable() -> Optional[str]: + return shutil.which("node") + + +def _jsii_import_failure_hint(stderr: str) -> str: + if "JSONDecodeError" not in stderr and "Expecting value" not in stderr: + return "" + return ( + "\nThis usually means the JSII Node.js helper process failed to start.\n" + "On the deployment machine, check:\n" + " 1. Node.js is installed and on PATH: node --version (LTS 18–22 recommended)\n" + " 2. Reinstall CDK Python deps: pip install --force-reinstall -r requirements.txt\n" + " 3. Test manually: python -c \"from aws_cdk import App; print('ok')\"\n" + " 4. If still failing, run: set JSII_DEBUG=1 (then retry synth for details)" + ) + + +def verify_node_for_jsii() -> str: + """Return the node executable path, or exit with guidance if missing.""" + node = _resolve_node_executable() + if not node: + raise SystemExit( + "Node.js not found on PATH. aws-cdk-lib (JSII) requires Node.js.\n" + "Install Node.js LTS from https://nodejs.org/ and ensure 'node' is on PATH." + ) + try: + result = subprocess.run( + [node, "--version"], + capture_output=True, + text=True, + timeout=10, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + raise SystemExit(f"Could not run Node.js ({node}): {exc}") from exc + if result.returncode != 0: + raise SystemExit( + f"Node.js ({node}) returned exit code {result.returncode}.\n" + f"{(result.stderr or result.stdout or '').strip()}" + ) + version = (result.stdout or result.stderr or "").strip() + print(f"Node.js for JSII: {node} ({version})") + return node + + +def _python_has_aws_cdk(python_exe: Path) -> Tuple[bool, str]: + if not python_exe.is_file(): + return False, "not found" + cmd = [ + str(python_exe), + "-c", + "import aws_cdk; print(getattr(aws_cdk, '__version__', 'unknown'))", + ] + try: + result = subprocess.run( + cmd, + cwd=str(CDK_DIR), + capture_output=True, + text=True, + timeout=30, + check=False, + ) + if result.returncode == 0: + version = (result.stdout or "").strip().splitlines()[-1] + return True, version + detail = (result.stderr or result.stdout or "import failed").strip() + hint = _jsii_import_failure_hint(detail) + if hint: + detail = f"{detail[:400]}{hint}" + return False, detail[:800] + except (OSError, subprocess.TimeoutExpired) as exc: + return False, str(exc) + + +def discover_python_candidates() -> List[Tuple[Path, str]]: + found: List[Tuple[Path, str]] = [] + for candidate in _venv_python_paths(): + ok, detail = _python_has_aws_cdk(candidate) + if ok: + found.append((candidate.resolve(), detail)) + return found + + +def resolve_python_executable( + override: Optional[str] = None, + interactive: bool = True, + assume_yes: bool = False, +) -> Path: + if override: + path = Path(override).resolve() + ok, detail = _python_has_aws_cdk(path) + if not ok: + raise SystemExit( + f"Python at {path} cannot import aws-cdk-lib: {detail}\n" + "Install CDK deps: pip install -r requirements.txt" + ) + print(f"Using Python: {path} (aws-cdk-lib {detail})") + return path + + candidates = discover_python_candidates() + if not candidates: + raise SystemExit( + "No Python interpreter with aws-cdk-lib found.\n" + "Activate your venv and run: pip install -r requirements.txt\n" + "Or pass --python /path/to/python" + ) + + if len(candidates) == 1 or assume_yes or not interactive: + path, version = candidates[0] + print(f"Using Python: {path} (aws-cdk-lib {version})") + return path + + print("Python interpreters with aws-cdk-lib:") + labels = [f"{p} (aws-cdk-lib {v})" for p, v in candidates] + idx = ask_choice("Select Python for cdk.json", labels, default_index=0) + path, version = candidates[idx] + print(f"Selected: {path} (aws-cdk-lib {version})") + return path + + +def format_cdk_app_command(python_exe: Path) -> str: + """Format the cdk.json ``app`` command for the current OS.""" + py = str(python_exe) + if os.name == "nt": + py = py.replace("/", "\\") + return f"{py} app.py" + + +def load_cdk_json_context(existing_path: Path) -> Dict[str, Any]: + if existing_path.is_file(): + try: + data = json.loads(existing_path.read_text(encoding="utf-8")) + ctx = data.get("context") + if isinstance(ctx, dict): + return dict(ctx) + except json.JSONDecodeError: + pass + if CDK_JSON_EXAMPLE.is_file(): + try: + data = json.loads(CDK_JSON_EXAMPLE.read_text(encoding="utf-8")) + ctx = data.get("context") + if isinstance(ctx, dict): + return dict(ctx) + except json.JSONDecodeError: + pass + return dict(DEFAULT_CDK_JSON_CONTEXT) + + +def write_cdk_json( + python_exe: Path, + *, + force: bool = False, + skip: bool = False, +) -> Optional[Path]: + if skip: + if not CDK_JSON_PATH.is_file(): + raise SystemExit( + f"{CDK_JSON_PATH.name} not found; cannot use --skip-cdk-json." + ) + print(f"Keeping existing {CDK_JSON_PATH.name} (--skip-cdk-json)") + return CDK_JSON_PATH + + if CDK_JSON_PATH.is_file() and force: + backup_file(CDK_JSON_PATH) + elif not CDK_JSON_PATH.is_file(): + pass + else: + # Update app only — backup first + backup_file(CDK_JSON_PATH) + + context = load_cdk_json_context(CDK_JSON_PATH) + payload = { + "app": format_cdk_app_command(python_exe), + "output": "cdk.out", + "context": context, + } + CDK_JSON_PATH.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + print(f"Wrote {CDK_JSON_PATH}") + return CDK_JSON_PATH + + +def build_cdk_subprocess_env( + values: Optional[Dict[str, str]] = None, +) -> Dict[str, str]: + """ + Environment for child Python / CDK processes. + + Merges config/cdk_config.env over the current process environment so stale + empty defaults from an earlier cdk_config import during the wizard do not + block load_dotenv (override=False) in app.py. + """ + env = os.environ.copy() + env["CDK_CONFIG_PATH"] = str(ENV_PATH) + cdk_folder = "" + if values: + cdk_folder = (values.get("CDK_FOLDER") or "").strip() + if not cdk_folder and ENV_PATH.is_file(): + cdk_folder = (read_env_file(ENV_PATH).get("CDK_FOLDER") or "").strip() + if not cdk_folder: + cdk_folder = str(CDK_DIR).replace("\\", "/") + "/" + env["CDK_FOLDER"] = cdk_folder + + file_values: Dict[str, str] = {} + if ENV_PATH.is_file(): + file_values = read_env_file(ENV_PATH) + elif values: + file_values = dict(values) + + for key, val in file_values.items(): + env[key] = str(val) + return env + + +def apply_cdk_runtime_env(values: Dict[str, str]) -> None: + """Expose written config to app.py imports (smoke test, cdk synth/deploy).""" + os.environ.update(build_cdk_subprocess_env(values)) + + +def smoke_test_python_app(python_exe: Path) -> None: + verify_node_for_jsii() + cmd = [str(python_exe), "-c", "import aws_cdk; import app"] + print("Smoke test: import aws_cdk and app ...") + result = subprocess.run( + cmd, + cwd=str(CDK_DIR), + env=build_cdk_subprocess_env(), + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + output = (result.stderr or result.stdout or "").strip() + print(output) + hint = _jsii_import_failure_hint(output) + if hint: + print(hint) + raise SystemExit( + "Python smoke test failed. Fix CDK dependencies on this machine before deploying." + ) + + +def run_smoke_test_if_needed( + python_exe: Optional[Path], + args: argparse.Namespace, +) -> None: + if args.config_only or args.skip_cdk_json or not python_exe: + return + smoke_test_python_app(python_exe) + + +# --------------------------------------------------------------------------- +# AWS discovery +# --------------------------------------------------------------------------- + + +def get_aws_identity(region: Optional[str] = None) -> Tuple[str, str]: + import boto3 + + session = boto3.Session(region_name=region) + sts = session.client("sts") + ident = sts.get_caller_identity() + account = ident["Account"] + resolved_region = ( + region + or session.region_name + or os.environ.get("AWS_REGION") + or os.environ.get("AWS_DEFAULT_REGION") + or "eu-west-2" + ) + return account, resolved_region + + +def resolve_cdk_executable() -> str: + """Return the AWS CDK CLI on PATH (cdk.cmd on Windows).""" + for name in ("cdk", "cdk.cmd", "cdk.exe"): + found = shutil.which(name) + if found: + return found + raise SystemExit( + "AWS CDK CLI not found on PATH.\n" + "Install: npm install -g aws-cdk (or use npx aws-cdk)" + ) + + +def check_cdk_cli() -> str: + cdk_exe = resolve_cdk_executable() + result = subprocess.run( + [cdk_exe, "--version"], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + raise SystemExit( + "AWS CDK CLI not found on PATH.\n" + "Install: npm install -g aws-cdk (or use npx aws-cdk)" + ) + version = (result.stdout or result.stderr or "").strip() + if version.startswith("Python "): + raise SystemExit( + "CDK CLI check failed: 'cdk' resolved to the Python interpreter instead " + "of the AWS CDK CLI. Install aws-cdk globally (npm install -g aws-cdk) " + "and ensure it is on PATH before python.exe." + ) + return version + + +def cdk_bootstrap_needed(account: str, region: str) -> bool: + import boto3 + from botocore.exceptions import ClientError + + cfn = boto3.client("cloudformation", region_name=region) + try: + cfn.describe_stacks(StackName="CDKToolkit") + return False + except ClientError as exc: + if "does not exist" in str(exc): + return True + raise + + +def list_vpcs(region: str) -> List[Dict[str, str]]: + import boto3 + + ec2 = boto3.client("ec2", region_name=region) + response = ec2.describe_vpcs() + vpcs: List[Dict[str, str]] = [] + for vpc in response.get("Vpcs", []): + name = "" + for tag in vpc.get("Tags", []): + if tag.get("Key") == "Name": + name = tag.get("Value", "") + break + cidrs = vpc_cidr_blocks_from_describe(vpc) + vpcs.append( + { + "id": vpc["VpcId"], + "name": name or vpc["VpcId"], + "cidr": cidrs[0] if cidrs else "", + "cidrs": ",".join(cidrs), + } + ) + return sorted(vpcs, key=lambda x: x["name"]) + + +def vpc_cidr_blocks_from_describe(vpc: Dict[str, Any]) -> List[str]: + """Primary and associated IPv4 CIDR blocks for one ``describe_vpcs`` entry.""" + blocks: List[str] = [] + seen: set[str] = set() + primary = (vpc.get("CidrBlock") or "").strip() + if primary and primary not in seen: + seen.add(primary) + blocks.append(primary) + for assoc in vpc.get("CidrBlockAssociationSet") or []: + cidr = (assoc.get("CidrBlock") or "").strip() + if cidr and cidr not in seen: + seen.add(cidr) + blocks.append(cidr) + return blocks + + +def list_vpc_cidr_blocks_in_region(region: str) -> List[str]: + """All IPv4 VPC CIDR blocks in the account/region (primary + associations).""" + import boto3 + + ec2 = boto3.client("ec2", region_name=region) + blocks: List[str] = [] + seen: set[str] = set() + for vpc in ec2.describe_vpcs().get("Vpcs", []): + for cidr in vpc_cidr_blocks_from_describe(vpc): + if cidr not in seen: + seen.add(cidr) + blocks.append(cidr) + return blocks + + +VPC_CIDR_PREFIX_LEN = 24 +SUBNET_CIDR_PREFIX_LEN = 28 +# ECS Express provisions a managed ALB in each public subnet (8+ free IPs required). +# /28 subnets (~11 usable IPs) exhaust quickly with VPC interface endpoints and tasks. +EXPRESS_PUBLIC_SUBNET_CIDR_PREFIX_LEN = 27 +PUBLIC_SUBNET_CIDR_PREFIX_LEN = 26 +PRIVATE_SUBNET_CIDR_PREFIX_LEN = 28 +EXPRESS_ALB_MIN_SUBNET_PREFIX_LEN = 27 +_VPC_CIDR_SEARCH_SUPERNETS = ("10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16") + + +def parse_ipv4_networks(cidrs: Sequence[str]) -> List[ipaddress.IPv4Network]: + networks: List[ipaddress.IPv4Network] = [] + for cidr in cidrs: + raw = (cidr or "").strip() + if not raw: + continue + try: + networks.append(ipaddress.ip_network(raw, strict=False)) + except ValueError: + continue + return networks + + +def suggest_vpc_cidr_block( + existing_vpc_cidrs: Sequence[str], + *, + prefix_len: int = VPC_CIDR_PREFIX_LEN, + search_supernets: Sequence[str] = _VPC_CIDR_SEARCH_SUPERNETS, +) -> str: + """Return the lowest non-overlapping ``prefix_len`` block in RFC1918 space.""" + occupied = parse_ipv4_networks(existing_vpc_cidrs) + for supernet_cidr in search_supernets: + supernet = ipaddress.ip_network(supernet_cidr, strict=False) + for candidate in supernet.subnets(new_prefix=prefix_len): + if any(candidate.overlaps(block) for block in occupied): + continue + return str(candidate) + raise ValueError( + f"No unused /{prefix_len} VPC CIDR found in region " + f"({len(occupied)} existing block(s))." + ) + + +def validate_new_vpc_cidr_format(cidr: str) -> Optional[str]: + """Return an error when ``cidr`` is not a valid AWS VPC IPv4 block.""" + raw = (cidr or "").strip() + if not raw: + return "NEW_VPC_CIDR is required when creating a new VPC." + try: + network = ipaddress.ip_network(raw, strict=False) + except ValueError: + return f"NEW_VPC_CIDR={raw!r} is not a valid IPv4 CIDR block." + if network.version != 4: + return f"NEW_VPC_CIDR must be IPv4 (got {raw!r})." + if network.prefixlen < 16 or network.prefixlen > 28: + return ( + f"NEW_VPC_CIDR={raw!r} must use a prefix length between /16 and /28 " + f"(got /{network.prefixlen})." + ) + if not network.is_private: + return ( + f"NEW_VPC_CIDR={raw!r} should be in RFC1918 private space " + "(10.0.0.0/8, 172.16.0.0/12, or 192.168.0.0/16)." + ) + return None + + +def validate_new_vpc_cidr_no_overlap( + cidr: str, + existing_vpc_cidrs: Sequence[str], +) -> Optional[str]: + """Return an error when ``cidr`` overlaps an existing VPC block in the region.""" + format_error = validate_new_vpc_cidr_format(cidr) + if format_error: + return format_error + candidate = ipaddress.ip_network(cidr.strip(), strict=False) + for block in parse_ipv4_networks(existing_vpc_cidrs): + if candidate.overlaps(block): + return ( + f"NEW_VPC_CIDR={candidate!s} overlaps existing VPC CIDR {block!s} " + "in this account/region." + ) + return None + + +def validate_new_vpc_cidr( + cidr: str, + existing_vpc_cidrs: Sequence[str], +) -> Optional[str]: + """Validate format and regional non-overlap for a new VPC CIDR.""" + return validate_new_vpc_cidr_no_overlap(cidr, existing_vpc_cidrs) + + +def canonicalize_vpc_cidr(cidr: str) -> str: + """Normalize a VPC CIDR string (e.g. 10.0.0.0/24).""" + return str(ipaddress.ip_network(cidr.strip(), strict=False)) + + +def subnet_cidr_prefix_len_for_tier(answers: "InstallAnswers", tier: str) -> int: + """Return installer subnet sizing for public/private tiers and deployment mode.""" + if tier == "public": + if answers_use_public_subnets_only(answers): + return EXPRESS_PUBLIC_SUBNET_CIDR_PREFIX_LEN + return PUBLIC_SUBNET_CIDR_PREFIX_LEN + return PRIVATE_SUBNET_CIDR_PREFIX_LEN + + +def validate_public_subnet_cidr_for_express(cidr: str) -> Optional[str]: + """Return an error when a public subnet is too small for Express managed ALB.""" + raw = (cidr or "").strip() + if not raw: + return None + try: + network = ipaddress.ip_network(raw, strict=False) + except ValueError: + return f"Public subnet CIDR {raw!r} is not valid." + if network.prefixlen > EXPRESS_ALB_MIN_SUBNET_PREFIX_LEN: + return ( + f"Public subnet {network!s} is too small for ECS Express Mode: the managed " + f"ALB needs at least 8 free IP addresses per subnet. Use /" + f"{EXPRESS_ALB_MIN_SUBNET_PREFIX_LEN} or larger (e.g. /27 in a /24 VPC)." + ) + return None + + +def validate_new_vpc_cidr_env_values(values: Dict[str, str]) -> List[str]: + """Live regional overlap check for NEW_VPC_CIDR before deploy.""" + new_cidr = (values.get("NEW_VPC_CIDR") or "").strip() + if not new_cidr or (values.get("VPC_NAME") or "").strip(): + return [] + region = (values.get("AWS_REGION") or "").strip() + if not region: + return ["AWS_REGION is required to validate NEW_VPC_CIDR."] + occupied = list_vpc_cidr_blocks_in_region(region) + err = validate_new_vpc_cidr(new_cidr, occupied) + return [err] if err else [] + + +def prompt_new_vpc_cidr(answers: "InstallAnswers", *, interactive: bool) -> None: + """Fill ``answers.new_vpc_cidr`` via auto-suggest and/or manual wizard prompt.""" + region = answers.aws_region + occupied = list_vpc_cidr_blocks_in_region(region) + + if answers.new_vpc_cidr: + err = validate_new_vpc_cidr(answers.new_vpc_cidr, occupied) + if err: + raise SystemExit(err) + answers.new_vpc_cidr = canonicalize_vpc_cidr(answers.new_vpc_cidr) + return + + auto_prompt = ( + f"Auto-select lowest available /{VPC_CIDR_PREFIX_LEN} VPC CIDR " f"in {region}?" + ) + + def _select_suggested() -> str: + suggested = suggest_vpc_cidr_block(occupied) + print(f"Selected VPC CIDR: {suggested}") + return suggested + + def _prompt_manual(default_cidr: str) -> str: + while True: + raw = ask("New VPC CIDR", default_cidr) + err = validate_new_vpc_cidr(raw, occupied) + if err: + print(err) + continue + return canonicalize_vpc_cidr(raw) + + if interactive: + try: + default_cidr = suggest_vpc_cidr_block(occupied) + can_auto_select = True + except ValueError as exc: + print(f"Warning: {exc}") + default_cidr = "" + can_auto_select = False + + if can_auto_select and ask_yes_no(auto_prompt, default=True): + answers.new_vpc_cidr = _select_suggested() + else: + if not default_cidr: + print( + f"Enter a non-overlapping RFC1918 VPC CIDR (/16–/28) for {region}." + ) + answers.new_vpc_cidr = _prompt_manual(default_cidr) + return + + try: + answers.new_vpc_cidr = _select_suggested() + except ValueError as exc: + raise SystemExit( + f"Could not auto-select a VPC CIDR in {region}. " + f"Pass --new-vpc-cidr with a non-overlapping RFC1918 /16–/28 block. {exc}" + ) from exc + + +def suggest_subnet_cidr_blocks( + vpc_cidr: str, + existing_subnet_cidrs: Sequence[str], + count: int, + *, + prefix_len: int = SUBNET_CIDR_PREFIX_LEN, + reserved_cidrs: Optional[Sequence[str]] = None, +) -> List[str]: + """Return ``count`` lowest non-overlapping ``prefix_len`` blocks inside ``vpc_cidr``.""" + if count < 1: + return [] + vpc_net = ipaddress.ip_network(vpc_cidr.strip(), strict=False) + occupied = parse_ipv4_networks(existing_subnet_cidrs) + if reserved_cidrs: + occupied.extend(parse_ipv4_networks(reserved_cidrs)) + + suggestions: List[str] = [] + for candidate in vpc_net.subnets(new_prefix=prefix_len): + if any(candidate.overlaps(block) for block in occupied): + continue + suggestions.append(str(candidate)) + occupied.append(candidate) + if len(suggestions) >= count: + return suggestions + + raise ValueError( + f"Only {len(suggestions)} available /{prefix_len} subnet block(s) in " + f"{vpc_cidr}; need {count}." + ) + + +def list_availability_zones(region: str) -> List[str]: + import boto3 + + ec2 = boto3.client("ec2", region_name=region) + zones = ec2.describe_availability_zones( + Filters=[{"Name": "state", "Values": ["available"]}] + ) + return [z["ZoneName"] for z in zones.get("AvailabilityZones", [])] + + +def list_subnets_in_vpc(vpc_id: str, region: str) -> List[Dict[str, str]]: + sys.path.insert(0, str(CDK_DIR)) + from cdk_functions import _get_existing_subnets_in_vpc + + data = _get_existing_subnets_in_vpc(vpc_id) + subnets: List[Dict[str, str]] = [] + for name, info in data.get("by_name", {}).items(): + subnets.append( + { + "name": name, + "id": info.get("id", ""), + "cidr": info.get("cidr", ""), + "az": info.get("az", ""), + } + ) + return sorted(subnets, key=lambda x: x["name"]) + + +def list_acm_certificates(region: str) -> List[Dict[str, str]]: + import boto3 + + acm = boto3.client("acm", region_name=region) + certs: List[Dict[str, str]] = [] + paginator = acm.get_paginator("list_certificates") + for page in paginator.paginate(CertificateStatuses=["ISSUED"]): + for summary in page.get("CertificateSummaryList", []): + arn = summary["CertificateArn"] + detail = acm.describe_certificate(CertificateArn=arn)["Certificate"] + domain = detail.get("DomainName", "") + sans = detail.get("SubjectAlternativeNames", []) + label = domain + if sans and sans[0] != domain: + label = f"{domain} (+{len(sans) - 1} SANs)" + certs.append({"arn": arn, "domain": domain, "label": label}) + return certs + + +def list_igws_for_vpc(vpc_id: str, region: str) -> List[Dict[str, str]]: + import boto3 + + ec2 = boto3.client("ec2", region_name=region) + response = ec2.describe_internet_gateways( + Filters=[{"Name": "attachment.vpc-id", "Values": [vpc_id]}] + ) + return [ + {"id": igw["InternetGatewayId"]} for igw in response.get("InternetGateways", []) + ] + + +def list_albs_in_vpc(vpc_id: str, region: str) -> List[Dict[str, str]]: + import boto3 + + elbv2 = boto3.client("elbv2", region_name=region) + response = elbv2.describe_load_balancers() + albs: List[Dict[str, str]] = [] + for lb in response.get("LoadBalancers", []): + if lb.get("VpcId") == vpc_id: + albs.append( + { + "arn": lb["LoadBalancerArn"], + "dns": lb["DNSName"], + "name": lb["LoadBalancerName"], + } + ) + return albs + + +# --------------------------------------------------------------------------- +# Env builder / writer +# --------------------------------------------------------------------------- + + +def format_list_env(values: Sequence[str], *, use_single_quotes: bool = False) -> str: + if not values: + return "[]" + if use_single_quotes: + inner = ", ".join(f"'{v}'" for v in values) + else: + inner = ", ".join(f'"{v}"' for v in values) + return f"[{inner}]" + + +SUBNET_TIER_MODES = ("auto", "existing", "create") + + +def parse_subnet_name_csv(raw: str) -> List[str]: + return [part.strip() for part in raw.split(",") if part.strip()] + + +def resolve_subnet_tier_modes(args: argparse.Namespace) -> Tuple[str, str]: + """CLI: --subnet-mode sets both tiers; per-tier flags override individually.""" + default = getattr(args, "subnet_mode", None) or "auto" + public = getattr(args, "public_subnet_mode", None) or default + private = getattr(args, "private_subnet_mode", None) or default + return public, private + + +def apply_subnet_cli_flags(args: argparse.Namespace, answers: "InstallAnswers") -> None: + if getattr(args, "public_subnet_names", None): + answers.public_subnet_names = parse_subnet_name_csv(args.public_subnet_names) + if getattr(args, "private_subnet_names", None): + answers.private_subnet_names = parse_subnet_name_csv(args.private_subnet_names) + + +def lookup_subnets_by_name( + vpc_subnets: Sequence[Dict[str, str]], +) -> Dict[str, Dict[str, str]]: + return {subnet["name"]: subnet for subnet in vpc_subnets if subnet.get("name")} + + +def fill_existing_subnet_tier_metadata( + names: Sequence[str], + lookup: Dict[str, Dict[str, str]], + *, + cidrs_out: List[str], + azs_out: List[str], +) -> List[str]: + """Resolve CIDR/AZ for existing subnet names (order preserved).""" + errors: List[str] = [] + cidrs_out.clear() + azs_out.clear() + for name in names: + info = lookup.get(name) + if not info: + errors.append(f"Subnet '{name}' was not found in the VPC.") + continue + cidr = (info.get("cidr") or "").strip() + az = (info.get("az") or "").strip() + if not az: + errors.append(f"Subnet '{name}' has no availability zone in AWS.") + continue + cidrs_out.append(cidr) + azs_out.append(az) + return errors + + +def enrich_existing_subnet_details_from_aws(answers: "InstallAnswers") -> List[str]: + """Look up CIDR/AZ for tiers using existing named subnets (wizard or CLI).""" + if answers.vpc_mode != "existing" or not answers.vpc_name: + return [] + + tiers: List[Tuple[str, str, List[str], List[str], List[str]]] = [] + if answers.public_subnet_mode == "existing" and answers.public_subnet_names: + tiers.append( + ( + "Public", + answers.public_subnet_mode, + answers.public_subnet_names, + answers.public_subnet_cidrs, + answers.public_subnet_azs, + ) + ) + if answers.private_subnet_mode == "existing" and answers.private_subnet_names: + tiers.append( + ( + "Private", + answers.private_subnet_mode, + answers.private_subnet_names, + answers.private_subnet_cidrs, + answers.private_subnet_azs, + ) + ) + if not tiers: + return [] + + vpc_list = list_vpcs(answers.aws_region) + vpc_id = next( + (vpc["id"] for vpc in vpc_list if vpc["name"] == answers.vpc_name), + None, + ) + if not vpc_id: + return [f"VPC '{answers.vpc_name}' not found in {answers.aws_region}."] + + lookup = lookup_subnets_by_name(list_subnets_in_vpc(vpc_id, answers.aws_region)) + errors: List[str] = [] + for _label, _mode, names, cidrs_out, azs_out in tiers: + if len(cidrs_out) == len(names) == len(azs_out) and names: + continue + errors.extend( + fill_existing_subnet_tier_metadata( + names, + lookup, + cidrs_out=cidrs_out, + azs_out=azs_out, + ) + ) + return errors + + +def apply_subnet_tier_env( + values: Dict[str, str], + answers: "InstallAnswers", + *, + tier: str, + mode: str, +) -> None: + """Write PUBLIC_* or PRIVATE_* subnet keys for one tier (auto | existing | create).""" + prefix = "PUBLIC" if tier == "public" else "PRIVATE" + names = ( + answers.public_subnet_names + if tier == "public" + else answers.private_subnet_names + ) + cidrs = ( + answers.public_subnet_cidrs + if tier == "public" + else answers.private_subnet_cidrs + ) + azs = answers.public_subnet_azs if tier == "public" else answers.private_subnet_azs + + if mode == "auto": + values[f"{prefix}_SUBNETS_TO_USE"] = "" + values[f"{prefix}_SUBNET_CIDR_BLOCKS"] = "" + values[f"{prefix}_SUBNET_AVAILABILITY_ZONES"] = "" + elif mode == "existing": + values[f"{prefix}_SUBNETS_TO_USE"] = format_list_env(names) + if names and len(cidrs) == len(names) == len(azs): + values[f"{prefix}_SUBNET_CIDR_BLOCKS"] = format_list_env( + cidrs, use_single_quotes=True + ) + values[f"{prefix}_SUBNET_AVAILABILITY_ZONES"] = format_list_env( + azs, use_single_quotes=True + ) + else: + values[f"{prefix}_SUBNET_CIDR_BLOCKS"] = "" + values[f"{prefix}_SUBNET_AVAILABILITY_ZONES"] = "" + else: + values[f"{prefix}_SUBNETS_TO_USE"] = format_list_env(names) + values[f"{prefix}_SUBNET_CIDR_BLOCKS"] = format_list_env( + cidrs, use_single_quotes=True + ) + values[f"{prefix}_SUBNET_AVAILABILITY_ZONES"] = format_list_env( + azs, use_single_quotes=True + ) + + +def validate_subnet_answers(answers: "InstallAnswers") -> List[str]: + errors: List[str] = [] + if answers.vpc_mode != "existing": + return errors + tiers: List[tuple] = [ + ( + "Public", + answers.public_subnet_mode, + answers.public_subnet_names, + answers.public_subnet_cidrs, + ), + ] + if not answers_use_express_mode(answers) and not answers_use_public_subnets_only( + answers + ): + tiers.append( + ( + "Private", + answers.private_subnet_mode, + answers.private_subnet_names, + answers.private_subnet_cidrs, + ) + ) + for label, mode, names, cidrs in tiers: + if mode == "existing" and not names: + errors.append( + f"{label} subnets: provide at least one existing subnet name." + ) + if mode == "create" and not names: + errors.append(f"{label} subnets: create mode requires subnet names.") + if mode == "create" and names and cidrs and len(names) != len(cidrs): + errors.append(f"{label} subnets: CIDR count must match subnet name count.") + if label == "Public" and answers_use_express_mode(answers) and cidrs: + for cidr in cidrs: + subnet_error = validate_public_subnet_cidr_for_express(cidr) + if subnet_error: + errors.append(subnet_error) + return errors + + +def ask_subnet_tier_mode(tier_label: str) -> str: + idx = ask_choice( + f"{tier_label} subnets", + [ + "Auto-discover suitable subnets", + "Use existing named subnets", + "Create new stack-specific subnets (name + CIDR + AZ)", + ], + default_index=0, + ) + return SUBNET_TIER_MODES[idx] + + +def configure_subnet_tier( + answers: "InstallAnswers", + tier: str, + mode: str, + vpc_subnets: List[Dict[str, str]], + azs: Sequence[str], + *, + interactive: bool, + vpc_cidr: str = "", + reserved_subnet_cidrs: Optional[Sequence[str]] = None, +) -> None: + """Collect subnet names/CIDRs for one tier in the wizard.""" + is_public = tier == "public" + label = "Public" if is_public else "Private" + prefix_label = "Public" if is_public else "Private" + + if mode == "auto": + return + + if mode == "existing": + names = ( + answers.public_subnet_names if is_public else answers.private_subnet_names + ) + if interactive and not names: + print( + f"Enter {label.lower()} subnet names (comma-separated), or pick from:" + ) + for subnet in vpc_subnets: + print(f" - {subnet['name']} ({subnet['cidr']}, {subnet['az']})") + names = parse_subnet_name_csv(ask(f"{label} subnet names", "")) + if is_public: + answers.public_subnet_names = names + cidrs_out = answers.public_subnet_cidrs + azs_out = answers.public_subnet_azs + else: + answers.private_subnet_names = names + cidrs_out = answers.private_subnet_cidrs + azs_out = answers.private_subnet_azs + if names: + lookup = lookup_subnets_by_name(vpc_subnets) + missing = fill_existing_subnet_tier_metadata( + names, + lookup, + cidrs_out=cidrs_out, + azs_out=azs_out, + ) + for err in missing: + print(f"Warning: {err}") + return + + n_az = len(azs) + prefix_label = "Public" if tier == "public" else "Private" + prefix_len = subnet_cidr_prefix_len_for_tier(answers, tier) + names = [f"{answers.cdk_prefix}{prefix_label}Subnet{i + 1}" for i in range(n_az)] + if is_public: + answers.public_subnet_names = names + cidrs_list = answers.public_subnet_cidrs + azs_list = answers.public_subnet_azs + cidr_base = 0 + else: + answers.private_subnet_names = names + cidrs_list = answers.private_subnet_cidrs + azs_list = answers.private_subnet_azs + cidr_base = 10 + + if interactive: + print(f"Suggested AZs for {label.lower()} subnets: {', '.join(azs)}") + auto_assign = ask_yes_no( + f"Auto-assign lowest available /{prefix_len} CIDR blocks " + f"for {label.lower()} subnets?", + default=True, + ) + else: + auto_assign = bool(vpc_cidr) + + existing_subnet_cidrs = [ + subnet.get("cidr", "") for subnet in vpc_subnets if subnet.get("cidr") + ] + suggested: Optional[List[str]] = None + if auto_assign and vpc_cidr: + try: + suggested = suggest_subnet_cidr_blocks( + vpc_cidr, + existing_subnet_cidrs, + len(names), + prefix_len=prefix_len, + reserved_cidrs=reserved_subnet_cidrs, + ) + except ValueError as exc: + print(f"Warning: {exc}") + if not interactive: + raise SystemExit(str(exc)) from exc + + cidrs_list.clear() + azs_list.clear() + for i, name in enumerate(names): + if suggested: + cidr = suggested[i] + if interactive: + print(f" {name}: {cidr}") + elif interactive: + default_cidr = f"10.0.{cidr_base + i}.0/{prefix_len}" + if vpc_cidr: + try: + default_cidr = suggest_subnet_cidr_blocks( + vpc_cidr, + existing_subnet_cidrs + cidrs_list, + 1, + prefix_len=prefix_len, + reserved_cidrs=reserved_subnet_cidrs, + )[0] + except ValueError: + pass + cidr = ask(f"CIDR for {label.lower()} {name}", default_cidr) + else: + cidr = f"10.0.{cidr_base + i}.0/{prefix_len}" + cidrs_list.append(cidr) + azs_list.append(azs[i % len(azs)]) + + +@dataclass +class InstallAnswers: + profile: str = "demo" + aws_account_id: str = "" + aws_region: str = "" + cdk_prefix: str = "" + cognito_domain_prefix: str = "" + s3_log_bucket_name: str = "" + s3_output_bucket_name: str = "" + github_branch: str = "main" + vpc_mode: str = "existing" # new | existing + vpc_name: str = "" + new_vpc_cidr: str = "" + public_subnet_mode: str = "auto" # auto | existing | create + private_subnet_mode: str = "auto" # auto | existing | create + public_subnet_names: List[str] = field(default_factory=list) + private_subnet_names: List[str] = field(default_factory=list) + public_subnet_cidrs: List[str] = field(default_factory=list) + private_subnet_cidrs: List[str] = field(default_factory=list) + public_subnet_azs: List[str] = field(default_factory=list) + private_subnet_azs: List[str] = field(default_factory=list) + existing_igw_id: str = "" + existing_alb_arn: str = "" + existing_alb_dns: str = "" + acm_cert_arn: str = "" + ssl_domain: str = "" + cloudfront_geo: str = "" + enable_pi_express: bool = False + enable_pi_legacy: bool = False + enable_service_connect: bool = False + enable_s3_batch: bool = False + enable_headless: bool = False + enable_headless_output_notifications: bool = False + headless_output_notify_email: str = "" + headless_output_iam_user_name: str = "" + enable_dynamo_usage_log_export: bool = False + dynamo_export_schedule_time: str = "06:00" + dynamo_export_schedule_days: str = "daily" + dynamo_export_s3_key: str = "" + dynamo_export_date_attribute: str = "timestamp" + ecs_memory: str = "8192" + pi_alb_routing: str = "path" + pi_alb_path_prefix: str = "/agent" + pi_alb_host_header: str = "" + pi_alb_listener_rule_priority: str = "" + pi_gradio_port: str = "7862" + sc_discovery_name: str = "llm-topic" + pi_default_provider: str = "amazon-bedrock" + write_pi_agent_env: bool = True + overwrite_pi_agent_env: bool = False + write_app_config_env: bool = True + overwrite_app_config_env: bool = False + custom_overrides: Dict[str, str] = field(default_factory=dict) + python_path: Optional[str] = None + + @property + def pi_enabled(self) -> bool: + return self.enable_pi_express or self.enable_pi_legacy + + +# Cognito hosted-UI prefix domains cannot contain these substrings (AWS docs). +COGNITO_DOMAIN_RESERVED_SUBSTRINGS = ("amazon", "cognito", "aws") + + +def sanitize_cognito_domain_prefix(raw: str, *, max_length: int = 63) -> str: + """Normalize a Cognito domain prefix: allowed chars, no reserved words.""" + prefix = re.sub(r"[^a-z0-9-]", "-", (raw or "").lower()) + for reserved in COGNITO_DOMAIN_RESERVED_SUBSTRINGS: + prefix = prefix.replace(reserved, "") + prefix = re.sub(r"-{2,}", "-", prefix).strip("-") + if max_length > 0: + prefix = prefix[:max_length].strip("-") + if not prefix or not re.match(r"^[a-z0-9]", prefix): + return "llm-topic-app" + prefix = prefix.rstrip("-") + if not prefix: + return "llm-topic-app" + for reserved in COGNITO_DOMAIN_RESERVED_SUBSTRINGS: + if reserved in prefix: + return "llm-topic-app" + return prefix + + +def default_cognito_domain_prefix_from_cdk_prefix(cdk_prefix: str) -> str: + """Derive a short installer default from CDK_PREFIX (globally unique hint).""" + slug = re.sub(r"[^a-z0-9-]", "-", (cdk_prefix or "").lower()).strip("-") + return sanitize_cognito_domain_prefix(slug, max_length=20) + + +def validate_cognito_domain_prefix(prefix: str) -> Optional[str]: + """Return an error message when prefix violates Cognito domain rules.""" + cleaned = (prefix or "").strip().lower() + if not cleaned: + return "COGNITO_USER_POOL_DOMAIN_PREFIX is required." + reserved = [word for word in COGNITO_DOMAIN_RESERVED_SUBSTRINGS if word in cleaned] + if reserved: + return ( + "COGNITO_USER_POOL_DOMAIN_PREFIX cannot contain reserved words: " + + ", ".join(reserved) + + "." + ) + if not re.match(r"^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$", cleaned): + return ( + "COGNITO_USER_POOL_DOMAIN_PREFIX must use lowercase letters, " + "numbers, and hyphens only." + ) + return None + + +def normalize_pi_path_prefix(raw: str) -> str: + segment = (raw or "agent").strip().strip("/") + return f"/{segment}" if segment else "/agent" + + +def default_pi_listener_priority(use_cloudfront: bool = False) -> str: + """Default Pi path/host rule priority (1–2 reserved for CloudFront / Express rules).""" + del use_cloudfront # kept for call-site compatibility + return "3" + + +def derive_ecs_resource_names(cdk_prefix: str) -> Dict[str, str]: + """ECS/CodeBuild resource names from CDK_PREFIX (matches cdk_config.py defaults).""" + prefix = (cdk_prefix or "").strip() + if not prefix: + return {} + return { + "CLUSTER_NAME": f"{prefix}Cluster", + "ECS_SERVICE_NAME": f"{prefix}ECSService", + "ECS_EXPRESS_SERVICE_NAME": f"{prefix}ECSService", + "ECS_PI_EXPRESS_SERVICE_NAME": f"{prefix}PiExpressService", + "ECS_PI_SERVICE_NAME": f"{prefix}PiAgentService", + "COGNITO_USER_POOL_CLIENT_SECRET_NAME": f"{prefix}ParamCognitoSecret", + "CODEBUILD_PROJECT_NAME": f"{prefix}CodeBuildProject", + "CODEBUILD_PI_PROJECT_NAME": f"{prefix}CodeBuildPiProject", + } + + +def derive_s3_bucket_names(cdk_prefix: str) -> Dict[str, str]: + """S3 bucket names from CDK_PREFIX (must be written to cdk_config.env explicitly).""" + prefix = (cdk_prefix or "").strip().lower() + if not prefix: + return {} + return { + "S3_LOG_CONFIG_BUCKET_NAME": f"{prefix}s3-logs", + "S3_OUTPUT_BUCKET_NAME": f"{prefix}s3-output", + } + + +S3_BUCKET_NAME_MAX_LEN = 63 + + +def normalize_s3_bucket_name(raw: str) -> str: + """Lowercase S3 bucket name with allowed characters only (3–63 chars).""" + name = re.sub(r"[^a-z0-9.-]", "-", (raw or "").lower()) + name = re.sub(r"-{2,}", "-", name).strip("-") + if len(name) < 3: + name = f"{name}-bucket".strip("-") + return name[:S3_BUCKET_NAME_MAX_LEN] + + +def suggest_available_s3_bucket_name( + preferred: str, + account_id: str, + *, + s3_client: Any = None, + max_attempts: int = 8, +) -> str: + """Return a globally available S3 bucket name, preferring ``preferred``.""" + from cdk_functions import resolve_s3_bucket_availability + + account = re.sub(r"[^0-9]", "", account_id or "") + bases = [preferred] + if account: + bases.extend( + [ + f"{account}-{preferred}", + f"{preferred}-{account}", + ] + ) + seen: set[str] = set() + for base in bases: + candidate = normalize_s3_bucket_name(base) + if not candidate or candidate in seen: + continue + seen.add(candidate) + status, _ = resolve_s3_bucket_availability(candidate, s3_client=s3_client) + if status == "available": + return candidate + for _ in range(max_attempts): + suffix = secrets.token_hex(3) + candidate = normalize_s3_bucket_name(f"{account}-{preferred}-{suffix}") + if candidate in seen: + continue + seen.add(candidate) + status, _ = resolve_s3_bucket_availability(candidate, s3_client=s3_client) + if status == "available": + return candidate + raise RuntimeError( + f"Could not find an available S3 bucket name near {preferred!r}. " + "Set S3_LOG_CONFIG_BUCKET_NAME / S3_OUTPUT_BUCKET_NAME manually." + ) + + +def suggest_available_cognito_domain_prefix( + preferred: str, + account_id: str, + region: str, + *, + cognito_client: Any = None, + max_attempts: int = 8, +) -> str: + """Return an available Cognito hosted UI domain prefix in ``region``.""" + from cdk_functions import resolve_cognito_domain_prefix_availability + + account = re.sub(r"[^0-9]", "", account_id or "") + preferred_clean = sanitize_cognito_domain_prefix(preferred, max_length=63) + candidates = [preferred_clean] + if account: + candidates.extend( + [ + sanitize_cognito_domain_prefix( + f"{account}-{preferred_clean}", max_length=63 + ), + sanitize_cognito_domain_prefix( + f"{preferred_clean}-{account[-8:]}", max_length=63 + ), + ] + ) + seen: set[str] = set() + for candidate in candidates: + if not candidate or candidate in seen: + continue + seen.add(candidate) + if ( + resolve_cognito_domain_prefix_availability( + candidate, region_name=region, cognito_client=cognito_client + ) + == "available" + ): + return candidate + for _ in range(max_attempts): + suffix = secrets.token_hex(2) + candidate = sanitize_cognito_domain_prefix( + f"{account}-{preferred_clean}-{suffix}", max_length=63 + ) + if not candidate or candidate in seen: + continue + seen.add(candidate) + if ( + resolve_cognito_domain_prefix_availability( + candidate, region_name=region, cognito_client=cognito_client + ) + == "available" + ): + return candidate + raise RuntimeError( + f"Could not find an available Cognito domain prefix near {preferred!r}. " + "Set COGNITO_USER_POOL_DOMAIN_PREFIX manually." + ) + + +def _prompt_globally_unique_s3_bucket( + label: str, + preferred: str, + account_id: str, + *, + interactive: bool, + assume_yes: bool, + cli_override: str = "", + s3_client: Any = None, +) -> str: + from cdk_functions import resolve_s3_bucket_availability + + if cli_override: + name = normalize_s3_bucket_name(cli_override) + status, _ = resolve_s3_bucket_availability(name, s3_client=s3_client) + if status == "globally_taken": + raise SystemExit( + f"{label}: bucket name {name!r} is taken globally by another AWS " + "account. Choose a different --s3-log-bucket / --s3-output-bucket." + ) + return name + + preferred = normalize_s3_bucket_name(preferred) + status, _ = resolve_s3_bucket_availability(preferred, s3_client=s3_client) + if status == "owned": + print(f"{label}: using existing bucket in this account: {preferred}") + return preferred + if status == "available": + print(f"{label}: bucket name available: {preferred}") + return preferred + + suggested = suggest_available_s3_bucket_name( + preferred, account_id, s3_client=s3_client + ) + print( + f"{label}: S3 bucket name {preferred!r} is taken globally by another AWS " + f"account (S3 names are unique worldwide)." + ) + if not interactive: + if assume_yes: + print(f"{label}: using suggested name: {suggested}") + return suggested + raise SystemExit( + f"{label}: bucket name {preferred!r} is taken globally. " + f"Re-run with --yes to accept {suggested!r}, or pass an explicit bucket name." + ) + if ask_yes_no(f"Use suggested bucket name {suggested!r}?", default=True): + return suggested + while True: + custom = ask(f"{label}: enter a globally unique bucket name", suggested) + custom = normalize_s3_bucket_name(custom) + custom_status, _ = resolve_s3_bucket_availability(custom, s3_client=s3_client) + if custom_status == "available": + return custom + if custom_status == "owned": + print(f"Bucket {custom!r} already exists in this account; will import it.") + return custom + print(f"Name {custom!r} is still taken globally. Try another name.") + + +def _prompt_globally_unique_cognito_prefix( + preferred: str, + account_id: str, + region: str, + *, + interactive: bool, + assume_yes: bool, + cli_override: str = "", + cognito_client: Any = None, +) -> str: + from cdk_functions import resolve_cognito_domain_prefix_availability + + if cli_override: + override = sanitize_cognito_domain_prefix(cli_override) + validation_error = validate_cognito_domain_prefix(override) + if validation_error: + raise SystemExit(validation_error) + if ( + resolve_cognito_domain_prefix_availability( + override, region_name=region, cognito_client=cognito_client + ) + != "available" + ): + raise SystemExit( + f"Cognito domain prefix {override!r} is not available in {region} " + "(likely taken by another AWS account). Pass a different " + "--cognito-prefix." + ) + return override + + preferred = sanitize_cognito_domain_prefix(preferred) + availability = resolve_cognito_domain_prefix_availability( + preferred, region_name=region, cognito_client=cognito_client + ) + + if availability == "taken": + suggested = suggest_available_cognito_domain_prefix( + preferred, account_id, region, cognito_client=cognito_client + ) + print( + f"Cognito hosted UI domain prefix {preferred!r} is not available in " + f"{region} (likely taken by another AWS account)." + ) + if not interactive: + if assume_yes: + print(f"Using suggested Cognito domain prefix: {suggested}") + return suggested + raise SystemExit( + f"Cognito domain prefix {preferred!r} is not available in {region}. " + f"Re-run with --yes to accept {suggested!r}, or pass --cognito-prefix." + ) + if ask_yes_no( + f"Use suggested Cognito domain prefix {suggested!r}?", default=True + ): + return suggested + preferred = suggested + + if not interactive: + print(f"Cognito domain prefix available: {preferred}") + return preferred + + while True: + raw = ask( + "Cognito hosted UI domain prefix " + "(must be globally unique in this region; " + "cannot contain aws, amazon, or cognito)", + preferred, + ) + choice = sanitize_cognito_domain_prefix(raw) + validation_error = validate_cognito_domain_prefix(choice) + if validation_error: + print(validation_error) + continue + if ( + resolve_cognito_domain_prefix_availability( + choice, region_name=region, cognito_client=cognito_client + ) + == "available" + ): + if choice != preferred: + print(f"Using Cognito domain prefix: {choice}") + else: + print(f"Cognito domain prefix available: {choice}") + return choice + print( + f"Prefix {choice!r} is not available in {region} " + "(likely taken by another AWS account). Try another prefix." + ) + + +def resolve_globally_unique_install_names( + answers: InstallAnswers, + *, + interactive: bool, + assume_yes: bool, + args: Optional[argparse.Namespace] = None, +) -> None: + """Check S3 bucket and Cognito domain names; prompt for alternatives when taken.""" + import boto3 + + args = args or argparse.Namespace() + region = answers.aws_region + account_id = answers.aws_account_id + s3_client = boto3.client("s3", region_name=region) + cognito_client = boto3.client("cognito-idp", region_name=region) + + defaults = derive_s3_bucket_names(answers.cdk_prefix) + answers.s3_log_bucket_name = _prompt_globally_unique_s3_bucket( + "S3 log/config bucket", + defaults.get("S3_LOG_CONFIG_BUCKET_NAME", ""), + account_id, + interactive=interactive, + assume_yes=assume_yes, + cli_override=getattr(args, "s3_log_bucket", "") or "", + s3_client=s3_client, + ) + answers.s3_output_bucket_name = _prompt_globally_unique_s3_bucket( + "S3 output bucket", + defaults.get("S3_OUTPUT_BUCKET_NAME", ""), + account_id, + interactive=interactive, + assume_yes=assume_yes, + cli_override=getattr(args, "s3_output_bucket", "") or "", + s3_client=s3_client, + ) + + if answers_use_headless(answers): + answers.cognito_domain_prefix = "" + return + + preferred_cognito = default_cognito_domain_prefix_from_cdk_prefix( + answers.cdk_prefix + ) + answers.cognito_domain_prefix = _prompt_globally_unique_cognito_prefix( + preferred_cognito, + account_id, + region, + interactive=interactive, + assume_yes=assume_yes, + cli_override=getattr(args, "cognito_prefix", "") or "", + cognito_client=cognito_client, + ) + + +def validate_globally_unique_env_values(values: Dict[str, str]) -> List[str]: + """Live AWS checks for globally unique resource names before deploy.""" + from cdk_functions import ( + resolve_cognito_domain_prefix_availability, + resolve_s3_bucket_availability, + ) + + errors: List[str] = [] + region = (values.get("AWS_REGION") or "").strip() + for env_key in ("S3_LOG_CONFIG_BUCKET_NAME", "S3_OUTPUT_BUCKET_NAME"): + bucket_name = (values.get(env_key) or "").strip().lower() + if not bucket_name: + continue + status, _ = resolve_s3_bucket_availability(bucket_name) + if status == "globally_taken": + errors.append( + f"{env_key}={bucket_name!r} is taken globally by another AWS account. " + "Re-run cdk_install.py to pick a unique name." + ) + if values.get("ENABLE_HEADLESS_DEPLOYMENT") != "True": + cognito_prefix = (values.get("COGNITO_USER_POOL_DOMAIN_PREFIX") or "").strip() + if cognito_prefix and region: + if ( + resolve_cognito_domain_prefix_availability( + cognito_prefix, region_name=region + ) + == "taken" + ): + errors.append( + f"COGNITO_USER_POOL_DOMAIN_PREFIX={cognito_prefix!r} is not " + f"available in {region} (taken by another AWS account or existing " + "pool). Re-run cdk_install.py to pick a unique prefix." + ) + return errors + + +def resolve_fixup_env_values(values: Dict[str, str]) -> Dict[str, str]: + """Fill missing ECS cluster/service keys from CDK_PREFIX for post-deploy boto3 calls.""" + resolved = dict(values) + for key, default in derive_ecs_resource_names( + resolved.get("CDK_PREFIX", "") + ).items(): + if not (resolved.get(key) or "").strip(): + resolved[key] = default + for key, default in derive_s3_bucket_names(resolved.get("CDK_PREFIX", "")).items(): + if not (resolved.get(key) or "").strip(): + resolved[key] = default + return resolved + + +def merge_preset( + profile: str, overrides: Optional[Dict[str, str]] = None +) -> Dict[str, str]: + base: Dict[str, str] = {} + if profile == "demo": + base.update(DEMO_PRESET) + elif profile == "production": + base.update(PRODUCTION_PRESET) + elif profile == "headless": + base.update(HEADLESS_PRESET) + if overrides: + base.update(overrides) + return base + + +def answers_preset_profile(answers: "InstallAnswers") -> str: + """Preset profile for env defaults (legacy ``headless`` maps to demo).""" + if answers.profile == "headless": + return "demo" + return answers.profile + + +def answers_use_headless(answers: "InstallAnswers") -> bool: + return answers.profile == "headless" or answers.enable_headless + + +def profile_allows_headless_add_on(answers: "InstallAnswers") -> bool: + """Headless batch is incompatible with the Demonstration (Express) route.""" + if answers.profile == "demo": + return False + if answers.profile in ("production", "headless"): + return True + if answers.profile == "custom": + return answers.custom_overrides.get("USE_ECS_EXPRESS_MODE") != "True" + return False + + +def headless_profile_error(answers: "InstallAnswers") -> Optional[str]: + if not answers_use_headless(answers): + return None + if answers.profile == "headless": + return None + if answers.profile == "demo": + return ( + "Headless batch mode is not available with the Demonstration (Express) " + "profile. Use the Headless shortcut profile, Production with headless, " + "or Custom without ECS Express." + ) + if ( + answers.profile == "custom" + and answers.custom_overrides.get("USE_ECS_EXPRESS_MODE") == "True" + ): + return ( + "Headless batch mode requires USE_ECS_EXPRESS_MODE=False. " + "Disable ECS Express in the custom profile or omit --headless." + ) + return None + + +def validate_notify_email(email: str) -> Optional[str]: + """Return an error message when email is not a plausible notification address.""" + cleaned = (email or "").strip() + if not cleaned: + return "Notification email address is required." + if not re.match(r"^[^@\s]+@[^@\s]+\.[^@\s]+$", cleaned): + return f"Invalid notification email address: {cleaned!r}" + return None + + +DYNAMO_EXPORT_SCHEDULE_DAY_PRESETS = { + "daily": "*", + "weekdays": "MON-FRI", + "weekends": "SAT-SUN", +} + + +def default_dynamo_export_s3_key() -> str: + return "reports/dynamodb-usage/dynamodb_logs_export.csv" + + +def parse_schedule_time_hhmm(raw: str) -> Tuple[int, int]: + match = re.match(r"^(\d{1,2}):(\d{2})$", (raw or "").strip()) + if not match: + raise ValueError( + "Schedule time must use HH:MM in 24-hour UTC format, e.g. 06:00." + ) + hour, minute = int(match.group(1)), int(match.group(2)) + if hour > 23 or minute > 59: + raise ValueError("Schedule time hour must be 0-23 and minute 0-59.") + return hour, minute + + +def validate_schedule_time_hhmm(raw: str) -> Optional[str]: + try: + parse_schedule_time_hhmm(raw) + return None + except ValueError as exc: + return str(exc) + + +def build_dynamo_export_cron_expression( + hour: int, minute: int, days_preset: str +) -> str: + """Build an EventBridge cron expression (UTC).""" + dow = DYNAMO_EXPORT_SCHEDULE_DAY_PRESETS.get( + (days_preset or "").strip().lower(), (days_preset or "*").strip().upper() + ) + return f"cron({minute} {hour} ? * {dow} *)" + + +def prompt_dynamo_usage_log_export_options( + answers: "InstallAnswers", + args: argparse.Namespace, + *, + interactive: bool, +) -> None: + if getattr(args, "dynamo_usage_log_export", False): + answers.enable_dynamo_usage_log_export = True + elif getattr(args, "no_dynamo_usage_log_export", False): + answers.enable_dynamo_usage_log_export = False + elif interactive: + answers.enable_dynamo_usage_log_export = ask_yes_no( + "Enable scheduled DynamoDB usage log export to S3 (EventBridge)?", + default=False, + ) + else: + return + + if not answers.enable_dynamo_usage_log_export: + return + + if getattr(args, "dynamo_export_schedule_time", "").strip(): + answers.dynamo_export_schedule_time = args.dynamo_export_schedule_time.strip() + elif interactive: + while True: + raw = ask( + "Export schedule time (UTC, HH:MM)", + answers.dynamo_export_schedule_time, + ) + err = validate_schedule_time_hhmm(raw) + if not err: + answers.dynamo_export_schedule_time = raw + break + print(err) + + if getattr(args, "dynamo_export_schedule_days", "").strip(): + answers.dynamo_export_schedule_days = ( + args.dynamo_export_schedule_days.strip().lower() + ) + elif interactive: + idx = ask_choice( + "Run export on which days?", + ["Every day", "Weekdays (Mon-Fri)", "Weekends (Sat-Sun)"], + default_index=0, + ) + answers.dynamo_export_schedule_days = ["daily", "weekdays", "weekends"][idx] + + default_key = answers.dynamo_export_s3_key or default_dynamo_export_s3_key() + if getattr(args, "dynamo_export_s3_key", "").strip(): + answers.dynamo_export_s3_key = args.dynamo_export_s3_key.strip() + elif interactive: + answers.dynamo_export_s3_key = ( + ask("S3 object key for the CSV export", default_key).strip() or default_key + ) + + if getattr(args, "dynamo_export_date_attribute", "").strip(): + answers.dynamo_export_date_attribute = args.dynamo_export_date_attribute.strip() + elif interactive: + raw = ask( + "DynamoDB date attribute for filtering", + answers.dynamo_export_date_attribute, + ) + if raw.strip(): + answers.dynamo_export_date_attribute = raw.strip() + + +def validate_install_answers(answers: "InstallAnswers") -> List[str]: + errors: List[str] = [] + msg = headless_profile_error(answers) + if msg: + errors.append(msg) + if answers.enable_headless_output_notifications: + if not answers_use_headless(answers): + errors.append( + "Headless output notifications require headless batch deployment." + ) + email_error = validate_notify_email(answers.headless_output_notify_email) + if email_error: + errors.append(email_error) + if answers.enable_dynamo_usage_log_export: + time_error = validate_schedule_time_hhmm(answers.dynamo_export_schedule_time) + if time_error: + errors.append(time_error) + days = (answers.dynamo_export_schedule_days or "").strip().lower() + if days not in DYNAMO_EXPORT_SCHEDULE_DAY_PRESETS: + errors.append( + "DynamoDB export schedule days must be one of: " + + ", ".join(DYNAMO_EXPORT_SCHEDULE_DAY_PRESETS) + ) + if not (answers.dynamo_export_s3_key.strip() or default_dynamo_export_s3_key()): + errors.append("DynamoDB export S3 object key is required.") + return errors + + +def answers_use_express_mode(answers: "InstallAnswers") -> bool: + """True when the selected install profile/options enable ECS Express ingress.""" + if answers_use_headless(answers): + return False + preset = merge_preset(answers.profile, answers.custom_overrides) + return preset.get("USE_ECS_EXPRESS_MODE") == "True" + + +def answers_use_public_subnets_only(answers: "InstallAnswers") -> bool: + """Express ingress or demo-style headless batch (public subnets, no private install).""" + if answers_use_express_mode(answers): + return True + if not answers_use_headless(answers): + return False + return answers_preset_profile(answers) in ("headless", "demo", "custom") + + +def build_env_values(answers: InstallAnswers) -> Dict[str, str]: + cdk_folder = str(CDK_DIR).replace("\\", "/") + if not cdk_folder.endswith("/"): + cdk_folder += "/" + + s3_bucket_names = derive_s3_bucket_names(answers.cdk_prefix) + if answers.s3_log_bucket_name: + s3_bucket_names["S3_LOG_CONFIG_BUCKET_NAME"] = answers.s3_log_bucket_name + if answers.s3_output_bucket_name: + s3_bucket_names["S3_OUTPUT_BUCKET_NAME"] = answers.s3_output_bucket_name + + values: Dict[str, str] = merge_preset( + answers_preset_profile(answers), answers.custom_overrides + ) + values.update( + { + "CDK_PREFIX": answers.cdk_prefix, + "AWS_REGION": answers.aws_region, + "AWS_ACCOUNT_ID": answers.aws_account_id, + "CDK_FOLDER": cdk_folder, + **derive_ecs_resource_names(answers.cdk_prefix), + **s3_bucket_names, + "CONTEXT_FILE": "precheck.context.json", + "COGNITO_USER_POOL_DOMAIN_PREFIX": sanitize_cognito_domain_prefix( + answers.cognito_domain_prefix + ), + "GITHUB_REPO_BRANCH": answers.github_branch, + "ECS_TASK_MEMORY_SIZE": answers.ecs_memory, + "EXISTING_IGW_ID": answers.existing_igw_id, + "EXISTING_LOAD_BALANCER_ARN": answers.existing_alb_arn, + "EXISTING_LOAD_BALANCER_DNS": answers.existing_alb_dns + or "placeholder_load_balancer_dns.net", + "ENABLE_PI_AGENT_EXPRESS_SERVICE": ( + "True" if answers.enable_pi_express else "False" + ), + "ENABLE_PI_AGENT_ECS_SERVICE": ( + "True" if answers.enable_pi_legacy else "False" + ), + "ENABLE_ECS_SERVICE_CONNECT": ( + "True" if answers.enable_service_connect else "False" + ), + "ENABLE_S3_BATCH_ECS_TRIGGER": ( + "True" + if answers.enable_s3_batch or answers_use_headless(answers) + else "False" + ), + "ENABLE_HEADLESS_DEPLOYMENT": ( + "True" if answers_use_headless(answers) else "False" + ), + } + ) + + if answers_use_headless(answers): + values.update( + { + "ENABLE_HEADLESS_DEPLOYMENT": "True", + "ENABLE_S3_BATCH_ECS_TRIGGER": "True", + "USE_ECS_EXPRESS_MODE": "False", + "USE_CLOUDFRONT": "False", + "RUN_USEAST_STACK": "False", + "COGNITO_AUTH": "False", + "ENABLE_PI_AGENT_EXPRESS_SERVICE": "False", + "ENABLE_PI_AGENT_ECS_SERVICE": "False", + "ENABLE_ECS_SERVICE_CONNECT": "False", + } + ) + if answers.enable_headless_output_notifications: + iam_user = ( + answers.headless_output_iam_user_name.strip() + or f"{answers.cdk_prefix}s3-output-reader" + ) + values.update( + { + "ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS": "True", + "HEADLESS_OUTPUT_NOTIFY_EMAIL": ( + answers.headless_output_notify_email.strip() + ), + "HEADLESS_OUTPUT_IAM_USER_NAME": iam_user, + } + ) + + if answers.enable_dynamo_usage_log_export: + hour, minute = parse_schedule_time_hhmm(answers.dynamo_export_schedule_time) + s3_key = answers.dynamo_export_s3_key.strip() or default_dynamo_export_s3_key() + values.update( + { + "ENABLE_DYNAMODB_USAGE_LOG_EXPORT": "True", + "DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE": build_dynamo_export_cron_expression( + hour, minute, answers.dynamo_export_schedule_days + ), + "DYNAMODB_USAGE_LOG_EXPORT_S3_KEY": s3_key, + "DYNAMODB_USAGE_LOG_EXPORT_DATE_ATTRIBUTE": ( + answers.dynamo_export_date_attribute.strip() or "timestamp" + ), + } + ) + + use_cloudfront = values.get("USE_CLOUDFRONT") == "True" + if answers.pi_enabled: + values.update( + { + "PI_GRADIO_PORT": answers.pi_gradio_port, + "ECS_SERVICE_CONNECT_DISCOVERY_NAME": answers.sc_discovery_name, + } + ) + if answers.enable_pi_express: + values["ECS_EXPRESS_SC_PORT_NAME"] = "port-7860" + values["ECS_PI_EXPRESS_SC_PORT_NAME"] = f"port-{answers.pi_gradio_port}" + else: + routing = answers.pi_alb_routing.strip().lower() + path_prefix = normalize_pi_path_prefix(answers.pi_alb_path_prefix) + priority = ( + answers.pi_alb_listener_rule_priority.strip() + or default_pi_listener_priority(use_cloudfront) + ) + values.update( + { + "PI_ALB_ROUTING": routing, + "PI_ALB_PATH_PREFIX": path_prefix, + "PI_ALB_LISTENER_RULE_PRIORITY": priority, + } + ) + if routing in ("host", "both"): + values["PI_ALB_HOST_HEADER"] = answers.pi_alb_host_header.strip() + + if answers.profile == "production": + values["ACM_SSL_CERTIFICATE_ARN"] = answers.acm_cert_arn + values["SSL_CERTIFICATE_DOMAIN"] = answers.ssl_domain + values["CLOUDFRONT_DOMAIN"] = "cloudfront_placeholder.net" + if answers.cloudfront_geo: + values["CLOUDFRONT_GEO_RESTRICTION"] = answers.cloudfront_geo + + if answers.vpc_mode == "new": + values["VPC_NAME"] = "" + values["NEW_VPC_CIDR"] = answers.new_vpc_cidr + values["NEW_VPC_DEFAULT_NAME"] = f"{answers.cdk_prefix}vpc" + values["PUBLIC_SUBNETS_TO_USE"] = "" + values["PRIVATE_SUBNETS_TO_USE"] = "" + values["PUBLIC_SUBNET_CIDR_BLOCKS"] = "" + values["PRIVATE_SUBNET_CIDR_BLOCKS"] = "" + values["PUBLIC_SUBNET_AVAILABILITY_ZONES"] = "" + values["PRIVATE_SUBNET_AVAILABILITY_ZONES"] = "" + else: + values["VPC_NAME"] = answers.vpc_name + values["NEW_VPC_CIDR"] = "" + apply_subnet_tier_env( + values, answers, tier="public", mode=answers.public_subnet_mode + ) + if answers_use_public_subnets_only(answers): + values["ECS_EXPRESS_USE_PUBLIC_SUBNETS"] = "True" + values["PRIVATE_SUBNETS_TO_USE"] = "" + values["PRIVATE_SUBNET_CIDR_BLOCKS"] = "" + values["PRIVATE_SUBNET_AVAILABILITY_ZONES"] = "" + else: + apply_subnet_tier_env( + values, answers, tier="private", mode=answers.private_subnet_mode + ) + + if values.get("ENABLE_APPREGISTRY") == "True": + values["APPREGISTRY_STACK_NAME"] = ( + f"{answers.cdk_prefix}{APPREGISTRY_STACK_SUFFIX}" + ) + + # Pi agent is not supported in llm_topic_modeller yet (stack code retained for future). + values["ENABLE_PI_AGENT_EXPRESS_SERVICE"] = "False" + values["ENABLE_PI_AGENT_ECS_SERVICE"] = "False" + + return values + + +def validate_env_values(values: Dict[str, str]) -> List[str]: + errors: List[str] = [] + + express = values.get("USE_ECS_EXPRESS_MODE") == "True" + acm = (values.get("ACM_SSL_CERTIFICATE_ARN") or "").strip() + if express and acm: + errors.append( + "USE_ECS_EXPRESS_MODE=True cannot be used with ACM_SSL_CERTIFICATE_ARN." + ) + + if values.get("ENABLE_ECS_SERVICE_CONNECT") == "True" and express: + errors.append( + "ENABLE_ECS_SERVICE_CONNECT requires legacy Fargate (not Express)." + ) + + if values.get("ENABLE_S3_BATCH_ECS_TRIGGER") == "True" and express: + errors.append( + "ENABLE_S3_BATCH_ECS_TRIGGER requires legacy Fargate (not Express)." + ) + + headless = values.get("ENABLE_HEADLESS_DEPLOYMENT") == "True" + if headless: + if values.get("ENABLE_S3_BATCH_ECS_TRIGGER") != "True": + errors.append( + "ENABLE_HEADLESS_DEPLOYMENT requires ENABLE_S3_BATCH_ECS_TRIGGER=True." + ) + if express: + errors.append( + "ENABLE_HEADLESS_DEPLOYMENT cannot be combined with USE_ECS_EXPRESS_MODE=True " + "(use --profile headless, production --headless, or custom without Express)." + ) + if values.get("USE_CLOUDFRONT") == "True": + errors.append( + "ENABLE_HEADLESS_DEPLOYMENT is incompatible with USE_CLOUDFRONT=True." + ) + + pi_ecs = values.get("ENABLE_PI_AGENT_ECS_SERVICE") == "True" + pi_express = values.get("ENABLE_PI_AGENT_EXPRESS_SERVICE") == "True" + if pi_ecs and pi_express: + errors.append("Enable at most one agent mode (ECS or Express).") + if pi_express and not express: + errors.append( + "ENABLE_PI_AGENT_EXPRESS_SERVICE requires USE_ECS_EXPRESS_MODE=True." + ) + if pi_ecs and express: + errors.append( + "ENABLE_PI_AGENT_ECS_SERVICE requires USE_ECS_EXPRESS_MODE=False." + ) + + vpc_name = (values.get("VPC_NAME") or "").strip() + new_cidr = (values.get("NEW_VPC_CIDR") or "").strip() + if not vpc_name and not new_cidr: + errors.append("Set VPC_NAME (existing VPC) or NEW_VPC_CIDR (new VPC).") + elif new_cidr and not vpc_name: + format_error = validate_new_vpc_cidr_format(new_cidr) + if format_error: + errors.append(format_error) + + if express: + public_cidrs_raw = (values.get("PUBLIC_SUBNET_CIDR_BLOCKS") or "").strip() + if public_cidrs_raw: + try: + import ast + + parsed = ast.literal_eval(public_cidrs_raw) + public_cidrs = parsed if isinstance(parsed, list) else [str(parsed)] + except (SyntaxError, ValueError): + public_cidrs = [ + part.strip().strip("'\"") + for part in public_cidrs_raw.split(",") + if part.strip() + ] + for cidr in public_cidrs: + subnet_error = validate_public_subnet_cidr_for_express(str(cidr)) + if subnet_error: + errors.append(subnet_error) + + if values.get("USE_CLOUDFRONT") == "True" and not express: + if not acm: + errors.append( + "Production CloudFront path requires ACM_SSL_CERTIFICATE_ARN." + ) + if not (values.get("SSL_CERTIFICATE_DOMAIN") or "").strip(): + errors.append("Production CloudFront path requires SSL_CERTIFICATE_DOMAIN.") + + if not (values.get("CDK_PREFIX") or "").strip(): + errors.append("CDK_PREFIX is required.") + prefix_lower = (values.get("CDK_PREFIX") or "").strip().lower() + for env_key, bare_suffix in ( + ("S3_LOG_CONFIG_BUCKET_NAME", "s3-logs"), + ("S3_OUTPUT_BUCKET_NAME", "s3-output"), + ): + bucket_name = (values.get(env_key) or "").strip().lower() + if prefix_lower and bucket_name == bare_suffix: + errors.append( + f"{env_key} must include CDK_PREFIX (got bare '{bare_suffix}'; " + f"expected '{prefix_lower}{bare_suffix}'). Re-run cdk_install or set " + f"{env_key}={prefix_lower}{bare_suffix} in config/cdk_config.env." + ) + if values.get("ENABLE_HEADLESS_DEPLOYMENT") != "True": + cognito_prefix_error = validate_cognito_domain_prefix( + values.get("COGNITO_USER_POOL_DOMAIN_PREFIX", "") + ) + if cognito_prefix_error: + errors.append(cognito_prefix_error) + + if pi_ecs and values.get("ENABLE_ECS_SERVICE_CONNECT") != "True": + errors.append( + "ENABLE_PI_AGENT_ECS_SERVICE=True requires ENABLE_ECS_SERVICE_CONNECT=True." + ) + + if headless and ( + pi_ecs or pi_express or values.get("ENABLE_ECS_SERVICE_CONNECT") == "True" + ): + errors.append( + "ENABLE_HEADLESS_DEPLOYMENT is incompatible with agent mode or Service Connect." + ) + + if values.get("ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS") == "True": + if not headless: + errors.append( + "ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS requires " + "ENABLE_HEADLESS_DEPLOYMENT=True." + ) + email_error = validate_notify_email( + values.get("HEADLESS_OUTPUT_NOTIFY_EMAIL", "") + ) + if email_error: + errors.append(email_error) + + if values.get("ENABLE_DYNAMODB_USAGE_LOG_EXPORT") == "True": + if values.get("SAVE_LOGS_TO_DYNAMODB") != "True": + errors.append( + "ENABLE_DYNAMODB_USAGE_LOG_EXPORT requires SAVE_LOGS_TO_DYNAMODB=True." + ) + schedule = (values.get("DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE") or "").strip() + if not schedule.startswith("cron(") and not schedule.startswith("rate("): + errors.append( + "DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE must be an EventBridge " + f"cron(...) or rate(...) expression; got {schedule!r}." + ) + if not (values.get("DYNAMODB_USAGE_LOG_EXPORT_S3_KEY") or "").strip(): + errors.append( + "DYNAMODB_USAGE_LOG_EXPORT_S3_KEY is required when " + "ENABLE_DYNAMODB_USAGE_LOG_EXPORT=True." + ) + + pi_routing = (values.get("PI_ALB_ROUTING") or "").strip().lower() + if pi_ecs: + if pi_routing not in PI_ALB_ROUTING_MODES: + errors.append( + f"PI_ALB_ROUTING must be one of {list(PI_ALB_ROUTING_MODES)}; got '{pi_routing}'." + ) + elif ( + pi_routing in ("host", "both") + and not (values.get("PI_ALB_HOST_HEADER") or "").strip() + ): + errors.append( + "PI_ALB_HOST_HEADER is required when PI_ALB_ROUTING is 'host' or 'both'." + ) + + return errors + + +def build_app_config_env_values(values: Dict[str, str]) -> Dict[str, str]: + """AWS deployment keys merged into config/app_config.env for ECS tasks.""" + prefix_lower = (values.get("CDK_PREFIX") or "").lower() + headless = values.get("ENABLE_HEADLESS_DEPLOYMENT") == "True" + pi_express = values.get("ENABLE_PI_AGENT_EXPRESS_SERVICE") == "True" + + def _name(env_key: str, suffix: str) -> str: + return (values.get(env_key) or "").strip() or f"{prefix_lower}{suffix}" + + output_bucket = _name("S3_OUTPUT_BUCKET_NAME", "s3-output") + + return { + "COGNITO_AUTH": "False" if headless or pi_express else "True", + "RUN_AWS_FUNCTIONS": "True", + "RUN_AWS_BEDROCK_MODELS": "True", + "RUN_LOCAL_MODEL": "False", + "RUN_GEMINI_MODELS": "False", + "RUN_AZURE_MODELS": "False", + "PRIORITISE_SSO_OVER_AWS_ENV_ACCESS_KEYS": "True", + "DISPLAY_FILE_NAMES_IN_LOGS": "False", + "SESSION_OUTPUT_FOLDER": "True", + "SAVE_LOGS_TO_CSV": "True", + "SAVE_LOGS_TO_DYNAMODB": "True", + "SHOW_COSTS": "True", + "S3_LOG_BUCKET": _name("S3_LOG_CONFIG_BUCKET_NAME", "s3-logs"), + "S3_OUTPUTS_BUCKET": output_bucket if headless else "", + "S3_OUTPUTS_FOLDER": "output/" if headless else "", + "SAVE_OUTPUTS_TO_S3": "True" if headless else "False", + "ACCESS_LOG_DYNAMODB_TABLE_NAME": _name( + "ACCESS_LOG_DYNAMODB_TABLE_NAME", "dynamodb-access-logs" + ), + "FEEDBACK_LOG_DYNAMODB_TABLE_NAME": _name( + "FEEDBACK_LOG_DYNAMODB_TABLE_NAME", "dynamodb-feedback-logs" + ), + "USAGE_LOG_DYNAMODB_TABLE_NAME": _name( + "USAGE_LOG_DYNAMODB_TABLE_NAME", "dynamodb-usage-logs" + ), + } + + +def write_app_config_env_file( + answers: InstallAnswers, + values: Dict[str, str], + *, + overwrite: bool = False, +) -> Optional[Path]: + if not answers.write_app_config_env: + return None + + updates = build_app_config_env_values(values) + APP_CONFIG_ENV_PATH.parent.mkdir(parents=True, exist_ok=True) + + if APP_CONFIG_ENV_PATH.is_file() and not overwrite: + existing = read_env_file(APP_CONFIG_ENV_PATH) + existing.update(updates) + backup_file(APP_CONFIG_ENV_PATH) + lines = [f"{k}={v}" for k, v in existing.items()] + APP_CONFIG_ENV_PATH.write_text("\n".join(lines) + "\n", encoding="utf-8") + print(f"Updated {APP_CONFIG_ENV_PATH} (AWS deployment keys merged)") + return APP_CONFIG_ENV_PATH + + if APP_CONFIG_ENV_EXAMPLE.is_file() and not overwrite: + if APP_CONFIG_ENV_PATH.is_file(): + backup_file(APP_CONFIG_ENV_PATH) + shutil.copy2(APP_CONFIG_ENV_EXAMPLE, APP_CONFIG_ENV_PATH) + existing = read_env_file(APP_CONFIG_ENV_PATH) + existing.update(updates) + lines = [f"{k}={v}" for k, v in existing.items()] + APP_CONFIG_ENV_PATH.write_text("\n".join(lines) + "\n", encoding="utf-8") + print(f"Created {APP_CONFIG_ENV_PATH} from example + AWS defaults") + return APP_CONFIG_ENV_PATH + + if APP_CONFIG_ENV_PATH.is_file(): + backup_file(APP_CONFIG_ENV_PATH) + lines = [ + "# Generated by cdk_install.py — app runtime config for AWS ECS", + *[f"{k}={v}" for k, v in updates.items()], + ] + APP_CONFIG_ENV_PATH.write_text("\n".join(lines) + "\n", encoding="utf-8") + print(f"Wrote {APP_CONFIG_ENV_PATH}") + return APP_CONFIG_ENV_PATH + + +def build_pi_agent_env_values(answers: InstallAnswers) -> Dict[str, str]: + """Runtime settings for the Pi agent Gradio app (uploaded to S3 as pi_agent.env).""" + values = { + "PI_DEPLOYMENT_PROFILE": "aws-ecs", + "PI_DEFAULT_PROVIDER": answers.pi_default_provider, + "DOC_SUMMARISATION_GRADIO_URL": f"http://{answers.sc_discovery_name}:7860", + "RUN_AWS_FUNCTIONS": "True", + "AWS_REGION": answers.aws_region, + "PI_GRADIO_PORT": answers.pi_gradio_port, + "PI_DEFAULT_OCR_METHOD": "AWS Textract service - all PDF types", + "PI_DEFAULT_PII_METHOD": "AWS Comprehend", + } + if answers.enable_pi_express: + values["RUN_FASTAPI"] = "True" + return values + path_prefix = normalize_pi_path_prefix(answers.pi_alb_path_prefix) + if answers.pi_alb_routing.strip().lower() in ("path", "both"): + values["PI_ROOT_PATH"] = path_prefix + values["ROOT_PATH"] = path_prefix + values["FASTAPI_ROOT_PATH"] = path_prefix + return values + + +def write_pi_agent_env_file( + answers: InstallAnswers, + *, + overwrite: bool = False, +) -> Optional[Path]: + if not answers.pi_enabled or not answers.write_pi_agent_env: + return None + + updates = build_pi_agent_env_values(answers) + PI_AGENT_ENV_PATH.parent.mkdir(parents=True, exist_ok=True) + + if PI_AGENT_ENV_PATH.is_file() and not overwrite: + existing = read_env_file(PI_AGENT_ENV_PATH) + existing.update(updates) + backup_file(PI_AGENT_ENV_PATH) + lines = [f"{k}={v}" for k, v in existing.items()] + PI_AGENT_ENV_PATH.write_text("\n".join(lines) + "\n", encoding="utf-8") + print(f"Updated {PI_AGENT_ENV_PATH} (AWS/Pi agent keys merged)") + return PI_AGENT_ENV_PATH + + if PI_AGENT_ENV_EXAMPLE.is_file() and not overwrite: + backup_file(PI_AGENT_ENV_PATH) if PI_AGENT_ENV_PATH.is_file() else None + shutil.copy2(PI_AGENT_ENV_EXAMPLE, PI_AGENT_ENV_PATH) + existing = read_env_file(PI_AGENT_ENV_PATH) + existing.update(updates) + lines = [f"{k}={v}" for k, v in existing.items()] + PI_AGENT_ENV_PATH.write_text("\n".join(lines) + "\n", encoding="utf-8") + print(f"Created {PI_AGENT_ENV_PATH} from example + AWS defaults") + return PI_AGENT_ENV_PATH + + backup_file(PI_AGENT_ENV_PATH) if PI_AGENT_ENV_PATH.is_file() else None + lines = [ + "# Generated by cdk_install.py — Pi agent runtime config for AWS ECS", + *[f"{k}={v}" for k, v in updates.items()], + ] + PI_AGENT_ENV_PATH.write_text("\n".join(lines) + "\n", encoding="utf-8") + print(f"Wrote {PI_AGENT_ENV_PATH}") + return PI_AGENT_ENV_PATH + + +def write_env_file(path: Path, values: Dict[str, str]) -> Path: + backup_file(path) + lines = [f"{key}={val}" for key, val in values.items()] + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + print(f"Wrote {path}") + return path + + +def read_env_file(path: Path) -> Dict[str, str]: + values: Dict[str, str] = {} + if not path.is_file(): + return values + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + if "=" in line: + key, _, val = line.partition("=") + values[key.strip()] = val.strip() + return values + + +def patch_env_file(path: Path, updates: Dict[str, str]) -> None: + values = read_env_file(path) + values.update(updates) + write_env_file(path, values) + + +def print_summary(values: Dict[str, str], python_exe: Optional[Path] = None) -> None: + keys = [ + "CDK_PREFIX", + "AWS_REGION", + "AWS_ACCOUNT_ID", + "USE_ECS_EXPRESS_MODE", + "USE_CLOUDFRONT", + "ENABLE_HEADLESS_DEPLOYMENT", + "ENABLE_S3_BATCH_ECS_TRIGGER", + "ENABLE_RESOURCE_DELETE_PROTECTION", + "VPC_NAME", + "NEW_VPC_CIDR", + "ACM_SSL_CERTIFICATE_ARN", + "SSL_CERTIFICATE_DOMAIN", + "ENABLE_PI_AGENT_EXPRESS_SERVICE", + "ENABLE_PI_AGENT_ECS_SERVICE", + ] + print("\n--- Configuration summary ---") + for key in keys: + if key in values: + print(f" {key}={values[key]}") + if python_exe: + print(f" cdk.json app={format_cdk_app_command(python_exe)}") + print("----------------------------\n") + + +# --------------------------------------------------------------------------- +# Deploy runner +# --------------------------------------------------------------------------- + + +def _deploy_env() -> Dict[str, str]: + return build_cdk_subprocess_env() + + +def run_cdk_command( + args: List[str], *, check: bool = True +) -> subprocess.CompletedProcess: + cdk_exe = resolve_cdk_executable() + cmd = [cdk_exe, *args] + print(f"Running: {' '.join(cmd)}") + return subprocess.run( + cmd, + cwd=str(CDK_DIR), + env=_deploy_env(), + check=check, + ) + + +@dataclass(frozen=True) +class ExistingStack: + name: str + region: str + status: str + termination_protection: bool = False + + +def _should_check_cloudfront_stack( + config_values: Optional[Dict[str, str]] = None, +) -> bool: + """Skip us-east-1 CloudFront stack lookup when config disables that path.""" + if not config_values: + return True + use_cloudfront = config_values.get("USE_CLOUDFRONT") + run_useast = config_values.get("RUN_USEAST_STACK") + if use_cloudfront is not None and use_cloudfront != "True": + return False + if run_useast is not None and run_useast != "True": + return False + return True + + +def _stack_check_skippable_error(exc: Exception) -> bool: + """Permission / SCP errors in one region must not block checks in others.""" + from botocore.exceptions import ClientError + + if not isinstance(exc, ClientError): + return False + code = exc.response.get("Error", {}).get("Code", "") + return code in ( + "AccessDenied", + "UnauthorizedOperation", + "AuthorizationError", + "AccessDeniedException", + ) + + +def derived_appregistry_stack_name( + config_values: Optional[Dict[str, str]] = None, +) -> Optional[str]: + """Resolve AppRegistry stack name from config (explicit or CDK_PREFIX default).""" + if not config_values: + return None + explicit = (config_values.get("APPREGISTRY_STACK_NAME") or "").strip() + if explicit: + return explicit + prefix = (config_values.get("CDK_PREFIX") or "").strip() + if prefix: + return f"{prefix}{APPREGISTRY_STACK_SUFFIX}" + return None + + +def list_regional_appregistry_stack_names(region: str) -> List[str]: + """List active AppRegistry stacks in the regional account (suffix match).""" + import boto3 + from botocore.exceptions import ClientError + + cfn = boto3.client("cloudformation", region_name=region) + skip_statuses = {"DELETE_COMPLETE", "DELETE_IN_PROGRESS"} + names: List[str] = [] + try: + paginator = cfn.get_paginator("list_stacks") + for page in paginator.paginate(): + for summary in page.get("StackSummaries", []): + name = summary.get("StackName", "") + status = summary.get("StackStatus", "") + if not name.endswith(APPREGISTRY_STACK_SUFFIX): + continue + if status in skip_statuses: + continue + names.append(name) + except ClientError as exc: + if _stack_check_skippable_error(exc): + code = exc.response.get("Error", {}).get("Code", "ClientError") + print( + f"Warning: could not list AppRegistry stacks in {region} " + f"({code}). Continuing with configured stack names only." + ) + return [] + raise + return sorted(dict.fromkeys(names)) + + +def appregistry_stack_names_to_check( + config_values: Optional[Dict[str, str]] = None, + *, + discovered_names: Optional[Sequence[str]] = None, +) -> List[str]: + """Merge configured and discovered AppRegistry stack names (deduped).""" + names: List[str] = [] + derived = derived_appregistry_stack_name(config_values) + if derived: + names.append(derived) + for name in discovered_names or (): + if name not in names: + names.append(name) + return names + + +def stacks_to_check( + regional_region: str, + config_values: Optional[Dict[str, str]] = None, + *, + discovered_appregistry_names: Optional[Sequence[str]] = None, +) -> List[Tuple[str, str]]: + """Return (stack_name, region) pairs in safe deletion order.""" + checks: List[Tuple[str, str]] = [] + if _should_check_cloudfront_stack(config_values): + checks.append((CLOUDFRONT_STACK, CLOUDFRONT_STACK_REGION)) + for appregistry_name in appregistry_stack_names_to_check( + config_values, + discovered_names=discovered_appregistry_names, + ): + checks.append((appregistry_name, regional_region)) + checks.append((REGIONAL_STACK, regional_region)) + return checks + + +def describe_existing_stack(stack_name: str, region: str) -> Optional[ExistingStack]: + import boto3 + from botocore.exceptions import ClientError + + cfn = boto3.client("cloudformation", region_name=region) + try: + response = cfn.describe_stacks(StackName=stack_name) + except ClientError as exc: + code = exc.response.get("Error", {}).get("Code", "") + if code in ("ValidationError", "ResourceNotFoundException") or ( + "does not exist" in str(exc) + ): + return None + raise + stack = response["Stacks"][0] + status = stack.get("StackStatus", "") + if status == "DELETE_COMPLETE": + return None + return ExistingStack( + name=stack_name, + region=region, + status=status, + termination_protection=bool(stack.get("EnableTerminationProtection", False)), + ) + + +def discover_existing_llm_topic_stacks( + regional_region: str, + config_values: Optional[Dict[str, str]] = None, +) -> List[ExistingStack]: + """Find deployed llm_topic_modeller CloudFormation stacks in the target account.""" + from botocore.exceptions import ClientError + + discovered_appregistry = list_regional_appregistry_stack_names(regional_region) + checks = stacks_to_check( + regional_region, + config_values, + discovered_appregistry_names=discovered_appregistry, + ) + ordered: List[ExistingStack] = [] + for stack_name, region in checks: + try: + info = describe_existing_stack(stack_name, region) + except ClientError as exc: + if _stack_check_skippable_error(exc): + code = exc.response.get("Error", {}).get("Code", "ClientError") + print( + f"Warning: could not check stack {stack_name!r} in {region} " + f"({code}). Continuing with other regions." + ) + continue + raise + if info: + ordered.append(info) + return ordered + + +def discover_existing_doc_summarisation_stacks( + regional_region: str, + config_values: Optional[Dict[str, str]] = None, +) -> List[ExistingStack]: + """Deprecated alias for discover_existing_llm_topic_stacks.""" + return discover_existing_llm_topic_stacks(regional_region, config_values) + + +def force_delete_cloudformation_stacks( + stacks: Sequence[ExistingStack], + *, + wait: bool = True, +) -> None: + """Delete stacks via CloudFormation (disables termination protection first).""" + import boto3 + from botocore.exceptions import ClientError + + clients: Dict[str, Any] = {} + + def cfn_for(region: str): + if region not in clients: + clients[region] = boto3.client("cloudformation", region_name=region) + return clients[region] + + for stack in stacks: + cfn = cfn_for(stack.region) + if stack.termination_protection: + print(f"Disabling termination protection on {stack.name} ({stack.region})") + cfn.update_termination_protection( + EnableTerminationProtection=False, + StackName=stack.name, + ) + if stack.status == "DELETE_IN_PROGRESS": + print(f"{stack.name} ({stack.region}) is already deleting.") + else: + print(f"Deleting stack {stack.name} ({stack.region}) ...") + try: + cfn.delete_stack(StackName=stack.name) + except ClientError as exc: + code = exc.response.get("Error", {}).get("Code", "") + if code == "ValidationError" and "DELETE_IN_PROGRESS" in str(exc): + print(f"{stack.name} is already deleting.") + else: + raise + + if not wait: + return + + for stack in stacks: + cfn = cfn_for(stack.region) + print(f"Waiting for {stack.name} ({stack.region}) to finish deleting ...") + try: + cfn.get_waiter("stack_delete_complete").wait( + StackName=stack.name, + WaiterConfig={"Delay": 15, "MaxAttempts": 120}, + ) + except ClientError: + detail = describe_existing_stack(stack.name, stack.region) + status = detail.status if detail else "unknown" + raise SystemExit( + f"Stack {stack.name} ({stack.region}) did not delete cleanly " + f"(status={status}). Resolve in the CloudFormation console and retry." + ) from None + + +def handle_existing_stacks_at_start( + args: argparse.Namespace, + regional_region: str, + *, + config_values: Optional[Dict[str, str]] = None, +) -> None: + """At wizard start: report existing stacks and optionally force-delete them.""" + if getattr(args, "skip_stack_check", False) or args.config_only or args.synth_only: + return + + if config_values is None and ENV_PATH.is_file(): + config_values = read_env_file(ENV_PATH) + + try: + existing = discover_existing_llm_topic_stacks(regional_region, config_values) + except Exception as exc: + print(f"Existing stack check skipped: {exc}") + return + + if not existing: + return + + print("\n--- Existing llm_topic_modeller CloudFormation stacks ---") + for stack in existing: + line = f" {stack.name} ({stack.region}): {stack.status}" + if stack.termination_protection: + line += " [termination protection ON]" + print(line) + + if args.force_delete_stacks: + should_delete = True + elif args.yes: + print( + "\nStacks already exist in this account/region. " + "Pass --force-delete-stacks to remove them before deploy, " + "or omit it to update in place." + ) + return + else: + should_delete = ask_yes_no( + "Force-delete these stacks before continuing? " + "(disables termination protection and deletes all stack resources)", + default=False, + ) + + if not should_delete: + print("Keeping existing stacks (deploy will update them in place).\n") + return + + if not args.force_delete_stacks and not args.yes: + if not ask_yes_no( + "This permanently deletes AWS resources in these stacks. Proceed?", + default=False, + ): + print("Stack deletion cancelled.\n") + return + + force_delete_cloudformation_stacks(existing) + print("Existing stacks deleted.\n") + + +def fetch_stack_output( + stack_name: str, + output_key: str, + region: str, +) -> Optional[str]: + import boto3 + from botocore.exceptions import ClientError + + cfn = boto3.client("cloudformation", region_name=region) + try: + response = cfn.describe_stacks(StackName=stack_name) + except ClientError: + return None + for stack in response.get("Stacks", []): + for output in stack.get("Outputs", []): + if output.get("OutputKey") == output_key: + return output.get("OutputValue") + return None + + +def apply_post_deploy_fixup(values: Dict[str, str], assume_yes: bool) -> bool: + """Return True if post-deploy fixup changed env and/or Cognito callbacks.""" + from cdk_post_deploy import ( + apply_cognito_alb_callback_fixup, + cognito_alb_callbacks_need_update, + ) + + if values.get("ENABLE_HEADLESS_DEPLOYMENT") == "True": + return False + + region = values.get("AWS_REGION", "") + express = values.get("USE_ECS_EXPRESS_MODE") == "True" + cloudfront = values.get("USE_CLOUDFRONT") == "True" + fixup_applied = False + + pool_id = fetch_stack_output(REGIONAL_STACK, "CognitoPoolId", region) + client_id = fetch_stack_output(REGIONAL_STACK, "CognitoAppClientId", region) + if not pool_id or not client_id: + print( + "CognitoPoolId or CognitoAppClientId stack output missing; " + "skipping Cognito callback fixup." + ) + return False + + if express and not cloudfront: + from cdk_config import normalize_https_redirect_url + + endpoint = fetch_stack_output(REGIONAL_STACK, "ExpressServiceEndpoint", region) + if not endpoint: + print("No ExpressServiceEndpoint output found; skipping Cognito URL fixup.") + return False + endpoint = normalize_https_redirect_url(endpoint) + current = (values.get("ECS_EXPRESS_COGNITO_REDIRECT_BASE") or "").strip() + env_needs_patch = current != endpoint + cognito_needs_update = cognito_alb_callbacks_need_update( + pool_id, client_id, endpoint, aws_region=region + ) + if not env_needs_patch and not cognito_needs_update: + print("Express Cognito redirect base and callback URLs already set.") + else: + if env_needs_patch: + print(f"Setting ECS_EXPRESS_COGNITO_REDIRECT_BASE={endpoint}") + patch_env_file( + ENV_PATH, + { + "ECS_EXPRESS_COGNITO_REDIRECT_BASE": endpoint, + "COGNITO_REDIRECTION_URL": endpoint, + }, + ) + fixup_applied = True + if cognito_needs_update and ( + assume_yes + or ask_yes_no( + "Update Cognito app client callback URLs for the Express endpoint " + "(no CDK redeploy)?", + True, + ) + ): + if apply_cognito_alb_callback_fixup( + user_pool_id=pool_id, + client_id=client_id, + redirect_base=endpoint, + aws_region=region, + ): + fixup_applied = True + elif env_needs_patch and cognito_needs_update: + print( + "Env updated with Express URL; Cognito callbacks unchanged. " + "ALB login will fail until callback URLs are updated." + ) + + elif cloudfront: + cf_domain = fetch_stack_output( + CLOUDFRONT_STACK, "CloudFrontDistributionURL", CLOUDFRONT_STACK_REGION + ) + if not cf_domain: + print( + "No CloudFrontDistributionURL output found; skipping CloudFront fixup." + ) + return fixup_applied + redirect_base = f"https://{cf_domain.strip()}" + current = (values.get("CLOUDFRONT_DOMAIN") or "").strip() + env_needs_patch = current != cf_domain.strip() + cognito_needs_update = cognito_alb_callbacks_need_update( + pool_id, client_id, redirect_base, aws_region=region + ) + if not env_needs_patch and not cognito_needs_update: + print("CloudFront domain and Cognito callback URLs already set.") + return fixup_applied + if env_needs_patch: + print(f"Setting CLOUDFRONT_DOMAIN={cf_domain}") + patch_env_file( + ENV_PATH, + { + "CLOUDFRONT_DOMAIN": cf_domain, + "COGNITO_REDIRECTION_URL": redirect_base, + }, + ) + fixup_applied = True + if cognito_needs_update and ( + assume_yes + or ask_yes_no( + "Update Cognito app client callback URLs for the CloudFront domain " + "(no CDK redeploy)?", + True, + ) + ): + if apply_cognito_alb_callback_fixup( + user_pool_id=pool_id, + client_id=client_id, + redirect_base=redirect_base, + aws_region=region, + ): + fixup_applied = True + elif env_needs_patch and cognito_needs_update: + print( + "Env updated with CloudFront domain; Cognito callbacks unchanged. " + "ALB login will fail until callback URLs are updated." + ) + + if express: + from cdk_post_deploy import apply_cognito_secret_fixup_from_stack + + resolved = resolve_fixup_env_values(values) + cluster_name = resolved["CLUSTER_NAME"] + main_service = resolved["ECS_EXPRESS_SERVICE_NAME"] + secret_name = resolved.get("COGNITO_USER_POOL_CLIENT_SECRET_NAME", "") + if not main_service: + print( + "ECS express service name could not be resolved; " + "skipping Express Cognito secret sync." + ) + return fixup_applied + try: + if pool_id and client_id: + if apply_cognito_secret_fixup_from_stack( + stack_name=REGIONAL_STACK, + secret_name=secret_name, + cluster_name=cluster_name, + main_service_name=main_service, + aws_region=region, + recycle_tasks=False, + ): + fixup_applied = True + except Exception as exc: + print(f"Warning: could not sync Express Cognito secret: {exc}") + + return fixup_applied + + +def run_quickstart(python_exe: Path) -> None: + if not QUICKSTART_SCRIPT.is_file(): + raise SystemExit(f"Quickstart script not found: {QUICKSTART_SCRIPT}") + cmd = [str(python_exe), str(QUICKSTART_SCRIPT)] + print(f"Running: {' '.join(cmd)}") + subprocess.run(cmd, cwd=str(CDK_DIR), env=_deploy_env(), check=True) + + +# --------------------------------------------------------------------------- +# Pi agent configuration (disabled in installer; stack code retained for future) +# --------------------------------------------------------------------------- + +_DEPRECATED_PI_CLI_FLAGS = ( + "--enable-pi", + "--enable-pi-express", + "--enable-pi-legacy", + "--pi-alb-routing", + "--pi-path-prefix", + "--pi-host-header", + "--pi-listener-priority", + "--pi-gradio-port", + "--sc-discovery-name", + "--pi-provider", + "--skip-pi-agent-env", +) + + +def warn_deprecated_pi_cli_flags(args: argparse.Namespace) -> None: + """Ignore Pi agent CLI flags — not supported in llm_topic_modeller yet.""" + argv = getattr(args, "_argv", ()) or () + if any( + token == flag or token.startswith(f"{flag}=") + for flag in _DEPRECATED_PI_CLI_FLAGS + for token in argv + ): + print( + "Note: Pi agent is not supported in llm_topic_modeller yet; " + "ignoring Pi-related flags." + ) + answers_pi = ( + getattr(args, "enable_pi", False) + or getattr(args, "enable_pi_express", False) + or getattr(args, "enable_pi_legacy", False) + ) + if answers_pi: + print( + "Note: Pi agent is not supported in llm_topic_modeller yet; " + "ignoring Pi-related flags." + ) + + +def apply_pi_cli_flags(args: argparse.Namespace, answers: InstallAnswers) -> None: + if getattr(args, "enable_pi", False): + preset = merge_preset(answers.profile, answers.custom_overrides) + if preset.get("USE_ECS_EXPRESS_MODE") == "True": + answers.enable_pi_express = True + else: + answers.enable_pi_legacy = True + answers.enable_service_connect = True + if getattr(args, "enable_pi_express", False): + answers.enable_pi_express = True + if getattr(args, "enable_pi_legacy", False): + answers.enable_pi_legacy = True + answers.enable_service_connect = True + if args.pi_alb_routing: + answers.pi_alb_routing = args.pi_alb_routing + if args.pi_path_prefix: + answers.pi_alb_path_prefix = args.pi_path_prefix + if args.pi_host_header: + answers.pi_alb_host_header = args.pi_host_header + if args.pi_listener_priority: + answers.pi_alb_listener_rule_priority = args.pi_listener_priority + if args.pi_gradio_port: + answers.pi_gradio_port = args.pi_gradio_port + if args.sc_discovery_name: + answers.sc_discovery_name = args.sc_discovery_name + if args.pi_provider: + answers.pi_default_provider = args.pi_provider + if args.skip_pi_agent_env: + answers.write_pi_agent_env = False + + +def configure_app_config_options( + answers: InstallAnswers, + args: argparse.Namespace, + *, + interactive: bool, +) -> None: + if getattr(args, "skip_app_config_env", False): + answers.write_app_config_env = False + return + + if getattr(args, "overwrite_app_config_env", False): + answers.write_app_config_env = True + answers.overwrite_app_config_env = True + return + + if not interactive: + answers.write_app_config_env = True + return + + example_label = ( + APP_CONFIG_ENV_EXAMPLE.name if APP_CONFIG_ENV_EXAMPLE.is_file() else "defaults" + ) + if ask_yes_no( + f"Write/update {APP_CONFIG_ENV_PATH.name} for the deployed app " + f"(from {example_label} + AWS resource names)?", + default=True, + ): + answers.write_app_config_env = True + if APP_CONFIG_ENV_PATH.is_file(): + answers.overwrite_app_config_env = ask_yes_no( + f"{APP_CONFIG_ENV_PATH.name} exists — replace file from example template?", + default=False, + ) + else: + answers.write_app_config_env = False + + +def configure_pi_options( + answers: InstallAnswers, + args: argparse.Namespace, + *, + interactive: bool, + assume_yes: bool, +) -> None: + preset = merge_preset(answers.profile, answers.custom_overrides) + use_express = preset.get("USE_ECS_EXPRESS_MODE") == "True" + use_cloudfront = preset.get("USE_CLOUDFRONT") == "True" + + apply_pi_cli_flags(args, answers) + + if not answers.pi_enabled and interactive: + if use_express: + label = ( + "Deploy agent mode (second Gradio app on Express, dedicated HTTPS URL)?" + ) + answers.enable_pi_express = ask_yes_no(label, default=False) + else: + label = "Deploy agent mode (second Gradio app on legacy Fargate + Service Connect)?" + if ask_yes_no(label, default=False): + answers.enable_pi_legacy = True + answers.enable_service_connect = True + + if not answers.pi_enabled: + return + + if not answers.pi_alb_listener_rule_priority and not answers.enable_pi_express: + answers.pi_alb_listener_rule_priority = default_pi_listener_priority( + use_cloudfront + ) + + if answers.enable_pi_express: + if interactive and not assume_yes: + print("\n--- Agent mode (ECS Express) ---") + answers.sc_discovery_name = ask( + "Service Connect discovery name for main app", + answers.sc_discovery_name, + ) + answers.pi_gradio_port = ask("Agent Gradio port", answers.pi_gradio_port) + if ask_yes_no( + f"Write/update {PI_AGENT_ENV_PATH.name} for AWS ECS (DOC_SUMMARISATION_GRADIO_URL, etc.)?", + default=True, + ): + answers.write_pi_agent_env = True + if PI_AGENT_ENV_PATH.is_file(): + answers.overwrite_pi_agent_env = ask_yes_no( + "pi_agent.env exists — replace file from example template?", + default=False, + ) + else: + answers.write_pi_agent_env = False + print( + f"Agent mode: Express (dedicated HTTPS endpoint per service); " + f"Service Connect discovery={answers.sc_discovery_name}" + ) + return + + if interactive and not ( + args.pi_alb_routing or args.pi_path_prefix or args.pi_host_header or assume_yes + ): + print("\n--- Agent mode ALB routing ---") + ridx = ask_choice( + "How should the shared ALB route traffic to the Agent UI?", + [ + "Path prefix (default /agent/ — e.g. https://host/agent/)", + "Dedicated hostname (PI_ALB_HOST_HEADER)", + "Both path prefix and hostname", + ], + default_index=0, + ) + answers.pi_alb_routing = PI_ALB_ROUTING_MODES[ridx] + + if answers.pi_alb_routing in ("path", "both"): + answers.pi_alb_path_prefix = ask( + "Agent path prefix", + answers.pi_alb_path_prefix, + ) + if answers.pi_alb_routing in ("host", "both"): + default_host = "" + if answers.ssl_domain: + default_host = f"agent.{answers.ssl_domain}" + answers.pi_alb_host_header = ask( + "Agent ALB host header (DNS CNAME to CloudFront/ALB)", + default_host, + ) + + if use_cloudfront: + default_pri = default_pi_listener_priority(True) + answers.pi_alb_listener_rule_priority = ask( + "Agent ALB listener rule priority (default 3; priorities 1–2 reserved)", + default_pri, + ) + + answers.sc_discovery_name = ask( + "Service Connect discovery name for main app", + answers.sc_discovery_name, + ) + answers.pi_gradio_port = ask("Agent Gradio port", answers.pi_gradio_port) + + if ask_yes_no( + f"Write/update {PI_AGENT_ENV_PATH.name} for AWS ECS (DOC_SUMMARISATION_GRADIO_URL, etc.)?", + default=True, + ): + answers.write_pi_agent_env = True + if PI_AGENT_ENV_PATH.is_file(): + answers.overwrite_pi_agent_env = ask_yes_no( + "pi_agent.env exists — replace file from example template?", + default=False, + ) + else: + answers.write_pi_agent_env = False + + print( + f"Agent mode: " + f"{'Express' if answers.enable_pi_express else 'legacy Fargate'}, " + f"routing={answers.pi_alb_routing}, " + f"prefix={normalize_pi_path_prefix(answers.pi_alb_path_prefix)}" + ) + + +# --------------------------------------------------------------------------- +# Interactive wizard +# --------------------------------------------------------------------------- + + +def run_wizard(args: argparse.Namespace) -> InstallAnswers: + answers = InstallAnswers() + interactive = not args.yes + assume_yes = args.yes + + if args.profile: + answers.profile = args.profile + elif interactive: + idx = ask_choice( + "Deployment profile", + [ + "Demonstration (Express, no CloudFront, no delete protection)", + "Production (ACM cert, CloudFront, delete protection)", + "Headless batch only (demo-style networking, no Express web UI)", + "Custom (configure individual toggles)", + ], + ) + answers.profile = ("demo", "production", "headless", "custom")[idx] + else: + answers.profile = "demo" + + if answers.profile == "headless": + answers.enable_headless = True + elif getattr(args, "headless", False): + if answers.profile == "demo": + raise SystemExit( + "Headless batch mode is not available with --profile demo (Express). " + "Use --profile headless, --profile production --headless, " + "or --profile custom without ECS Express." + ) + answers.enable_headless = True + elif interactive and answers.profile == "production": + answers.enable_headless = ask_yes_no( + "Enable headless batch-only deployment (S3 → Lambda → one-shot ECS, " + "no always-on web UI)?", + default=False, + ) + + if answers.profile == "custom" and interactive: + answers.custom_overrides["USE_ECS_EXPRESS_MODE"] = ( + "True" if ask_yes_no("Use ECS Express Mode?", False) else "False" + ) + answers.custom_overrides["USE_CLOUDFRONT"] = ( + "True" if ask_yes_no("Use CloudFront?", True) else "False" + ) + answers.custom_overrides["RUN_USEAST_STACK"] = ( + "True" + if answers.custom_overrides.get("USE_CLOUDFRONT") == "True" + and ask_yes_no( + "Deploy us-east-1 CloudFront stack (RUN_USEAST_STACK)?", True + ) + else "False" + ) + answers.custom_overrides["ENABLE_RESOURCE_DELETE_PROTECTION"] = ( + "True" if ask_yes_no("Enable delete protection?", True) else "False" + ) + answers.custom_overrides["ENABLE_APPREGISTRY"] = ( + "True" if ask_yes_no("Enable AppRegistry?", True) else "False" + ) + if getattr(args, "headless", False): + if answers.custom_overrides.get("USE_ECS_EXPRESS_MODE") == "True": + raise SystemExit( + "Headless batch mode requires USE_ECS_EXPRESS_MODE=False in " + "the custom profile." + ) + answers.enable_headless = True + elif profile_allows_headless_add_on(answers): + answers.enable_headless = ask_yes_no( + "Enable headless batch-only deployment (S3 → Lambda → one-shot ECS, " + "no always-on web UI)?", + default=False, + ) + + headless_err = headless_profile_error(answers) + if headless_err: + raise SystemExit(headless_err) + + try: + account, region = get_aws_identity(args.region) + except Exception as exc: + raise SystemExit(f"AWS credentials error: {exc}") from exc + + answers.aws_account_id = args.account or account + answers.aws_region = args.region or region + + if interactive and not args.account: + answers.aws_account_id = ask("AWS account ID", answers.aws_account_id) + if interactive and not args.region: + answers.aws_region = ask("AWS region", answers.aws_region) + + default_prefix = f"{answers.profile.capitalize()}-Summarisation-" + answers.cdk_prefix = args.cdk_prefix or ( + ask("CDK resource prefix (e.g. MyOrg-Summarisation-)", default_prefix) + if interactive + else default_prefix + ) + if not answers.cdk_prefix.endswith("-"): + answers.cdk_prefix += "-" + + if answers_use_headless(answers): + answers.cognito_domain_prefix = "" + print("Headless batch mode: skipping Cognito hosted UI domain (no web login).") + + resolve_globally_unique_install_names( + answers, + interactive=interactive, + assume_yes=assume_yes, + args=args, + ) + + # VPC + if args.new_vpc_cidr: + answers.vpc_mode = "new" + answers.new_vpc_cidr = args.new_vpc_cidr + elif args.vpc_name: + answers.vpc_mode = "existing" + answers.vpc_name = args.vpc_name + elif interactive: + vpc_idx = ask_choice( + "VPC", + ["Create new VPC", "Use existing VPC"], + ) + answers.vpc_mode = "new" if vpc_idx == 0 else "existing" + else: + answers.vpc_mode = "existing" + + if answers.vpc_mode == "new": + prompt_new_vpc_cidr(answers, interactive=interactive) + else: + if not answers.vpc_name: + vpcs = list_vpcs(answers.aws_region) + if not vpcs: + raise SystemExit("No VPCs found in region.") + if interactive: + labels = [f"{v['name']} ({v['cidr']})" for v in vpcs] + vidx = ask_choice("Select VPC", labels) + answers.vpc_name = vpcs[vidx]["name"] + else: + raise SystemExit( + "--vpc-name required for existing VPC in non-interactive mode." + ) + vpc_list = list_vpcs(answers.aws_region) + vpc_id = next( + (v["id"] for v in vpc_list if v["name"] == answers.vpc_name), + None, + ) + + # Subnets (public only for Express or demo-style headless) + apply_subnet_cli_flags(args, answers) + public_subnets_only = answers_use_public_subnets_only(answers) + if public_subnets_only: + if interactive: + answers.public_subnet_mode = ask_subnet_tier_mode("Public") + else: + answers.public_subnet_mode, _ = resolve_subnet_tier_modes(args) + answers.private_subnet_mode = "auto" + answers.private_subnet_names = [] + answers.private_subnet_cidrs = [] + answers.private_subnet_azs = [] + elif interactive: + layout_idx = ask_choice( + "Subnets in existing VPC", + [ + "Auto-discover public and private subnets", + "Use existing named subnets (both tiers)", + "Create new stack-specific subnets (both tiers)", + "Configure public and private separately", + ], + default_index=0, + ) + if layout_idx < 3: + shared_mode = SUBNET_TIER_MODES[layout_idx] + answers.public_subnet_mode = shared_mode + answers.private_subnet_mode = shared_mode + else: + answers.public_subnet_mode = ask_subnet_tier_mode("Public") + answers.private_subnet_mode = ask_subnet_tier_mode("Private") + else: + answers.public_subnet_mode, answers.private_subnet_mode = ( + resolve_subnet_tier_modes(args) + ) + + azs = list_availability_zones(answers.aws_region)[:3] + vpc_cidr = next( + (v["cidr"] for v in vpc_list if v["name"] == answers.vpc_name), + "", + ) + if vpc_id: + subnets = list_subnets_in_vpc(vpc_id, answers.aws_region) + configure_subnet_tier( + answers, + "public", + answers.public_subnet_mode, + subnets, + azs, + interactive=interactive, + vpc_cidr=vpc_cidr, + ) + if not public_subnets_only: + reserved_public = ( + list(answers.public_subnet_cidrs) + if answers.public_subnet_mode == "create" + else [] + ) + configure_subnet_tier( + answers, + "private", + answers.private_subnet_mode, + subnets, + azs, + interactive=interactive, + vpc_cidr=vpc_cidr, + reserved_subnet_cidrs=reserved_public, + ) + + # Optional infra reuse + if ( + vpc_id + and interactive + and ask_yes_no("Configure optional existing infra (IGW/ALB)?", False) + ): + igws = list_igws_for_vpc(vpc_id, answers.aws_region) + if igws: + # igws_ids = [g["id"] for g in igws] + print("Internet gateways:", ", ".join(g["id"] for g in igws)) + answers.existing_igw_id = ask("EXISTING_IGW_ID (blank=new)", "") + albs = list_albs_in_vpc(vpc_id, answers.aws_region) + if albs: + labels = [f"{a['name']} ({a['dns']})" for a in albs] + if ask_yes_no("Reuse an existing ALB?", False): + aidx = ask_choice("Select ALB", labels) + answers.existing_alb_arn = albs[aidx]["arn"] + answers.existing_alb_dns = albs[aidx]["dns"] + + # Production TLS + preset = merge_preset(answers_preset_profile(answers), answers.custom_overrides) + use_express = preset.get("USE_ECS_EXPRESS_MODE") == "True" + if ( + preset.get("USE_CLOUDFRONT") == "True" + and preset.get("USE_ECS_EXPRESS_MODE") != "True" + and not answers_use_headless(answers) + ): + if args.cert_arn: + answers.acm_cert_arn = args.cert_arn + answers.ssl_domain = args.domain or "" + elif interactive: + certs = list_acm_certificates(answers.aws_region) + if not certs: + raise SystemExit(f"No ISSUED ACM certificates in {answers.aws_region}.") + labels = [f"{c['label']} — {c['arn']}" for c in certs] + cidx = ask_choice("Select ACM certificate", labels) + answers.acm_cert_arn = certs[cidx]["arn"] + answers.ssl_domain = ask( + "SSL certificate domain (Cognito/CloudFront URL)", + certs[cidx]["domain"], + ) + else: + raise SystemExit( + "--cert-arn and --domain required for production in non-interactive mode." + ) + if interactive: + answers.cloudfront_geo = ask( + "CloudFront geo restriction (e.g. GB, blank=none)", "" + ) + + if not answers_use_headless(answers): + warn_deprecated_pi_cli_flags(args) + + # Advanced add-ons (non-Pi agent) + is_express = use_express and not answers_use_headless(answers) + if answers_use_headless(answers) and interactive: + mem = ask("ECS task memory (MB)", answers.ecs_memory) + if mem: + answers.ecs_memory = mem + if getattr(args, "headless_output_notifications", False): + answers.enable_headless_output_notifications = True + elif getattr(args, "no_headless_output_notifications", False): + answers.enable_headless_output_notifications = False + elif ask_yes_no( + "Enable email notifications when new analysis outputs are uploaded to S3?", + default=True, + ): + answers.enable_headless_output_notifications = True + if answers.enable_headless_output_notifications: + default_email = getattr(args, "headless_notify_email", "") or "" + while True: + answers.headless_output_notify_email = ask( + "Notification email address (SNS subscription; confirm via AWS email)", + default_email, + ).strip() + email_error = validate_notify_email( + answers.headless_output_notify_email + ) + if not email_error: + break + print(email_error) + default_iam_user = ( + getattr(args, "headless_output_iam_user", "").strip() + or f"{answers.cdk_prefix}s3-output-reader" + ) + iam_user = ask( + "IAM user name for programmatic S3 output access", + default_iam_user, + ).strip() + answers.headless_output_iam_user_name = iam_user or default_iam_user + elif answers_use_headless(answers): + if getattr(args, "headless_output_notifications", False): + answers.enable_headless_output_notifications = True + answers.headless_output_notify_email = ( + getattr(args, "headless_notify_email", "") or "" + ).strip() + answers.headless_output_iam_user_name = ( + getattr(args, "headless_output_iam_user", "").strip() + or f"{answers.cdk_prefix}s3-output-reader" + ) + email_error = validate_notify_email(answers.headless_output_notify_email) + if email_error: + raise SystemExit(email_error) + elif interactive and ask_yes_no("Configure other optional add-ons?", False): + if not is_express: + if not answers.enable_service_connect: + answers.enable_service_connect = ask_yes_no( + "Enable ECS Service Connect (without agent mode)?", False + ) + answers.enable_s3_batch = ask_yes_no("Enable S3 batch ECS trigger?", False) + mem = ask("ECS task memory (MB)", answers.ecs_memory) + if mem: + answers.ecs_memory = mem + + prompt_dynamo_usage_log_export_options(answers, args, interactive=interactive) + + configure_app_config_options(answers, args, interactive=interactive) + + answers.python_path = args.python + return answers + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + description="Interactive CDK installer for llm_topic_modeller", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + p.add_argument( + "--profile", + choices=("demo", "production", "headless", "custom"), + help="Base profile; use --profile headless or production --headless for batch-only", + ) + p.add_argument( + "--yes", action="store_true", help="Accept defaults / skip confirmations" + ) + p.add_argument("--config-only", action="store_true", help="Write config files only") + p.add_argument( + "--deploy-only", + action="store_true", + help="Skip wizard; use existing cdk_config.env", + ) + p.add_argument( + "--synth-only", action="store_true", help="Run cdk synth only (no deploy)" + ) + p.add_argument("--skip-deploy", action="store_true", help="Skip cdk deploy") + p.add_argument( + "--skip-quickstart", + action="store_true", + help="Skip post_cdk_build_quickstart.py", + ) + p.add_argument( + "--skip-stack-check", + action="store_true", + help="Do not check for existing CloudFormation stacks at startup", + ) + p.add_argument( + "--force-delete-stacks", + action="store_true", + help="If llm_topic_modeller stacks already exist, delete them before continuing " + "(disables termination protection; implies consent in non-interactive mode)", + ) + p.add_argument( + "--skip-cdk-json", action="store_true", help="Do not update cdk.json" + ) + p.add_argument( + "--refresh-cdk-json", action="store_true", help="Force rewrite cdk.json app key" + ) + p.add_argument("--python", help="Python executable for cdk.json") + p.add_argument("--region", help="AWS region") + p.add_argument("--account", help="AWS account ID") + p.add_argument("--cdk-prefix", help="CDK resource prefix") + p.add_argument("--cognito-prefix", help="Cognito user pool domain prefix") + p.add_argument( + "--s3-log-bucket", + help="S3 log/config bucket name (globally unique; skips wizard suggestion)", + ) + p.add_argument( + "--s3-output-bucket", + help="S3 output bucket name (globally unique; skips wizard suggestion)", + ) + p.add_argument("--vpc-name", help="Existing VPC Name tag") + p.add_argument("--new-vpc-cidr", help="CIDR for new VPC") + p.add_argument( + "--subnet-mode", + choices=SUBNET_TIER_MODES, + help="Subnet mode for both public and private tiers (overridden by per-tier flags)", + ) + p.add_argument( + "--public-subnet-mode", + choices=SUBNET_TIER_MODES, + help="Public subnet tier: auto, existing, or create", + ) + p.add_argument( + "--private-subnet-mode", + choices=SUBNET_TIER_MODES, + help="Private subnet tier: auto, existing, or create", + ) + p.add_argument( + "--public-subnet-names", + help="Comma-separated existing public subnet names (with --public-subnet-mode existing)", + ) + p.add_argument( + "--private-subnet-names", + help="Comma-separated existing private subnet names (with --private-subnet-mode existing)", + ) + p.add_argument( + "--headless", + action="store_true", + help="Batch-only add-on (production/custom without Express only; not valid with demo)", + ) + p.add_argument( + "--headless-output-notifications", + action="store_true", + help="Enable S3 output PutRequests alarm, SNS email, and IAM reader user (headless)", + ) + p.add_argument( + "--no-headless-output-notifications", + action="store_true", + help="Disable headless output notifications in interactive headless installs", + ) + p.add_argument( + "--headless-notify-email", + help="Email for SNS notifications when headless outputs are uploaded", + ) + p.add_argument( + "--headless-output-iam-user", + help="IAM user name for programmatic download of headless S3 outputs", + ) + p.add_argument( + "--dynamo-usage-log-export", + action="store_true", + help="Enable scheduled DynamoDB usage log export Lambda (EventBridge)", + ) + p.add_argument( + "--no-dynamo-usage-log-export", + action="store_true", + help="Disable DynamoDB usage log export in interactive installs", + ) + p.add_argument( + "--dynamo-export-schedule-time", + metavar="HH:MM", + help="UTC schedule time for DynamoDB usage log export (default: 06:00)", + ) + p.add_argument( + "--dynamo-export-schedule-days", + choices=tuple(DYNAMO_EXPORT_SCHEDULE_DAY_PRESETS), + help="Days to run DynamoDB usage log export (default: daily)", + ) + p.add_argument( + "--dynamo-export-s3-key", + help="S3 object key for the exported CSV (default: reports/dynamodb-usage/...)", + ) + p.add_argument( + "--dynamo-export-date-attribute", + help="DynamoDB attribute used for date filtering (default: timestamp)", + ) + p.add_argument( + "--create-headless-output-access-key", + action="store_true", + help="After deploy, create IAM access key for the headless output reader user", + ) + p.add_argument( + "--skip-headless-output-access-key", + action="store_true", + help="Do not offer to create IAM access key after headless deploy", + ) + p.add_argument("--cert-arn", help="ACM certificate ARN (production)") + p.add_argument("--domain", help="SSL certificate domain (production)") + p.add_argument( + "--skip-app-config-env", + action="store_true", + help="Do not write cdk/config/app_config.env from app_config.env.example", + ) + p.add_argument( + "--overwrite-app-config-env", + action="store_true", + help="Replace existing config/app_config.env from app_config.env.example (non-interactive)", + ) + return p + + +def main(argv: Optional[Sequence[str]] = None) -> int: + argv = list(argv) if argv is not None else None + args = build_arg_parser().parse_args(argv) + args._argv = tuple(argv if argv is not None else sys.argv[1:]) + CONFIG_DIR.mkdir(parents=True, exist_ok=True) + + print("llm_topic_modeller CDK installer\n") + + # Prerequisites + try: + cdk_version = check_cdk_cli() + print(f"CDK CLI: {cdk_version}") + except SystemExit: + if not args.config_only: + raise + print("Warning: CDK CLI not found; config-only mode continuing.") + + if not args.config_only and not args.synth_only: + try: + _stack_account, stack_region = get_aws_identity(args.region) + handle_existing_stacks_at_start(args, stack_region) + except SystemExit: + raise + except Exception as exc: + print(f"Existing stack check skipped: {exc}") + + python_exe: Optional[Path] = None + if not args.skip_cdk_json: + python_exe = resolve_python_executable( + override=args.python, + interactive=not args.yes and not args.deploy_only, + assume_yes=args.yes or args.deploy_only, + ) + write_cdk_json( + python_exe, + force=args.refresh_cdk_json or not CDK_JSON_PATH.is_file(), + skip=args.skip_cdk_json, + ) + elif CDK_JSON_PATH.is_file(): + # Best-effort parse python from existing cdk.json + try: + data = json.loads(CDK_JSON_PATH.read_text(encoding="utf-8")) + app_cmd = data.get("app", "") + if app_cmd: + python_exe = Path(app_cmd.split()[0]) + except (json.JSONDecodeError, IndexError): + pass + + values: Dict[str, str] + if args.deploy_only: + if not ENV_PATH.is_file(): + raise SystemExit( + f"No config at {ENV_PATH}. Run without --deploy-only first." + ) + values = read_env_file(ENV_PATH) + pre_deploy_errors = validate_globally_unique_env_values( + values + ) + validate_new_vpc_cidr_env_values(values) + if pre_deploy_errors: + print("Pre-deploy configuration conflicts:") + for err in pre_deploy_errors: + print(f" - {err}") + return 1 + if not python_exe or not python_exe.is_file(): + python_exe = resolve_python_executable( + override=args.python, interactive=False, assume_yes=True + ) + else: + answers = run_wizard(args) + errors = validate_install_answers(answers) + errors.extend(validate_subnet_answers(answers)) + errors.extend(enrich_existing_subnet_details_from_aws(answers)) + values = build_env_values(answers) + errors.extend(validate_env_values(values)) + if errors: + print("Configuration errors:") + for err in errors: + print(f" - {err}") + return 1 + + if not python_exe: + python_exe = resolve_python_executable( + override=answers.python_path or args.python, + interactive=not args.yes, + assume_yes=args.yes, + ) + if not args.skip_cdk_json: + write_cdk_json(python_exe, force=args.refresh_cdk_json) + + print_summary(values, python_exe) + if not args.yes and not ask_yes_no("Write config/cdk_config.env?", True): + print("Aborted.") + return 0 + + write_env_file(ENV_PATH, values) + if answers.write_app_config_env: + write_app_config_env_file( + answers, + values, + overwrite=answers.overwrite_app_config_env, + ) + + apply_cdk_runtime_env(values) + + if args.config_only: + print("Config written. Exiting (--config-only).") + return 0 + + run_smoke_test_if_needed(python_exe, args) + + # Bootstrap prompt + account = values.get("AWS_ACCOUNT_ID", "") + region = values.get("AWS_REGION", "") + if account and region: + try: + if cdk_bootstrap_needed(account, region): + print(f"CDK bootstrap not found in {region}.") + if args.yes or ask_yes_no( + f"Run cdk bootstrap aws://{account}/{region}?", True + ): + run_cdk_command( + ["bootstrap", f"aws://{account}/{region}"], + check=True, + ) + except Exception as exc: + print(f"Bootstrap check skipped: {exc}") + + verify_node_for_jsii() + run_cdk_command(["synth"], check=True) + + if args.synth_only or args.skip_deploy: + print("Synth complete. Skipping deploy.") + return 0 + + if not args.yes and not ask_yes_no("Run cdk deploy --all?", True): + print("Deploy skipped.") + return 0 + + run_cdk_command(["deploy", "--all", "--require-approval", "broadening"], check=True) + + values = read_env_file(ENV_PATH) + apply_post_deploy_fixup(values, assume_yes=args.yes) + + run_qs = False + if args.skip_quickstart: + print("Skipping post-deploy quickstart.") + else: + is_headless = values.get("ENABLE_HEADLESS_DEPLOYMENT") == "True" + is_demo = values.get("USE_ECS_EXPRESS_MODE") == "True" + if args.yes: + run_qs = True + elif is_headless: + run_qs = ask_yes_no( + "Run post_cdk_build_quickstart.py (CodeBuild image only; no ECS service)?", + default=True, + ) + else: + run_qs = ask_yes_no( + "Run post_cdk_build_quickstart.py (CodeBuild + scale ECS)?", + default=is_demo, + ) + if run_qs: + if not python_exe: + python_exe = resolve_python_executable( + assume_yes=True, interactive=False + ) + run_quickstart(python_exe) + + values = read_env_file(ENV_PATH) + is_headless = values.get("ENABLE_HEADLESS_DEPLOYMENT") == "True" + if values.get("USE_ECS_EXPRESS_MODE") == "True" and not is_headless: + from cdk_post_deploy import print_express_mode_next_steps + + if args.skip_quickstart or not run_qs: + print_express_mode_next_steps(values) + return 0 + + if is_headless: + from cdk_post_deploy import ( + print_headless_deployment_next_steps, + print_headless_output_notification_steps, + provision_headless_output_reader_access_key, + ) + + if args.skip_quickstart or not run_qs: + print_headless_deployment_next_steps(values) + if values.get("ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS") == "True": + print_headless_output_notification_steps(values) + create_key = False + if args.skip_headless_output_access_key: + create_key = False + elif args.create_headless_output_access_key: + create_key = True + elif not args.yes and not args.deploy_only: + create_key = ask_yes_no( + "Create IAM access key for the headless output reader user now?", + default=True, + ) + elif args.yes: + create_key = True + if create_key: + provision_headless_output_reader_access_key(values) + return 0 + + print("\nDone. Next steps:") + if values.get("USE_CLOUDFRONT") == "True": + domain = values.get("SSL_CERTIFICATE_DOMAIN", "") + cf = values.get("CLOUDFRONT_DOMAIN", "") + if domain and cf and cf != "cloudfront_placeholder.net": + print(f" - Point DNS CNAME {domain} -> {cf}") + elif values.get("USE_ECS_EXPRESS_MODE") == "True": + ep = values.get("ECS_EXPRESS_COGNITO_REDIRECT_BASE", "") + if ep: + from cdk_config import normalize_https_redirect_url + + print(f" - Express endpoint: {normalize_https_redirect_url(ep)}") + if APP_CONFIG_ENV_PATH.is_file(): + print(f" - App runtime config: {APP_CONFIG_ENV_PATH} (uploaded by quickstart)") + print(f" - Config: {ENV_PATH}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cdk/cdk_post_deploy.py b/cdk/cdk_post_deploy.py new file mode 100644 index 0000000000000000000000000000000000000000..1c556eb1348026c82fe30a3b820c3f538967f239 --- /dev/null +++ b/cdk/cdk_post_deploy.py @@ -0,0 +1,1408 @@ +""" +Post-deploy helpers (boto3 only). + +Use this module from post_cdk_build_quickstart.py so you do not need Node.js or +aws-cdk-lib installed to start CodeBuild / ECS after deployment. +""" + +from __future__ import annotations + +import copy +import json +import os +import re +import time +from typing import Any, Dict, List, Optional, Union + +import boto3 +from cdk_config import ( + AWS_REGION, + CLOUDFRONT_DOMAIN, +) + +_TASK_DEF_REGISTER_KEYS = ( + "family", + "taskRoleArn", + "executionRoleArn", + "networkMode", + "containerDefinitions", + "volumes", + "placementConstraints", + "requiresCompatibilities", + "cpu", + "memory", + "pidMode", + "ipcMode", + "proxyConfiguration", + "inferenceAccelerators", + "ephemeralStorage", + "runtimePlatform", +) + +_CONTAINER_REGISTER_OMIT_KEYS = frozenset( + { + "containerArn", + "taskDefinitionArn", + "status", + "lastStatus", + "managedAgents", + "networkInterfaces", + "healthStatus", + "cpu", + "memory", + "gpu", + } +) + + +def start_codebuild_build(project_name: str, aws_region: str = AWS_REGION) -> None: + """Start an existing CodeBuild project build.""" + client = boto3.client("codebuild", region_name=aws_region) + + try: + print(f"Attempting to start build for project: {project_name}") + response = client.start_build(projectName=project_name) + build_id = response["build"]["id"] + print(f"Successfully started build with ID: {build_id}") + print(f"Build ARN: {response['build']['arn']}") + print( + f"https://{aws_region}.console.aws.amazon.com/codesuite/codebuild/projects/" + f"{project_name}/build/{build_id.split(':')[-1]}/detail" + ) + except client.exceptions.ResourceNotFoundException: + print(f"Error: Project '{project_name}' not found in region '{aws_region}'.") + except Exception as e: + print(f"An unexpected error occurred: {e}") + + +def upload_file_to_s3( + local_file_paths: Union[str, List[str]], + s3_key: str, + s3_bucket: str, + run_aws_functions: str = "1", + aws_region: str = AWS_REGION, +) -> str: + """Upload local file(s) to S3.""" + final_out_message: List[str] = [] + final_out_message_str = "" + + if run_aws_functions != "1": + return "App not set to run AWS functions" + + try: + if not (s3_bucket and local_file_paths): + return "At least one essential variable is empty, could not upload to S3" + + s3_client = boto3.client("s3", region_name=aws_region) + paths = ( + [local_file_paths] + if isinstance(local_file_paths, str) + else list(local_file_paths) + ) + + for file_path in paths: + try: + file_name = os.path.basename(file_path) + s3_key_full = s3_key + file_name + print("S3 key: ", s3_key_full) + s3_client.upload_file(file_path, s3_bucket, s3_key_full) + out_message = f"File {file_name} uploaded successfully!" + print(out_message) + except Exception as e: + out_message = f"Error uploading file(s): {e}" + print(out_message) + final_out_message.append(out_message) + + final_out_message_str = "\n".join(final_out_message) + except Exception as e: + final_out_message_str = "Could not upload files to S3 due to: " + str(e) + print(final_out_message_str) + + return final_out_message_str + + +def start_ecs_task( + cluster_name: str, + service_name: str, + aws_region: str = AWS_REGION, +) -> dict: + """Scale a legacy Fargate ECS service to one running task.""" + ecs_client = boto3.client("ecs", region_name=aws_region) + + try: + ecs_client.update_service( + cluster=cluster_name, service=service_name, desiredCount=1 + ) + return { + "statusCode": 200, + "body": ( + f"Service {service_name} in cluster {cluster_name} " + "has been updated to 1 task." + ), + } + except Exception as e: + return {"statusCode": 500, "body": f"Error updating service: {str(e)}"} + + +EXPRESS_GATEWAY_ACTIVE_SCALING_TARGET = { + "minTaskCount": 1, + "maxTaskCount": 1, + "autoScalingMetric": "AVERAGE_CPU", + "autoScalingTargetValue": 60, +} + + +def resolve_express_gateway_service_arn( + cluster_name: str, + service_name: str, + aws_region: str = AWS_REGION, +) -> str: + """Look up an Express gateway service ARN by cluster and service name.""" + ecs_client = boto3.client("ecs", region_name=aws_region) + paginator = ecs_client.get_paginator("list_services") + for page in paginator.paginate(cluster=cluster_name): + for arn in page.get("serviceArns", []): + if arn.rstrip("/").split("/")[-1] == service_name: + return arn + raise ValueError( + f"Express gateway service '{service_name}' not found in cluster " + f"'{cluster_name}'." + ) + + +def _task_definition_has_port_name( + task_definition: Dict[str, Any], port_name: str +) -> bool: + for container in task_definition.get("containerDefinitions", []): + for mapping in container.get("portMappings") or []: + if mapping.get("name") == port_name: + return True + return False + + +def _container_definitions_with_named_port( + container_definitions: List[Dict[str, Any]], + *, + port_name: str, + container_port: int, +) -> List[Dict[str, Any]]: + updated: List[Dict[str, Any]] = [] + has_matching_port = any( + mapping.get("containerPort") == container_port + for container in container_definitions + for mapping in container.get("portMappings") or [] + ) + for index, container in enumerate(container_definitions): + container = { + key: value + for key, value in container.items() + if key not in _CONTAINER_REGISTER_OMIT_KEYS + } + port_mappings = [ + dict(mapping) for mapping in container.get("portMappings") or [] + ] + matched = False + for mapping in port_mappings: + if mapping.get("containerPort") == container_port: + matched = True + mapping["name"] = port_name + mapping.setdefault("protocol", "tcp") + if not matched and not has_matching_port and index == 0: + port_mappings.append( + { + "name": port_name, + "containerPort": container_port, + "protocol": "tcp", + } + ) + container["portMappings"] = port_mappings + updated.append(container) + return updated + + +def resolve_service_task_definition_arn( + cluster_name: str, + service_name: str, + aws_region: str = AWS_REGION, +) -> str: + """ + Resolve the task definition ARN for a Fargate or Express gateway ECS service. + + Express gateway services omit ``taskDefinition`` on ``describe_services``; use the + active service revision from ``describe_express_gateway_service`` instead. + """ + ecs_client = boto3.client("ecs", region_name=aws_region) + services = ecs_client.describe_services( + cluster=cluster_name, services=[service_name] + ).get("services", []) + if services: + task_definition_arn = services[0].get("taskDefinition") + if task_definition_arn: + return task_definition_arn + service_arn = services[0].get("serviceArn") + else: + service_arn = None + + if not service_arn: + service_arn = resolve_express_gateway_service_arn( + cluster_name, service_name, aws_region + ) + + express = ecs_client.describe_express_gateway_service(serviceArn=service_arn) + active_configs = (express.get("service") or {}).get("activeConfigurations") or [] + if not active_configs: + raise ValueError( + f"Could not resolve task definition for service '{service_name}' in " + f"cluster '{cluster_name}' (no active Express gateway configuration)." + ) + revision_arn = active_configs[0].get("serviceRevisionArn") + if not revision_arn: + raise ValueError( + f"Could not resolve task definition for service '{service_name}' " + "(active Express configuration has no serviceRevisionArn)." + ) + revisions = ecs_client.describe_service_revisions( + serviceRevisionArns=[revision_arn] + ).get("serviceRevisions", []) + if not revisions: + raise ValueError( + f"Service revision '{revision_arn}' not found for service " + f"'{service_name}'." + ) + task_definition_arn = revisions[0].get("taskDefinition") + if not task_definition_arn: + raise ValueError( + f"Service revision '{revision_arn}' has no taskDefinition for service " + f"'{service_name}'." + ) + return task_definition_arn + + +def ensure_ecs_service_port_mapping_name( + cluster_name: str, + service_name: str, + port_name: str, + container_port: int, + aws_region: str = AWS_REGION, +) -> str: + """ + Service Connect requires a named portMapping in the task definition. + Express gateway services only set containerPort at create time. + """ + ecs_client = boto3.client("ecs", region_name=aws_region) + task_definition_arn = resolve_service_task_definition_arn( + cluster_name, service_name, aws_region + ) + task_definition = ecs_client.describe_task_definition( + taskDefinition=task_definition_arn + )["taskDefinition"] + if _task_definition_has_port_name(task_definition, port_name): + return task_definition_arn + + new_containers = _container_definitions_with_named_port( + task_definition["containerDefinitions"], + port_name=port_name, + container_port=container_port, + ) + register_kwargs = { + key: copy.deepcopy(task_definition[key]) + for key in _TASK_DEF_REGISTER_KEYS + if key in task_definition + } + register_kwargs["containerDefinitions"] = new_containers + if task_definition.get("tags"): + register_kwargs["tags"] = [ + {"key": tag["key"], "value": tag["value"]} + for tag in task_definition["tags"] + ] + + new_task_definition = ecs_client.register_task_definition(**register_kwargs)[ + "taskDefinition" + ] + new_arn = new_task_definition["taskDefinitionArn"] + ecs_client.update_service( + cluster=cluster_name, + service=service_name, + taskDefinition=new_arn, + forceNewDeployment=True, + ) + print( + f"Registered task definition {new_arn} with Service Connect port " + f"name {port_name!r} on container port {container_port}." + ) + return new_arn + + +def apply_ecs_service_connect( + cluster_name: str, + service_name: str, + service_connect_configuration: Dict[str, Any], + aws_region: str = AWS_REGION, +) -> None: + ecs_client = boto3.client("ecs", region_name=aws_region) + ecs_client.update_service( + cluster=cluster_name, + service=service_name, + serviceConnectConfiguration=service_connect_configuration, + forceNewDeployment=True, + ) + print(f"Applied Service Connect to {service_name} in cluster {cluster_name}.") + + +def configure_express_pi_service_connect( + cluster_name: str, + main_service_name: str, + pi_service_name: str, + namespace: str, + main_port_name: str, + discovery_name: str, + main_port: int, + aws_region: str = AWS_REGION, +) -> None: + """Enable Service Connect for Pi Express -> main Express (post image build).""" + ensure_ecs_service_port_mapping_name( + cluster_name, + main_service_name, + main_port_name, + main_port, + aws_region=aws_region, + ) + apply_ecs_service_connect( + cluster_name, + main_service_name, + { + "enabled": True, + "namespace": namespace, + "services": [ + { + "portName": main_port_name, + "discoveryName": discovery_name, + "clientAliases": [ + {"port": int(main_port), "dnsName": discovery_name} + ], + } + ], + }, + aws_region=aws_region, + ) + apply_ecs_service_connect( + cluster_name, + pi_service_name, + {"enabled": True, "namespace": namespace}, + aws_region=aws_region, + ) + + +def start_express_gateway_service( + cluster_name: str, + service_name: str, + aws_region: str = AWS_REGION, +) -> dict: + """Scale an ECS Express gateway service to one running task after image build.""" + ecs_client = boto3.client("ecs", region_name=aws_region) + + try: + service_arn = resolve_express_gateway_service_arn( + cluster_name, service_name, aws_region=aws_region + ) + ecs_client.update_express_gateway_service( + serviceArn=service_arn, + scalingTarget=EXPRESS_GATEWAY_ACTIVE_SCALING_TARGET, + ) + return { + "statusCode": 200, + "body": ( + f"Express service {service_name} in cluster {cluster_name} " + "has been updated to run 1 task." + ), + } + except Exception as e: + return { + "statusCode": 500, + "body": f"Error updating Express gateway service: {str(e)}", + } + + +_ALB_COGNITO_CALLBACK_SUFFIX = "/oauth2/idpresponse" + +# Fields preserved from describe_user_pool_client when updating CallbackURLs only. +_USER_POOL_CLIENT_UPDATE_PASSTHROUGH_KEYS = ( + "ClientName", + "RefreshTokenValidity", + "AccessTokenValidity", + "IdTokenValidity", + "TokenValidityUnits", + "ReadAttributes", + "WriteAttributes", + "ExplicitAuthFlows", + "SupportedIdentityProviders", + "DefaultRedirectURI", + "AllowedOAuthFlows", + "AllowedOAuthScopes", + "AllowedOAuthFlowsUserPoolClient", + "AnalyticsConfiguration", + "PreventUserExistenceErrors", + "EnableTokenRevocation", + "EnablePropagateAdditionalUserContextData", + "AuthSessionValidity", + "RefreshTokenRotation", +) + + +def cognito_https_callback_urls(redirect_base: str) -> List[str]: + """ + ALB authenticate-cognito requires the app URL and ``/oauth2/idpresponse``. + """ + base = (redirect_base or "").strip().rstrip("/") + if not base: + raise ValueError("redirect_base is required for Cognito callback URLs") + if not base.startswith("https://"): + base = f"https://{base.lstrip('/')}" + return [base, f"{base}{_ALB_COGNITO_CALLBACK_SUFFIX}"] + + +def cognito_callback_urls_match( + existing_callbacks: List[str], + desired_callbacks: List[str], +) -> bool: + return set(existing_callbacks or []) == set(desired_callbacks) + + +def get_user_pool_client_callback_urls( + user_pool_id: str, + client_id: str, + *, + aws_region: str = AWS_REGION, +) -> List[str]: + cognito_client = boto3.client("cognito-idp", region_name=aws_region) + existing = cognito_client.describe_user_pool_client( + UserPoolId=user_pool_id, + ClientId=client_id, + )["UserPoolClient"] + return list(existing.get("CallbackURLs") or []) + + +def cognito_alb_callbacks_need_update( + user_pool_id: str, + client_id: str, + redirect_base: str, + *, + aws_region: str = AWS_REGION, +) -> bool: + desired = cognito_https_callback_urls(redirect_base) + current = get_user_pool_client_callback_urls( + user_pool_id, client_id, aws_region=aws_region + ) + return not cognito_callback_urls_match(current, desired) + + +def update_user_pool_client_callback_urls( + user_pool_id: str, + client_id: str, + callback_urls: List[str], + *, + aws_region: str = AWS_REGION, +) -> None: + """ + Set Cognito app client callback URLs without a CDK redeploy. + + Merges existing client settings from ``describe_user_pool_client`` so OAuth + flows/scopes and token validity are not reset. + """ + cognito_client = boto3.client("cognito-idp", region_name=aws_region) + existing = cognito_client.describe_user_pool_client( + UserPoolId=user_pool_id, + ClientId=client_id, + )["UserPoolClient"] + + update_kwargs: Dict[str, Any] = { + "UserPoolId": user_pool_id, + "ClientId": client_id, + "CallbackURLs": callback_urls, + } + for key in _USER_POOL_CLIENT_UPDATE_PASSTHROUGH_KEYS: + value = existing.get(key) + if value is not None: + update_kwargs[key] = value + logout_urls = existing.get("LogoutURLs") + if logout_urls: + update_kwargs["LogoutURLs"] = logout_urls + + cognito_client.update_user_pool_client(**update_kwargs) + print("Updated Cognito app client callback URLs: " + ", ".join(callback_urls)) + + +def apply_cognito_alb_callback_fixup( + *, + user_pool_id: str, + client_id: str, + redirect_base: str, + aws_region: str = AWS_REGION, +) -> bool: + """ + Update Cognito callbacks when they differ from ``redirect_base``. + + Returns True if URLs were updated, False if already correct. + """ + desired = cognito_https_callback_urls(redirect_base) + cognito_client = boto3.client("cognito-idp", region_name=aws_region) + existing = cognito_client.describe_user_pool_client( + UserPoolId=user_pool_id, + ClientId=client_id, + )["UserPoolClient"] + current = existing.get("CallbackURLs") or [] + if cognito_callback_urls_match(current, desired): + print("Cognito app client callback URLs already match the target endpoint.") + return False + update_user_pool_client_callback_urls( + user_pool_id, + client_id, + desired, + aws_region=aws_region, + ) + return True + + +_TARGET_GROUP_REGISTER_EVENT = re.compile( + r"target-group (arn:aws:elasticloadbalancing:[^\s)]+)", + re.IGNORECASE, +) + +EXPRESS_TARGET_GROUP_RESOLVE_MAX_WAIT_SECONDS = 600 +EXPRESS_TARGET_GROUP_RESOLVE_POLL_INTERVAL_SECONDS = 15 + + +def target_group_arn_from_ecs_register_event(message: str) -> Optional[str]: + """Parse target group ARN from ECS ``registered N targets in (target-group ...)``.""" + if "registered" not in (message or "").lower(): + return None + match = _TARGET_GROUP_REGISTER_EVENT.search(message) + return match.group(1) if match else None + + +def _target_group_arn_from_service_events( + events: List[Dict[str, Any]], +) -> Optional[str]: + for event in events: + target_group_arn = target_group_arn_from_ecs_register_event( + event.get("message", "") + ) + if target_group_arn: + return target_group_arn + return None + + +def _task_private_ipv4_addresses( + ecs_client: Any, + cluster_name: str, + service_name: str, +) -> set[str]: + task_arns = ecs_client.list_tasks( + cluster=cluster_name, serviceName=service_name + ).get("taskArns", []) + if not task_arns: + return set() + tasks = ecs_client.describe_tasks(cluster=cluster_name, tasks=task_arns).get( + "tasks", [] + ) + addresses: set[str] = set() + for task in tasks: + for attachment in task.get("attachments", []): + for detail in attachment.get("details", []): + if detail.get("name") == "privateIPv4Address" and detail.get("value"): + addresses.add(detail["value"]) + return addresses + + +def _target_group_arn_from_task_ips( + elbv2_client: Any, + load_balancer_arn: str, + task_ips: set[str], +) -> Optional[str]: + if not task_ips: + return None + target_groups = elbv2_client.describe_target_groups( + LoadBalancerArn=load_balancer_arn + ).get("TargetGroups", []) + for target_group in target_groups: + health = elbv2_client.describe_target_health( + TargetGroupArn=target_group["TargetGroupArn"] + ).get("TargetHealthDescriptions", []) + for description in health: + target_id = (description.get("Target") or {}).get("Id") + if target_id in task_ips: + return target_group["TargetGroupArn"] + return None + + +def resolve_express_service_target_group_arn( + cluster_name: str, + service_name: str, + *, + aws_region: str = AWS_REGION, + load_balancer_arn: Optional[str] = None, + max_wait_seconds: int = EXPRESS_TARGET_GROUP_RESOLVE_MAX_WAIT_SECONDS, + poll_interval_seconds: int = EXPRESS_TARGET_GROUP_RESOLVE_POLL_INTERVAL_SECONDS, +) -> str: + """ + Target group where Express most recently registered tasks. + + After post-deploy scaling, this ARN can differ from the TG baked into the CDK + Cognito listener custom resource at deploy time. Polls until a registration + event appears or running tasks are visible in an Express ALB target group. + """ + ecs_client = boto3.client("ecs", region_name=aws_region) + elbv2_client = ( + boto3.client("elbv2", region_name=aws_region) if load_balancer_arn else None + ) + deadline = time.monotonic() + max_wait_seconds + attempt = 0 + while True: + attempt += 1 + services = ecs_client.describe_services( + cluster=cluster_name, services=[service_name] + ).get("services", []) + if not services: + raise ValueError( + f"ECS service '{service_name}' not found in cluster '{cluster_name}'." + ) + service = services[0] + target_group_arn = _target_group_arn_from_service_events( + service.get("events", []) + ) + if target_group_arn: + if attempt > 1: + print( + f"Resolved target group for '{service_name}' " + f"after {attempt} poll(s)." + ) + return target_group_arn + + if elbv2_client and load_balancer_arn: + task_ips = _task_private_ipv4_addresses( + ecs_client, cluster_name, service_name + ) + target_group_arn = _target_group_arn_from_task_ips( + elbv2_client, load_balancer_arn, task_ips + ) + if target_group_arn: + print( + f"Resolved target group for '{service_name}' from running task IPs." + ) + return target_group_arn + + if time.monotonic() >= deadline: + running = service.get("runningCount", 0) + desired = service.get("desiredCount", 0) + raise ValueError( + f"No target group registration event found for service " + f"'{service_name}' after {max_wait_seconds}s " + f"(running {running}/{desired} tasks). " + "Ensure CodeBuild finished and the service scaled to at least one task." + ) + + if attempt == 1: + print( + f"Waiting for target group registration for '{service_name}' " + f"(up to {max_wait_seconds}s)..." + ) + time.sleep(poll_interval_seconds) + + +def find_express_gateway_https_listener( + *, + aws_region: str = AWS_REGION, +) -> Dict[str, str]: + """Return Express-managed ALB HTTPS listener metadata.""" + elbv2 = boto3.client("elbv2", region_name=aws_region) + for load_balancer in elbv2.describe_load_balancers().get("LoadBalancers", []): + if not load_balancer["LoadBalancerName"].startswith("ecs-express-gateway-alb"): + continue + listeners = elbv2.describe_listeners( + LoadBalancerArn=load_balancer["LoadBalancerArn"] + ).get("Listeners", []) + https_listener = next( + (listener for listener in listeners if listener.get("Port") == 443), + None, + ) + if https_listener: + return { + "load_balancer_arn": load_balancer["LoadBalancerArn"], + "listener_arn": https_listener["ListenerArn"], + "dns_name": load_balancer["DNSName"], + } + raise ValueError( + "Express gateway ALB (ecs-express-gateway-alb-*) with HTTPS listener not found." + ) + + +def listener_actions_with_target_group( + existing_actions: List[Dict[str, Any]], + target_group_arn: str, +) -> List[Dict[str, Any]]: + """Copy listener/rule actions, replacing the forward target group ARN.""" + updated_actions: List[Dict[str, Any]] = [] + for action in sorted(existing_actions, key=lambda item: item.get("Order", 0)): + action_copy = copy.deepcopy(action) + if action_copy.get("Type") == "forward": + action_copy["TargetGroupArn"] = target_group_arn + forward_config = action_copy.setdefault("ForwardConfig", {}) + forward_config["TargetGroups"] = [ + {"TargetGroupArn": target_group_arn, "Weight": 1} + ] + updated_actions.append(action_copy) + return updated_actions + + +def apply_express_alb_listener_target_group_fixup( + *, + cluster_name: str, + main_service_name: str, + pi_service_name: Optional[str] = None, + pi_path_prefixes: Optional[List[str]] = None, + aws_region: str = AWS_REGION, +) -> bool: + """ + Point ALB Cognito listener actions at the target groups Express tasks use. + + Express creates fresh target groups when a service scales up after deploy; the + CDK custom resource may still forward authenticated traffic to an empty TG. + """ + ingress = find_express_gateway_https_listener(aws_region=aws_region) + load_balancer_arn = ingress["load_balancer_arn"] + main_target_group_arn = resolve_express_service_target_group_arn( + cluster_name, + main_service_name, + aws_region=aws_region, + load_balancer_arn=load_balancer_arn, + ) + pi_target_group_arn = None + if pi_service_name: + try: + pi_target_group_arn = resolve_express_service_target_group_arn( + cluster_name, + pi_service_name, + aws_region=aws_region, + load_balancer_arn=load_balancer_arn, + ) + except ValueError as exc: + print(f"Note: skipping Pi listener rule TG fixup: {exc}") + + elbv2 = boto3.client("elbv2", region_name=aws_region) + listener_arn = ingress["listener_arn"] + listener = elbv2.describe_listeners(ListenerArns=[listener_arn])["Listeners"][0] + current_default = listener.get("DefaultActions", []) + current_forward_arn = next( + ( + action.get("TargetGroupArn") + for action in current_default + if action.get("Type") == "forward" + ), + None, + ) + changed = current_forward_arn != main_target_group_arn + + if changed: + elbv2.modify_listener( + ListenerArn=listener_arn, + DefaultActions=listener_actions_with_target_group( + current_default, main_target_group_arn + ), + ) + print( + "Updated Express ALB default listener forward target group to " + f"{main_target_group_arn}." + ) + else: + print( + "Express ALB default listener already forwards to the active target group." + ) + + if pi_target_group_arn and pi_path_prefixes: + rules = elbv2.describe_rules(ListenerArn=listener_arn).get("Rules", []) + prefixes = {prefix.rstrip("/") for prefix in pi_path_prefixes} + for rule in rules: + if rule.get("IsDefault"): + continue + path_values = [] + for condition in rule.get("Conditions", []): + if condition.get("Field") == "path-pattern": + path_values.extend(condition.get("Values", [])) + if not prefixes.intersection({value.rstrip("/") for value in path_values}): + continue + current_actions = rule.get("Actions", []) + current_pi_forward = next( + ( + action.get("TargetGroupArn") + for action in current_actions + if action.get("Type") == "forward" + ), + None, + ) + if current_pi_forward == pi_target_group_arn: + continue + elbv2.modify_rule( + RuleArn=rule["RuleArn"], + Actions=listener_actions_with_target_group( + current_actions, pi_target_group_arn + ), + ) + print( + "Updated Pi ALB listener rule forward target group to " + f"{pi_target_group_arn}." + ) + changed = True + + return changed + + +def listener_rule_has_cognito_auth(actions: List[Dict[str, Any]]) -> bool: + return any(action.get("Type") == "authenticate-cognito" for action in actions) + + +def _listener_rule_is_cloudfront_bypass_without_cognito( + rule: Dict[str, Any], + *, + cloudfront_host_header: str, +) -> bool: + """True for legacy forward-only host-header rules matching the CloudFront domain.""" + if listener_rule_has_cognito_auth(rule.get("Actions", [])): + return False + host = (cloudfront_host_header or "").strip() + if not host or host == "cloudfront_placeholder.net": + return False + for condition in rule.get("Conditions") or []: + if condition.get("Field") != "host-header": + continue + values = (condition.get("HostHeaderConfig") or {}).get("Values") or [] + if host in values: + return True + return False + + +def remove_express_listener_rules_without_cognito( + *, + aws_region: str = AWS_REGION, + cloudfront_host_header: Optional[str] = None, +) -> bool: + """ + Delete legacy Express ALB rules that forward CloudFront host traffic without Cognito. + + Only removes host-header rules whose value matches ``CLOUDFRONT_DOMAIN`` (or the + supplied override). ECS Express-managed ``*.ecs.*.on.aws`` host rules are left + intact. + """ + cloudfront_host = (cloudfront_host_header or CLOUDFRONT_DOMAIN or "").strip() + ingress = find_express_gateway_https_listener(aws_region=aws_region) + elbv2 = boto3.client("elbv2", region_name=aws_region) + listener_arn = ingress["listener_arn"] + changed = False + for rule in elbv2.describe_rules(ListenerArn=listener_arn).get("Rules", []): + if rule.get("IsDefault"): + continue + if not _listener_rule_is_cloudfront_bypass_without_cognito( + rule, cloudfront_host_header=cloudfront_host + ): + continue + rule_arn = rule["RuleArn"] + elbv2.delete_rule(RuleArn=rule_arn) + print( + "Removed legacy Express CloudFront bypass ALB listener rule " + f"(host {cloudfront_host!r}): {rule_arn}" + ) + changed = True + return changed + + +def build_cognito_secret_payload( + user_pool_id: str, + client_id: str, + *, + aws_region: str = AWS_REGION, +) -> Dict[str, str]: + """Build Secrets Manager JSON for SUMMARISATION_* Cognito keys.""" + cognito_client = boto3.client("cognito-idp", region_name=aws_region) + client = cognito_client.describe_user_pool_client( + UserPoolId=user_pool_id, + ClientId=client_id, + )["UserPoolClient"] + client_secret = client.get("ClientSecret") or "" + return { + "SUMMARISATION_USER_POOL_ID": user_pool_id, + "SUMMARISATION_CLIENT_ID": client_id, + "SUMMARISATION_CLIENT_SECRET": client_secret, + } + + +def cognito_secret_payload_matches( + existing_secret_string: str, + desired_payload: Dict[str, str], +) -> bool: + try: + current = json.loads(existing_secret_string or "{}") + except json.JSONDecodeError: + return False + return all(current.get(key) == value for key, value in desired_payload.items()) + + +def apply_cognito_secret_fixup( + *, + secret_name: str, + user_pool_id: str, + client_id: str, + aws_region: str = AWS_REGION, + recycle_express_service: Optional[Dict[str, str]] = None, +) -> bool: + """ + Sync imported Secrets Manager JSON with the stack's Cognito pool and app client. + + Express tasks read ``AWS_USER_POOL_ID`` / ``AWS_CLIENT_*`` from this secret. + When the secret predates a redeploy, values can reference a deleted user pool. + """ + desired_payload = build_cognito_secret_payload( + user_pool_id, client_id, aws_region=aws_region + ) + secrets_client = boto3.client("secretsmanager", region_name=aws_region) + current = secrets_client.get_secret_value(SecretId=secret_name) + current_string = current.get("SecretString") or "" + if cognito_secret_payload_matches(current_string, desired_payload): + print( + "Cognito secret already matches Cognito user pool ID " + "and Cognito app client ID." + ) + return False + + secrets_client.put_secret_value( + SecretId=secret_name, + SecretString=json.dumps(desired_payload), + ) + print( + "Updated Cognito secret for Cognito user pool ID " "and Cognito app client ID." + ) + if recycle_express_service: + recycle_express_gateway_tasks( + recycle_express_service["cluster_name"], + recycle_express_service["service_name"], + aws_region=aws_region, + ) + return True + + +def recycle_express_gateway_tasks( + cluster_name: str, + service_name: str, + *, + aws_region: str = AWS_REGION, +) -> None: + """Stop running Express tasks so replacements pick up updated secrets/env.""" + ecs_client = boto3.client("ecs", region_name=aws_region) + task_arns = ecs_client.list_tasks( + cluster=cluster_name, + serviceName=service_name, + ).get("taskArns", []) + for task_arn in task_arns: + ecs_client.stop_task( + cluster=cluster_name, + task=task_arn, + reason="Recycle task after Cognito secret/config sync", + ) + if task_arns: + print( + f"Stopped {len(task_arns)} task(s) for {service_name} to pick up Cognito updates." + ) + + +def apply_express_disable_in_app_cognito_auth( + cluster_name: str, + service_name: str, + *, + aws_region: str = AWS_REGION, +) -> bool: + """ + Set ``COGNITO_AUTH=False`` on a running Express service revision. + + ALB ``authenticate-cognito`` already gates access; in-app Gradio login is redundant + and fails when Secrets Manager still references an old user pool. + """ + ecs_client = boto3.client("ecs", region_name=aws_region) + service_arn = resolve_express_gateway_service_arn( + cluster_name, service_name, aws_region=aws_region + ) + express = ecs_client.describe_express_gateway_service(serviceArn=service_arn)[ + "service" + ] + active_configs = express.get("activeConfigurations") or [] + if not active_configs: + raise ValueError( + f"No active configuration for Express service '{service_name}'." + ) + active = active_configs[0] + primary = copy.deepcopy(active.get("primaryContainer") or {}) + environment = { + item["name"]: item["value"] + for item in primary.get("environment") or [] + if item.get("name") + } + if environment.get("COGNITO_AUTH") == "False": + print(f"{service_name} already has COGNITO_AUTH=False.") + return False + environment["COGNITO_AUTH"] = "False" + primary["environment"] = [ + {"name": name, "value": value} for name, value in sorted(environment.items()) + ] + update_kwargs: Dict[str, Any] = { + "serviceArn": service_arn, + "primaryContainer": primary, + } + for key in ( + "executionRoleArn", + "taskRoleArn", + "cpu", + "memory", + "healthCheckPath", + "networkConfiguration", + ): + value = active.get(key) + if value is not None: + update_kwargs[key] = value + scaling = express.get("scalingTarget") or active.get("scalingTarget") + if scaling is not None: + update_kwargs["scalingTarget"] = scaling + ecs_client.update_express_gateway_service(**update_kwargs) + print(f"Set COGNITO_AUTH=False on Express service {service_name}.") + return True + + +def apply_cognito_secret_fixup_from_stack( + *, + stack_name: str, + secret_name: str, + cluster_name: str, + main_service_name: str, + aws_region: str = AWS_REGION, + recycle_tasks: bool = True, +) -> bool: + """Read Cognito outputs from CloudFormation and sync the app client secret.""" + cfn_client = boto3.client("cloudformation", region_name=aws_region) + stacks = cfn_client.describe_stacks(StackName=stack_name).get("Stacks", []) + outputs = { + item["OutputKey"]: item["OutputValue"] + for item in (stacks[0].get("Outputs") or []) + } + user_pool_id = outputs.get("CognitoPoolId") + client_id = outputs.get("CognitoAppClientId") + if not user_pool_id or not client_id: + raise ValueError( + f"Stack '{stack_name}' is missing CognitoPoolId or CognitoAppClientId outputs." + ) + return apply_cognito_secret_fixup( + secret_name=secret_name, + user_pool_id=user_pool_id, + client_id=client_id, + aws_region=aws_region, + recycle_express_service=( + {"cluster_name": cluster_name, "service_name": main_service_name} + if recycle_tasks + else None + ), + ) + + +def get_stack_output( + stack_name: str, + output_key: str, + region: str, +) -> Optional[str]: + """Return a CloudFormation stack output value, or None if missing.""" + from botocore.exceptions import ClientError + + cfn = boto3.client("cloudformation", region_name=region) + try: + response = cfn.describe_stacks(StackName=stack_name) + except ClientError: + return None + for stack in response.get("Stacks", []): + for output in stack.get("Outputs", []): + if output.get("OutputKey") == output_key: + return output.get("OutputValue") + return None + + +def print_express_mode_next_steps( + values: Dict[str, str], + *, + stack_name: str = "SummarisationStack", + region: Optional[str] = None, +) -> None: + """Print user-facing next steps after Express mode deploy + quickstart.""" + from cdk_config import normalize_https_redirect_url + from cdk_functions import format_express_pi_public_url + + aws_region = region or values.get("AWS_REGION") or AWS_REGION + + main_raw = (values.get("ECS_EXPRESS_COGNITO_REDIRECT_BASE") or "").strip() + if not main_raw: + main_raw = ( + get_stack_output(stack_name, "ExpressServiceEndpoint", aws_region) or "" + ) + main_url = normalize_https_redirect_url(main_raw) if main_raw else "" + + print("\nDone. Next steps:") + print(" - Wait 10 minutes for app deployment to finish.") + pi_express = values.get("ENABLE_PI_AGENT_EXPRESS_SERVICE") == "True" + if pi_express: + print( + "- Register a Cognito user in AWS Console, change password at the login page URL " + "(available in the Cognito AWS console -> App Clients -> Login pages -> View login page " + "and sign in at the Pi agent URL below. The main topic-modelling backend runs without " + "in-app login so the Pi agent can call it over Service Connect. You can disable " + "Cognito login by setting COGNITO_AUTH to False in the ECS task definition / " + "ECS service options." + ) + else: + print( + " - If you have enabled Cognito login in the app, register a new user " + "to your Cognito user pool in AWS Console and complete sign up with the app " + "client login. If you do not want Cognito login, then set COGNITO_AUTH to " + "False in the ECS task definition / ECS service options." + ) + if main_url: + print(f" - The main topic modelling app can be accessed at {main_url}") + else: + print( + " - The main topic modelling app URL: see ExpressServiceEndpoint stack output" + ) + + if values.get("ENABLE_PI_AGENT_EXPRESS_SERVICE") == "True": + pi_raw = get_stack_output(stack_name, "PiExpressEndpoint", aws_region) or "" + pi_url = ( + format_express_pi_public_url(normalize_https_redirect_url(pi_raw)) + if pi_raw + else "" + ) + if pi_url: + print(f" - The Pi agent app can be accessed at {pi_url}") + else: + print(" - The Pi agent app URL: see PiExpressEndpoint stack output") + + +def _s3_object_exists(s3_client, bucket: str, key: str) -> bool: + from botocore.exceptions import ClientError + + try: + s3_client.head_object(Bucket=bucket, Key=key) + return True + except ClientError as exc: + code = exc.response.get("Error", {}).get("Code", "") + if code in ("404", "NoSuchKey", "NotFound"): + return False + raise + + +def seed_headless_batch_s3_layout( + output_bucket: str, + *, + input_prefix: str = "input/", + env_prefix: str = "input/config/", + example_env_local_path: str, + example_env_basename: str = "example_headless_env_file.env", + aws_region: str = AWS_REGION, +) -> None: + """ + Ensure headless batch prefixes and example job .env exist on the output bucket. + + Idempotent: existing keys are left unchanged (safe to run from quickstart). + """ + if not output_bucket: + print("Skipping headless S3 layout seed: output bucket name is empty.") + return + + input_prefix_norm = ( + input_prefix if input_prefix.endswith("/") else f"{input_prefix}/" + ) + env_prefix_norm = env_prefix if env_prefix.endswith("/") else f"{env_prefix}/" + if not env_prefix_norm.startswith(input_prefix_norm): + env_prefix_norm = f"{input_prefix_norm}{env_prefix_norm.lstrip('/')}" + if not env_prefix_norm.endswith("/"): + env_prefix_norm += "/" + + s3_client = boto3.client("s3", region_name=aws_region) + markers = ( + input_prefix_norm, + env_prefix_norm, + ) + for key in markers: + if _s3_object_exists(s3_client, output_bucket, key): + print(f"S3 prefix marker already present: s3://{output_bucket}/{key}") + continue + s3_client.put_object(Bucket=output_bucket, Key=key, Body=b"") + print(f"Created S3 prefix marker: s3://{output_bucket}/{key}") + + example_key = f"{env_prefix_norm}{example_env_basename}" + if _s3_object_exists(s3_client, output_bucket, example_key): + print(f"Example job .env already present: s3://{output_bucket}/{example_key}") + return + + if not os.path.isfile(example_env_local_path): + print(f"Skipping example job .env upload: {example_env_local_path} not found.") + return + + with open(example_env_local_path, "rb") as handle: + s3_client.put_object(Bucket=output_bucket, Key=example_key, Body=handle.read()) + print(f"Uploaded example job .env to s3://{output_bucket}/{example_key}") + + +def print_headless_deployment_next_steps( + values: Dict[str, str], + *, + stack_name: str = "SummarisationStack", + region: Optional[str] = None, +) -> None: + """Print user-facing next steps after headless deploy + quickstart.""" + aws_region = region or values.get("AWS_REGION") or AWS_REGION + output_bucket = (values.get("S3_OUTPUT_BUCKET_NAME") or "").strip() + input_prefix = (values.get("S3_BATCH_INPUT_PREFIX") or "input/").strip("/") + config_prefix = (values.get("S3_BATCH_ENV_PREFIX") or "input/config/").strip("/") + + lambda_name = (values.get("S3_BATCH_LAMBDA_FUNCTION_NAME") or "").strip() + if not lambda_name: + lambda_name = ( + get_stack_output(stack_name, "BatchEcsTriggerLambdaName", aws_region) or "" + ).strip() + if not lambda_name: + prefix = (values.get("CDK_PREFIX") or "").strip() + lambda_name = f"{prefix}S3BatchEcsTrigger" if prefix else "S3BatchEcsTrigger" + + input_uri = ( + f"s3://{output_bucket}/{input_prefix}/" + if output_bucket + else f"/{input_prefix}/" + ) + config_uri = ( + f"s3://{output_bucket}/{config_prefix}/" + if output_bucket + else f"/{config_prefix}/" + ) + output_uri = ( + f"s3://{output_bucket}/output//" + if output_bucket + else "/output//" + ) + example_name = "example_headless_env_file.env" + + print("\nDone. Next steps:") + print( + f" - Upload a consultation spreadsheet (.xlsx) to {input_uri}. " + "You can use dummy_consultation_response.xlsx as a placeholder filename " + "in the example job .env below." + ) + print( + f" - Create a job .env file for submitting a task to {config_uri}. " + f"You can copy the example at cdk/config/headless_s3_seed/input/config/{example_name} " + f"(also seeded at {config_uri}{example_name}) and adjust DIRECT_MODE_INPUT_FILE " + "and DIRECT_MODE_TEXT_COLUMN. Further DIRECT_MODE_* variables are documented in " + "tools/config.py and cli_topics.py." + ) + print(f" - Upload your job .env file to {config_uri}") + print( + f" - AWS Lambda function {lambda_name} will start an ECS task to run " + "topic modelling according to your job .env (RUN_DIRECT_MODE=1)." + ) + print( + f" - Outputs are written to {output_uri} when SAVE_OUTPUTS_TO_S3=True " + "in your job .env (see the example file)." + ) + log_group = (values.get("ECS_LOG_GROUP_NAME") or "").strip() + if log_group: + print( + f" - Batch task logs: CloudWatch log group {log_group} " + "(streams appear once the container starts)." + ) + + +def print_headless_output_notification_steps(values: Dict[str, str]) -> None: + """Print follow-on steps for S3 output alarms, SNS, and IAM reader user.""" + email = (values.get("HEADLESS_OUTPUT_NOTIFY_EMAIL") or "").strip() + iam_user = (values.get("HEADLESS_OUTPUT_IAM_USER_NAME") or "").strip() + output_bucket = (values.get("S3_OUTPUT_BUCKET_NAME") or "").strip() + output_prefix = (values.get("HEADLESS_OUTPUT_S3_PREFIX") or "output/").strip() + alarm_name = (values.get("HEADLESS_OUTPUT_ALARM_NAME") or "").strip() + secret_name = (values.get("HEADLESS_OUTPUT_IAM_SECRET_NAME") or "").strip() + + print("\nHeadless output notifications:") + if email: + print( + f" - Confirm the SNS email subscription sent to {email} " + "(check spam; status shows Pending until confirmed)." + ) + if output_bucket: + print( + f" - CloudWatch alarm fires when objects are uploaded under " + f"s3://{output_bucket}/{output_prefix.rstrip('/')}/" + ) + if alarm_name: + print(f" - Alarm name: {alarm_name}") + if iam_user: + print( + f" - IAM user {iam_user} can list/get/put/delete objects in the " + f"output bucket for programmatic result download." + ) + if secret_name: + print( + " - If an access key was created, credentials are stored in " + "the relevantSecrets Manager secret." + ) + + +def provision_headless_output_reader_access_key( + values: Dict[str, str], + *, + region: Optional[str] = None, +) -> Optional[str]: + """ + Create an IAM access key for the headless output reader user. + + Stores JSON credentials in Secrets Manager when HEADLESS_OUTPUT_IAM_SECRET_NAME + is set in cdk_config.env. + """ + from botocore.exceptions import ClientError + + iam_user = (values.get("HEADLESS_OUTPUT_IAM_USER_NAME") or "").strip() + if not iam_user: + print("Warning: HEADLESS_OUTPUT_IAM_USER_NAME not set; skipping access key.") + return None + + aws_region = region or values.get("AWS_REGION") or AWS_REGION + iam_client = boto3.client("iam", region_name=aws_region) + secret_name = (values.get("HEADLESS_OUTPUT_IAM_SECRET_NAME") or "").strip() + + try: + response = iam_client.create_access_key(UserName=iam_user) + except ClientError as exc: + print(f"Could not create IAM access key for {iam_user}: {exc}") + return None + + access_key = response["AccessKey"] + payload = json.dumps( + { + "aws_access_key_id": access_key["AccessKeyId"], + "aws_secret_access_key": access_key["SecretAccessKey"], + } + ) + + if secret_name: + sm_client = boto3.client("secretsmanager", region_name=aws_region) + try: + sm_client.create_secret(Name=secret_name, SecretString=payload) + print("Stored output-reader credentials in Secrets Manager secret") + except ClientError as exc: + if exc.response["Error"]["Code"] == "ResourceExistsException": + sm_client.put_secret_value(SecretId=secret_name, SecretString=payload) + print("Updated output-reader credentials in Secrets Manager secret") + else: + print(f"Warning: could not store credentials in Secrets Manager: {exc}") + print("IAM access key created (view in IAM console if needed):") + else: + print("IAM access key created (view in IAM console if needed):") + + return access_key["AccessKeyId"] diff --git a/cdk/cdk_stack.py b/cdk/cdk_stack.py new file mode 100644 index 0000000000000000000000000000000000000000..479a20e8ef46bf39f49572d1f61a0db2bcc934c9 --- /dev/null +++ b/cdk/cdk_stack.py @@ -0,0 +1,3180 @@ +import json # You might still need json if loading task_definition.json +import os +from typing import Any, Dict, List + +from aws_cdk import ( + CfnOutput, # <-- Import CfnOutput directly + Duration, + SecretValue, + Stack, +) +from aws_cdk import aws_cloudfront as cloudfront +from aws_cdk import aws_cloudfront_origins as origins +from aws_cdk import aws_codebuild as codebuild +from aws_cdk import aws_cognito as cognito +from aws_cdk import aws_dynamodb as dynamodb # Import the DynamoDB module +from aws_cdk import aws_ec2 as ec2 +from aws_cdk import aws_ecr as ecr +from aws_cdk import aws_ecs as ecs +from aws_cdk import aws_elasticloadbalancingv2 as elbv2 +from aws_cdk import aws_iam as iam +from aws_cdk import aws_kms as kms +from aws_cdk import aws_logs as logs +from aws_cdk import aws_s3 as s3 +from aws_cdk import aws_secretsmanager as secretsmanager +from aws_cdk import aws_wafv2 as wafv2 +from cdk_cloudfront_headers import ( + create_secure_cloudfront_response_headers_policy, + resolve_cloudfront_csp_urls, +) +from cdk_config import ( + ACCESS_LOG_DYNAMODB_TABLE_NAME, + ACM_SSL_CERTIFICATE_ARN, + ALB_NAME, + ALB_NAME_SECURITY_GROUP_NAME, + ALB_TARGET_GROUP_NAME, + APP_CONFIG_ENV_BASENAME, + APP_CONFIG_ENV_FILE, + AWS_ACCOUNT_ID, + AWS_MANAGED_TASK_ROLES_LIST, + AWS_REGION, + CDK_FOLDER, + CDK_PREFIX, + CLOUDFRONT_DISTRIBUTION_NAME, + CLOUDFRONT_DOMAIN, + CLOUDFRONT_ENABLE_SECURE_RESPONSE_HEADERS, + CLOUDFRONT_GEO_RESTRICTION, + CLOUDFRONT_PREFIX_LIST_ID, + CLUSTER_NAME, + CODEBUILD_PI_PROJECT_NAME, + CODEBUILD_PROJECT_NAME, + CODEBUILD_ROLE_NAME, + COGNITO_ACCESS_TOKEN_VALIDITY, + COGNITO_ID_TOKEN_VALIDITY, + COGNITO_REDIRECTION_URL, + COGNITO_REFRESH_TOKEN_VALIDITY, + COGNITO_USER_POOL_CLIENT_NAME, + COGNITO_USER_POOL_CLIENT_SECRET_NAME, + COGNITO_USER_POOL_DOMAIN_PREFIX, + COGNITO_USER_POOL_LOGIN_URL, + COGNITO_USER_POOL_NAME, + CUSTOM_HEADER, + CUSTOM_HEADER_VALUE, + CUSTOM_KMS_KEY_NAME, + DAYS_TO_DISPLAY_WHOLE_DOCUMENT_JOBS, + DYNAMODB_USAGE_LOG_EXPORT_DATE_ATTRIBUTE, + DYNAMODB_USAGE_LOG_EXPORT_LAMBDA_NAME, + DYNAMODB_USAGE_LOG_EXPORT_OUTPUT_FILENAME, + DYNAMODB_USAGE_LOG_EXPORT_S3_KEY, + DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE, + ECR_CDK_REPO_NAME, + ECR_PI_REPO_NAME, + ECS_AVAILABILITY_ZONE_REBALANCING, + ECS_EXECUTION_ROLE_MANAGED_POLICIES, + ECS_EXECUTION_ROLE_POLICY_ARNS, + ECS_EXECUTION_ROLE_POLICY_FILES, + ECS_EXPRESS_HEALTH_CHECK_PATH, + ECS_EXPRESS_INFRASTRUCTURE_ROLE_NAME, + ECS_EXPRESS_SERVICE_NAME, + ECS_EXPRESS_USE_PUBLIC_SUBNETS, + ECS_LOG_GROUP_NAME, + ECS_PI_EXPRESS_HEALTH_CHECK_PATH, + ECS_PI_EXPRESS_SECURITY_GROUP_NAME, + ECS_PI_EXPRESS_SERVICE_NAME, + ECS_PI_LOG_GROUP_NAME, + ECS_PI_SECURITY_GROUP_NAME, + ECS_PI_SERVICE_NAME, + ECS_PI_TASK_CPU_SIZE, + ECS_PI_TASK_DEFINITION_NAME, + ECS_PI_TASK_MEMORY_SIZE, + ECS_READ_ONLY_FILE_SYSTEM, + ECS_SECURITY_GROUP_NAME, + ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_IDS_LIST, + ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES_TO_LOOKUP, + ECS_SERVICE_CONNECT_CLIENT_SG_NAME_SUFFIX, + ECS_SERVICE_CONNECT_DISCOVERY_NAME, + ECS_SERVICE_CONNECT_DNS_NAME, + ECS_SERVICE_CONNECT_NAMESPACE, + ECS_SERVICE_CONNECT_PORT_MAPPING_NAME, + ECS_SERVICE_NAME, + ECS_TASK_CPU_SIZE, + ECS_TASK_EXECUTION_ROLE_NAME, + ECS_TASK_MEMORY_SIZE, + ECS_TASK_ROLE_NAME, + ECS_USE_FARGATE_SPOT, + ENABLE_DYNAMODB_USAGE_LOG_EXPORT, + ENABLE_ECS_SERVICE_CONNECT, + ENABLE_ECS_VPC_INTERFACE_ENDPOINTS, + ENABLE_HEADLESS_DEPLOYMENT, + ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS, + ENABLE_PI_AGENT_ECS_SERVICE, + ENABLE_PI_AGENT_EXPRESS_SERVICE, + ENABLE_S3_BATCH_ECS_TRIGGER, + EXISTING_IGW_ID, + EXISTING_LOAD_BALANCER_ARN, + EXISTING_LOAD_BALANCER_DNS, + FARGATE_TASK_DEFINITION_NAME, + FEEDBACK_LOG_DYNAMODB_TABLE_NAME, + GITHUB_REPO_BRANCH, + GITHUB_REPO_NAME, + GITHUB_REPO_USERNAME, + GRADIO_SERVER_PORT, + HEADLESS_OUTPUT_ALARM_NAME, + HEADLESS_OUTPUT_IAM_USER_NAME, + HEADLESS_OUTPUT_NOTIFY_EMAIL, + HEADLESS_OUTPUT_S3_METRIC_FILTER_ID, + HEADLESS_OUTPUT_S3_PREFIX, + HEADLESS_OUTPUT_SNS_TOPIC_NAME, + LOAD_BALANCER_WEB_ACL_NAME, + NAT_GATEWAY_NAME, + NEW_VPC_CIDR, + NEW_VPC_DEFAULT_NAME, + PI_AGENT_ENV_S3_KEY, + PI_ALB_HOST_HEADER, + PI_ALB_LISTENER_RULE_PRIORITY, + PI_ALB_PATH_PREFIX_NORMALIZED, + PI_ALB_ROUTING, + PI_ALB_TARGET_GROUP_NAME, + PI_GRADIO_PORT, + POLICY_FILE_ARNS, + POLICY_FILE_LOCATIONS, + PRIVATE_SUBNET_AVAILABILITY_ZONES, + PRIVATE_SUBNET_CIDR_BLOCKS, + PRIVATE_SUBNETS_TO_USE, + PUBLIC_SUBNET_AVAILABILITY_ZONES, + PUBLIC_SUBNET_CIDR_BLOCKS, + PUBLIC_SUBNETS_TO_USE, + S3_BATCH_CONFIG_PREFIX, + S3_BATCH_DEFAULT_PARAMS_KEY, + S3_BATCH_ENV_PREFIX, + S3_BATCH_ENV_SUFFIX, + S3_BATCH_GENERAL_ENV_PREFIX, + S3_BATCH_INPUT_PREFIX, + S3_BATCH_LAMBDA_FUNCTION_NAME, + S3_LOG_CONFIG_BUCKET_NAME, + S3_OUTPUT_BUCKET_NAME, + SAVE_LOGS_TO_DYNAMODB, + SINGLE_NAT_GATEWAY_ID, + SSL_CERTIFICATE_DOMAIN, + TASK_DEFINITION_FILE_LOCATION, + USAGE_LOG_DYNAMODB_TABLE_NAME, + USE_CLOUDFRONT, + USE_CUSTOM_KMS_KEY, + USE_ECS_EXPRESS_MODE, + VPC_NAME, + WEB_ACL_NAME, +) +from cdk_functions import ( # Only keep CDK-native functions + add_alb_https_listener_with_cert, + add_custom_policies, + add_s3_enforce_ssl_policy, + allow_express_load_balancer_to_ecs_security_group, + attach_managed_policy_arns, + attach_pi_agent_to_shared_alb, + build_ecs_execution_role_kms_policy, + build_ecs_task_role_inline_policy, + build_express_gateway_primary_container, + build_express_pi_primary_container, + build_pi_express_container_environment, + configure_public_github_codebuild_source, + create_dynamo_usage_log_export_lambda, + create_ecs_express_infrastructure_role, + create_ecs_vpc_endpoints_for_private_subnets, + create_express_gateway_service, + create_headless_output_notifications, + create_headless_s3_batch_seed, + create_nat_gateway, + create_pi_agent_ecs_resources, + create_s3_batch_ecs_trigger_lambda, + create_subnets, + create_web_acl_with_common_rules, + default_secrets_manager_kms_key_arn, + ecr_empty_on_delete, + ecs_availability_zone_rebalancing, + express_ingress_first_load_balancer_security_group, + express_ingress_load_balancer_arn, + format_express_pi_public_url, + format_pi_public_urls, + load_app_config_env_for_express, + managed_resource_removal_policy, + pi_alb_root_path_for_container, + pi_listener_rule_count, + public_github_codebuild_source, + resolve_ecs_s3_gateway_subnet_selection, + resolve_ecs_vpc_endpoint_subnet_selection, + resolve_policy_file_paths, + resolve_service_connect_client_security_group_ids, + resource_deletion_protection_flag, + s3_auto_delete_objects_on_stack_destroy, + wire_public_subnet_internet_access, +) +from constructs import Construct + + +def _get_env_list(env_var_name: str) -> List[str]: + """Parses a comma-separated environment variable into a list of strings.""" + value = env_var_name[1:-1].strip().replace('"', "").replace("'", "") + if not value: + return [] + # Split by comma and filter out any empty strings that might result from extra commas + return [s.strip() for s in value.split(",") if s.strip()] + + +# 1. Try to load CIDR/AZs from environment variables +if PUBLIC_SUBNETS_TO_USE: + PUBLIC_SUBNETS_TO_USE = _get_env_list(PUBLIC_SUBNETS_TO_USE) +if PRIVATE_SUBNETS_TO_USE: + PRIVATE_SUBNETS_TO_USE = _get_env_list(PRIVATE_SUBNETS_TO_USE) + +if PUBLIC_SUBNET_CIDR_BLOCKS: + PUBLIC_SUBNET_CIDR_BLOCKS = _get_env_list("PUBLIC_SUBNET_CIDR_BLOCKS") +if PUBLIC_SUBNET_AVAILABILITY_ZONES: + PUBLIC_SUBNET_AVAILABILITY_ZONES = _get_env_list("PUBLIC_SUBNET_AVAILABILITY_ZONES") +if PRIVATE_SUBNET_CIDR_BLOCKS: + PRIVATE_SUBNET_CIDR_BLOCKS = _get_env_list("PRIVATE_SUBNET_CIDR_BLOCKS") +if PRIVATE_SUBNET_AVAILABILITY_ZONES: + PRIVATE_SUBNET_AVAILABILITY_ZONES = _get_env_list( + "PRIVATE_SUBNET_AVAILABILITY_ZONES" + ) + +# AWS_MANAGED_TASK_ROLES_LIST and POLICY_* lists are parsed in cdk_config.py. + + +class CdkStack(Stack): + + def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None: + super().__init__(scope, construct_id, **kwargs) + + # --- Helper to get context values --- + def get_context_bool(key: str, default: bool = False) -> bool: + value = self.node.try_get_context(key) + if value is None: + return default + if isinstance(value, bool): + return value + if isinstance(value, str): + return value.lower() in ("true", "1", "yes") + return bool(value) + + def get_context_str(key: str, default: str = None) -> str: + return self.node.try_get_context(key) or default + + def get_context_dict(key: str, default: dict = None) -> dict: + return self.node.try_get_context(key) or default + + def get_context_list_of_dicts(key: str) -> List[Dict[str, Any]]: + ctx_value = self.node.try_get_context(key) + if not isinstance(ctx_value, list): + print( + f"Warning: Context key '{key}' not found or not a list. Returning empty list." + ) + return [] + # Optional: Add validation that all items in the list are dicts + return ctx_value + + resource_removal_policy = managed_resource_removal_policy() + resource_delete_protection = resource_deletion_protection_flag() + s3_auto_delete_objects = s3_auto_delete_objects_on_stack_destroy() + + self.template_options.description = "Deployment of the llm_topic_modeller Gradio app for LLM-based topic modelling. Git repo: https://github.com/seanpedrick-case/llm_topic_modeller." + + use_express_ingress = ( + not ACM_SSL_CERTIFICATE_ARN and USE_ECS_EXPRESS_MODE == "True" + ) + enable_headless = ENABLE_HEADLESS_DEPLOYMENT == "True" + express_public_subnets_only = ( + use_express_ingress and ECS_EXPRESS_USE_PUBLIC_SUBNETS == "True" + ) or ( + enable_headless + and not use_express_ingress + and ECS_EXPRESS_USE_PUBLIC_SUBNETS == "True" + ) + deploy_web_ingress = not use_express_ingress and not enable_headless + enable_service_connect = ( + ENABLE_ECS_SERVICE_CONNECT == "True" and not use_express_ingress + ) + enable_pi_agent = ( + ENABLE_PI_AGENT_ECS_SERVICE == "True" and not use_express_ingress + ) + enable_pi_express = ( + ENABLE_PI_AGENT_EXPRESS_SERVICE == "True" and use_express_ingress + ) + enable_pi_build = enable_pi_agent or enable_pi_express + if enable_headless: + print( + "ENABLE_HEADLESS_DEPLOYMENT=True: S3 batch trigger + one-shot Fargate " + "tasks only (no ALB, CloudFront, or always-on ECS service)." + ) + elif use_express_ingress: + print( + "USE_ECS_EXPRESS_MODE=True: using ECS Express Mode for HTTPS ingress " + "(no manual ALB/Fargate service)." + ) + if express_public_subnets_only: + print( + "ECS_EXPRESS_USE_PUBLIC_SUBNETS=True: Express tasks and VPC " + "endpoints use public subnets only (no private subnet install)." + ) + elif enable_headless: + print( + "ENABLE_HEADLESS_DEPLOYMENT=True: batch Fargate tasks use " + "legacy private subnets (or public if configured)." + ) + service_connect_client_sg_ids: List[str] = [] + + if enable_service_connect: + if ( + not ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_IDS_LIST + and not ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES_TO_LOOKUP + and not enable_pi_agent + ): + raise ValueError( + "ENABLE_ECS_SERVICE_CONNECT=True requires at least one of " + "ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_IDS, " + "ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES, or " + "ECS_SERVICE_CONNECT_CLIENT_CDK_PREFIXES (other apps' CDK_PREFIX " + f"values, resolved to {{prefix}}{ECS_SERVICE_CONNECT_CLIENT_SG_NAME_SUFFIX} " + "in this VPC), unless ENABLE_PI_AGENT_ECS_SERVICE=True (Pi SG is wired in-stack)." + ) + service_connect_client_sg_ids = ( + resolve_service_connect_client_security_group_ids( + ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_IDS_LIST, + ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES_TO_LOOKUP, + get_context_str, + ) + ) + print( + "ENABLE_ECS_SERVICE_CONNECT=True: advertising Fargate service on " + f"Service Connect as {ECS_SERVICE_CONNECT_DISCOVERY_NAME}; " + f"client SGs: {', '.join(service_connect_client_sg_ids)}" + ) + + # --- VPC and Subnets (Assuming VPC is always lookup, Subnets are created/returned by create_subnets) --- + new_vpc_created = False + imported_vpc_cidr_block = None + imported_vpc_cidr_blocks: List[str] = [] + if VPC_NAME: + vpc_id = get_context_str("vpc_id") + if not vpc_id: + raise ValueError( + f"VPC '{VPC_NAME}' was not resolved during pre-check (missing " + "'vpc_id' in context). Re-run from the cdk/ directory so " + "precheck.context.json is generated." + ) + availability_zones = list( + dict.fromkeys( + (PUBLIC_SUBNET_AVAILABILITY_ZONES or []) + + (PRIVATE_SUBNET_AVAILABILITY_ZONES or []) + ) + ) + if not availability_zones: + raise ValueError( + "vpc_id is in context but no subnet availability zones are " + "configured. Set PUBLIC_SUBNET_AVAILABILITY_ZONES and/or " + "PRIVATE_SUBNET_AVAILABILITY_ZONES in cdk_config.env." + ) + vpc_cidr_block = get_context_str("vpc_cidr_block") + imported_vpc_cidr_block = vpc_cidr_block + imported_vpc_cidr_blocks = list( + self.node.try_get_context("vpc_cidr_blocks") or [] + ) + if ( + imported_vpc_cidr_block + and imported_vpc_cidr_block not in imported_vpc_cidr_blocks + ): + imported_vpc_cidr_blocks.insert(0, imported_vpc_cidr_block) + vpc_attrs = { + "vpc_id": vpc_id, + "availability_zones": availability_zones, + } + if vpc_cidr_block: + vpc_attrs["vpc_cidr_block"] = vpc_cidr_block + vpc = ec2.Vpc.from_vpc_attributes(self, "VPC", **vpc_attrs) + cidr_log = ( + ", ".join(imported_vpc_cidr_blocks) + if imported_vpc_cidr_blocks + else vpc_cidr_block + ) + print( + f"Using VPC from pre-check context: {vpc_id}" + + (f" (CIDR(s) {cidr_log})" if cidr_log else "") + ) + + elif NEW_VPC_DEFAULT_NAME and not VPC_NAME: + new_vpc_created = True + print( + f"NEW_VPC_DEFAULT_NAME ('{NEW_VPC_DEFAULT_NAME}') is set. Creating a new VPC." + ) + + # Configuration for the new VPC + # You can make these configurable via context as well, e.g., + # new_vpc_cidr = self.node.try_get_context("new_vpc_cidr") or "10.0.0.0/24" + # new_vpc_max_azs = self.node.try_get_context("new_vpc_max_azs") or 2 # Use 2 AZs by default for HA + # new_vpc_nat_gateways = self.node.try_get_context("new_vpc_nat_gateways") or new_vpc_max_azs # One NAT GW per AZ for HA + # or 1 for cost savings if acceptable + if not NEW_VPC_CIDR: + raise Exception( + "App has been instructed to create a new VPC but not VPC CDR range provided to variable NEW_VPC_CIDR" + ) + + print("Provided NEW_VPC_CIDR range:", NEW_VPC_CIDR) + + new_vpc_cidr = NEW_VPC_CIDR + new_vpc_max_azs = 2 # Creates resources in 2 AZs. Adjust as needed. + + # For "a NAT gateway", you can set nat_gateways=1. + # For resilience (NAT GW per AZ), set nat_gateways=new_vpc_max_azs. + # The Vpc construct will create NAT Gateway(s) if subnet_type PRIVATE_WITH_EGRESS is used + # and nat_gateways > 0. + if express_public_subnets_only: + new_vpc_nat_gateways = 0 + new_vpc_subnet_configuration = [ + ec2.SubnetConfiguration( + name="Public", + subnet_type=ec2.SubnetType.PUBLIC, + # /27 (~27 usable IPs): Express managed ALB needs 8+ free IPs per + # subnet alongside VPC interface endpoints and task ENIs. + cidr_mask=27, + ), + ] + else: + new_vpc_nat_gateways = ( + 1 # Creates a single NAT Gateway for cost-effectiveness. + ) + new_vpc_subnet_configuration = [ + ec2.SubnetConfiguration( + name="Public", # Name prefix for public subnets + subnet_type=ec2.SubnetType.PUBLIC, + cidr_mask=26, + ), + ec2.SubnetConfiguration( + name="Private", # Name prefix for private subnets + subnet_type=ec2.SubnetType.PRIVATE_WITH_EGRESS, # Ensures these subnets have NAT Gateway access + cidr_mask=28, + ), + # You could also add ec2.SubnetType.PRIVATE_ISOLATED if needed + ] + # If you need one NAT GW per AZ for higher availability, set nat_gateways to new_vpc_max_azs. + + vpc = ec2.Vpc( + self, + "MyNewLogicalVpc", # This is the CDK construct ID + vpc_name=NEW_VPC_DEFAULT_NAME, + ip_addresses=ec2.IpAddresses.cidr(new_vpc_cidr), + max_azs=new_vpc_max_azs, + nat_gateways=new_vpc_nat_gateways, # Number of NAT gateways to create + subnet_configuration=new_vpc_subnet_configuration, + # Internet Gateway is created and configured automatically for PUBLIC subnets. + # Route tables for public subnets will point to the IGW. + # Route tables for PRIVATE_WITH_EGRESS subnets will point to the NAT Gateway(s). + ) + print( + f"Successfully created new VPC: {vpc.vpc_id} with name '{NEW_VPC_DEFAULT_NAME}'" + ) + # If nat_gateways > 0, vpc.nat_gateway_ips will contain EIPs if Vpc created them. + # vpc.public_subnets, vpc.private_subnets, vpc.isolated_subnets are populated. + + else: + raise Exception( + "VPC_NAME for current VPC not found, and NEW_VPC_DEFAULT_NAME not found to create a new VPC" + ) + + # --- Subnet Handling (Check Context and Create/Import) --- + # Initialize lists to hold ISubnet objects (L2) and CfnSubnet/CfnRouteTable (L1) + # We will store ISubnet for consistency, as CfnSubnet has a .subnet_id property + self.public_subnets: List[ec2.ISubnet] = [] + self.private_subnets: List[ec2.ISubnet] = [] + # Store L1 CfnRouteTables explicitly if you need to reference them later + self.private_route_tables_cfn: List[ec2.CfnRouteTable] = [] + self.public_route_tables_cfn: List[ec2.CfnRouteTable] = ( + [] + ) # New: to store public RTs + + names_to_create_private = [] + names_to_create_public = [] + + if not PUBLIC_SUBNETS_TO_USE and not PRIVATE_SUBNETS_TO_USE: + if express_public_subnets_only: + print( + "Express public-subnet mode: auto-selecting public subnets only " + "(private subnets are not installed)." + ) + selected_public_subnets = vpc.select_subnets( + subnet_type=ec2.SubnetType.PUBLIC, one_per_az=True + ) + if len(selected_public_subnets.subnet_ids) < 2: + raise Exception( + "Express mode needs at least two public subnets in different " + "availability zones." + ) + self.public_subnets = selected_public_subnets.subnets + self.private_subnets = [] + print( + f"Selected {len(self.public_subnets)} public subnets for Express." + ) + else: + print( + "Warning: No public or private subnets specified in *_SUBNETS_TO_USE. Attempting to select from existing VPC subnets." + ) + + print("vpc.public_subnets:", vpc.public_subnets) + print("vpc.private_subnets:", vpc.private_subnets) + + if ( + vpc.public_subnets + ): # These are already one_per_az if max_azs was used and Vpc created them + self.public_subnets.extend(vpc.public_subnets) + else: + self.node.add_warning("No public subnets found in the VPC.") + + # Get private subnets with egress specifically + # selected_private_subnets_with_egress = vpc.select_subnets(subnet_type=ec2.SubnetType.PRIVATE_WITH_EGRESS) + + print( + f"Selected from VPC: {len(self.public_subnets)} public, {len(self.private_subnets)} private_with_egress subnets." + ) + + if ( + len(self.public_subnets) < 1 or len(self.private_subnets) < 1 + ): # Simplified check for new VPC + # If new_vpc_max_azs was 1, you'd have 1 of each. If 2, then 2 of each. + # The original check ' < 2' might be too strict if new_vpc_max_azs=1 + pass # For new VPC, allow single AZ setups if configured that way. The VPC construct ensures one per AZ up to max_azs. + + if not self.public_subnets and not self.private_subnets: + print( + "Error: No public or private subnets could be found in the VPC for automatic selection. " + "You must either specify subnets in *_SUBNETS_TO_USE or ensure the VPC has discoverable subnets." + ) + raise RuntimeError( + "No suitable subnets found for automatic selection." + ) + else: + print( + f"Automatically selected {len(self.public_subnets)} public and {len(self.private_subnets)} private subnets based on VPC properties." + ) + + selected_public_subnets = vpc.select_subnets( + subnet_type=ec2.SubnetType.PUBLIC, one_per_az=True + ) + private_subnets_egress = vpc.select_subnets( + subnet_type=ec2.SubnetType.PRIVATE_WITH_EGRESS, one_per_az=True + ) + + if private_subnets_egress.subnets: + self.private_subnets.extend(private_subnets_egress.subnets) + else: + self.node.add_warning( + "No PRIVATE_WITH_EGRESS subnets found in the VPC." + ) + + try: + private_subnets_isolated = vpc.select_subnets( + subnet_type=ec2.SubnetType.PRIVATE_ISOLATED, one_per_az=True + ) + except Exception as e: + private_subnets_isolated = [] + print("Could not find any isolated subnets due to:", e) + + ### + combined_subnet_objects = [] + + if private_subnets_isolated: + if private_subnets_egress.subnets: + # Add the first PRIVATE_WITH_EGRESS subnet + combined_subnet_objects.append( + private_subnets_egress.subnets[0] + ) + elif not private_subnets_isolated: + if private_subnets_egress.subnets: + # Add the first PRIVATE_WITH_EGRESS subnet + combined_subnet_objects.extend(private_subnets_egress.subnets) + else: + self.node.add_warning( + "No PRIVATE_WITH_EGRESS subnets found to select the first one." + ) + + # Add all PRIVATE_ISOLATED subnets *except* the first one (if they exist) + try: + if len(private_subnets_isolated.subnets) > 1: + combined_subnet_objects.extend( + private_subnets_isolated.subnets[1:] + ) + elif ( + private_subnets_isolated.subnets + ): # Only 1 isolated subnet, add a warning if [1:] was desired + self.node.add_warning( + "Only one PRIVATE_ISOLATED subnet found, private_subnets_isolated.subnets[1:] will be empty." + ) + else: + self.node.add_warning("No PRIVATE_ISOLATED subnets found.") + except Exception as e: + print("Could not identify private isolated subnets due to:", e) + + # Create an ec2.SelectedSubnets object from the combined private subnet list. + selected_private_subnets = vpc.select_subnets( + subnets=combined_subnet_objects + ) + + print("selected_public_subnets:", selected_public_subnets) + print("selected_private_subnets:", selected_private_subnets) + + if ( + len(selected_public_subnets.subnet_ids) < 2 + or len(selected_private_subnets.subnet_ids) < 2 + ): + raise Exception( + "Need at least two public or private subnets in different availability zones" + ) + + if not selected_public_subnets and not selected_private_subnets: + # If no subnets could be found even with automatic selection, raise an error. + # This ensures the stack doesn't proceed if it absolutely needs subnets. + print( + "Error: No existing public or private subnets could be found in the VPC for automatic selection. " + "You must either specify subnets in *_SUBNETS_TO_USE or ensure the VPC has discoverable subnets." + ) + raise RuntimeError( + "No suitable subnets found for automatic selection." + ) + else: + self.public_subnets = selected_public_subnets.subnets + self.private_subnets = selected_private_subnets.subnets + print( + f"Automatically selected {len(self.public_subnets)} public and {len(self.private_subnets)} private subnets based on VPC discovery." + ) + + print("self.public_subnets:", self.public_subnets) + print("self.private_subnets:", self.private_subnets) + # Since subnets are now assigned, we can exit this processing block. + # The rest of the original code (which iterates *_SUBNETS_TO_USE) will be skipped. + + checked_public_subnets_ctx = get_context_dict("checked_public_subnets") + checked_private_subnets_ctx = get_context_dict("checked_private_subnets") + + public_subnets_data_for_creation_ctx = get_context_list_of_dicts( + "public_subnets_to_create" + ) + private_subnets_data_for_creation_ctx = get_context_list_of_dicts( + "private_subnets_to_create" + ) + + # --- 3. Process Public Subnets --- + print("\n--- Processing Public Subnets ---") + public_internet_gateway_attachment = None + if not new_vpc_created: + resolved_igw_id = ( + get_context_str("internet_gateway_id") or EXISTING_IGW_ID or "" + ).strip() + if resolved_igw_id and ( + PUBLIC_SUBNETS_TO_USE + or public_subnets_data_for_creation_ctx + or get_context_list_of_dicts("public_subnets_needing_igw_route") + ): + public_internet_gateway_attachment = wire_public_subnet_internet_access( + self, + "PublicSubnetInternet", + vpc_id=vpc.vpc_id, + internet_gateway_id=resolved_igw_id, + needs_igw_vpc_attachment=get_context_bool( + "internet_gateway_needs_vpc_attachment", False + ), + subnets_needing_route=get_context_list_of_dicts( + "public_subnets_needing_igw_route" + ), + ) + + # Import existing public subnets + if checked_public_subnets_ctx: + for i, subnet_name in enumerate(PUBLIC_SUBNETS_TO_USE): + subnet_info = checked_public_subnets_ctx.get(subnet_name) + if subnet_info and subnet_info.get("exists"): + subnet_id = subnet_info.get("id") + if not subnet_id: + raise RuntimeError( + f"Context for existing public subnet '{subnet_name}' is missing 'id'." + ) + subnet_az = subnet_info.get("az") + if ( + not subnet_az + and PUBLIC_SUBNET_AVAILABILITY_ZONES + and i < len(PUBLIC_SUBNET_AVAILABILITY_ZONES) + ): + subnet_az = PUBLIC_SUBNET_AVAILABILITY_ZONES[i] + if not subnet_az: + raise RuntimeError( + f"Context for existing public subnet '{subnet_name}' is missing 'az'." + ) + subnet_attrs = { + "subnet_id": subnet_id, + "availability_zone": subnet_az, + } + route_table_id = subnet_info.get("route_table_id") + if route_table_id: + subnet_attrs["route_table_id"] = route_table_id + try: + imported_subnet = ec2.Subnet.from_subnet_attributes( + self, + f"ImportedPublicSubnet{subnet_name.replace('-', '')}{i}", + **subnet_attrs, + ) + self.public_subnets.append(imported_subnet) + print( + f"Imported existing public subnet: {subnet_name} (ID: {subnet_id})" + ) + except Exception as e: + raise RuntimeError( + f"Failed to import public subnet '{subnet_name}' with ID '{subnet_id}'. Error: {e}" + ) + + # Create new public subnets based on public_subnets_data_for_creation_ctx + if public_subnets_data_for_creation_ctx: + names_to_create_public = [ + s["name"] for s in public_subnets_data_for_creation_ctx + ] + cidrs_to_create_public = [ + s["cidr"] for s in public_subnets_data_for_creation_ctx + ] + azs_to_create_public = [ + s["az"] for s in public_subnets_data_for_creation_ctx + ] + + if names_to_create_public: + print( + f"Attempting to create {len(names_to_create_public)} new public subnets: {names_to_create_public}" + ) + igw_for_new_subnets = ( + get_context_str("internet_gateway_id") or EXISTING_IGW_ID + ) + newly_created_public_subnets, newly_created_public_rts_cfn = ( + create_subnets( + self, + vpc, + CDK_PREFIX, + names_to_create_public, + cidrs_to_create_public, + azs_to_create_public, + is_public=True, + internet_gateway_id=igw_for_new_subnets, + internet_gateway_attachment=public_internet_gateway_attachment, + ) + ) + self.public_subnets.extend(newly_created_public_subnets) + self.public_route_tables_cfn.extend(newly_created_public_rts_cfn) + + if ( + not self.public_subnets + and not names_to_create_public + and not PUBLIC_SUBNETS_TO_USE + ): + raise Exception("No public subnets found or created, exiting.") + + # --- NAT Gateway Creation/Lookup --- + self.single_nat_gateway_id = None + if express_public_subnets_only: + print( + "Express public-subnet mode: skipping NAT Gateway install " + "(not required for public Express tasks)." + ) + else: + print("Creating NAT gateway/located existing") + + nat_gw_id_from_context = SINGLE_NAT_GATEWAY_ID or get_context_str( + "id:NatGateway" + ) + + if nat_gw_id_from_context: + print( + f"Using existing NAT Gateway ID from context: {nat_gw_id_from_context}" + ) + self.single_nat_gateway_id = nat_gw_id_from_context + + elif ( + new_vpc_created + and new_vpc_nat_gateways > 0 + and hasattr(vpc, "nat_gateways") + and vpc.nat_gateways + ): + self.single_nat_gateway_id = vpc.nat_gateways[0].gateway_id + print( + f"Using NAT Gateway {self.single_nat_gateway_id} created by the new VPC construct." + ) + + if not self.single_nat_gateway_id: + print("Creating a new NAT gateway") + + if hasattr(vpc, "nat_gateways") and vpc.nat_gateways: + print("Existing NAT gateway found in vpc") + pass + + # If not in context, create a new one, but only if we have a public subnet. + elif self.public_subnets: + print("NAT Gateway ID not found in context. Creating a new one.") + # Place the NAT GW in the first available public subnet + first_public_subnet = self.public_subnets[0] + + self.single_nat_gateway_id = create_nat_gateway( + self, + first_public_subnet, + nat_gateway_name=NAT_GATEWAY_NAME, + nat_gateway_id_context_key=SINGLE_NAT_GATEWAY_ID, + ) + else: + print( + "WARNING: No public subnets available and NAT gateway not found in existing VPC. Cannot create a NAT Gateway." + ) + + # --- 4. Process Private Subnets --- + if express_public_subnets_only: + if PRIVATE_SUBNETS_TO_USE or private_subnets_data_for_creation_ctx: + print( + "Note: PRIVATE_* subnet settings are ignored in Express public-subnet mode." + ) + else: + print("\n--- Processing Private Subnets ---") + if checked_private_subnets_ctx: + for i, subnet_name in enumerate(PRIVATE_SUBNETS_TO_USE): + subnet_info = checked_private_subnets_ctx.get(subnet_name) + if subnet_info and subnet_info.get("exists"): + subnet_id = subnet_info.get("id") + if not subnet_id: + raise RuntimeError( + f"Context for existing private subnet '{subnet_name}' is missing 'id'." + ) + subnet_az = subnet_info.get("az") + if ( + not subnet_az + and PRIVATE_SUBNET_AVAILABILITY_ZONES + and i < len(PRIVATE_SUBNET_AVAILABILITY_ZONES) + ): + subnet_az = PRIVATE_SUBNET_AVAILABILITY_ZONES[i] + if not subnet_az: + raise RuntimeError( + f"Context for existing private subnet '{subnet_name}' is missing 'az'." + ) + subnet_attrs = { + "subnet_id": subnet_id, + "availability_zone": subnet_az, + } + route_table_id = subnet_info.get("route_table_id") + if route_table_id: + subnet_attrs["route_table_id"] = route_table_id + try: + imported_subnet = ec2.Subnet.from_subnet_attributes( + self, + f"ImportedPrivateSubnet{subnet_name.replace('-', '')}{i}", + **subnet_attrs, + ) + self.private_subnets.append(imported_subnet) + print( + f"Imported existing private subnet: {subnet_name} (ID: {subnet_id})" + ) + except Exception as e: + raise RuntimeError( + f"Failed to import private subnet '{subnet_name}' with ID '{subnet_id}'. Error: {e}" + ) + + # Create new private subnets + if private_subnets_data_for_creation_ctx: + names_to_create_private = [ + s["name"] for s in private_subnets_data_for_creation_ctx + ] + cidrs_to_create_private = [ + s["cidr"] for s in private_subnets_data_for_creation_ctx + ] + azs_to_create_private = [ + s["az"] for s in private_subnets_data_for_creation_ctx + ] + + if names_to_create_private: + print( + f"Attempting to create {len(names_to_create_private)} new private subnets: {names_to_create_private}" + ) + # --- CALL THE NEW CREATE_SUBNETS FUNCTION FOR PRIVATE --- + # Ensure self.single_nat_gateway_id is available before this call + if not self.single_nat_gateway_id: + raise ValueError( + "A single NAT Gateway ID is required for private subnets but was not resolved." + ) + + newly_created_private_subnets_cfn, newly_created_private_rts_cfn = ( + create_subnets( + self, + vpc, + CDK_PREFIX, + names_to_create_private, + cidrs_to_create_private, + azs_to_create_private, + is_public=False, + single_nat_gateway_id=self.single_nat_gateway_id, # Pass the single NAT Gateway ID + ) + ) + self.private_subnets.extend(newly_created_private_subnets_cfn) + self.private_route_tables_cfn.extend(newly_created_private_rts_cfn) + print( + f"Successfully defined {len(newly_created_private_subnets_cfn)} new private subnets and their route tables for creation." + ) + else: + print( + "No private subnets specified for creation in context ('private_subnets_to_create')." + ) + + # if not self.private_subnets: + # raise Exception("No private subnets found or created, exiting.") + + if ( + not self.private_subnets + and not names_to_create_private + and not PRIVATE_SUBNETS_TO_USE + ): + # This condition might need adjustment for new VPCs. + raise Exception("No private subnets found or created, exiting.") + + # --- 5. Sanity Check and Output --- + # Output the single NAT Gateway ID for verification + if self.single_nat_gateway_id: + CfnOutput( + self, + "SingleNatGatewayId", + value=self.single_nat_gateway_id, + description="ID of the single NAT Gateway resolved or created.", + ) + elif express_public_subnets_only: + print( + "INFO: Express public-subnet mode — NAT Gateway not installed or required." + ) + elif ( + NEW_VPC_DEFAULT_NAME + and (self.node.try_get_context("new_vpc_nat_gateways") or 1) > 0 + ): + print( + "INFO: A new VPC was created with NAT Gateway(s). Their routing is handled by the VPC construct. No single_nat_gateway_id was explicitly set for separate output." + ) + else: + out_message = "WARNING: No single NAT Gateway was resolved or created explicitly by the script's logic after VPC setup." + print(out_message) + raise Exception(out_message) + + # --- Outputs for other stacks/regions --- + # These are crucial for cross-stack, cross-region referencing + + self.params = dict() + self.params["vpc_id"] = vpc.vpc_id + self.params["private_subnets"] = self.private_subnets + self.params["private_route_tables"] = self.private_route_tables_cfn + self.params["public_subnets"] = self.public_subnets + self.params["public_route_tables"] = self.public_route_tables_cfn + + private_subnet_selection = ec2.SubnetSelection(subnets=self.private_subnets) + public_subnet_selection = ec2.SubnetSelection(subnets=self.public_subnets) + + for sub in private_subnet_selection.subnets: + print( + "private subnet:", + sub.subnet_id, + "is in availability zone:", + sub.availability_zone, + ) + + for sub in public_subnet_selection.subnets: + print( + "public subnet:", + sub.subnet_id, + "is in availability zone:", + sub.availability_zone, + ) + + print("Private subnet route tables:", self.private_route_tables_cfn) + + CfnOutput( + self, + "VpcIdOutput", + value=vpc.vpc_id, + description="The ID of the VPC used by this stack.", + ) + + # --- IAM Roles --- + cognito_secret_name = COGNITO_USER_POOL_CLIENT_SECRET_NAME + secret_kms_key_arn_from_context = get_context_str( + f"kms_key_arn:{cognito_secret_name}" + ) + + if USE_CUSTOM_KMS_KEY == "1": + kms_key = kms.Key( + self, + "SummarisationSharedKmsKey", + alias=CUSTOM_KMS_KEY_NAME, + removal_policy=resource_removal_policy, + ) + shared_kms_key_arn = kms_key.key_arn + secret_kms_key_arn = secret_kms_key_arn_from_context or kms_key.key_arn + else: + kms_key = None + shared_kms_key_arn = None + secret_kms_key_arn = ( + secret_kms_key_arn_from_context + or default_secrets_manager_kms_key_arn(AWS_REGION, AWS_ACCOUNT_ID) + ) + + task_role_inline_policy = json.dumps( + build_ecs_task_role_inline_policy( + output_bucket_name=S3_OUTPUT_BUCKET_NAME, + log_config_bucket_name=S3_LOG_CONFIG_BUCKET_NAME, + shared_kms_key_arn=shared_kms_key_arn, + ), + indent=4, + ) + if enable_headless: + execution_role_kms_policy = json.dumps( + { + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "STSCallerIdentity", + "Effect": "Allow", + "Action": ["sts:GetCallerIdentity"], + "Resource": "*", + } + ], + }, + indent=4, + ) + else: + execution_role_kms_policy = json.dumps( + build_ecs_execution_role_kms_policy( + secret_kms_key_arn=secret_kms_key_arn, + ), + indent=4, + ) + + try: + codebuild_role_name = CODEBUILD_ROLE_NAME + + if get_context_bool(f"exists:{codebuild_role_name}"): + # If exists, lookup/import the role using ARN from context + role_arn = get_context_str(f"arn:{codebuild_role_name}") + if not role_arn: + raise ValueError( + f"Context value 'arn:{codebuild_role_name}' is required if role exists." + ) + codebuild_role = iam.Role.from_role_arn( + self, "CodeBuildRole", role_arn=role_arn + ) + print("Using existing CodeBuild role") + else: + # If not exists, create the role + codebuild_role = iam.Role( + self, + "CodeBuildRole", # Logical ID + role_name=codebuild_role_name, # Explicit resource name + assumed_by=iam.ServicePrincipal("codebuild.amazonaws.com"), + ) + codebuild_role.add_managed_policy( + iam.ManagedPolicy.from_aws_managed_policy_name( + "EC2InstanceProfileForImageBuilderECRContainerBuilds" + ) + ) + print("Successfully created new CodeBuild role") + + task_role_name = ECS_TASK_ROLE_NAME + if get_context_bool(f"exists:{task_role_name}"): + role_arn = get_context_str(f"arn:{task_role_name}") + if not role_arn: + raise ValueError( + f"Context value 'arn:{task_role_name}' is required if role exists." + ) + task_role = iam.Role.from_role_arn(self, "TaskRole", role_arn=role_arn) + print("Using existing ECS task role") + else: + task_role = iam.Role( + self, + "TaskRole", # Logical ID + role_name=task_role_name, # Explicit resource name + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + for role in AWS_MANAGED_TASK_ROLES_LIST: + print(f"Adding {role} to policy") + task_role.add_managed_policy( + iam.ManagedPolicy.from_aws_managed_policy_name(f"{role}") + ) + attach_managed_policy_arns(task_role, POLICY_FILE_ARNS) + print("Successfully created new ECS task role") + task_role = add_custom_policies( + self, + task_role, + policy_file_locations=resolve_policy_file_paths( + POLICY_FILE_LOCATIONS, cdk_folder=CDK_FOLDER + ), + custom_policy_text=task_role_inline_policy, + ) + + execution_role_name = ECS_TASK_EXECUTION_ROLE_NAME + if get_context_bool(f"exists:{execution_role_name}"): + role_arn = get_context_str(f"arn:{execution_role_name}") + if not role_arn: + raise ValueError( + f"Context value 'arn:{execution_role_name}' is required if role exists." + ) + execution_role = iam.Role.from_role_arn( + self, "ExecutionRole", role_arn=role_arn + ) + print("Using existing ECS execution role") + else: + execution_role = iam.Role( + self, + "ExecutionRole", # Logical ID + role_name=execution_role_name, # Explicit resource name + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + for role in ECS_EXECUTION_ROLE_MANAGED_POLICIES: + print(f"Adding {role} to execution role") + execution_role.add_managed_policy( + iam.ManagedPolicy.from_aws_managed_policy_name(f"{role}") + ) + attach_managed_policy_arns( + execution_role, ECS_EXECUTION_ROLE_POLICY_ARNS + ) + print("Successfully created new ECS execution role") + execution_role = add_custom_policies( + self, + execution_role, + policy_file_locations=resolve_policy_file_paths( + ECS_EXECUTION_ROLE_POLICY_FILES, cdk_folder=CDK_FOLDER + ), + custom_policy_text=execution_role_kms_policy, + ) + + except Exception as e: + raise Exception("Failed at IAM role step due to:", e) + + # --- S3 Buckets --- + try: + log_bucket_name = S3_LOG_CONFIG_BUCKET_NAME + if get_context_bool(f"globally_taken:{log_bucket_name}"): + raise ValueError( + f"S3 bucket name {log_bucket_name!r} is taken globally by another " + "AWS account. Set S3_LOG_CONFIG_BUCKET_NAME in cdk/config/cdk_config.env " + "to a unique name (re-run cdk_install.py or check_resources.py)." + ) + if get_context_bool(f"exists:{log_bucket_name}"): + bucket = s3.Bucket.from_bucket_name( + self, "LogConfigBucket", bucket_name=log_bucket_name + ) + print("Using existing S3 bucket", log_bucket_name) + else: + log_bucket_lifecycle = [ + s3.LifecycleRule( + abort_incomplete_multipart_upload_after=Duration.days(7) + ) + ] + if USE_CUSTOM_KMS_KEY == "1" and isinstance(kms_key, kms.Key): + bucket = s3.Bucket( + self, + "LogConfigBucket", + bucket_name=log_bucket_name, + lifecycle_rules=log_bucket_lifecycle, + versioned=False, + removal_policy=resource_removal_policy, + auto_delete_objects=s3_auto_delete_objects, + encryption=s3.BucketEncryption.KMS, + encryption_key=kms_key, + ) + else: + bucket = s3.Bucket( + self, + "LogConfigBucket", + bucket_name=log_bucket_name, + lifecycle_rules=log_bucket_lifecycle, + versioned=False, + removal_policy=resource_removal_policy, + auto_delete_objects=s3_auto_delete_objects, + ) + + print("Created S3 bucket", log_bucket_name) + + # Add policies - this will apply to both created and imported buckets + # CDK handles idempotent policy additions + bucket.add_to_resource_policy( + iam.PolicyStatement( + effect=iam.Effect.ALLOW, + principals=[task_role], # Pass the role object directly + actions=["s3:GetObject", "s3:PutObject"], + resources=[f"{bucket.bucket_arn}/*"], + ) + ) + bucket.add_to_resource_policy( + iam.PolicyStatement( + effect=iam.Effect.ALLOW, + principals=[task_role], + actions=["s3:ListBucket"], + resources=[bucket.bucket_arn], + ) + ) + + output_bucket_name = S3_OUTPUT_BUCKET_NAME + if get_context_bool(f"globally_taken:{output_bucket_name}"): + raise ValueError( + f"S3 bucket name {output_bucket_name!r} is taken globally by another " + "AWS account. Set S3_OUTPUT_BUCKET_NAME in cdk/config/cdk_config.env " + "to a unique name (re-run cdk_install.py or check_resources.py)." + ) + if get_context_bool(f"exists:{output_bucket_name}"): + output_bucket = s3.Bucket.from_bucket_name( + self, "OutputBucket", bucket_name=output_bucket_name + ) + print("Using existing Output bucket", output_bucket_name) + else: + if USE_CUSTOM_KMS_KEY == "1" and isinstance(kms_key, kms.Key): + output_bucket = s3.Bucket( + self, + "OutputBucket", + bucket_name=output_bucket_name, + lifecycle_rules=[ + s3.LifecycleRule( + expiration=Duration.days( + int(DAYS_TO_DISPLAY_WHOLE_DOCUMENT_JOBS) + ) + ) + ], + versioned=False, + removal_policy=resource_removal_policy, + auto_delete_objects=s3_auto_delete_objects, + encryption=s3.BucketEncryption.KMS, + encryption_key=kms_key, + ) + else: + output_bucket = s3.Bucket( + self, + "OutputBucket", + bucket_name=output_bucket_name, + lifecycle_rules=[ + s3.LifecycleRule( + expiration=Duration.days( + int(DAYS_TO_DISPLAY_WHOLE_DOCUMENT_JOBS) + ) + ) + ], + versioned=False, + removal_policy=resource_removal_policy, + auto_delete_objects=s3_auto_delete_objects, + ) + + print("Created Output bucket:", output_bucket_name) + + add_s3_enforce_ssl_policy(bucket) + add_s3_enforce_ssl_policy(output_bucket) + + # Add policies to output bucket + output_bucket.add_to_resource_policy( + iam.PolicyStatement( + effect=iam.Effect.ALLOW, + principals=[task_role], + actions=["s3:GetObject", "s3:PutObject"], + resources=[f"{output_bucket.bucket_arn}/*"], + ) + ) + output_bucket.add_to_resource_policy( + iam.PolicyStatement( + effect=iam.Effect.ALLOW, + principals=[task_role], + actions=["s3:ListBucket"], + resources=[output_bucket.bucket_arn], + ) + ) + # Identity-based S3 access is scoped via build_ecs_task_role_inline_policy on + # task_role (output + log/config buckets). Bucket policies above remain for + # imported buckets and org policies that expect explicit bucket principals. + + except Exception as e: + raise Exception("Could not handle S3 buckets due to:", e) + + # --- Elastic Container Registry --- + try: + full_ecr_repo_name = ECR_CDK_REPO_NAME + if get_context_bool(f"exists:{full_ecr_repo_name}"): + ecr_repo = ecr.Repository.from_repository_name( + self, "ECRRepo", repository_name=full_ecr_repo_name + ) + print("Using existing ECR repository") + else: + ecr_repo = ecr.Repository( + self, + "ECRRepo", + repository_name=full_ecr_repo_name, + removal_policy=resource_removal_policy, + empty_on_delete=ecr_empty_on_delete(), + ) # Explicitly set repository_name + print("Created ECR repository", full_ecr_repo_name) + + ecr_image_loc = ecr_repo.repository_uri + except Exception as e: + raise Exception("Could not handle ECR repo due to:", e) + + pi_ecr_image_loc = ecr_image_loc + + # --- CODEBUILD --- + try: + codebuild_project_name = CODEBUILD_PROJECT_NAME + if get_context_bool(f"exists:{codebuild_project_name}"): + # Lookup CodeBuild project by ARN from context + project_arn = get_context_str(f"arn:{codebuild_project_name}") + if not project_arn: + raise ValueError( + f"Context value 'arn:{codebuild_project_name}' is required if project exists." + ) + codebuild.Project.from_project_arn( + self, "CodeBuildProject", project_arn=project_arn + ) + print( + "Using existing CodeBuild project " + "(public GitHub source is applied in post_cdk_build_quickstart)." + ) + else: + main_codebuild_project = codebuild.Project( + self, + "CodeBuildProject", # Logical ID + project_name=codebuild_project_name, # Explicit resource name + role=codebuild_role, + source=public_github_codebuild_source( + owner=GITHUB_REPO_USERNAME, + repo=GITHUB_REPO_NAME, + branch_or_ref=GITHUB_REPO_BRANCH, + ), + environment=codebuild.BuildEnvironment( + build_image=codebuild.LinuxBuildImage.STANDARD_7_0, + privileged=True, + environment_variables={ + "ECR_REPO_NAME": codebuild.BuildEnvironmentVariable( + value=full_ecr_repo_name + ), + "AWS_DEFAULT_REGION": codebuild.BuildEnvironmentVariable( + value=AWS_REGION + ), + "AWS_ACCOUNT_ID": codebuild.BuildEnvironmentVariable( + value=AWS_ACCOUNT_ID + ), + "APP_MODE": codebuild.BuildEnvironmentVariable( + value="gradio" + ), + }, + ), + build_spec=codebuild.BuildSpec.from_object( + { + "version": "0.2", + "phases": { + "pre_build": { + "commands": [ + "echo Logging in to Amazon ECR", + "aws ecr get-login-password --region $AWS_DEFAULT_REGION | docker login --username AWS --password-stdin $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com", + ] + }, + "build": { + "commands": [ + "echo Building the Docker image", + "docker build --build-arg APP_MODE=$APP_MODE --target $APP_MODE -t $ECR_REPO_NAME:latest .", + "docker tag $ECR_REPO_NAME:latest $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/$ECR_REPO_NAME:latest", + ] + }, + "post_build": { + "commands": [ + "echo Pushing the Docker image", + "docker push $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/$ECR_REPO_NAME:latest", + ] + }, + }, + } + ), + ) + configure_public_github_codebuild_source( + main_codebuild_project, + GITHUB_REPO_USERNAME, + GITHUB_REPO_NAME, + GITHUB_REPO_BRANCH, + ) + print("Successfully created CodeBuild project", codebuild_project_name) + + # Imported projects have role=undefined in CDK; use the actual service + # role from context (existing project) or the managed codebuild_role (new). + if get_context_bool(f"exists:{codebuild_project_name}"): + project_service_role_arn = get_context_str( + f"service_role_arn:{codebuild_project_name}" + ) + if project_service_role_arn: + ecr_grantee = iam.Role.from_role_arn( + self, + "CodeBuildProjectServiceRole", + role_arn=project_service_role_arn, + mutable=True, + ) + else: + ecr_grantee = codebuild_role + else: + ecr_grantee = codebuild_role + ecr_repo.grant_pull_push(ecr_grantee) + + if enable_pi_build: + pi_codebuild_name = CODEBUILD_PI_PROJECT_NAME + if get_context_bool(f"exists:{pi_codebuild_name}"): + project_arn = get_context_str(f"arn:{pi_codebuild_name}") + if project_arn: + codebuild.Project.from_project_arn( + self, "CodeBuildPiProject", project_arn=project_arn + ) + print("Using existing Pi agent CodeBuild project") + else: + pi_codebuild_project = codebuild.Project( + self, + "CodeBuildPiProject", + project_name=pi_codebuild_name, + role=codebuild_role, + source=public_github_codebuild_source( + owner=GITHUB_REPO_USERNAME, + repo=GITHUB_REPO_NAME, + branch_or_ref=GITHUB_REPO_BRANCH, + ), + environment=codebuild.BuildEnvironment( + build_image=codebuild.LinuxBuildImage.STANDARD_7_0, + privileged=True, + environment_variables={ + "ECR_REPO_NAME": codebuild.BuildEnvironmentVariable( + value=ECR_PI_REPO_NAME + ), + "AWS_DEFAULT_REGION": codebuild.BuildEnvironmentVariable( + value=AWS_REGION + ), + "AWS_ACCOUNT_ID": codebuild.BuildEnvironmentVariable( + value=AWS_ACCOUNT_ID + ), + }, + ), + build_spec=codebuild.BuildSpec.from_object( + { + "version": "0.2", + "phases": { + "pre_build": { + "commands": [ + "echo Logging in to Amazon ECR", + "aws ecr get-login-password --region $AWS_DEFAULT_REGION | docker login --username AWS --password-stdin $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com", + "test -f config/pi_agent.env.example", + "test -f agent-redact/pi-agent/Dockerfile", + ] + }, + "build": { + "commands": [ + "docker build -f agent-redact/pi-agent/Dockerfile -t $ECR_REPO_NAME:latest .", + "docker tag $ECR_REPO_NAME:latest $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/$ECR_REPO_NAME:latest", + ] + }, + "post_build": { + "commands": [ + "docker push $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/$ECR_REPO_NAME:latest", + ] + }, + }, + } + ), + ) + configure_public_github_codebuild_source( + pi_codebuild_project, + GITHUB_REPO_USERNAME, + GITHUB_REPO_NAME, + GITHUB_REPO_BRANCH, + ) + print("Created Pi agent CodeBuild project", pi_codebuild_name) + + pi_ecr_repo_name = ECR_PI_REPO_NAME + if get_context_bool(f"exists:{pi_ecr_repo_name}"): + pi_ecr_repo = ecr.Repository.from_repository_name( + self, "ECRPiRepo", repository_name=pi_ecr_repo_name + ) + else: + pi_ecr_repo = ecr.Repository( + self, + "ECRPiRepo", + repository_name=pi_ecr_repo_name, + removal_policy=resource_removal_policy, + empty_on_delete=ecr_empty_on_delete(), + ) + pi_ecr_image_loc = pi_ecr_repo.repository_uri + pi_ecr_repo.grant_pull_push(ecr_grantee) + CfnOutput(self, "ECRPiRepoUri", value=pi_ecr_repo.repository_uri) + + except Exception as e: + raise Exception("Could not handle Codebuild project due to:", e) + + pi_ecs_service = None + pi_ecs_security_group = None + + # --- Security Groups --- + try: + ecs_security_group_name = ECS_SECURITY_GROUP_NAME + + try: + ecs_security_group = ec2.SecurityGroup( + self, + "ECSSecurityGroup", # Logical ID + security_group_name=ecs_security_group_name, # Explicit resource name + vpc=vpc, + ) + print(f"Created Security Group: {ecs_security_group_name}") + except Exception as e: # If lookup fails, create + print("Failed to create ECS security group due to:", e) + + ec2_port_gradio_server_port = ec2.Port.tcp(int(GRADIO_SERVER_PORT)) + + if deploy_web_ingress: + alb_security_group_name = ALB_NAME_SECURITY_GROUP_NAME + + try: + alb_security_group = ec2.SecurityGroup( + self, + "ALBSecurityGroup", # Logical ID + security_group_name=alb_security_group_name, + vpc=vpc, + ) + print(f"Created Security Group: {alb_security_group_name}") + except Exception as e: + print("Failed to create ALB security group due to:", e) + + ecs_security_group.add_ingress_rule( + peer=alb_security_group, + connection=ec2_port_gradio_server_port, + description="ALB traffic", + ) + + alb_security_group.add_ingress_rule( + peer=ec2.Peer.prefix_list(CLOUDFRONT_PREFIX_LIST_ID), + connection=ec2.Port.all_traffic(), + description="CloudFront traffic", + ) + else: + alb_security_group = None + if USE_CLOUDFRONT == "True": + ecs_security_group.add_ingress_rule( + peer=ec2.Peer.prefix_list(CLOUDFRONT_PREFIX_LIST_ID), + connection=ec2_port_gradio_server_port, + description="CloudFront to ECS (Express Mode)", + ) + if enable_service_connect: + for index, client_sg_id in enumerate(service_connect_client_sg_ids): + client_sg = ec2.SecurityGroup.from_security_group_id( + self, + f"ServiceConnectClientSg{index}", + security_group_id=client_sg_id, + ) + ecs_security_group.add_ingress_rule( + peer=client_sg, + connection=ec2_port_gradio_server_port, + description=( + f"Service Connect client {client_sg_id} to app port" + ), + ) + print( + "Service Connect ingress allowed from security groups: " + + ", ".join(service_connect_client_sg_ids) + ) + + except Exception as e: + raise Exception("Could not handle security groups due to:", e) + + endpoint_subnet_selection = resolve_ecs_vpc_endpoint_subnet_selection( + use_express_ingress=use_express_ingress, + express_use_public_subnets=ECS_EXPRESS_USE_PUBLIC_SUBNETS == "True", + public_subnets=self.public_subnets, + private_subnets=self.private_subnets, + ) + s3_gateway_subnet_selection = resolve_ecs_s3_gateway_subnet_selection( + public_subnets=self.public_subnets, + private_subnets=self.private_subnets, + ) + + if ENABLE_ECS_VPC_INTERFACE_ENDPOINTS == "True" and ( + endpoint_subnet_selection or s3_gateway_subnet_selection + ): + if ( + VPC_NAME + and not imported_vpc_cidr_block + and not imported_vpc_cidr_blocks + ): + raise ValueError( + "vpc_cidr_block / vpc_cidr_blocks missing from precheck.context.json. " + "Re-run check_resources.py from the cdk/ directory so the VPC " + "CIDR(s) are stored for VPC endpoints and security groups." + ) + existing_endpoint_services = frozenset( + self.node.try_get_context("existing_vpc_endpoint_service_names") or [] + ) + if VPC_NAME and not existing_endpoint_services: + print( + "Note: existing_vpc_endpoint_service_names not in precheck context; " + "re-run check_resources.py to skip duplicate endpoints in shared VPCs." + ) + try: + endpoint_tier = ( + "public" + if use_express_ingress and ECS_EXPRESS_USE_PUBLIC_SUBNETS == "True" + else "private" + ) + create_ecs_vpc_endpoints_for_private_subnets( + self, + vpc=vpc, + subnets=endpoint_subnet_selection, + s3_gateway_subnets=s3_gateway_subnet_selection, + logical_id_prefix="SummarisationEcs", + include_secrets_and_kms=True, + vpc_cidr_block=imported_vpc_cidr_block, + vpc_cidr_blocks=imported_vpc_cidr_blocks or None, + skip_service_names=existing_endpoint_services, + aws_region=AWS_REGION, + ) + s3_subnet_count = len( + (s3_gateway_subnet_selection.subnets or []) + if s3_gateway_subnet_selection + else [] + ) + print( + "Defined ECS VPC interface endpoints (ECR, Logs, Secrets Manager, " + f"KMS) for {endpoint_tier} subnets where not already present; " + f"S3 gateway for {s3_subnet_count} stack subnet(s) (public + " + "private) where not already present." + ) + except Exception as e: + raise Exception( + "Could not create ECS VPC interface endpoints for ECS task subnets. " + "If this VPC already has them, re-run check_resources.py (auto-skip) " + "or set ENABLE_ECS_VPC_INTERFACE_ENDPOINTS=False in cdk_config.env " + "and ensure task subnets reach ECR (NAT, IGW, or existing endpoints).", + e, + ) from e + + # --- DynamoDB tables for logs (optional) --- + usage_log_table = None + + if SAVE_LOGS_TO_DYNAMODB == "True": + try: + print("Creating DynamoDB tables for logs") + + dynamodb.Table( + self, + "SummarisationAccessDataTable", + table_name=ACCESS_LOG_DYNAMODB_TABLE_NAME, + partition_key=dynamodb.Attribute( + name="id", type=dynamodb.AttributeType.STRING + ), + billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST, + deletion_protection=resource_delete_protection, + removal_policy=resource_removal_policy, + ) + + dynamodb.Table( + self, + "SummarisationFeedbackDataTable", + table_name=FEEDBACK_LOG_DYNAMODB_TABLE_NAME, + partition_key=dynamodb.Attribute( + name="id", type=dynamodb.AttributeType.STRING + ), + billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST, + deletion_protection=resource_delete_protection, + removal_policy=resource_removal_policy, + ) + + usage_log_table = dynamodb.Table( + self, + "SummarisationUsageDataTable", + table_name=USAGE_LOG_DYNAMODB_TABLE_NAME, + partition_key=dynamodb.Attribute( + name="id", type=dynamodb.AttributeType.STRING + ), + billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST, + deletion_protection=resource_delete_protection, + removal_policy=resource_removal_policy, + ) + + except Exception as e: + raise Exception("Could not create DynamoDB tables due to:", e) + + if ENABLE_DYNAMODB_USAGE_LOG_EXPORT == "True": + try: + if usage_log_table is None: + raise ValueError( + "ENABLE_DYNAMODB_USAGE_LOG_EXPORT=True requires " + "SAVE_LOGS_TO_DYNAMODB=True and the usage log table." + ) + lambda_asset_dir = os.path.join( + os.path.dirname(__file__), "lambda_dynamo_logs_export" + ) + create_dynamo_usage_log_export_lambda( + self, + "DynamoUsageLogExport", + function_name=DYNAMODB_USAGE_LOG_EXPORT_LAMBDA_NAME or None, + lambda_asset_path=lambda_asset_dir, + dynamodb_table=usage_log_table, + output_bucket=output_bucket, + s3_output_key=DYNAMODB_USAGE_LOG_EXPORT_S3_KEY, + schedule_expression=DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE, + dynamodb_table_name=USAGE_LOG_DYNAMODB_TABLE_NAME, + date_attribute=DYNAMODB_USAGE_LOG_EXPORT_DATE_ATTRIBUTE, + output_filename=DYNAMODB_USAGE_LOG_EXPORT_OUTPUT_FILENAME, + shared_kms_key_arn=shared_kms_key_arn, + ) + print( + "Scheduled DynamoDB usage log export Lambda defined " + f"({DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE})." + ) + except Exception as e: + raise Exception( + "Could not handle DynamoDB usage log export Lambda due to:", e + ) from e + + alb = None + load_balancer_name = ALB_NAME + if len(load_balancer_name) > 32: + load_balancer_name = load_balancer_name[-32:] + + if deploy_web_ingress: + # --- ALB (legacy path) --- + try: + alb_arn = get_context_str(f"arn:{load_balancer_name}") or ( + EXISTING_LOAD_BALANCER_ARN or None + ) + alb_dns_name = get_context_str(f"dns:{load_balancer_name}") or ( + EXISTING_LOAD_BALANCER_DNS or None + ) + if alb_arn and alb_dns_name: + alb_security_group_id = ( + get_context_str(f"security_group_id:{load_balancer_name}") + or alb_security_group.security_group_id + ) + alb_attrs = { + "load_balancer_arn": alb_arn, + "load_balancer_dns_name": alb_dns_name, + "security_group_id": alb_security_group_id, + "vpc": vpc, + } + alb_canonical_zone_id = get_context_str( + f"canonical_hosted_zone_id:{load_balancer_name}" + ) + if alb_canonical_zone_id: + alb_attrs["load_balancer_canonical_hosted_zone_id"] = ( + alb_canonical_zone_id + ) + alb = elbv2.ApplicationLoadBalancer.from_application_load_balancer_attributes( + self, + "ALB", + **alb_attrs, + ) + print( + f"Using existing Application Load Balancer {load_balancer_name}." + ) + else: + alb = elbv2.ApplicationLoadBalancer( + self, + "ALB", + load_balancer_name=load_balancer_name, + vpc=vpc, + internet_facing=True, + security_group=alb_security_group, + vpc_subnets=public_subnet_selection, + drop_invalid_header_fields=True, + deletion_protection=resource_delete_protection, + ) + print("Successfully created new Application Load Balancer") + except Exception as e: + raise Exception("Could not handle application load balancer due to:", e) + + # --- Cognito User Pool (web login; skipped for headless batch-only) --- + user_pool = None + user_pool_client = None + user_pool_domain = None + secret = None + if enable_headless: + print( + "ENABLE_HEADLESS_DEPLOYMENT=True: skipping Cognito user pool, " + "hosted UI domain, and client secret (no web login for batch tasks)." + ) + else: + try: + if get_context_bool(f"exists:{COGNITO_USER_POOL_NAME}"): + # Lookup by ID from context + user_pool_id = get_context_str(f"id:{COGNITO_USER_POOL_NAME}") + if not user_pool_id: + raise ValueError( + f"Context value 'id:{COGNITO_USER_POOL_NAME}' is required if User Pool exists." + ) + user_pool = cognito.UserPool.from_user_pool_id( + self, "UserPool", user_pool_id=user_pool_id + ) + print(f"Using existing user pool {user_pool_id}.") + else: + user_pool = cognito.UserPool( + self, + "UserPool", + user_pool_name=COGNITO_USER_POOL_NAME, + mfa=cognito.Mfa.OFF, # Adjust as needed + sign_in_aliases=cognito.SignInAliases(email=True), + deletion_protection=resource_delete_protection, + removal_policy=resource_removal_policy, + ) # Adjust as needed + print(f"Created new user pool {user_pool.user_pool_id}.") + + # HTTPS ALB (ACM cert or Express Mode) needs oauth2/idpresponse callback URLs. + if ACM_SSL_CERTIFICATE_ARN or use_express_ingress: + redirect_uris = [ + COGNITO_REDIRECTION_URL, + COGNITO_REDIRECTION_URL + "/oauth2/idpresponse", + ] + else: + redirect_uris = [COGNITO_REDIRECTION_URL] + + user_pool_client_name = COGNITO_USER_POOL_CLIENT_NAME + if get_context_bool(f"exists:{user_pool_client_name}"): + # Lookup by ID from context (requires User Pool object) + user_pool_client_id = get_context_str(f"id:{user_pool_client_name}") + if not user_pool_client_id: + raise ValueError( + f"Context value 'id:{user_pool_client_name}' is required if User Pool Client exists." + ) + user_pool_client = cognito.UserPoolClient.from_user_pool_client_id( + self, "UserPoolClient", user_pool_client_id=user_pool_client_id + ) + print(f"Using existing user pool client {user_pool_client_id}.") + else: + user_pool_client = cognito.UserPoolClient( + self, + "UserPoolClient", + auth_flows=cognito.AuthFlow( + user_srp=True, user_password=True + ), # Example: enable SRP for secure sign-in + user_pool=user_pool, + generate_secret=True, + user_pool_client_name=user_pool_client_name, + supported_identity_providers=[ + cognito.UserPoolClientIdentityProvider.COGNITO + ], + o_auth=cognito.OAuthSettings( + flows=cognito.OAuthFlows(authorization_code_grant=True), + scopes=[ + cognito.OAuthScope.OPENID, + cognito.OAuthScope.EMAIL, + cognito.OAuthScope.PROFILE, + ], + callback_urls=redirect_uris, + ), + refresh_token_validity=Duration.minutes( + COGNITO_REFRESH_TOKEN_VALIDITY + ), + id_token_validity=Duration.minutes(COGNITO_ID_TOKEN_VALIDITY), + access_token_validity=Duration.minutes( + COGNITO_ACCESS_TOKEN_VALIDITY + ), + ) + + CfnOutput( + self, + "CognitoAppClientId", + value=user_pool_client.user_pool_client_id, + ) + + print( + f"Created new user pool client {user_pool_client.user_pool_client_id}." + ) + + # Add a domain to the User Pool (crucial for ALB integration) + domain_prefix = (COGNITO_USER_POOL_DOMAIN_PREFIX or "").strip().lower() + if get_context_bool(f"cognito_domain_taken:{domain_prefix}"): + raise ValueError( + f"Cognito hosted UI domain prefix {domain_prefix!r} is not " + f"available in this region (taken by another AWS account or " + "an existing pool). Set COGNITO_USER_POOL_DOMAIN_PREFIX in " + "cdk/config/cdk_config.env to a unique value and re-run " + "cdk_install.py / check_resources.py." + ) + user_pool_domain = user_pool.add_domain( + "UserPoolDomain", + cognito_domain=cognito.CognitoDomainOptions( + domain_prefix=COGNITO_USER_POOL_DOMAIN_PREFIX + ), + ) + + # Apply removal_policy to the created UserPoolDomain construct + user_pool_domain.apply_removal_policy(policy=resource_removal_policy) + + CfnOutput( + self, "CognitoUserPoolLoginUrl", value=user_pool_domain.base_url() + ) + + except Exception as e: + raise Exception("Could not handle Cognito resources due to:", e) + + # --- Secrets Manager Secret --- + try: + secret_name = COGNITO_USER_POOL_CLIENT_SECRET_NAME + if get_context_bool(f"exists:{secret_name}"): + secret_arn = get_context_str(f"arn:{secret_name}") + if secret_arn: + secret = secretsmanager.Secret.from_secret_complete_arn( + self, + "CognitoSecret", + secret_complete_arn=secret_arn, + ) + print("Using existing Secret (ARN from precheck context).") + else: + secret = secretsmanager.Secret.from_secret_name_v2( + self, "CognitoSecret", secret_name=secret_name + ) + print( + "Using existing Secret by name (IAM grants use ARN wildcard " + "suffix; re-run precheck to pin the full ARN)." + ) + else: + if USE_CUSTOM_KMS_KEY == "1" and isinstance(kms_key, kms.Key): + secret = secretsmanager.Secret( + self, + "CognitoSecret", # Logical ID + secret_name=secret_name, # Explicit resource name + secret_object_value={ + "SUMMARISATION_USER_POOL_ID": SecretValue.unsafe_plain_text( + user_pool.user_pool_id + ), # Use the CDK attribute + "SUMMARISATION_CLIENT_ID": SecretValue.unsafe_plain_text( + user_pool_client.user_pool_client_id + ), # Use the CDK attribute + "SUMMARISATION_CLIENT_SECRET": user_pool_client.user_pool_client_secret, # Use the CDK attribute + }, + encryption_key=kms_key, + removal_policy=resource_removal_policy, + ) + else: + secret = secretsmanager.Secret( + self, + "CognitoSecret", # Logical ID + secret_name=secret_name, # Explicit resource name + secret_object_value={ + "SUMMARISATION_USER_POOL_ID": SecretValue.unsafe_plain_text( + user_pool.user_pool_id + ), # Use the CDK attribute + "SUMMARISATION_CLIENT_ID": SecretValue.unsafe_plain_text( + user_pool_client.user_pool_client_id + ), # Use the CDK attribute + "SUMMARISATION_CLIENT_SECRET": user_pool_client.user_pool_client_secret, # Use the CDK attribute + }, + removal_policy=resource_removal_policy, + ) + + print( + "Created new secret in Secrets Manager for Cognito user pool and related details." + ) + + except Exception as e: + raise Exception("Could not handle Secrets Manager secret due to:", e) + + try: + secret.grant_read(task_role) + secret.grant_read(execution_role) + except Exception as e: + raise Exception("Could not grant access to Secrets Manager due to:", e) + + try: + # ECS environmentFiles (app_config.env) are fetched by the execution role at task start. + bucket.grant_read(execution_role, APP_CONFIG_ENV_BASENAME) + # KMS: task role uses shared S3 CMK via build_ecs_task_role_inline_policy; + # execution role uses the secret's CMK via build_ecs_execution_role_kms_policy. + except Exception as e: + raise Exception("Could not grant bucket read to execution role due to:", e) + + # --- ECS Cluster (shared by legacy Fargate and Express paths) --- + try: + cluster_kwargs = { + "cluster_name": CLUSTER_NAME, + "enable_fargate_capacity_providers": True, + "vpc": vpc, + } + if enable_service_connect or enable_pi_express: + cluster_kwargs["default_cloud_map_namespace"] = ( + ecs.CloudMapNamespaceOptions( + name=ECS_SERVICE_CONNECT_NAMESPACE, + vpc=vpc, + ) + ) + cluster = ecs.Cluster(self, "ECSCluster", **cluster_kwargs) + print("Successfully created new ECS cluster") + except Exception as e: + raise Exception("Could not handle ECS cluster due to:", e) + + express_service = None + express_alb_security_group_id = None + + if use_express_ingress: + try: + express_log_group = logs.LogGroup( + self, + "ExpressTaskLogGroup", + log_group_name=f"/ecs/{ECS_EXPRESS_SERVICE_NAME}-logs".lower(), + retention=logs.RetentionDays.ONE_MONTH, + removal_policy=resource_removal_policy, + ) + express_log_group.grant_write(execution_role) + + express_infra_role = create_ecs_express_infrastructure_role( + self, + "ExpressInfrastructureRole", + ECS_EXPRESS_INFRASTRUCTURE_ROLE_NAME, + ) + + express_app_overrides: Dict[str, str] = {} + if ENABLE_HEADLESS_DEPLOYMENT == "True": + express_app_overrides["COGNITO_AUTH"] = "False" + elif enable_pi_express: + # Pi agent calls main over Service Connect; Gradio auth blocks + # gradio_client unless credentials are passed on every call. + express_app_overrides["COGNITO_AUTH"] = "False" + express_app_environment = load_app_config_env_for_express( + APP_CONFIG_ENV_FILE, + overrides=express_app_overrides or None, + ) + primary_container = build_express_gateway_primary_container( + image_uri=ecr_image_loc + ":latest", + container_port=int(GRADIO_SERVER_PORT), + log_group_name=express_log_group.log_group_name, + aws_region=AWS_REGION, + secret=secret, + environment=express_app_environment, + ) + + express_use_public_subnets = ECS_EXPRESS_USE_PUBLIC_SUBNETS == "True" + express_subnet_ids = [ + s.subnet_id + for s in ( + self.public_subnets + if express_use_public_subnets + else self.private_subnets + ) + ] + if not express_subnet_ids: + tier = "public" if express_use_public_subnets else "private" + raise ValueError( + f"No {tier} subnets available for ECS Express Mode. " + f"Set ECS_EXPRESS_USE_PUBLIC_SUBNETS=False to use private " + "subnets (internal ALB only), or create/import public subnets." + ) + if express_use_public_subnets: + print( + "ECS Express Mode using public subnets " + "(internet-facing managed ALB)." + ) + else: + print( + "ECS Express Mode using private subnets " + "(internal managed ALB)." + ) + + # MinTaskCount=0 until post_cdk_build_quickstart builds/pushes :latest. + express_service = create_express_gateway_service( + self, + "ExpressGatewayService", + service_name=ECS_EXPRESS_SERVICE_NAME, + cluster_name=CLUSTER_NAME, + execution_role_arn=execution_role.role_arn, + infrastructure_role_arn=express_infra_role.role_arn, + task_role_arn=task_role.role_arn, + cpu=str(ECS_TASK_CPU_SIZE), + memory=str(ECS_TASK_MEMORY_SIZE), + health_check_path=ECS_EXPRESS_HEALTH_CHECK_PATH, + primary_container=primary_container, + subnet_ids=express_subnet_ids, + security_group_ids=[ecs_security_group.security_group_id], + ) + express_service.node.add_dependency(cluster) + + allow_express_load_balancer_to_ecs_security_group( + self, + "ExpressAlbToEcsIngress", + express_service=express_service, + ecs_security_group=ecs_security_group, + container_port=int(GRADIO_SERVER_PORT), + ) + + express_alb_arn = express_ingress_load_balancer_arn(express_service) + express_alb_dns = express_service.attr_endpoint + express_alb_security_group_id = ( + express_ingress_first_load_balancer_security_group(express_service) + ) + + alb = elbv2.ApplicationLoadBalancer.from_application_load_balancer_attributes( + self, + "ALB", + load_balancer_arn=express_alb_arn, + load_balancer_dns_name=express_alb_dns, + security_group_id=express_alb_security_group_id, + vpc=vpc, + ) + + # Express Mode manages host-header listener rules (priorities 1, 2, …). + # Do not add ALB authenticate-cognito rules here; use in-app COGNITO_AUTH. + + CfnOutput( + self, + "ExpressServiceEndpoint", + value=express_service.attr_endpoint, + description="HTTPS URL for the ECS Express Mode service", + ) + CfnOutput( + self, + "ExpressServiceArn", + value=express_service.attr_service_arn, + ) + CfnOutput( + self, + "ExpressManagedCertificateArn", + value=express_service.attr_ecs_managed_resource_arns_ingress_path_certificate_arn, + ) + + if enable_pi_express: + try: + pi_express_log_group = logs.LogGroup( + self, + "ExpressPiTaskLogGroup", + log_group_name=f"/ecs/{ECS_PI_EXPRESS_SERVICE_NAME}-logs".lower(), + retention=logs.RetentionDays.ONE_MONTH, + removal_policy=resource_removal_policy, + ) + pi_express_log_group.grant_write(execution_role) + + pi_express_security_group = ec2.SecurityGroup( + self, + "ExpressPiSecurityGroup", + vpc=vpc, + security_group_name=ECS_PI_EXPRESS_SECURITY_GROUP_NAME, + description="Pi agent ECS Express tasks", + ) + + pi_express_environment = build_pi_express_container_environment( + service_connect_discovery_name=ECS_SERVICE_CONNECT_DISCOVERY_NAME, + main_app_port=int(GRADIO_SERVER_PORT), + pi_gradio_port=int(PI_GRADIO_PORT), + cognito_auth=ENABLE_HEADLESS_DEPLOYMENT != "True", + ) + pi_primary_container = build_express_pi_primary_container( + image_uri=pi_ecr_image_loc + ":latest", + container_port=int(PI_GRADIO_PORT), + log_group_name=pi_express_log_group.log_group_name, + aws_region=AWS_REGION, + environment=pi_express_environment, + secret=secret, + cognito_auth=ENABLE_HEADLESS_DEPLOYMENT != "True", + ) + + express_pi_service = create_express_gateway_service( + self, + "ExpressPiGatewayService", + service_name=ECS_PI_EXPRESS_SERVICE_NAME, + cluster_name=CLUSTER_NAME, + execution_role_arn=execution_role.role_arn, + infrastructure_role_arn=express_infra_role.role_arn, + task_role_arn=task_role.role_arn, + cpu=str(ECS_PI_TASK_CPU_SIZE), + memory=str(ECS_PI_TASK_MEMORY_SIZE), + health_check_path=ECS_PI_EXPRESS_HEALTH_CHECK_PATH, + primary_container=pi_primary_container, + subnet_ids=express_subnet_ids, + security_group_ids=[ + pi_express_security_group.security_group_id + ], + ) + express_pi_service.node.add_dependency(cluster) + express_pi_service.node.add_dependency(express_service) + + allow_express_load_balancer_to_ecs_security_group( + self, + "ExpressAlbToPiExpressIngress", + express_service=express_pi_service, + ecs_security_group=pi_express_security_group, + container_port=int(PI_GRADIO_PORT), + ) + + pi_express_security_group.add_egress_rule( + peer=ecs_security_group, + connection=ec2.Port.tcp(int(GRADIO_SERVER_PORT)), + description="Pi Express (Service Connect) to main summarisation app", + ) + ecs_security_group.add_ingress_rule( + peer=pi_express_security_group, + connection=ec2.Port.tcp(int(GRADIO_SERVER_PORT)), + description="Pi Express (Service Connect) to main summarisation app", + ) + + # Service Connect for Express is applied in post_cdk_build_quickstart.py + # after CodeBuild pushes :latest. Express primary containers do not + # define named portMappings at create time; CDK cannot enable SC here. + + pi_public_url = format_express_pi_public_url( + express_pi_service.attr_endpoint, + ) + sc_backend = ( + f"http://{ECS_SERVICE_CONNECT_DISCOVERY_NAME}:" + f"{GRADIO_SERVER_PORT}" + ) + CfnOutput( + self, + "PiExpressEndpoint", + value=express_pi_service.attr_endpoint, + description="HTTPS URL for the Pi ECS Express service (AWS-managed cert)", + ) + CfnOutput( + self, + "PiPublicUrl", + value=pi_public_url, + description="Public URL for Pi Express UI (managed HTTPS endpoint)", + ) + CfnOutput( + self, + "PiDocSummarisationBackendUrl", + value=sc_backend, + description="DOC_SUMMARISATION_GRADIO_URL on Pi Express (Service Connect, no Cognito)", + ) + CfnOutput( + self, + "PiExpressServiceName", + value=ECS_PI_EXPRESS_SERVICE_NAME, + ) + CfnOutput( + self, + "ServiceConnectNamespace", + value=ECS_SERVICE_CONNECT_NAMESPACE, + description="Cloud Map namespace for Express Service Connect", + ) + print( + "ECS Express Pi gateway service defined with Service Connect " + f"backend {sc_backend}; public URL: {pi_public_url}." + ) + except Exception as e: + raise Exception( + "Could not handle ECS Express Pi agent due to:", e + ) + + print("ECS Express Gateway service defined.") + except Exception as e: + raise Exception("Could not handle ECS Express Mode due to:", e) + + if not use_express_ingress: + # --- Fargate Task Definition --- + try: + fargate_task_definition_name = FARGATE_TASK_DEFINITION_NAME + + read_only_file_system = ECS_READ_ONLY_FILE_SYSTEM == "True" + + if os.path.exists(TASK_DEFINITION_FILE_LOCATION): + with open(TASK_DEFINITION_FILE_LOCATION) as f: # Use correct path + task_def_params = json.load(f) + # Need to ensure taskRoleArn and executionRoleArn in JSON are correct ARN strings + else: + epheremal_storage_volume_name = "appEphemeralVolume" + + task_def_params = {} + task_def_params["taskRoleArn"] = ( + task_role.role_arn + ) # Use CDK role object ARN + task_def_params["executionRoleArn"] = ( + execution_role.role_arn + ) # Use CDK role object ARN + task_def_params["memory"] = ECS_TASK_MEMORY_SIZE + task_def_params["cpu"] = ECS_TASK_CPU_SIZE + container_def = { + "name": full_ecr_repo_name, + "image": ecr_image_loc + ":latest", + "essential": True, + "portMappings": [ + { + "containerPort": int(GRADIO_SERVER_PORT), + "hostPort": int(GRADIO_SERVER_PORT), + "protocol": "tcp", + "appProtocol": "http", + } + ], + "logConfiguration": { + "logDriver": "awslogs", + "options": { + "awslogs-group": ECS_LOG_GROUP_NAME, + "awslogs-region": AWS_REGION, + "awslogs-stream-prefix": "ecs", + }, + }, + "environmentFiles": ( + [] + if enable_headless + else [ + { + "value": bucket.bucket_arn + + f"/{APP_CONFIG_ENV_BASENAME}", + "type": "s3", + } + ] + ), + "memoryReservation": int(task_def_params["memory"]) + - 512, # Reserve some memory for the container + "mountPoints": [ + { + "sourceVolume": epheremal_storage_volume_name, + "containerPath": "/home/user/app/logs", + "readOnly": False, + }, + { + "sourceVolume": epheremal_storage_volume_name, + "containerPath": "/home/user/app/feedback", + "readOnly": False, + }, + { + "sourceVolume": epheremal_storage_volume_name, + "containerPath": "/home/user/app/usage", + "readOnly": False, + }, + { + "sourceVolume": epheremal_storage_volume_name, + "containerPath": "/home/user/app/input", + "readOnly": False, + }, + { + "sourceVolume": epheremal_storage_volume_name, + "containerPath": "/home/user/app/output", + "readOnly": False, + }, + { + "sourceVolume": epheremal_storage_volume_name, + "containerPath": "/home/user/app/config", + "readOnly": False, + }, + { + "sourceVolume": epheremal_storage_volume_name, + "containerPath": "/tmp/matplotlib_cache", + "readOnly": False, + }, + { + "sourceVolume": epheremal_storage_volume_name, + "containerPath": "/tmp", + "readOnly": False, + }, + { + "sourceVolume": epheremal_storage_volume_name, + "containerPath": "/var/tmp", + "readOnly": False, + }, + { + "sourceVolume": epheremal_storage_volume_name, + "containerPath": "/tmp/gradio_tmp", + "readOnly": False, + }, + ], + "readonlyRootFilesystem": read_only_file_system, + "user": "1000", + } + task_def_params["containerDefinitions"] = [container_def] + + log_group_name_from_config = task_def_params["containerDefinitions"][0][ + "logConfiguration" + ]["options"]["awslogs-group"] + + cdk_managed_log_group = logs.LogGroup( + self, + "MyTaskLogGroup", # CDK Logical ID + log_group_name=log_group_name_from_config, + retention=logs.RetentionDays.ONE_MONTH, + removal_policy=resource_removal_policy, + ) + cdk_managed_log_group.grant_write(execution_role) + + epheremal_storage_volume_cdk_obj = ecs.Volume( + name=epheremal_storage_volume_name + ) + + fargate_task_definition = ecs.FargateTaskDefinition( + self, + "FargateTaskDefinition", # Logical ID + family=fargate_task_definition_name, + cpu=int(task_def_params["cpu"]), + memory_limit_mib=int(task_def_params["memory"]), + task_role=task_role, + execution_role=execution_role, + runtime_platform=ecs.RuntimePlatform( + cpu_architecture=ecs.CpuArchitecture.X86_64, + operating_system_family=ecs.OperatingSystemFamily.LINUX, + ), + ephemeral_storage_gib=21, # Minimum is 21 GiB + volumes=[epheremal_storage_volume_cdk_obj], + ) + print("Fargate task definition defined.") + + # Add container definitions to the task definition object + if task_def_params["containerDefinitions"]: + container_def_params = task_def_params["containerDefinitions"][0] + + env_files = [] + if container_def_params.get("environmentFiles"): + for env_file_param in container_def_params["environmentFiles"]: + # Need to parse the ARN to get the bucket object and key + env_file_arn_parts = env_file_param["value"].split(":::") + bucket_name_and_key = env_file_arn_parts[-1] + env_bucket_name, env_key = bucket_name_and_key.split("/", 1) + + env_file = ecs.EnvironmentFile.from_bucket(bucket, env_key) + + env_files.append(env_file) + + container_kwargs: Dict[str, Any] = { + "image": ecs.ContainerImage.from_registry( + container_def_params["image"] + ), + "logging": ecs.LogDriver.aws_logs( + stream_prefix=container_def_params["logConfiguration"][ + "options" + ]["awslogs-stream-prefix"], + log_group=cdk_managed_log_group, + ), + "environment_files": env_files if env_files else None, + "readonly_root_filesystem": read_only_file_system, + "user": container_def_params.get("user", "1000"), + } + if not enable_headless: + container_kwargs["secrets"] = { + "AWS_USER_POOL_ID": ecs.Secret.from_secrets_manager( + secret, "SUMMARISATION_USER_POOL_ID" + ), + "AWS_CLIENT_ID": ecs.Secret.from_secrets_manager( + secret, "SUMMARISATION_CLIENT_ID" + ), + "AWS_CLIENT_SECRET": ecs.Secret.from_secrets_manager( + secret, "SUMMARISATION_CLIENT_SECRET" + ), + } + container = fargate_task_definition.add_container( + container_def_params["name"], + **container_kwargs, + ) + + for port_mapping in container_def_params["portMappings"]: + container.add_port_mappings( + ecs.PortMapping( + container_port=int(port_mapping["containerPort"]), + host_port=int(port_mapping["hostPort"]), + name="port-" + str(port_mapping["containerPort"]), + app_protocol=ecs.AppProtocol.http, + protocol=ecs.Protocol.TCP, + ) + ) + + container.add_port_mappings( + ecs.PortMapping( + container_port=80, + host_port=80, + name="port-80", + app_protocol=ecs.AppProtocol.http, + protocol=ecs.Protocol.TCP, + ) + ) + + if container_def_params.get("mountPoints"): + mount_points = [] + for mount_point in container_def_params["mountPoints"]: + mount_points.append( + ecs.MountPoint( + container_path=mount_point["containerPath"], + read_only=mount_point["readOnly"], + source_volume=epheremal_storage_volume_name, + ) + ) + container.add_mount_points(*mount_points) + + except Exception as e: + raise Exception("Could not handle Fargate task definition due to:", e) + ecs_service = None + if deploy_web_ingress: + # --- ECS Service --- + try: + ecs_service_name = ECS_SERVICE_NAME + + if ECS_USE_FARGATE_SPOT == "True": + use_fargate_spot = "FARGATE_SPOT" + if ECS_USE_FARGATE_SPOT == "False": + use_fargate_spot = "FARGATE" + + # Check if service exists - from_service_arn or from_service_name (needs cluster) + try: + # from_service_name is useful if you have the cluster object + ecs_service = ecs.FargateService.from_service_attributes( + self, + "ECSService", # Logical ID + cluster=cluster, # Requires the cluster object + service_name=ecs_service_name, + ) + print(f"Using existing ECS service {ecs_service_name}.") + if enable_service_connect: + print( + "Warning: ENABLE_ECS_SERVICE_CONNECT=True but an existing " + "ECS service was imported; enable Service Connect on that " + "service in the ECS console or replace the service via CDK." + ) + except Exception: + service_connect_configuration = None + if enable_service_connect: + sc_dns_name = ( + ECS_SERVICE_CONNECT_DNS_NAME + or ECS_SERVICE_CONNECT_DISCOVERY_NAME + ) + service_connect_configuration = ecs.ServiceConnectProps( + namespace=ECS_SERVICE_CONNECT_NAMESPACE, + services=[ + ecs.ServiceConnectService( + port_mapping_name=ECS_SERVICE_CONNECT_PORT_MAPPING_NAME, + discovery_name=ECS_SERVICE_CONNECT_DISCOVERY_NAME, + dns_name=sc_dns_name, + port=int(GRADIO_SERVER_PORT), + ) + ], + ) + # Service will be created with a count of 0, because you haven't yet actually built the initial Docker container with CodeBuild + ecs_service = ecs.FargateService( + self, + "ECSService", # Logical ID + service_name=ecs_service_name, # Explicit resource name + platform_version=ecs.FargatePlatformVersion.LATEST, + capacity_provider_strategies=[ + ecs.CapacityProviderStrategy( + capacity_provider=use_fargate_spot, base=0, weight=1 + ) + ], + cluster=cluster, + task_definition=fargate_task_definition, # Link to TD + security_groups=[ecs_security_group], # Link to SG + vpc_subnets=ec2.SubnetSelection( + subnets=self.private_subnets + ), # Link to subnets + min_healthy_percent=0, + max_healthy_percent=600, + desired_count=0, + availability_zone_rebalancing=ecs_availability_zone_rebalancing( + ECS_AVAILABILITY_ZONE_REBALANCING + ), + service_connect_configuration=service_connect_configuration, + ) + print("Successfully created new ECS service") + + # Note: Auto-scaling setup would typically go here if needed for the service + + except Exception as e: + raise Exception("Could not handle ECS service due to:", e) + + if enable_pi_agent: + try: + pi_ecs_service, pi_ecs_security_group, _pi_task_def = ( + create_pi_agent_ecs_resources( + self, + "PiAgent", + vpc=vpc, + cluster=cluster, + private_subnets=self.private_subnets, + pi_ecr_image_uri=pi_ecr_image_loc, + container_name=ECR_PI_REPO_NAME, + task_role=task_role, + execution_role=execution_role, + config_bucket=bucket, + pi_agent_env_s3_key=PI_AGENT_ENV_S3_KEY, + service_name=ECS_PI_SERVICE_NAME, + task_family=ECS_PI_TASK_DEFINITION_NAME, + security_group_name=ECS_PI_SECURITY_GROUP_NAME, + log_group_name=ECS_PI_LOG_GROUP_NAME, + cpu=int(ECS_PI_TASK_CPU_SIZE), + memory_mib=int(ECS_PI_TASK_MEMORY_SIZE), + pi_gradio_port=int(PI_GRADIO_PORT), + service_connect_namespace=ECS_SERVICE_CONNECT_NAMESPACE, + service_connect_discovery_name=ECS_SERVICE_CONNECT_DISCOVERY_NAME, + main_app_port=int(GRADIO_SERVER_PORT), + use_fargate_spot=use_fargate_spot, + pi_root_path=pi_alb_root_path_for_container( + PI_ALB_PATH_PREFIX_NORMALIZED, PI_ALB_ROUTING + ), + ) + ) + ecs_security_group.add_ingress_rule( + peer=pi_ecs_security_group, + connection=ec2_port_gradio_server_port, + description="Pi agent (Service Connect) to main summarisation app", + ) + print("Pi agent ECS service defined.") + except Exception as e: + raise Exception("Could not handle Pi agent ECS service due to:", e) + + if ENABLE_S3_BATCH_ECS_TRIGGER == "True": + try: + batch_subnet_ids = [s.subnet_id for s in self.private_subnets] + if not batch_subnet_ids: + batch_subnet_ids = [s.subnet_id for s in self.public_subnets] + if not batch_subnet_ids: + raise ValueError( + "S3 batch ECS trigger requires at least one public or " + "private subnet." + ) + lambda_asset_dir = os.path.join( + os.path.dirname(__file__), "config", "lambda" + ) + batch_lambda = create_s3_batch_ecs_trigger_lambda( + self, + "S3BatchEcsTrigger", + function_name=S3_BATCH_LAMBDA_FUNCTION_NAME or None, + lambda_asset_path=lambda_asset_dir, + output_bucket=output_bucket, + config_bucket=bucket, + cluster_name=CLUSTER_NAME, + task_definition_arn=fargate_task_definition.task_definition_arn, + container_name=full_ecr_repo_name, + subnet_ids=batch_subnet_ids, + security_group_id=ecs_security_group.security_group_id, + execution_role=execution_role, + task_role=task_role, + env_prefix=S3_BATCH_ENV_PREFIX, + env_suffix=S3_BATCH_ENV_SUFFIX, + input_prefix=S3_BATCH_INPUT_PREFIX, + config_prefix=S3_BATCH_CONFIG_PREFIX, + default_params_key=S3_BATCH_DEFAULT_PARAMS_KEY, + general_env_prefix=S3_BATCH_GENERAL_ENV_PREFIX, + default_task_type="extract", + assign_public_ip=not bool(self.private_subnets), + ) + CfnOutput( + self, + "BatchEcsTriggerLambdaArn", + value=batch_lambda.function_arn, + description="Lambda ARN for S3-triggered batch ECS tasks", + ) + CfnOutput( + self, + "BatchJobEnvPrefix", + value=f"s3://{output_bucket.bucket_name}/{S3_BATCH_ENV_PREFIX}", + description="Upload job .env files here to start batch topic-modelling tasks", + ) + CfnOutput( + self, + "BatchInputPrefix", + value=f"s3://{output_bucket.bucket_name}/{S3_BATCH_INPUT_PREFIX}", + description="Upload consultation spreadsheets and other input files for batch jobs", + ) + CfnOutput( + self, + "BatchEcsTriggerLambdaName", + value=batch_lambda.function_name, + description="Lambda that starts ECS batch tasks on job .env upload", + ) + if enable_headless: + seed_asset_dir = os.path.join( + os.path.dirname(__file__), "config", "headless_s3_seed" + ) + create_headless_s3_batch_seed( + self, + "HeadlessBatchS3Seed", + destination_bucket=output_bucket, + seed_asset_directory=seed_asset_dir, + s3_outputs_bucket_name=output_bucket.bucket_name, + ) + if ( + enable_headless + and ENABLE_HEADLESS_OUTPUT_NOTIFICATIONS == "True" + ): + create_headless_output_notifications( + self, + "HeadlessOutputNotifications", + output_bucket=output_bucket, + output_prefix=HEADLESS_OUTPUT_S3_PREFIX, + notify_email=HEADLESS_OUTPUT_NOTIFY_EMAIL, + iam_user_name=HEADLESS_OUTPUT_IAM_USER_NAME, + metric_filter_id=HEADLESS_OUTPUT_S3_METRIC_FILTER_ID, + sns_topic_name=HEADLESS_OUTPUT_SNS_TOPIC_NAME, + alarm_name=HEADLESS_OUTPUT_ALARM_NAME, + kms_key_arn=shared_kms_key_arn, + ) + print( + "Headless output notifications enabled: S3 PutRequests " + f"alarm -> SNS ({HEADLESS_OUTPUT_NOTIFY_EMAIL})." + ) + print("S3 batch ECS trigger Lambda defined.") + except Exception as e: + raise Exception("Could not handle S3 batch ECS trigger due to:", e) + + if deploy_web_ingress: + # --- ALB TARGET GROUPS AND LISTENERS --- + # This section should primarily define the resources if they are managed by this stack. + # CDK handles adding/removing targets and actions on updates. + # If they might pre-exist outside the stack, you need lookups. + cookie_duration = Duration.hours(8) + target_group_name = ALB_TARGET_GROUP_NAME # Explicit resource name + cloudfront_distribution_url = "cloudfront_placeholder.net" # Need to replace this afterwards with the actual cloudfront_distribution.domain_name + cloudfront_http_rule_priority = ( + PI_ALB_LISTENER_RULE_PRIORITY + + (pi_listener_rule_count(PI_ALB_ROUTING) if enable_pi_agent else 0) + if enable_pi_agent + else 1 + ) + https_listener = None + + try: + # --- CREATING TARGET GROUPS AND ADDING THE CLOUDFRONT LISTENER RULE --- + + target_group = elbv2.ApplicationTargetGroup( + self, + "AppTargetGroup", # Logical ID + target_group_name=target_group_name, # Explicit resource name + port=int(GRADIO_SERVER_PORT), # Ensure port is int + protocol=elbv2.ApplicationProtocol.HTTP, + targets=[ecs_service], # Link to ECS Service + stickiness_cookie_duration=cookie_duration, + vpc=vpc, # Target Groups need VPC + ) + print(f"ALB target group {target_group_name} defined.") + + # First HTTP + listener_port = 80 + # Check if Listener exists - from_listener_arn or lookup by port/ALB + + http_listener = alb.add_listener( + "HttpListener", # Logical ID + port=listener_port, + open=False, # Be cautious with open=True, usually restrict source SG + ) + print(f"ALB listener on port {listener_port} defined.") + + if ACM_SSL_CERTIFICATE_ARN: + http_listener.add_action( + "DefaultAction", # Logical ID for the default action + action=elbv2.ListenerAction.redirect( + protocol="HTTPS", + host="#{host}", + port="443", + path="/#{path}", + query="#{query}", + ), + ) + else: + if USE_CLOUDFRONT == "True": + + # The following default action can be added for the listener after a host header rule is added to the listener manually in the Console as suggested in the above comments. + http_listener.add_action( + "DefaultAction", # Logical ID for the default action + action=elbv2.ListenerAction.fixed_response( + status_code=403, + content_type="text/plain", + message_body="Access denied", + ), + ) + + # Add the Listener Rule for the specific CloudFront Host Header + http_listener.add_action( + "CloudFrontHostHeaderRule", + action=elbv2.ListenerAction.forward( + target_groups=[target_group], + stickiness_duration=cookie_duration, + ), + priority=cloudfront_http_rule_priority, + conditions=[ + elbv2.ListenerCondition.host_headers( + [cloudfront_distribution_url] + ) # May have to redefine url in console afterwards if not specified in config file + ], + ) + + else: + # Add the Listener Rule for the specific CloudFront Host Header + http_listener.add_action( + "CloudFrontHostHeaderRule", + action=elbv2.ListenerAction.forward( + target_groups=[target_group], + stickiness_duration=cookie_duration, + ), + priority=cloudfront_http_rule_priority, + ) + + print("Added targets and actions to ALB HTTP listener.") + + # Now the same for HTTPS if you have an ACM certificate + if ACM_SSL_CERTIFICATE_ARN: + listener_port_https = 443 + # Check if Listener exists - from_listener_arn or lookup by port/ALB + + https_listener = add_alb_https_listener_with_cert( + self, + "MyHttpsListener", # Logical ID for the HTTPS listener + alb, + acm_certificate_arn=ACM_SSL_CERTIFICATE_ARN, + default_target_group=target_group, + enable_cognito_auth=True, + cognito_user_pool=user_pool, + cognito_user_pool_client=user_pool_client, + cognito_user_pool_domain=user_pool_domain, + listener_open_to_internet=True, + stickiness_cookie_duration=cookie_duration, + ) + + if https_listener: + CfnOutput( + self, + "HttpsListenerArn", + value=https_listener.listener_arn, + ) + + print(f"ALB listener on port {listener_port_https} defined.") + + # if USE_CLOUDFRONT == 'True': + # # Add default action to the listener + # https_listener.add_action( + # "DefaultAction", # Logical ID for the default action + # action=elbv2.ListenerAction.fixed_response( + # status_code=403, + # content_type="text/plain", + # message_body="Access denied", + # ), + # ) + + # # Add the Listener Rule for the specific CloudFront Host Header + # https_listener.add_action( + # "CloudFrontHostHeaderRuleHTTPS", + # action=elbv2.ListenerAction.forward(target_groups=[target_group],stickiness_duration=cookie_duration), + # priority=1, # Example priority. Adjust as needed. Lower is evaluated first. + # conditions=[ + # elbv2.ListenerCondition.host_headers([cloudfront_distribution_url]) + # ] + # ) + # else: + # https_listener.add_action( + # "CloudFrontHostHeaderRuleHTTPS", + # action=elbv2.ListenerAction.forward(target_groups=[target_group],stickiness_duration=cookie_duration)) + + print("Added targets and actions to ALB HTTPS listener.") + + if enable_pi_agent and pi_ecs_service and alb_security_group: + pi_tg_name = PI_ALB_TARGET_GROUP_NAME + if len(pi_tg_name) > 32: + pi_tg_name = pi_tg_name[-32:] + + _pi_public_urls = format_pi_public_urls( + routing_mode=PI_ALB_ROUTING, + path_prefix=PI_ALB_PATH_PREFIX_NORMALIZED, + host_header=PI_ALB_HOST_HEADER, + cloudfront_domain=( + CLOUDFRONT_DOMAIN if USE_CLOUDFRONT == "True" else "" + ), + use_https=bool(ACM_SSL_CERTIFICATE_ARN), + ) + attach_pi_agent_to_shared_alb( + self, + "PiAgent", + vpc=vpc, + alb_security_group=alb_security_group, + pi_security_group=pi_ecs_security_group, + pi_service=pi_ecs_service, + pi_port=int(PI_GRADIO_PORT), + routing_mode=PI_ALB_ROUTING, + path_prefix=PI_ALB_PATH_PREFIX_NORMALIZED, + pi_host_header=PI_ALB_HOST_HEADER.strip(), + listener_rule_priority=PI_ALB_LISTENER_RULE_PRIORITY, + target_group_name=pi_tg_name, + stickiness_cookie_duration=cookie_duration, + https_listener=https_listener, + http_listener=http_listener, + acm_certificate_arn=ACM_SSL_CERTIFICATE_ARN or "", + enable_cognito_auth=bool(ACM_SSL_CERTIFICATE_ARN), + cognito_user_pool=user_pool, + cognito_user_pool_client=user_pool_client, + cognito_user_pool_domain=user_pool_domain, + ) + pi_public_url = _pi_public_urls[0] if _pi_public_urls else "" + CfnOutput( + self, + "PiPublicUrl", + value=pi_public_url, + description="Primary public URL for Pi agent UI (path and/or host ALB rules)", + ) + if len(_pi_public_urls) > 1: + CfnOutput( + self, + "PiPublicUrls", + value=", ".join(_pi_public_urls), + description="All configured Pi UI entry URLs", + ) + CfnOutput( + self, + "PiAlbPathPrefix", + value=PI_ALB_PATH_PREFIX_NORMALIZED, + description="ALB path prefix for Pi when PI_ALB_ROUTING includes path", + ) + CfnOutput( + self, + "PiAgentServiceName", + value=ECS_PI_SERVICE_NAME, + ) + sc_backend = ( + f"http://{ECS_SERVICE_CONNECT_DISCOVERY_NAME}:" + f"{GRADIO_SERVER_PORT}" + ) + CfnOutput( + self, + "PiDocSummarisationBackendUrl", + value=sc_backend, + description="DOC_SUMMARISATION_GRADIO_URL set on Pi tasks (Service Connect)", + ) + print( + "Pi agent attached to shared ALB " + f"(routing={PI_ALB_ROUTING}, urls={', '.join(_pi_public_urls)})." + ) + + except Exception as e: + raise Exception( + "Could not handle ALB target groups and listeners due to:", e + ) + if not enable_headless: + # Create WAF to attach to load balancer + try: + web_acl_name = LOAD_BALANCER_WEB_ACL_NAME + if get_context_bool(f"exists:{web_acl_name}"): + # Lookup WAF ACL by ARN from context + web_acl_arn = get_context_str(f"arn:{web_acl_name}") + if not web_acl_arn: + raise ValueError( + f"Context value 'arn:{web_acl_name}' is required if Web ACL exists." + ) + + web_acl = create_web_acl_with_common_rules( + self, web_acl_name, waf_scope="REGIONAL" + ) # Assuming it takes scope and name + print(f"Handled ALB WAF web ACL {web_acl_name}.") + else: + web_acl = create_web_acl_with_common_rules( + self, web_acl_name, waf_scope="REGIONAL" + ) # Assuming it takes scope and name + print(f"Created ALB WAF web ACL {web_acl_name}.") + + wafv2.CfnWebACLAssociation( + self, + id="alb_waf_association", + resource_arn=alb.load_balancer_arn, + web_acl_arn=web_acl.attr_arn, + ) + + except Exception as e: + raise Exception("Could not handle create ALB WAF web ACL due to:", e) + + # --- Outputs for other stacks/regions --- + + self.params = dict() + self.params["alb_arn_output"] = alb.load_balancer_arn + if use_express_ingress: + self.params["alb_security_group_id"] = express_alb_security_group_id + else: + self.params["alb_security_group_id"] = ( + alb_security_group.security_group_id + ) + self.params["alb_dns_name"] = alb.load_balancer_dns_name + + CfnOutput( + self, + "AlbArnOutput", + value=alb.load_balancer_arn, + description="ARN of the Application Load Balancer", + export_name=f"{self.stack_name}-AlbArn", + ) # Export name must be unique within the account/region + + CfnOutput( + self, + "AlbSecurityGroupIdOutput", + value=( + express_alb_security_group_id + if use_express_ingress + else alb_security_group.security_group_id + ), + description="ID of the ALB's Security Group", + export_name=f"{self.stack_name}-AlbSgId", + ) + CfnOutput(self, "ALBName", value=load_balancer_name) + + CfnOutput(self, "RegionalAlbDnsName", value=alb.load_balancer_dns_name) + else: + self.params = dict() + CfnOutput( + self, + "HeadlessDeploymentMode", + value="True", + description="Stack deployed for S3-triggered direct-mode batch tasks only", + ) + CfnOutput( + self, + "ECSClusterName", + value=CLUSTER_NAME, + description="ECS cluster used for one-shot Fargate batch tasks", + ) + CfnOutput( + self, + "EcsBatchLogGroup", + value=ECS_LOG_GROUP_NAME, + description=( + "CloudWatch log group for batch tasks (streams appear only after " + "the container starts; init failures may have no stream)" + ), + ) + + if not enable_headless and user_pool is not None: + CfnOutput(self, "CognitoPoolId", value=user_pool.user_pool_id) + # Add other outputs if needed + + CfnOutput(self, "ECRRepoUri", value=ecr_repo.repository_uri) + + if enable_service_connect: + sc_host = ECS_SERVICE_CONNECT_DNS_NAME or ECS_SERVICE_CONNECT_DISCOVERY_NAME + sc_base = f"http://{sc_host}:{GRADIO_SERVER_PORT}" + CfnOutput( + self, + "ServiceConnectHttpBaseUrl", + value=sc_base, + description="Base URL for other ECS services in this cluster (Service Connect)", + ) + CfnOutput( + self, + "ServiceConnectAgentApiUrl", + value=f"{sc_base}/agent", + description="FastAPI Agent API prefix (when RUN_FASTAPI=True in app_config.env)", + ) + CfnOutput( + self, + "ServiceConnectNamespace", + value=ECS_SERVICE_CONNECT_NAMESPACE, + ) + + +# --- CLOUDFRONT DISTRIBUTION in separate stack (us-east-1 required) --- +class CdkStackCloudfront(Stack): + + def __init__( + self, + scope: Construct, + construct_id: str, + alb_arn: str, + alb_sec_group_id: str, + alb_dns_name: str, + **kwargs, + ) -> None: + super().__init__(scope, construct_id, **kwargs) + + # --- Helper to get context values --- + def get_context_bool(key: str, default: bool = False) -> bool: + return self.node.try_get_context(key) or default + + def get_context_str(key: str, default: str = None) -> str: + return self.node.try_get_context(key) or default + + def get_context_dict(scope: Construct, key: str, default: dict = None) -> dict: + return scope.node.try_get_context(key) or default + + resource_removal_policy = managed_resource_removal_policy() + + print(f"CloudFront Stack: Received ALB ARN: {alb_arn}") + print(f"CloudFront Stack: Received ALB Security Group ID: {alb_sec_group_id}") + + if not alb_arn: + raise ValueError("ALB ARN must be provided to CloudFront stack") + if not alb_sec_group_id: + raise ValueError( + "ALB Security Group ID must be provided to CloudFront stack" + ) + + # 2. Import the ALB using its ARN + # This imports an existing ALB as a construct in the CloudFront stack's context. + # CloudFormation will understand this reference at deploy time. + alb = elbv2.ApplicationLoadBalancer.from_application_load_balancer_attributes( + self, + "ImportedAlb", + load_balancer_arn=alb_arn, + security_group_id=alb_sec_group_id, + load_balancer_dns_name=alb_dns_name, + ) + + try: + web_acl_name = WEB_ACL_NAME + if get_context_bool(f"exists:{web_acl_name}"): + # Lookup WAF ACL by ARN from context + web_acl_arn = get_context_str(f"arn:{web_acl_name}") + if not web_acl_arn: + raise ValueError( + f"Context value 'arn:{web_acl_name}' is required if Web ACL exists." + ) + + web_acl = create_web_acl_with_common_rules( + self, web_acl_name + ) # Assuming it takes scope and name + print(f"Handled Cloudfront WAF web ACL {web_acl_name}.") + else: + web_acl = create_web_acl_with_common_rules( + self, web_acl_name + ) # Assuming it takes scope and name + print(f"Created Cloudfront WAF web ACL {web_acl_name}.") + + # Add ALB as CloudFront Origin + origin = origins.LoadBalancerV2Origin( + alb, # Use the created or looked-up ALB object + custom_headers={CUSTOM_HEADER: CUSTOM_HEADER_VALUE}, + origin_shield_enabled=False, + protocol_policy=cloudfront.OriginProtocolPolicy.HTTP_ONLY, + ) + + if CLOUDFRONT_GEO_RESTRICTION: + geo_restrict = cloudfront.GeoRestriction.allowlist( + CLOUDFRONT_GEO_RESTRICTION + ) + else: + geo_restrict = None + + response_headers_policy = None + if CLOUDFRONT_ENABLE_SECURE_RESPONSE_HEADERS == "True": + app_origin, cognito_login_url = resolve_cloudfront_csp_urls( + cognito_redirection_url=COGNITO_REDIRECTION_URL, + cloudfront_domain=CLOUDFRONT_DOMAIN, + cognito_user_pool_domain_prefix=COGNITO_USER_POOL_DOMAIN_PREFIX, + aws_region=AWS_REGION, + cognito_user_pool_login_url=COGNITO_USER_POOL_LOGIN_URL, + ssl_certificate_domain=SSL_CERTIFICATE_DOMAIN, + ) + policy_name = f"{CDK_PREFIX}SecureResponseHeaders"[:128] + response_headers_policy = ( + create_secure_cloudfront_response_headers_policy( + self, + "SecureResponseHeadersPolicy", + policy_name=policy_name, + app_origin=app_origin, + cognito_login_url=cognito_login_url, + ) + ) + print( + "CloudFront secure response headers: " + f"app_origin={app_origin}, cognito_login_url={cognito_login_url}" + ) + + default_behavior = cloudfront.BehaviorOptions( + origin=origin, + viewer_protocol_policy=cloudfront.ViewerProtocolPolicy.REDIRECT_TO_HTTPS, + allowed_methods=cloudfront.AllowedMethods.ALLOW_ALL, + cache_policy=cloudfront.CachePolicy.CACHING_DISABLED, + origin_request_policy=cloudfront.OriginRequestPolicy.ALL_VIEWER, + response_headers_policy=response_headers_policy, + ) + + cloudfront_distribution = cloudfront.Distribution( + self, + "CloudFrontDistribution", # Logical ID + comment=CLOUDFRONT_DISTRIBUTION_NAME, # Use name as comment for easier identification + geo_restriction=geo_restrict, + default_behavior=default_behavior, + web_acl_id=web_acl.attr_arn, + ) + cloudfront_distribution.apply_removal_policy(resource_removal_policy) + print(f"Cloudfront distribution {CLOUDFRONT_DISTRIBUTION_NAME} defined.") + + except Exception as e: + raise Exception("Could not handle Cloudfront distribution due to:", e) + + # --- Outputs --- + CfnOutput( + self, "CloudFrontDistributionURL", value=cloudfront_distribution.domain_name + ) diff --git a/cdk/check_resources.py b/cdk/check_resources.py new file mode 100644 index 0000000000000000000000000000000000000000..2f10f365bbc64dcf0a2ab113878b3cdc21890dcf --- /dev/null +++ b/cdk/check_resources.py @@ -0,0 +1,583 @@ +import json +import os +from typing import Dict, List + +from cdk_config import ( # Import necessary config + ALB_NAME, + AWS_REGION, + CDK_CONFIG_PATH, + CDK_FOLDER, + CODEBUILD_PROJECT_NAME, + CODEBUILD_ROLE_NAME, + COGNITO_USER_POOL_CLIENT_NAME, + COGNITO_USER_POOL_CLIENT_SECRET_NAME, + COGNITO_USER_POOL_DOMAIN_PREFIX, + COGNITO_USER_POOL_NAME, + CONTEXT_FILE, + ECR_CDK_REPO_NAME, + ECR_PI_REPO_NAME, + ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES_TO_LOOKUP, + ECS_TASK_EXECUTION_ROLE_NAME, + ECS_TASK_ROLE_NAME, + ENABLE_ECS_SERVICE_CONNECT, + ENABLE_HEADLESS_DEPLOYMENT, + ENABLE_PI_AGENT_ECS_SERVICE, + ENABLE_S3_BATCH_ECS_TRIGGER, + EXISTING_IGW_ID, + PRIVATE_SUBNET_AVAILABILITY_ZONES, + PRIVATE_SUBNET_CIDR_BLOCKS, + PRIVATE_SUBNETS_TO_USE, + PUBLIC_SUBNET_AVAILABILITY_ZONES, + PUBLIC_SUBNET_CIDR_BLOCKS, + PUBLIC_SUBNETS_TO_USE, + S3_LOG_CONFIG_BUCKET_NAME, + S3_OUTPUT_BUCKET_NAME, + USE_ECS_EXPRESS_MODE, + VPC_NAME, + WEB_ACL_NAME, +) +from cdk_functions import ( # Import your check functions (assuming they use Boto3) + _get_existing_subnets_in_vpc, + audit_public_subnet_internet_connectivity, + check_alb_exists, + check_codebuild_project_exists, + check_ecr_repo_exists, + check_for_existing_role, + check_for_existing_user_pool, + check_for_existing_user_pool_client, + check_for_secret, + check_subnet_exists_by_name, + check_web_acl_exists, + get_secret_kms_key_arn, + get_security_group_id_by_name, + get_vpc_id_by_name, + list_existing_vpc_endpoint_service_names, + resolve_cognito_domain_prefix_availability, + resolve_s3_bucket_availability, + validate_subnet_creation_parameters, + # Add other check functions as needed +) + +cdk_folder = CDK_FOLDER # + +# Full path needed to find config file +os.environ["CDK_CONFIG_PATH"] = cdk_folder + CDK_CONFIG_PATH + + +# --- Helper to parse environment variables into lists --- +def _get_env_list(env_var_name: str) -> List[str]: + """Parses a comma-separated environment variable into a list of strings.""" + value = env_var_name[1:-1].strip().replace('"', "").replace("'", "") + if not value: + return [] + # Split by comma and filter out any empty strings that might result from extra commas + return [s.strip() for s in value.split(",") if s.strip()] + + +if PUBLIC_SUBNETS_TO_USE and not isinstance(PUBLIC_SUBNETS_TO_USE, list): + PUBLIC_SUBNETS_TO_USE = _get_env_list(PUBLIC_SUBNETS_TO_USE) +if PRIVATE_SUBNETS_TO_USE and not isinstance(PRIVATE_SUBNETS_TO_USE, list): + PRIVATE_SUBNETS_TO_USE = _get_env_list(PRIVATE_SUBNETS_TO_USE) +if PUBLIC_SUBNET_CIDR_BLOCKS and not isinstance(PUBLIC_SUBNET_CIDR_BLOCKS, list): + PUBLIC_SUBNET_CIDR_BLOCKS = _get_env_list(PUBLIC_SUBNET_CIDR_BLOCKS) +if PUBLIC_SUBNET_AVAILABILITY_ZONES and not isinstance( + PUBLIC_SUBNET_AVAILABILITY_ZONES, list +): + PUBLIC_SUBNET_AVAILABILITY_ZONES = _get_env_list(PUBLIC_SUBNET_AVAILABILITY_ZONES) +if PRIVATE_SUBNET_CIDR_BLOCKS and not isinstance(PRIVATE_SUBNET_CIDR_BLOCKS, list): + PRIVATE_SUBNET_CIDR_BLOCKS = _get_env_list(PRIVATE_SUBNET_CIDR_BLOCKS) +if PRIVATE_SUBNET_AVAILABILITY_ZONES and not isinstance( + PRIVATE_SUBNET_AVAILABILITY_ZONES, list +): + PRIVATE_SUBNET_AVAILABILITY_ZONES = _get_env_list(PRIVATE_SUBNET_AVAILABILITY_ZONES) + +# Check for the existence of elements in your AWS environment to see if it's necessary to create new versions of the same + + +def check_and_set_context(): + context_data = {} + + # --- Find the VPC ID first --- + if VPC_NAME: + print("VPC_NAME:", VPC_NAME) + vpc_lookup = get_vpc_id_by_name(VPC_NAME) + if not vpc_lookup: + raise RuntimeError( + f"Required VPC '{VPC_NAME}' not found. Cannot proceed with subnet checks." + ) + vpc_id, nat_gateways, vpc_cidr_block, vpc_cidr_blocks = vpc_lookup + + # If you expect only one, or one per AZ and you're creating one per AZ in CDK: + if nat_gateways: + # For simplicity, let's just check if *any* NAT exists in the VPC + # A more robust check would match by subnet, AZ, or a specific tag. + context_data["exists:NatGateway"] = True + context_data["id:NatGateway"] = nat_gateways[0][ + "NatGatewayId" + ] # Store the ID of the first one found + else: + context_data["exists:NatGateway"] = False + context_data["id:NatGateway"] = None + + context_data["vpc_id"] = vpc_id # Store VPC ID in context + if vpc_cidr_block: + context_data["vpc_cidr_block"] = vpc_cidr_block + if vpc_cidr_blocks: + context_data["vpc_cidr_blocks"] = vpc_cidr_blocks + + existing_endpoint_services = sorted( + list_existing_vpc_endpoint_service_names(vpc_id, region_name=AWS_REGION) + ) + if existing_endpoint_services: + context_data["existing_vpc_endpoint_service_names"] = ( + existing_endpoint_services + ) + print( + "Existing VPC endpoints in target VPC (will be skipped on deploy): " + + ", ".join(existing_endpoint_services) + ) + + # SUBNET CHECKS + all_proposed_subnets_data: List[Dict[str, str]] = [] + + # Flag to indicate if full validation mode (with CIDR/AZs) is active + full_validation_mode = False + + # Determine if full validation mode is possible/desired + # It's 'desired' if CIDR/AZs are provided, and their lengths match the name lists. + public_ready_for_full_validation = ( + len(PUBLIC_SUBNETS_TO_USE) > 0 + and len(PUBLIC_SUBNET_CIDR_BLOCKS) == len(PUBLIC_SUBNETS_TO_USE) + and len(PUBLIC_SUBNET_AVAILABILITY_ZONES) == len(PUBLIC_SUBNETS_TO_USE) + ) + private_ready_for_full_validation = ( + len(PRIVATE_SUBNETS_TO_USE) > 0 + and len(PRIVATE_SUBNET_CIDR_BLOCKS) == len(PRIVATE_SUBNETS_TO_USE) + and len(PRIVATE_SUBNET_AVAILABILITY_ZONES) == len(PRIVATE_SUBNETS_TO_USE) + ) + + # Activate full validation if *any* type of subnet (public or private) has its full details provided. + # You might adjust this logic if you require ALL subnet types to have CIDRs, or NONE. + if public_ready_for_full_validation or private_ready_for_full_validation: + full_validation_mode = True + + # If some are ready but others aren't, print a warning or raise an error based on your strictness + if ( + public_ready_for_full_validation + and not private_ready_for_full_validation + and PRIVATE_SUBNETS_TO_USE + ): + print( + "Warning: Public subnets have CIDRs/AZs, but private subnets do not. Only public will be fully validated/created with CIDRs." + ) + if ( + private_ready_for_full_validation + and not public_ready_for_full_validation + and PUBLIC_SUBNETS_TO_USE + ): + print( + "Warning: Private subnets have CIDRs/AZs, but public subnets do not. Only private will be fully validated/created with CIDRs." + ) + + # Prepare data for validate_subnet_creation_parameters for all subnets that have full details + if public_ready_for_full_validation: + for i, name in enumerate(PUBLIC_SUBNETS_TO_USE): + all_proposed_subnets_data.append( + { + "name": name, + "cidr": PUBLIC_SUBNET_CIDR_BLOCKS[i], + "az": PUBLIC_SUBNET_AVAILABILITY_ZONES[i], + } + ) + if private_ready_for_full_validation: + for i, name in enumerate(PRIVATE_SUBNETS_TO_USE): + all_proposed_subnets_data.append( + { + "name": name, + "cidr": PRIVATE_SUBNET_CIDR_BLOCKS[i], + "az": PRIVATE_SUBNET_AVAILABILITY_ZONES[i], + } + ) + + print(f"Target VPC ID for Boto3 lookup: {vpc_id}") + + # Fetch all existing subnets in the target VPC once to avoid repeated API calls + try: + existing_aws_subnets = _get_existing_subnets_in_vpc(vpc_id) + except Exception as e: + print(f"Failed to fetch existing VPC subnets. Aborting. Error: {e}") + raise SystemExit(1) # Exit immediately if we can't get baseline data + + print("\n--- Running Name-Only Subnet Existence Check Mode ---") + # Fallback: check only by name using the existing data + checked_public_subnets = {} + if PUBLIC_SUBNETS_TO_USE: + for subnet_name in PUBLIC_SUBNETS_TO_USE: + print("subnet_name:", subnet_name) + exists, subnet_id = check_subnet_exists_by_name( + subnet_name, existing_aws_subnets + ) + checked_public_subnets[subnet_name] = { + "exists": exists, + "id": subnet_id, + "az": ( + existing_aws_subnets["by_name"].get(subnet_name, {}).get("az") + if exists + else None + ), + "route_table_id": ( + existing_aws_subnets["by_name"] + .get(subnet_name, {}) + .get("route_table_id") + if exists + else None + ), + } + + # If the subnet exists, remove it from the proposed subnets list + if checked_public_subnets[subnet_name]["exists"] is True: + all_proposed_subnets_data = [ + subnet + for subnet in all_proposed_subnets_data + if subnet["name"] != subnet_name + ] + + context_data["checked_public_subnets"] = checked_public_subnets + + checked_private_subnets = {} + if PRIVATE_SUBNETS_TO_USE: + for subnet_name in PRIVATE_SUBNETS_TO_USE: + print("subnet_name:", subnet_name) + exists, subnet_id = check_subnet_exists_by_name( + subnet_name, existing_aws_subnets + ) + checked_private_subnets[subnet_name] = { + "exists": exists, + "id": subnet_id, + "az": ( + existing_aws_subnets["by_name"].get(subnet_name, {}).get("az") + if exists + else None + ), + "route_table_id": ( + existing_aws_subnets["by_name"] + .get(subnet_name, {}) + .get("route_table_id") + if exists + else None + ), + } + + # If the subnet exists, remove it from the proposed subnets list + if checked_private_subnets[subnet_name]["exists"] is True: + all_proposed_subnets_data = [ + subnet + for subnet in all_proposed_subnets_data + if subnet["name"] != subnet_name + ] + + context_data["checked_private_subnets"] = checked_private_subnets + + # Internet Gateway + public subnet default routes (legacy ALB / NAT public subnets) + if PUBLIC_SUBNETS_TO_USE and vpc_id: + public_entries_for_igw_audit = [] + for subnet_name in PUBLIC_SUBNETS_TO_USE: + info = checked_public_subnets.get(subnet_name, {}) + if not info.get("exists"): + continue + public_entries_for_igw_audit.append( + { + "name": subnet_name, + "subnet_id": info.get("id"), + "route_table_id": info.get("route_table_id"), + } + ) + if public_entries_for_igw_audit or EXISTING_IGW_ID: + print("\n--- Auditing Internet Gateway and public subnet routes ---") + try: + igw_audit = audit_public_subnet_internet_connectivity( + vpc_id, + EXISTING_IGW_ID, + public_entries_for_igw_audit, + ) + context_data["internet_gateway_id"] = igw_audit[ + "internet_gateway_id" + ] + context_data["internet_gateway_needs_vpc_attachment"] = igw_audit[ + "internet_gateway_needs_vpc_attachment" + ] + context_data["public_subnets_needing_igw_route"] = igw_audit[ + "public_subnets_needing_igw_route" + ] + needing = igw_audit["public_subnets_needing_igw_route"] + if igw_audit["internet_gateway_needs_vpc_attachment"]: + print( + f"CDK will attach IGW '{igw_audit['internet_gateway_id']}' " + f"to VPC '{vpc_id}' on deploy." + ) + if needing: + print( + f"CDK will add default internet routes on deploy for: " + f"{', '.join(n['name'] for n in needing)}" + ) + else: + print( + "All audited public subnets already have 0.0.0.0/0 -> IGW routes." + ) + except ValueError as e: + print( + f"\nFATAL ERROR: Internet Gateway / public route audit failed: {e}\n" + ) + raise SystemExit(1) from e + + print("\nName-only existence subnet check complete.\n") + + if full_validation_mode: + print( + "\n--- Running in Full Subnet Validation Mode (CIDR/AZs provided) ---" + ) + try: + validate_subnet_creation_parameters( + vpc_id, all_proposed_subnets_data, existing_aws_subnets + ) + print("\nPre-synth validation successful. Proceeding with CDK synth.\n") + + # Populate context_data for downstream CDK construct creation. + # Skip subnets that already exist in AWS (imported in the stack). + context_data["public_subnets_to_create"] = [] + if public_ready_for_full_validation: + for i, name in enumerate(PUBLIC_SUBNETS_TO_USE): + if checked_public_subnets.get(name, {}).get("exists"): + continue + context_data["public_subnets_to_create"].append( + { + "name": name, + "cidr": PUBLIC_SUBNET_CIDR_BLOCKS[i], + "az": PUBLIC_SUBNET_AVAILABILITY_ZONES[i], + "is_public": True, + } + ) + context_data["private_subnets_to_create"] = [] + if private_ready_for_full_validation: + for i, name in enumerate(PRIVATE_SUBNETS_TO_USE): + if checked_private_subnets.get(name, {}).get("exists"): + continue + context_data["private_subnets_to_create"].append( + { + "name": name, + "cidr": PRIVATE_SUBNET_CIDR_BLOCKS[i], + "az": PRIVATE_SUBNET_AVAILABILITY_ZONES[i], + "is_public": False, + } + ) + + except (ValueError, Exception) as e: + print(f"\nFATAL ERROR: Subnet parameter validation failed: {e}\n") + raise SystemExit(1) # Exit if validation fails + + # Example checks and setting context values + # IAM Roles + role_name = CODEBUILD_ROLE_NAME + exists, role_arn, _ = check_for_existing_role(role_name) + context_data[f"exists:{role_name}"] = exists + if exists: + context_data[f"arn:{role_name}"] = role_arn + + role_name = ECS_TASK_ROLE_NAME + exists, role_arn, _ = check_for_existing_role(role_name) + context_data[f"exists:{role_name}"] = exists + if exists: + context_data[f"arn:{role_name}"] = role_arn + + role_name = ECS_TASK_EXECUTION_ROLE_NAME + exists, role_arn, _ = check_for_existing_role(role_name) + context_data[f"exists:{role_name}"] = exists + if exists: + context_data[f"arn:{role_name}"] = role_arn + + # S3 Buckets + def _record_s3_bucket_context(name: str) -> None: + status, _ = resolve_s3_bucket_availability(name) + context_data[f"exists:{name}"] = status == "owned" + context_data[f"globally_taken:{name}"] = status == "globally_taken" + + bucket_name = S3_LOG_CONFIG_BUCKET_NAME + _record_s3_bucket_context(bucket_name) + + output_bucket_name = S3_OUTPUT_BUCKET_NAME + _record_s3_bucket_context(output_bucket_name) + + # ECR Repositories + for repo_name in (ECR_CDK_REPO_NAME, ECR_PI_REPO_NAME): + exists, _ = check_ecr_repo_exists(repo_name) + context_data[f"exists:{repo_name}"] = exists + + # CodeBuild Project + project_name = CODEBUILD_PROJECT_NAME + exists, project_arn, service_role_arn = check_codebuild_project_exists(project_name) + context_data[f"exists:{project_name}"] = exists + if exists: + context_data[f"arn:{project_name}"] = project_arn + if service_role_arn: + context_data[f"service_role_arn:{project_name}"] = service_role_arn + + # ALB (by name lookup) — skipped when Express Mode will provision its own ALB + alb_name = ALB_NAME[-32:] if len(ALB_NAME) > 32 else ALB_NAME + if USE_ECS_EXPRESS_MODE == "True": + context_data[f"exists:{alb_name}"] = False + print( + "USE_ECS_EXPRESS_MODE=True: skipping ALB pre-check (Express provisions ALB)." + ) + elif ENABLE_HEADLESS_DEPLOYMENT == "True": + context_data[f"exists:{alb_name}"] = False + print( + "ENABLE_HEADLESS_DEPLOYMENT=True: skipping ALB pre-check (no web ingress)." + ) + if ENABLE_HEADLESS_DEPLOYMENT == "True": + print( + "ENABLE_HEADLESS_DEPLOYMENT=True: requires ENABLE_S3_BATCH_ECS_TRIGGER=True " + "and USE_ECS_EXPRESS_MODE=False." + ) + elif ENABLE_S3_BATCH_ECS_TRIGGER == "True": + print( + "ENABLE_S3_BATCH_ECS_TRIGGER=True: requires legacy Fargate (USE_ECS_EXPRESS_MODE=False)." + ) + if ENABLE_PI_AGENT_ECS_SERVICE == "True": + print( + "ENABLE_PI_AGENT_ECS_SERVICE=True: requires legacy Fargate, Service Connect, " + "and PI_ALB_ROUTING (default path=/pi on shared ALB; host mode needs PI_ALB_HOST_HEADER)." + ) + elif ENABLE_HEADLESS_DEPLOYMENT == "True": + print( + "ENABLE_HEADLESS_DEPLOYMENT=True: legacy Fargate task definition + " + "S3 batch Lambda only (no ALB pre-check)." + ) + else: + exists, alb_object = check_alb_exists(alb_name, region_name=AWS_REGION) + context_data[f"exists:{alb_name}"] = exists + if exists: + print("alb_object:", alb_object) + context_data[f"arn:{alb_name}"] = alb_object["LoadBalancerArn"] + context_data[f"dns:{alb_name}"] = alb_object["DNSName"] + context_data[f"canonical_hosted_zone_id:{alb_name}"] = alb_object[ + "CanonicalHostedZoneId" + ] + if alb_object.get("SecurityGroups"): + context_data[f"security_group_id:{alb_name}"] = alb_object[ + "SecurityGroups" + ][0] + + # Cognito (web login only; headless batch mode has no Gradio/ALB ingress) + if ENABLE_HEADLESS_DEPLOYMENT != "True": + domain_prefix = (COGNITO_USER_POOL_DOMAIN_PREFIX or "").strip().lower() + if domain_prefix: + availability = resolve_cognito_domain_prefix_availability( + domain_prefix, region_name=AWS_REGION + ) + context_data[f"cognito_domain_taken:{domain_prefix}"] = ( + availability == "taken" + ) + + user_pool_name = COGNITO_USER_POOL_NAME + exists, user_pool_id, _ = check_for_existing_user_pool(user_pool_name) + context_data[f"exists:{user_pool_name}"] = exists + if exists: + context_data[f"id:{user_pool_name}"] = user_pool_id + + # Cognito User Pool Client (by name and pool ID) - requires User Pool ID from check + if user_pool_id: + user_pool_id_for_client_check = user_pool_id # context_data.get(f"id:{user_pool_name}") # Use ID from context + user_pool_client_name = COGNITO_USER_POOL_CLIENT_NAME + if user_pool_id_for_client_check: + exists, client_id, _ = check_for_existing_user_pool_client( + user_pool_client_name, user_pool_id_for_client_check + ) + context_data[f"exists:{user_pool_client_name}"] = exists + if exists: + context_data[f"id:{user_pool_client_name}"] = client_id + else: + print( + f"User pool '{user_pool_name}' exists but app client " + f"'{user_pool_client_name}' does not; CDK will create a new client." + ) + + # Secrets Manager Secret (by name) + secret_name = COGNITO_USER_POOL_CLIENT_SECRET_NAME + exists, secret_response = check_for_secret(secret_name) + context_data[f"exists:{secret_name}"] = exists + if exists: + secret_arn = ( + secret_response.get("ARN") + if isinstance(secret_response, dict) + else None + ) + if secret_arn: + context_data[f"arn:{secret_name}"] = secret_arn + print("Cognito secret ARN recorded for IAM grants.") + secret_kms_key_arn = get_secret_kms_key_arn( + secret_name, region_name=AWS_REGION + ) + if secret_kms_key_arn: + context_data[f"kms_key_arn:{secret_name}"] = secret_kms_key_arn + print("Cognito secret KMS key recorded for execution role decrypt.") + else: + print( + "Warning: Cognito secret exists but ARN was not returned; " + "CDK will use a name-based ARN wildcard in IAM policies." + ) + else: + print( + "ENABLE_HEADLESS_DEPLOYMENT=True: skipping Cognito user pool and secret pre-checks." + ) + + # Service Connect client security groups (by name in VPC) + if ENABLE_ECS_SERVICE_CONNECT == "True": + vpc_id_for_sg = context_data.get("vpc_id") + if vpc_id_for_sg: + for sg_name in ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES_TO_LOOKUP: + exists, sg_id = get_security_group_id_by_name( + sg_name, vpc_id_for_sg, region_name=AWS_REGION + ) + context_data[f"exists:sg:{sg_name}"] = exists + if exists: + context_data[f"security_group_id:{sg_name}"] = sg_id + print(f"Service Connect client SG '{sg_name}' -> {sg_id}") + else: + print( + f"Warning: Service Connect client SG '{sg_name}' " + f"not found in VPC {vpc_id_for_sg}" + ) + else: + print( + "Warning: vpc_id missing from context; cannot resolve Service " + "Connect client security group names." + ) + + # WAF Web ACL (by name and scope) + web_acl_name = WEB_ACL_NAME + exists, existing_web_acl = check_web_acl_exists(web_acl_name, scope="CLOUDFRONT") + context_data[f"exists:{web_acl_name}"] = exists + if exists: + context_data[f"arn:{web_acl_name}"] = existing_web_acl["ARN"] + + # Write the context data to the file + with open(CONTEXT_FILE, "w") as f: + json.dump(context_data, f, indent=2) + + print(f"Context data written to {CONTEXT_FILE}") + + +if __name__ == "__main__": + print(f"Pre-check context file: {CONTEXT_FILE}") + print( + "Running AWS pre-check (requires credentials for the target account/region)..." + ) + try: + check_and_set_context() + except SystemExit: + raise + except Exception as exc: + raise SystemExit(f"Pre-check failed: {exc}") from exc + if not os.path.exists(CONTEXT_FILE): + raise SystemExit(f"Pre-check finished but {CONTEXT_FILE} was not created.") + print(f"Pre-check complete. Context written to {os.path.abspath(CONTEXT_FILE)}") diff --git a/cdk/config/app_config.env.example b/cdk/config/app_config.env.example new file mode 100644 index 0000000000000000000000000000000000000000..06628a72c6ca61056a5af5207e408acc8822d358 --- /dev/null +++ b/cdk/config/app_config.env.example @@ -0,0 +1,26 @@ +# AWS deployment app settings (copy to cdk/config/app_config.env — do not commit secrets). +# Same variable names as tools/config.py; loaded by ECS from S3 or Express env. +COGNITO_AUTH=True +RUN_AWS_FUNCTIONS=True +RUN_AWS_BEDROCK_MODELS=True +RUN_LOCAL_MODEL=False +RUN_GEMINI_MODELS=False +RUN_AZURE_MODELS=False +PRIORITISE_SSO_OVER_AWS_ENV_ACCESS_KEYS=True +DISPLAY_FILE_NAMES_IN_LOGS=False +SESSION_OUTPUT_FOLDER=True +SAVE_LOGS_TO_CSV=True +SAVE_LOGS_TO_DYNAMODB=True +SHOW_COSTS=True +S3_LOG_BUCKET=demo-llm-topic-s3-logs +S3_OUTPUTS_BUCKET= +S3_OUTPUTS_FOLDER= +SAVE_OUTPUTS_TO_S3=False +UPLOAD_USAGE_LOG_TO_S3_OUTPUTS=False +UPLOAD_PROMPT_RESPONSE_LOG_TO_S3_OUTPUTS=False +ACCESS_LOG_DYNAMODB_TABLE_NAME=demo-llm-topic-dynamodb-access-logs +FEEDBACK_LOG_DYNAMODB_TABLE_NAME=demo-llm-topic-dynamodb-feedback-logs +USAGE_LOG_DYNAMODB_TABLE_NAME=demo-llm-topic-dynamodb-usage-logs +LOG_FILE_NAME=log.csv +USAGE_LOG_FILE_NAME=usage_log.csv +FEEDBACK_LOG_FILE_NAME=feedback_log.csv diff --git a/cdk/config/headless_s3_seed/input/config/example_headless_env_file.env b/cdk/config/headless_s3_seed/input/config/example_headless_env_file.env new file mode 100644 index 0000000000000000000000000000000000000000..3b5bfc2da90380824b535b23ef9552ba409a34c4 --- /dev/null +++ b/cdk/config/headless_s3_seed/input/config/example_headless_env_file.env @@ -0,0 +1,13 @@ +# Example job file for headless batch topic modelling. +# 1. Upload your consultation spreadsheet to s3:///input/ +# 2. Copy this file, set DIRECT_MODE_INPUT_FILE to that filename, save as my-job.env +# 3. Upload my-job.env to s3:///input/config/ to start the ECS task. +# +# More DIRECT_MODE_* variables: tools/config.py and cli_topics.py in this repo. + +RUN_DIRECT_MODE=1 +DIRECT_MODE_TASK=extract +DIRECT_MODE_INPUT_FILE=dummy_consultation_response.xlsx +DIRECT_MODE_TEXT_COLUMN="Response Text" +SAVE_OUTPUTS_TO_S3=True +S3_OUTPUTS_FOLDER=output/ diff --git a/cdk/config/lambda/lambda_function.py b/cdk/config/lambda/lambda_function.py new file mode 100644 index 0000000000000000000000000000000000000000..8770f73a5ae58da8df48d4bd9137ef2fde6bf4bc --- /dev/null +++ b/cdk/config/lambda/lambda_function.py @@ -0,0 +1,187 @@ +import os +import urllib.parse + +import boto3 + +ecs = boto3.client("ecs") +s3 = boto3.client("s3") + +# Move static config to Lambda env vars for easier ops +BUCKET = os.environ.get("BUCKET", "lambeth-llm-topic-modelling") +INPUT_PREFIX = os.environ.get("INPUT_PREFIX", "input/") +ENV_PREFIX = os.environ.get("ENV_PREFIX", f"{INPUT_PREFIX}config/") +GENERAL_ENV_PREFIX = os.environ.get("GENERAL_ENV_PREFIX", "general-config/") +ENV_SUFFIX = os.environ.get("ENV_SUFFIX", ".env") +DEFAULT_PARAMS_KEY = os.environ.get( + "DEFAULT_PARAMS_KEY", f"{GENERAL_ENV_PREFIX}app_defaults{ENV_SUFFIX}" +) +CLUSTER = os.environ.get("ECS_CLUSTER", "analytics_fargate_cluster") +TASK_DEF = os.environ.get("ECS_TASK_DEF", "llm_topic_modelling_fargate_def:5") +SUBNETS = os.environ.get( + "SUBNETS", "subnet-083650a35f93a0ee5,subnet-0a41ba70470cabab8" +).split(",") +SECURITY_GROUPS = os.environ.get("SECURITY_GROUPS", "sg-051effb4c5e752df1").split(",") +ECS_ASSIGN_PUBLIC_IP = os.environ.get("ECS_ASSIGN_PUBLIC_IP", "DISABLED").upper() +DEFAULT_INPUT_S3_URI = os.environ.get( + "DEFAULT_INPUT_S3_URI", + f"s3://{BUCKET}/{INPUT_PREFIX}dummy_consultation_response.xlsx", +) +DEFAULT_TASK_TYPE = os.environ.get("DEFAULT_TASK_TYPE", "extract") +CONTAINER_NAME = os.environ.get("CONTAINER_NAME", "llm_topic_modelling_container") + + +def _key_matches(key: str) -> bool: + return key.startswith(ENV_PREFIX) and key.endswith(ENV_SUFFIX) + + +def _derive_runtime_params_from_key(key: str) -> dict: + """Optional convention: decide task type from file name.""" + basename = key.split("/")[-1].lower() + if "train" in basename: + task_type = "train" + elif "infer" in basename or "staging" in basename: + task_type = "infer" + else: + task_type = DEFAULT_TASK_TYPE + + return { + "DIRECT_MODE_TASK": task_type, + "DIRECT_MODE_INPUT_FILE": DEFAULT_INPUT_S3_URI, + } + + +def _parse_dotenv(dotenv_bytes: bytes, bucket: str, input_prefix: str) -> dict: + """ + Parse a basic .env and prepend S3 paths to specific keys if they have values. + """ + env = {} + text = dotenv_bytes.decode("utf-8", errors="replace") + + # 1. Parse the file (Existing Logic) + for raw_line in text.splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + if line.startswith("export "): + line = line[len("export ") :].strip() + if "=" not in line: + continue + key, value = line.split("=", 1) + key = key.strip() + value = value.strip() + if (value.startswith('"') and value.endswith('"')) or ( + value.startswith("'") and value.endswith("'") + ): + value = value[1:-1] + env[key] = value + + # 2. Apply S3 Prefix Logic (New Logic) + # We define the base S3 path string using the provided arguments + s3_base = f"s3://{bucket}/{input_prefix}" + + # List of keys that need modification + keys_to_modify = ["DIRECT_MODE_INPUT_FILE", "DIRECT_MODE_CANDIDATE_TOPICS"] + + for key in keys_to_modify: + # env.get(key) returns the value. + # The 'if' check ensures the key exists AND is not an empty string. + if env.get(key): + env[key] = s3_base + env[key] + + return env + + +def _build_environment_array(*env_dicts): + """Merge dictionaries left→right and produce ECS environment array format.""" + merged = {} + for d in env_dicts: + merged.update(d or {}) + # ECS override format: [{"name": "...","value":"..."}] + return [{"name": k, "value": v} for k, v in merged.items()] + + +def lambda_handler(event, context): + runs = [] + + # Parse default env file with default config + # Fetch the default .env file content from S3 + default_obj = s3.get_object(Bucket=BUCKET, Key=DEFAULT_PARAMS_KEY) + default_dotenv_bytes = default_obj["Body"].read() + + # Parse .env → dict of env vars + default_file_env = _parse_dotenv(default_dotenv_bytes, BUCKET, INPUT_PREFIX) + + for record in event.get("Records", []): + + # Parse env file from user + s3rec = record.get("s3", {}) + bucket = s3rec.get("bucket", {}).get("name") + raw_key = s3rec.get("object", {}).get("key") + if not bucket or not raw_key: + print(f"Object in bucket {bucket} and key {raw_key} not found, exiting.") + continue + + key = urllib.parse.unquote_plus(raw_key) + + # combined_key = ENV_PREFIX + key + combined_key = key + + print("combined_key:", combined_key) + + # Only process the intended prefix/suffix + if not _key_matches(combined_key): + print(f"Key does not match: {combined_key}, exiting") + continue + + # Fetch the .env file content from S3 + obj = s3.get_object(Bucket=bucket, Key=combined_key) + dotenv_bytes = obj["Body"].read() + + # Parse .env → dict of env vars + file_env = _parse_dotenv(dotenv_bytes, BUCKET, INPUT_PREFIX) + + # Optionally derive run-time params from naming + # run_params = _derive_runtime_params_from_key(key) + + # Merge: file env (dominant), overwrites default values. + environment = _build_environment_array(file_env, default_file_env) + + print("Combined environment variables:", environment) + + # Execute ECS Fargate task with fully dynamic environment overrides + response = ecs.run_task( + cluster=CLUSTER, + launchType="FARGATE", + taskDefinition=TASK_DEF, + networkConfiguration={ + "awsvpcConfiguration": { + "subnets": SUBNETS, + "securityGroups": SECURITY_GROUPS, + "assignPublicIp": ( + ECS_ASSIGN_PUBLIC_IP + if ECS_ASSIGN_PUBLIC_IP in ("ENABLED", "DISABLED") + else "DISABLED" + ), + } + }, + overrides={ + "containerOverrides": [ + { + "name": CONTAINER_NAME, + "environment": environment, + } + ] + }, + ) + + runs.append( + { + "bucket": bucket, + "key": key, + "taskArns": [t["taskArn"] for t in response.get("tasks", [])], + "failures": response.get("failures", []), + "envCount": len(environment), + } + ) + + return {"runs": runs} diff --git a/cdk/docs/express_pi_service_connect_spike.md b/cdk/docs/express_pi_service_connect_spike.md new file mode 100644 index 0000000000000000000000000000000000000000..1b247bcd5caf2a16dc1c96fa34d644f2bed14b0f --- /dev/null +++ b/cdk/docs/express_pi_service_connect_spike.md @@ -0,0 +1,71 @@ +# Phase 0: Express Service Connect validation (dev account) + +Service Connect for Express is applied in `post_cdk_build_quickstart.py` (not during +`cdk deploy`). Express primary containers do not define named `portMappings` at create +time, so the post-deploy step registers a task-definition revision with +`port-{port}` and then calls `ecs:UpdateService` with `serviceConnectConfiguration`. +Use this checklist to confirm behaviour in a dev account before relying on Pi Express demos. + +## Prerequisites + +- Main Express stack deployed (`USE_ECS_EXPRESS_MODE=True`). +- Optional: deploy with `ENABLE_PI_AGENT_EXPRESS_SERVICE=True` and run + `python post_cdk_build_quickstart.py` (Service Connect is configured there). + +## Manual validation (without Pi CDK flag) + +1. Note cluster name, main Express **service name**, and task security group from stack outputs. +2. Ensure the cluster has a default Cloud Map namespace (CDK creates one when Pi Express is enabled). +3. Update the main Express service (server): + +```bash +aws ecs update-service \ + --cluster \ + --service \ + --force-new-deployment \ + --service-connect-configuration '{ + "enabled": true, + "namespace": "", + "services": [{ + "portName": "port-7860", + "discoveryName": "summarisation", + "clientAliases": [{"port": 7860, "dnsName": "summarisation"}] + }] + }' +``` + +4. Deploy a second Express service (Pi image, port **7862**) or use CDK Pi Express. +5. Update the Pi Express service (client only): + +```bash +aws ecs update-service \ + --cluster \ + --service \ + --force-new-deployment \ + --service-connect-configuration '{ + "enabled": true, + "namespace": "" + }' +``` + +6. Exec into a Pi task (ECS Exec enabled on the service): + +```bash +curl -sS -o /dev/null -w "%{http_code}\n" http://summarisation:7860/ +``` + +Expect HTTP **200** (or Gradio redirect) without Cognito. + +7. Optional: run a minimal `gradio_client` predict against `/doc_redact` from the Pi task. + +## Exit criteria + +- Service Connect DNS `summarisation` resolves inside the Pi task network namespace. +- Gradio responds on port 7860 without ALB Cognito. +- If `update-service` fails or `portName` is rejected, stop and use legacy Fargate Pi + + `ENABLE_ECS_SERVICE_CONNECT` until AWS/CDK support is confirmed (plan Phase 2 hybrid). + +## CDK deploy path + +With `ENABLE_PI_AGENT_EXPRESS_SERVICE=True`, run `post_cdk_build_quickstart.py` after +`cdk deploy`. Check its output for Service Connect / task-definition errors. diff --git a/cdk/example_headless_env_file.env b/cdk/example_headless_env_file.env new file mode 100644 index 0000000000000000000000000000000000000000..6e94e445c8f5b7b09a0c77e7eb2c1df864bdb22e --- /dev/null +++ b/cdk/example_headless_env_file.env @@ -0,0 +1 @@ +DIRECT_MODE_INPUT_FILE=example_of_emails_sent_to_a_professor_before_applying.pdf \ No newline at end of file diff --git a/cdk/lambda_dynamo_logs_export/lambda_function.py b/cdk/lambda_dynamo_logs_export/lambda_function.py new file mode 100644 index 0000000000000000000000000000000000000000..45232bcceaa0ad7e57985b0c97aee3d93071e894 --- /dev/null +++ b/cdk/lambda_dynamo_logs_export/lambda_function.py @@ -0,0 +1,316 @@ +""" +Lambda handler to export DynamoDB usage log table to CSV and upload to S3. + +All inputs are read from environment variables (no argparse). +Intended to run as an AWS Lambda function; can also be invoked locally +by setting env vars and calling lambda_handler({}, None). + +Environment variables (same semantics as load_dynamo_logs.py CLI): + DYNAMODB_TABLE_NAME - DynamoDB table name (default: summarisation_usage) + AWS_REGION - AWS region (optional; if unset, uses AWS_DEFAULT_REGION, + then region from Lambda context ARN, then eu-west-2) + OUTPUT_FOLDER - Local output directory, e.g. /tmp (optional) + OUTPUT_FILENAME - Local output file name (default: dynamodb_logs_export.csv) + OUTPUT - Full local output path (overrides folder + filename if set). + In Lambda only /tmp is writable; relative paths are auto-resolved to /tmp. + FROM_DATE - Only include entries on/after this date YYYY-MM-DD (optional) + TO_DATE - Only include entries on/before this date YYYY-MM-DD (optional) + DATE_ATTRIBUTE - Attribute name for date filtering (default: timestamp) + S3_OUTPUT_BUCKET - S3 bucket for the output CSV (required for upload) + S3_OUTPUT_KEY - S3 object key/path for the output CSV (required for upload) +""" + +import csv +import datetime +import os +from decimal import Decimal +from io import StringIO + +import boto3 + + +def _get_region_from_context(context): + """Extract region from Lambda context invoked_function_arn (arn:aws:lambda:REGION:ACCOUNT:function:NAME).""" + if context is None: + return None + arn = getattr(context, "invoked_function_arn", None) + if not arn or not isinstance(arn, str): + return None + parts = arn.split(":") + if len(parts) >= 4: + return parts[3] # region is 4th segment + return None + + +def get_config_from_env(context=None): + """Read all settings from environment variables (same inputs as load_dynamo_logs.py). + When running in Lambda, context can be passed to derive region from the function ARN if env is not set. + """ + today = datetime.datetime.now().date() + one_year_ago = today - datetime.timedelta(days=365) + + table_name = os.environ.get("DYNAMODB_TABLE_NAME") or os.environ.get( + "USAGE_LOG_DYNAMODB_TABLE_NAME", "summarisation_usage" + ) + region = ( + os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "" + ).strip() + output = os.environ.get("OUTPUT") + output_folder = os.environ.get("OUTPUT_FOLDER", "output/") + output_filename = os.environ.get("OUTPUT_FILENAME", "dynamodb_logs_export.csv") + from_date_str = os.environ.get("FROM_DATE") + to_date_str = os.environ.get("TO_DATE") + date_attribute = os.environ.get("DATE_ATTRIBUTE", "timestamp") + s3_output_bucket = os.environ.get("S3_OUTPUT_BUCKET") + s3_output_key = os.environ.get("S3_OUTPUT_KEY") + + if output: + local_output_path = output + else: + folder = output_folder.rstrip("/").rstrip("\\") + local_output_path = os.path.join(folder, output_filename) + + # In AWS Lambda only /tmp is writable; resolve relative paths to /tmp to avoid read-only FS errors + if os.environ.get("AWS_LAMBDA_FUNCTION_NAME"): + resolved = os.path.abspath(local_output_path) + if not resolved.startswith("/tmp"): + local_output_path = os.path.join( + "/tmp", os.path.basename(local_output_path) + ) + + # Region: env (AWS_REGION / AWS_DEFAULT_REGION) → Lambda context ARN → hardcoded fallback + if not region and context is not None: + region = _get_region_from_context(context) or "" + if not region: + region = "eu-west-2" + + from_date = None + to_date = None + if from_date_str: + from_date = datetime.datetime.strptime(from_date_str, "%Y-%m-%d").date() + if to_date_str: + to_date = datetime.datetime.strptime(to_date_str, "%Y-%m-%d").date() + if from_date is None and to_date is None: + from_date = one_year_ago + to_date = today + elif from_date is None: + from_date = one_year_ago + elif to_date is None: + to_date = today + + return { + "table_name": table_name, + "region": region, + "local_output_path": local_output_path, + "from_date": from_date, + "to_date": to_date, + "date_attribute": date_attribute, + "s3_output_bucket": s3_output_bucket, + "s3_output_key": s3_output_key, + } + + +def convert_types(item): + new_item = {} + for key, value in item.items(): + if isinstance(value, Decimal): + new_item[key] = int(value) if value % 1 == 0 else float(value) + elif isinstance(value, str): + try: + dt_obj = datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + new_item[key] = dt_obj.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3] + except (ValueError, TypeError): + new_item[key] = value + else: + new_item[key] = value + return new_item + + +def _parse_item_date(value): + """Parse a DynamoDB attribute value to datetime for comparison. Returns None if unparseable.""" + if value is None: + return None + if isinstance(value, Decimal): + try: + return datetime.datetime.utcfromtimestamp(float(value)) + except (ValueError, OSError): + return None + if isinstance(value, (int, float)): + try: + return datetime.datetime.utcfromtimestamp(float(value)) + except (ValueError, OSError): + return None + if isinstance(value, str): + for fmt in ( + "%Y-%m-%d %H:%M:%S.%f", + "%Y-%m-%d %H:%M:%S", + "%Y-%m-%d", + "%Y-%m-%dT%H:%M:%S", + ): + try: + return datetime.datetime.strptime(value, fmt) + except (ValueError, TypeError): + continue + try: + return datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + except (ValueError, TypeError): + pass + return None + + +def filter_items_by_date(items, from_date, to_date, date_attribute: str): + """Return items whose date attribute falls within [from_date, to_date] (inclusive).""" + if from_date is None and to_date is None: + return items + start = datetime.datetime.combine(from_date, datetime.time.min) + end = datetime.datetime.combine(to_date, datetime.time.max) + filtered = [] + for item in items: + raw = item.get(date_attribute) + dt = _parse_item_date(raw) + if dt is None: + continue + if dt.tzinfo: + dt = dt.replace(tzinfo=None) + if start <= dt <= end: + filtered.append(item) + return filtered + + +def scan_table(table): + """Paginated scan of DynamoDB table.""" + items = [] + response = table.scan() + items.extend(response["Items"]) + while "LastEvaluatedKey" in response: + response = table.scan(ExclusiveStartKey=response["LastEvaluatedKey"]) + items.extend(response["Items"]) + return items + + +def export_to_csv_buffer(items, fields_to_drop=None): + """ + Write items to a CSV in memory; return (csv_string, fieldnames). + Use for uploading to S3 without writing to disk. + """ + if not items: + return "", [] + + drop_set = set(fields_to_drop or []) + all_keys = set() + for item in items: + all_keys.update(item.keys()) + fieldnames = sorted(list(all_keys - drop_set)) + + buf = StringIO() + writer = csv.DictWriter( + buf, fieldnames=fieldnames, extrasaction="ignore", restval="" + ) + writer.writeheader() + for item in items: + writer.writerow(convert_types(item)) + return buf.getvalue(), fieldnames + + +def export_to_csv_file(items, output_path, fields_to_drop=None): + """Write items to a CSV file (for optional /tmp or local path).""" + csv_string, _ = export_to_csv_buffer(items, fields_to_drop) + if not csv_string: + return + os.makedirs(os.path.dirname(os.path.abspath(output_path)) or ".", exist_ok=True) + with open(output_path, "w", newline="", encoding="utf-8-sig") as f: + f.write(csv_string) + + +def run_export(config): + """ + Run the full export: scan DynamoDB, filter by date, write CSV (buffer and/or file), upload to S3. + """ + table_name = config["table_name"] + region = config["region"] + local_output_path = config["local_output_path"] + from_date = config["from_date"] + to_date = config["to_date"] + date_attribute = config["date_attribute"] + s3_output_bucket = config["s3_output_bucket"] + s3_output_key = config["s3_output_key"] + + if from_date > to_date: + raise ValueError("FROM_DATE must be on or before TO_DATE") + + dynamodb = boto3.resource("dynamodb", region_name=region or None) + table = dynamodb.Table(table_name) + + items = scan_table(table) + items = filter_items_by_date(items, from_date, to_date, date_attribute) + + csv_string, fieldnames = export_to_csv_buffer(items, fields_to_drop=[]) + result = { + "item_count": len(items), + "from_date": str(from_date), + "to_date": str(to_date), + "columns": fieldnames, + } + + if csv_string: + try: + export_to_csv_file(items, local_output_path, fields_to_drop=[]) + result["local_path"] = local_output_path + except Exception as e: + result["local_write_error"] = str(e) + + if s3_output_bucket and s3_output_key: + s3 = boto3.client("s3", region_name=region or None) + s3.put_object( + Bucket=s3_output_bucket, + Key=s3_output_key, + Body=csv_string.encode("utf-8-sig"), + ContentType="text/csv; charset=utf-8", + ) + result["s3_uri"] = f"s3://{s3_output_bucket}/{s3_output_key}" + elif s3_output_bucket or s3_output_key: + result["s3_skip_reason"] = ( + "Both S3_OUTPUT_BUCKET and S3_OUTPUT_KEY must be set" + ) + + return result + + +def lambda_handler(event, context): + """ + AWS Lambda entrypoint. Config is read from environment variables. + + Event is not required for config; it can be used to override env vars + (e.g. pass table_name, from_date, to_date, s3_output_bucket, s3_output_key). + """ + config = get_config_from_env(context=context) + + if isinstance(event, dict): + if event.get("table_name"): + config["table_name"] = event["table_name"] + if event.get("region"): + config["region"] = event["region"] + if event.get("from_date"): + config["from_date"] = datetime.datetime.strptime( + event["from_date"], "%Y-%m-%d" + ).date() + if event.get("to_date"): + config["to_date"] = datetime.datetime.strptime( + event["to_date"], "%Y-%m-%d" + ).date() + if event.get("date_attribute"): + config["date_attribute"] = event["date_attribute"] + if event.get("s3_output_bucket"): + config["s3_output_bucket"] = event["s3_output_bucket"] + if event.get("s3_output_key"): + config["s3_output_key"] = event["s3_output_key"] + + result = run_export(config) + return {"statusCode": 200, "body": result} + + +if __name__ == "__main__": + import json + + result = lambda_handler({}, None) + print(json.dumps(result, indent=2)) diff --git a/cdk/lambda_elbv2_listener_rule_upsert/lambda_function.py b/cdk/lambda_elbv2_listener_rule_upsert/lambda_function.py new file mode 100644 index 0000000000000000000000000000000000000000..032406a899cb76ea38183d8b273bbf9e24c2d7c9 --- /dev/null +++ b/cdk/lambda_elbv2_listener_rule_upsert/lambda_function.py @@ -0,0 +1,187 @@ +""" +CloudFormation custom resource: create or update an ALB listener rule (upsert by priority). + +If a rule already exists at the requested priority with matching conditions, it is +modified in place. This supports stack retries after partial deploys left orphaned rules. +If priority is taken by a rule with different conditions, the operation fails with a +clear error so operators can choose another PI_ALB_LISTENER_RULE_PRIORITY. +""" + +from __future__ import annotations + +import json +import logging +from typing import Any + +import boto3 +from botocore.exceptions import ClientError + +LOGGER = logging.getLogger() +LOGGER.setLevel(logging.INFO) + +elbv2 = boto3.client("elbv2") + +_DELETE_NOT_FOUND = frozenset( + { + "RuleNotFound", + "ResourceNotFound", + "ValidationError", + } +) +# CloudFormation custom-resource properties are strings; ELBv2 expects typed fields. +_ELBV2_INT_KEYS = frozenset( + { + "Order", + "Priority", + "SessionTimeout", + "DurationSeconds", + "SessionCookieName", + } +) +_ELBV2_BOOL_KEYS = frozenset({"Enabled"}) + + +def _parse_json(value: Any) -> Any: + if isinstance(value, str): + return json.loads(value) + return value + + +def _coerce_elbv2_api_value(key: str, value: Any) -> Any: + """Coerce CloudFormation string properties to types boto3 elbv2 expects.""" + if isinstance(value, dict): + return {k: _coerce_elbv2_api_value(k, v) for k, v in value.items()} + if isinstance(value, list): + return [_coerce_elbv2_api_value(key, item) for item in value] + if isinstance(value, str): + if key in _ELBV2_INT_KEYS: + try: + return int(value) + except ValueError: + return value + if key in _ELBV2_BOOL_KEYS: + lowered = value.lower() + if lowered == "true": + return True + if lowered == "false": + return False + return value + + +def _normalize_listener_rule_payload( + conditions: list[dict[str, Any]], + actions: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + return ( + [_coerce_elbv2_api_value("", c) for c in conditions], + [_coerce_elbv2_api_value("", a) for a in actions], + ) + + +def _condition_values(condition: dict[str, Any]) -> list[str]: + field = condition.get("Field") + if field == "path-pattern": + cfg = condition.get("PathPatternConfig") or {} + return sorted(cfg.get("Values") or condition.get("Values") or []) + if field == "host-header": + cfg = condition.get("HostHeaderConfig") or {} + return sorted(cfg.get("Values") or condition.get("Values") or []) + return sorted(condition.get("Values") or []) + + +def _conditions_match( + rule_conditions: list[dict[str, Any]], expected: list[dict[str, Any]] +) -> bool: + if len(rule_conditions) != len(expected): + return False + for exp in expected: + exp_field = exp.get("Field") + exp_vals = _condition_values(exp) + if not any( + rc.get("Field") == exp_field and _condition_values(rc) == exp_vals + for rc in rule_conditions + ): + return False + return True + + +def _rule_at_priority( + rules: list[dict[str, Any]], priority: int +) -> dict[str, Any] | None: + target = str(priority) + for rule in rules: + if rule.get("Priority") == target: + return rule + return None + + +def _upsert_rule( + *, + listener_arn: str, + priority: int, + conditions: list[dict[str, Any]], + actions: list[dict[str, Any]], +) -> str: + described = elbv2.describe_rules(ListenerArn=listener_arn) + existing = _rule_at_priority(described.get("Rules", []), priority) + if existing: + if _conditions_match(existing.get("Conditions", []), conditions): + rule_arn = existing["RuleArn"] + LOGGER.info( + "Modifying existing listener rule %s at priority %s", rule_arn, priority + ) + elbv2.modify_rule(RuleArn=rule_arn, Conditions=conditions, Actions=actions) + return rule_arn + raise RuntimeError( + f"Listener rule priority {priority} is already in use by {existing['RuleArn']} " + "with different conditions. Delete that rule in the ELB console or set " + "PI_ALB_LISTENER_RULE_PRIORITY to a free value, then redeploy." + ) + LOGGER.info("Creating listener rule at priority %s on %s", priority, listener_arn) + created = elbv2.create_rule( + ListenerArn=listener_arn, + Priority=priority, + Conditions=conditions, + Actions=actions, + ) + return created["Rules"][0]["RuleArn"] + + +def _delete_rule(rule_arn: str) -> None: + if not rule_arn: + return + try: + elbv2.delete_rule(RuleArn=rule_arn) + except ClientError as exc: + code = exc.response.get("Error", {}).get("Code", "") + if code in _DELETE_NOT_FOUND: + LOGGER.info("Rule %s already deleted (%s)", rule_arn, code) + return + raise + + +def handler(event, context): + request_type = event["RequestType"] + props = event["ResourceProperties"] + listener_arn = props["ListenerArn"] + priority = int(props["Priority"]) + conditions, actions = _normalize_listener_rule_payload( + _parse_json(props["Conditions"]), + _parse_json(props["Actions"]), + ) + + if request_type == "Delete": + _delete_rule(event.get("PhysicalResourceId", "")) + return {"PhysicalResourceId": event.get("PhysicalResourceId", "deleted")} + + rule_arn = _upsert_rule( + listener_arn=listener_arn, + priority=priority, + conditions=conditions, + actions=actions, + ) + if request_type == "Update": + prior_arn = event.get("PhysicalResourceId", "") + if prior_arn and prior_arn != rule_arn: + _delete_rule(prior_arn) + return {"PhysicalResourceId": rule_arn} diff --git a/cdk/post_cdk_build_quickstart.py b/cdk/post_cdk_build_quickstart.py new file mode 100644 index 0000000000000000000000000000000000000000..f74e7308bcd670ed7a055e8ac6b6b5162b8ce98e --- /dev/null +++ b/cdk/post_cdk_build_quickstart.py @@ -0,0 +1,195 @@ +import os +import time + +from cdk_config import ( + AWS_REGION, + CDK_PREFIX, + CLUSTER_NAME, + CODEBUILD_PI_PROJECT_NAME, + CODEBUILD_PROJECT_NAME, + COGNITO_USER_POOL_CLIENT_SECRET_NAME, + ECS_EXPRESS_COGNITO_REDIRECT_BASE, + ECS_EXPRESS_SC_PORT_NAME, + ECS_EXPRESS_SERVICE_NAME, + ECS_LOG_GROUP_NAME, + ECS_PI_EXPRESS_SERVICE_NAME, + ECS_PI_SERVICE_NAME, + ECS_SERVICE_CONNECT_DISCOVERY_NAME, + ECS_SERVICE_CONNECT_NAMESPACE, + ECS_SERVICE_NAME, + ENABLE_HEADLESS_DEPLOYMENT, + ENABLE_PI_AGENT_ECS_SERVICE, + ENABLE_PI_AGENT_EXPRESS_SERVICE, + GITHUB_REPO_BRANCH, + GITHUB_REPO_NAME, + GITHUB_REPO_USERNAME, + GRADIO_SERVER_PORT, + PI_AGENT_ENV_S3_KEY, + S3_BATCH_ENV_PREFIX, + S3_BATCH_INPUT_PREFIX, + S3_BATCH_LAMBDA_FUNCTION_NAME, + S3_LOG_CONFIG_BUCKET_NAME, + S3_OUTPUT_BUCKET_NAME, + USE_ECS_EXPRESS_MODE, +) +from cdk_functions import create_basic_config_env, ensure_codebuild_public_github_source + +# boto3-only module (does not import aws-cdk / Node.js) +from cdk_post_deploy import ( + apply_cognito_secret_fixup_from_stack, + configure_express_pi_service_connect, + print_headless_deployment_next_steps, + seed_headless_batch_s3_layout, + start_codebuild_build, + start_ecs_task, + start_express_gateway_service, + upload_file_to_s3, +) +from tqdm import tqdm + +# Create basic app_config.env file that user can use to run the app later. Input is the folder it is saved into. +create_basic_config_env( + "config", + headless=ENABLE_HEADLESS_DEPLOYMENT == "True", + pi_express_backend=ENABLE_PI_AGENT_EXPRESS_SERVICE == "True", +) + +# Ensure CodeBuild clones the public GitHub repo (not CodeConnections/OAuth). +ensure_codebuild_public_github_source( + CODEBUILD_PROJECT_NAME, + GITHUB_REPO_USERNAME, + GITHUB_REPO_NAME, + GITHUB_REPO_BRANCH, + aws_region=AWS_REGION, +) + +# Start CodeBuild for the main app image +print("Starting main app CodeBuild project.") +start_codebuild_build(project_name=CODEBUILD_PROJECT_NAME) + +_enable_pi_image_build = ( + ENABLE_PI_AGENT_ECS_SERVICE == "True" or ENABLE_PI_AGENT_EXPRESS_SERVICE == "True" +) +if _enable_pi_image_build: + ensure_codebuild_public_github_source( + CODEBUILD_PI_PROJECT_NAME, + GITHUB_REPO_USERNAME, + GITHUB_REPO_NAME, + GITHUB_REPO_BRANCH, + aws_region=AWS_REGION, + ) + print("Starting Pi agent CodeBuild project.") + start_codebuild_build(project_name=CODEBUILD_PI_PROJECT_NAME) + +# Upload app_config.env file to S3 bucket +upload_file_to_s3( + local_file_paths="config/app_config.env", + s3_key="", + s3_bucket=S3_LOG_CONFIG_BUCKET_NAME, +) + +if _enable_pi_image_build: + pi_env_local = os.path.join("config", "pi_agent.env") + if os.path.isfile(pi_env_local): + print( + f"Uploading {pi_env_local} to s3://{S3_LOG_CONFIG_BUCKET_NAME}/{PI_AGENT_ENV_S3_KEY}" + ) + upload_file_to_s3( + local_file_paths=pi_env_local, + s3_key="", + s3_bucket=S3_LOG_CONFIG_BUCKET_NAME, + ) + else: + print( + f"Skipping Pi env upload: {pi_env_local} not found. " + f"Create it (from config/pi_agent.env.example) and upload to " + f"s3://{S3_LOG_CONFIG_BUCKET_NAME}/{PI_AGENT_ENV_S3_KEY} before scaling the Pi service." + ) + +total_seconds = 480 # 8 minutes +update_interval = 1 # Update every second + +print("Waiting 8 minutes for CodeBuild container image(s) to build.") + +# tqdm iterates over a range, and you perform a small sleep in each iteration +for i in tqdm(range(total_seconds), desc="Building container"): + time.sleep(update_interval) + +# Scale main ECS service to one task (skipped for headless batch-only deployments) +if ENABLE_HEADLESS_DEPLOYMENT != "True": + if USE_ECS_EXPRESS_MODE == "True": + print(f"Starting Express ECS service {ECS_EXPRESS_SERVICE_NAME}") + start_express_gateway_service( + cluster_name=CLUSTER_NAME, service_name=ECS_EXPRESS_SERVICE_NAME + ) + if ENABLE_PI_AGENT_EXPRESS_SERVICE == "True": + print( + "Configuring Service Connect for main Express (server) and " + "Pi Express (client)." + ) + try: + configure_express_pi_service_connect( + cluster_name=CLUSTER_NAME, + main_service_name=ECS_EXPRESS_SERVICE_NAME, + pi_service_name=ECS_PI_EXPRESS_SERVICE_NAME, + namespace=ECS_SERVICE_CONNECT_NAMESPACE, + main_port_name=ECS_EXPRESS_SC_PORT_NAME, + discovery_name=ECS_SERVICE_CONNECT_DISCOVERY_NAME, + main_port=int(GRADIO_SERVER_PORT), + ) + except Exception as exc: + print("Warning: could not configure Express Service Connect: " f"{exc}") + print("Syncing Cognito app client secret for in-app authentication.") + apply_cognito_secret_fixup_from_stack( + stack_name="SummarisationStack", + secret_name=COGNITO_USER_POOL_CLIENT_SECRET_NAME, + cluster_name=CLUSTER_NAME, + main_service_name=ECS_EXPRESS_SERVICE_NAME, + recycle_tasks=False, + ) + else: + print(f"Starting ECS service {ECS_SERVICE_NAME}") + start_ecs_task(cluster_name=CLUSTER_NAME, service_name=ECS_SERVICE_NAME) +else: + print( + "Headless deployment: skipping always-on ECS service start " + "(tasks are started by the S3 batch Lambda)." + ) + seed_headless_batch_s3_layout( + S3_OUTPUT_BUCKET_NAME, + input_prefix=S3_BATCH_INPUT_PREFIX, + env_prefix=S3_BATCH_ENV_PREFIX, + example_env_local_path=os.path.join("config", "example_headless_env_file.env"), + ) + print_headless_deployment_next_steps( + { + "AWS_REGION": AWS_REGION, + "S3_OUTPUT_BUCKET_NAME": S3_OUTPUT_BUCKET_NAME, + "S3_BATCH_INPUT_PREFIX": S3_BATCH_INPUT_PREFIX, + "S3_BATCH_ENV_PREFIX": S3_BATCH_ENV_PREFIX, + "ECS_LOG_GROUP_NAME": ECS_LOG_GROUP_NAME, + "S3_BATCH_LAMBDA_FUNCTION_NAME": S3_BATCH_LAMBDA_FUNCTION_NAME, + "CDK_PREFIX": CDK_PREFIX, + } + ) + +if ENABLE_PI_AGENT_ECS_SERVICE == "True": + print(f"Starting Pi agent ECS service {ECS_PI_SERVICE_NAME}") + start_ecs_task(cluster_name=CLUSTER_NAME, service_name=ECS_PI_SERVICE_NAME) + +if ENABLE_PI_AGENT_EXPRESS_SERVICE == "True": + print(f"Starting Pi Express ECS service {ECS_PI_EXPRESS_SERVICE_NAME}") + start_express_gateway_service( + cluster_name=CLUSTER_NAME, service_name=ECS_PI_EXPRESS_SERVICE_NAME + ) + +if USE_ECS_EXPRESS_MODE == "True" and ENABLE_HEADLESS_DEPLOYMENT != "True": + from cdk_post_deploy import print_express_mode_next_steps + + print_express_mode_next_steps( + { + "AWS_REGION": AWS_REGION, + "ECS_EXPRESS_COGNITO_REDIRECT_BASE": ECS_EXPRESS_COGNITO_REDIRECT_BASE, + "ENABLE_PI_AGENT_EXPRESS_SERVICE": ENABLE_PI_AGENT_EXPRESS_SERVICE, + } + ) diff --git a/cdk/requirements.txt b/cdk/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..d284ff9c1beaa8eb81f5caf01e7350b7d4fd5bce --- /dev/null +++ b/cdk/requirements.txt @@ -0,0 +1,6 @@ +aws-cdk-lib==2.257.0 +aws-cdk.aws-servicecatalogappregistry-alpha~=2.257.0a0 +boto3 +pandas +nodejs +python-dotenv \ No newline at end of file diff --git a/cdk/test/test_cdk_install.py b/cdk/test/test_cdk_install.py new file mode 100644 index 0000000000000000000000000000000000000000..d2d4db40a11ff5952e81e446399b2af87aa9d86b --- /dev/null +++ b/cdk/test/test_cdk_install.py @@ -0,0 +1,1010 @@ +"""Unit tests for cdk_install.py (no live AWS or cdk deploy).""" + +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + +import cdk_install as inst + + +def _demo_answers() -> inst.InstallAnswers: + return inst.InstallAnswers( + profile="demo", + aws_account_id="123456789012", + aws_region="eu-west-2", + cdk_prefix="Test-Summarisation-", + cognito_domain_prefix="test-summarisation", + vpc_mode="existing", + vpc_name="test-vpc", + public_subnet_mode="auto", + private_subnet_mode="auto", + ) + + +def _production_answers() -> inst.InstallAnswers: + a = _demo_answers() + a.profile = "production" + a.acm_cert_arn = "arn:aws:acm:eu-west-2:123:certificate/abc" + a.ssl_domain = "summarisation.example.com" + a.cloudfront_geo = "GB" + return a + + +def _headless_answers() -> inst.InstallAnswers: + return inst.InstallAnswers( + profile="headless", + aws_account_id="123456789012", + aws_region="eu-west-2", + cdk_prefix="Headless-Summarisation-", + cognito_domain_prefix="headless-summarisation", + vpc_mode="existing", + vpc_name="test-vpc", + public_subnet_mode="auto", + private_subnet_mode="auto", + enable_s3_batch=True, + ) + + +def test_default_cognito_domain_prefix_strips_reserved_words_from_cdk_prefix(): + prefix = inst.default_cognito_domain_prefix_from_cdk_prefix( + "Lambeth-AWS-SharedServices-Dev-Summarisation-" + ) + assert prefix == "lambeth-sharedservic" + assert "aws" not in prefix + assert "amazon" not in prefix + assert "cognito" not in prefix + + +def test_sanitize_cognito_domain_prefix_removes_reserved_substrings(): + assert ( + inst.sanitize_cognito_domain_prefix("lambeth-aws-sharedse") + == "lambeth-sharedse" + ) + assert inst.sanitize_cognito_domain_prefix("my-cognito-login") == "my-login" + assert ( + inst.sanitize_cognito_domain_prefix("amazon-summarisation-dev") + == "summarisation-dev" + ) + assert inst.sanitize_cognito_domain_prefix("aws") == "llm-topic-app" + + +def test_validate_cognito_domain_prefix_rejects_reserved_words(): + assert inst.validate_cognito_domain_prefix("lambeth-aws-sharedse") is not None + assert inst.validate_cognito_domain_prefix("lambeth-sharedse-dev") is None + + +def test_build_env_values_demo(): + values = inst.build_env_values(_demo_answers()) + assert values["USE_ECS_EXPRESS_MODE"] == "True" + assert values["ECS_EXPRESS_USE_PUBLIC_SUBNETS"] == "True" + assert values["PRIVATE_SUBNETS_TO_USE"] == "" + assert values["USE_CLOUDFRONT"] == "False" + assert values["ENABLE_RESOURCE_DELETE_PROTECTION"] == "False" + assert values["VPC_NAME"] == "test-vpc" + assert values["CONTEXT_FILE"] == "precheck.context.json" + assert values["CDK_FOLDER"].endswith("/cdk/") + + +def test_build_env_values_production(): + values = inst.build_env_values(_production_answers()) + assert values["USE_ECS_EXPRESS_MODE"] == "False" + assert values["USE_CLOUDFRONT"] == "True" + assert values["RUN_USEAST_STACK"] == "True" + assert values["ENABLE_RESOURCE_DELETE_PROTECTION"] == "True" + assert values["ACM_SSL_CERTIFICATE_ARN"].startswith("arn:aws:acm:") + assert values["SSL_CERTIFICATE_DOMAIN"] == "summarisation.example.com" + assert values["APPREGISTRY_STACK_NAME"] == "Test-Summarisation-AppRegistryStack" + + +def test_build_env_values_uses_custom_s3_bucket_names(): + answers = _demo_answers() + answers.s3_log_bucket_name = "063418083240-demo-summarisation-s3-logs" + answers.s3_output_bucket_name = "063418083240-demo-summarisation-s3-output" + values = inst.build_env_values(answers) + assert ( + values["S3_LOG_CONFIG_BUCKET_NAME"] == "063418083240-demo-summarisation-s3-logs" + ) + assert ( + values["S3_OUTPUT_BUCKET_NAME"] == "063418083240-demo-summarisation-s3-output" + ) + + +def test_normalize_s3_bucket_name_truncates_and_sanitizes(): + long_name = "A" * 80 + "/underscores" + normalized = inst.normalize_s3_bucket_name(long_name) + assert normalized == normalized.lower() + assert len(normalized) <= inst.S3_BUCKET_NAME_MAX_LEN + assert "_" not in normalized + + +def test_suggest_available_s3_bucket_name_prefers_account_prefix(monkeypatch): + calls: list[str] = [] + + def fake_resolve(name: str, **kwargs): + calls.append(name) + if name.startswith("123456789012-"): + return ("available", name) + return ("globally_taken", name) + + monkeypatch.setattr("cdk_functions.resolve_s3_bucket_availability", fake_resolve) + suggested = inst.suggest_available_s3_bucket_name( + "demo-summarisation-s3-logs", "123456789012" + ) + assert suggested == "123456789012-demo-summarisation-s3-logs" + assert calls[0] == "demo-summarisation-s3-logs" + + +def test_validate_globally_unique_env_values_flags_taken_bucket(monkeypatch): + monkeypatch.setattr( + "cdk_functions.resolve_s3_bucket_availability", + lambda name, **kwargs: ("globally_taken", name), + ) + monkeypatch.setattr( + "cdk_functions.resolve_cognito_domain_prefix_availability", + lambda prefix, **kwargs: "available", + ) + values = inst.build_env_values(_demo_answers()) + errors = inst.validate_globally_unique_env_values(values) + assert any("taken globally" in err for err in errors) + + +def test_prompt_globally_unique_cognito_prefix_suggests_when_taken(monkeypatch): + monkeypatch.setattr( + "cdk_functions.resolve_cognito_domain_prefix_availability", + lambda prefix, **kwargs: ( + "taken" if prefix == "demo-summarisation" else "available" + ), + ) + monkeypatch.setattr( + inst, + "suggest_available_cognito_domain_prefix", + lambda preferred, account_id, region, **kwargs: "063418083240-demo-summarisation", + ) + monkeypatch.setattr(inst, "ask_yes_no", lambda *args, **kwargs: True) + result = inst._prompt_globally_unique_cognito_prefix( + "demo-summarisation", + "063418083240", + "eu-west-2", + interactive=True, + assume_yes=False, + ) + assert result == "063418083240-demo-summarisation" + + +def test_prompt_globally_unique_cognito_prefix_rejects_taken_cli_override(monkeypatch): + monkeypatch.setattr( + "cdk_functions.resolve_cognito_domain_prefix_availability", + lambda prefix, **kwargs: "taken", + ) + with pytest.raises(SystemExit): + inst._prompt_globally_unique_cognito_prefix( + "demo-summarisation", + "063418083240", + "eu-west-2", + interactive=False, + assume_yes=False, + cli_override="demo-summarisation", + ) + + +def test_subnet_cidr_prefix_len_for_express_demo(): + answers = _demo_answers() + assert inst.subnet_cidr_prefix_len_for_tier(answers, "public") == 27 + assert inst.subnet_cidr_prefix_len_for_tier(answers, "private") == 28 + + +def test_validate_public_subnet_cidr_for_express_rejects_small_blocks(): + assert inst.validate_public_subnet_cidr_for_express("10.0.0.0/28") is not None + assert inst.validate_public_subnet_cidr_for_express("10.0.0.0/27") is None + assert inst.validate_public_subnet_cidr_for_express("10.0.0.0/26") is None + + +def test_validate_env_values_rejects_express_public_subnets_too_small(): + values = inst.build_env_values(_demo_answers()) + values["PUBLIC_SUBNET_CIDR_BLOCKS"] = "['10.0.0.0/28', '10.0.0.16/28']" + errors = inst.validate_env_values(values) + assert any("too small for ECS Express" in err for err in errors) + + +def test_build_env_values_headless(): + values = inst.build_env_values(_headless_answers()) + assert values["USE_ECS_EXPRESS_MODE"] == "False" + assert values["USE_CLOUDFRONT"] == "False" + assert values["RUN_USEAST_STACK"] == "False" + assert values["ENABLE_HEADLESS_DEPLOYMENT"] == "True" + assert values["ENABLE_S3_BATCH_ECS_TRIGGER"] == "True" + assert values["COGNITO_AUTH"] == "False" + assert values["ECS_EXPRESS_USE_PUBLIC_SUBNETS"] == "True" + assert values["PRIVATE_SUBNETS_TO_USE"] == "" + assert values["S3_LOG_CONFIG_BUCKET_NAME"] == "headless-summarisation-s3-logs" + assert values["S3_OUTPUT_BUCKET_NAME"] == "headless-summarisation-s3-output" + + +def test_validate_env_values_rejects_bare_s3_bucket_names(): + values = inst.build_env_values(_headless_answers()) + values["S3_LOG_CONFIG_BUCKET_NAME"] = "s3-logs" + errors = inst.validate_env_values(values) + assert any("S3_LOG_CONFIG_BUCKET_NAME" in e for e in errors) + + +def test_headless_profile_uses_public_subnets_only(): + assert inst.answers_use_public_subnets_only(_headless_answers()) is True + + +def test_build_env_values_production_headless(): + answers = _production_answers() + answers.enable_headless = True + values = inst.build_env_values(answers) + assert values["USE_ECS_EXPRESS_MODE"] == "False" + assert values["USE_CLOUDFRONT"] == "False" + assert values["ENABLE_HEADLESS_DEPLOYMENT"] == "True" + assert values["ENABLE_RESOURCE_DELETE_PROTECTION"] == "True" + assert values.get("ECS_EXPRESS_USE_PUBLIC_SUBNETS") != "True" + assert "PRIVATE_SUBNET_CIDR_BLOCKS" in values + + +def test_validate_install_answers_rejects_demo_headless(): + answers = _demo_answers() + answers.enable_headless = True + errors = inst.validate_install_answers(answers) + assert any("Demonstration" in err for err in errors) + + +def test_validate_install_answers_rejects_custom_express_headless(): + answers = inst.InstallAnswers(profile="custom", enable_headless=True) + answers.custom_overrides["USE_ECS_EXPRESS_MODE"] = "True" + errors = inst.validate_install_answers(answers) + assert any("USE_ECS_EXPRESS_MODE=False" in err for err in errors) + + +def test_profile_allows_headless_add_on(): + assert inst.profile_allows_headless_add_on(_demo_answers()) is False + assert inst.profile_allows_headless_add_on(_production_answers()) is True + assert inst.profile_allows_headless_add_on(_headless_answers()) is True + custom = inst.InstallAnswers(profile="custom") + assert inst.profile_allows_headless_add_on(custom) is True + custom.custom_overrides["USE_ECS_EXPRESS_MODE"] = "True" + assert inst.profile_allows_headless_add_on(custom) is False + + +def test_validate_install_answers_skips_cognito_prefix_for_headless(): + values = inst.build_env_values(_headless_answers()) + values["COGNITO_USER_POOL_DOMAIN_PREFIX"] = "" + assert inst.validate_env_values(values) == [] + + +def test_validate_headless_rejects_pi(): + values = inst.build_env_values(_headless_answers()) + values["ENABLE_ECS_SERVICE_CONNECT"] = "True" + errors = inst.validate_env_values(values) + assert any("HEADLESS" in e for e in errors) + + +def test_validate_rejects_express_with_acm(): + values = inst.build_env_values(_demo_answers()) + values["ACM_SSL_CERTIFICATE_ARN"] = "arn:aws:acm:eu-west-2:123:certificate/x" + errors = inst.validate_env_values(values) + assert any("ACM_SSL_CERTIFICATE_ARN" in e for e in errors) + + +def test_build_env_values_mixed_subnet_tiers(): + answers = _production_answers() + answers.public_subnet_mode = "existing" + answers.private_subnet_mode = "create" + answers.public_subnet_names = ["existing-public-a", "existing-public-b"] + answers.public_subnet_cidrs = ["10.244.1.0/28", "10.244.2.0/28"] + answers.public_subnet_azs = ["eu-west-2a", "eu-west-2b"] + answers.private_subnet_names = ["Demo-Summarisation-PrivateSubnet1"] + answers.private_subnet_cidrs = ["10.0.10.0/28"] + answers.private_subnet_azs = ["eu-west-2a"] + values = inst.build_env_values(answers) + assert ( + values["PUBLIC_SUBNETS_TO_USE"] == '["existing-public-a", "existing-public-b"]' + ) + assert values["PUBLIC_SUBNET_CIDR_BLOCKS"] == "['10.244.1.0/28', '10.244.2.0/28']" + assert values["PUBLIC_SUBNET_AVAILABILITY_ZONES"] == "['eu-west-2a', 'eu-west-2b']" + assert values["PRIVATE_SUBNETS_TO_USE"] == '["Demo-Summarisation-PrivateSubnet1"]' + assert values["PRIVATE_SUBNET_CIDR_BLOCKS"] == "['10.0.10.0/28']" + assert values["PRIVATE_SUBNET_AVAILABILITY_ZONES"] == "['eu-west-2a']" + + +def test_enrich_existing_subnet_details_from_aws(monkeypatch): + answers = _demo_answers() + answers.vpc_name = "test-vpc" + answers.public_subnet_mode = "existing" + answers.public_subnet_names = ["pub-a", "pub-b"] + answers.private_subnet_mode = "create" + + monkeypatch.setattr( + inst, + "list_vpcs", + lambda _region: [{"name": "test-vpc", "id": "vpc-123", "cidr": "10.0.0.0/16"}], + ) + monkeypatch.setattr( + inst, + "list_subnets_in_vpc", + lambda _vpc_id, _region: [ + {"name": "pub-a", "cidr": "10.244.1.0/28", "az": "eu-west-2a"}, + {"name": "pub-b", "cidr": "10.244.2.0/28", "az": "eu-west-2b"}, + ], + ) + + errors = inst.enrich_existing_subnet_details_from_aws(answers) + assert errors == [] + assert answers.public_subnet_cidrs == ["10.244.1.0/28", "10.244.2.0/28"] + assert answers.public_subnet_azs == ["eu-west-2a", "eu-west-2b"] + + +def test_validate_subnet_answers_mixed_requires_public_names(): + answers = _demo_answers() + answers.public_subnet_mode = "existing" + answers.private_subnet_mode = "create" + answers.private_subnet_names = ["new-private"] + errors = inst.validate_subnet_answers(answers) + assert any("Public subnets" in err for err in errors) + + +def test_resolve_subnet_tier_modes_per_tier_override(): + args = argparse.Namespace( + subnet_mode="auto", + public_subnet_mode="existing", + private_subnet_mode="create", + ) + public, private = inst.resolve_subnet_tier_modes(args) + assert public == "existing" + assert private == "create" + + +def test_suggest_vpc_cidr_block_empty_region(): + assert inst.suggest_vpc_cidr_block([]) == "10.0.0.0/24" + + +def test_suggest_vpc_cidr_block_skips_overlaps(): + assert inst.suggest_vpc_cidr_block(["10.0.0.0/24"]) == "10.0.1.0/24" + assert inst.suggest_vpc_cidr_block(["10.0.0.0/16"]) == "10.1.0.0/24" + + +def test_validate_new_vpc_cidr_rejects_overlap_and_public_space(): + err = inst.validate_new_vpc_cidr("10.0.0.0/24", ["10.0.0.0/16"]) + assert err is not None + assert "overlaps" in err + err = inst.validate_new_vpc_cidr("8.8.8.0/24", []) + assert err is not None + assert "RFC1918" in err + + +def test_validate_new_vpc_cidr_accepts_available_block(): + assert inst.validate_new_vpc_cidr("10.2.0.0/24", ["10.0.0.0/16"]) is None + + +def test_prompt_new_vpc_cidr_validates_cli_override(monkeypatch): + answers = inst.InstallAnswers(aws_region="eu-west-2", new_vpc_cidr="10.0.0.0/24") + monkeypatch.setattr( + inst, "list_vpc_cidr_blocks_in_region", lambda _r: ["10.0.0.0/16"] + ) + with pytest.raises(SystemExit): + inst.prompt_new_vpc_cidr(answers, interactive=False) + + +def test_prompt_new_vpc_cidr_auto_select_noninteractive(monkeypatch): + answers = inst.InstallAnswers(aws_region="eu-west-2") + monkeypatch.setattr(inst, "list_vpc_cidr_blocks_in_region", lambda _r: []) + inst.prompt_new_vpc_cidr(answers, interactive=False) + assert answers.new_vpc_cidr == "10.0.0.0/24" + + +def test_suggest_subnet_cidr_blocks_lowest_available(): + blocks = inst.suggest_subnet_cidr_blocks("10.0.0.0/24", ["10.0.0.0/28"], 2) + assert blocks == ["10.0.0.16/28", "10.0.0.32/28"] + + +def test_suggest_subnet_cidr_blocks_respects_reserved(): + blocks = inst.suggest_subnet_cidr_blocks( + "10.0.0.0/24", + [], + 1, + reserved_cidrs=["10.0.0.0/28"], + ) + assert blocks == ["10.0.0.16/28"] + + +def test_vpc_cidr_blocks_from_describe_includes_associations(): + vpc = { + "CidrBlock": "10.0.0.0/16", + "CidrBlockAssociationSet": [{"CidrBlock": "10.1.0.0/16"}], + } + assert inst.vpc_cidr_blocks_from_describe(vpc) == [ + "10.0.0.0/16", + "10.1.0.0/16", + ] + + +def test_build_app_config_env_values_express_uses_in_app_cognito(): + values = inst.build_env_values(_demo_answers()) + updates = inst.build_app_config_env_values(values) + assert updates["COGNITO_AUTH"] == "True" + assert updates["RUN_AWS_FUNCTIONS"] == "True" + assert updates["S3_LOG_BUCKET"].endswith("s3-logs") + + +def test_build_app_config_env_values_express_pi_disables_main_cognito(): + values = inst.build_env_values(_demo_answers()) + values["ENABLE_PI_AGENT_EXPRESS_SERVICE"] = "True" + updates = inst.build_app_config_env_values(values) + assert updates["COGNITO_AUTH"] == "False" + + +def test_build_app_config_env_values_headless(): + answers = _headless_answers() + values = inst.build_env_values(answers) + updates = inst.build_app_config_env_values(values) + assert updates["COGNITO_AUTH"] == "False" + + +def test_write_app_config_env_from_example(tmp_path, monkeypatch): + example = tmp_path / "app_config.env.example" + example.write_text( + "RUN_AWS_FUNCTIONS=True\nS3_LOG_BUCKET=placeholder\n", + encoding="utf-8", + ) + target = tmp_path / "app_config.env" + monkeypatch.setattr(inst, "APP_CONFIG_ENV_EXAMPLE", example) + monkeypatch.setattr(inst, "APP_CONFIG_ENV_PATH", target) + + answers = _demo_answers() + answers.write_app_config_env = True + values = inst.build_env_values(answers) + inst.write_app_config_env_file(answers, values) + + written = inst.read_env_file(target) + assert written["RUN_AWS_FUNCTIONS"] == "True" + assert written["S3_LOG_BUCKET"].endswith("s3-logs") + assert written["COGNITO_AUTH"] == "True" + + +def test_format_list_env(): + assert inst.format_list_env(["a", "b"]) == '["a", "b"]' + assert ( + inst.format_list_env(["10.0.0.0/28"], use_single_quotes=True) + == "['10.0.0.0/28']" + ) + assert inst.format_list_env([]) == "[]" + + +def test_write_env_file_backs_up(tmp_path): + env_path = tmp_path / "cdk_config.env" + env_path.write_text("OLD=1\n", encoding="utf-8") + inst.write_env_file(env_path, {"NEW": "2"}) + backups = list(tmp_path.glob("cdk_config.env.bak.*")) + assert len(backups) == 1 + assert "NEW=2" in env_path.read_text(encoding="utf-8") + + +def test_write_cdk_json_preserves_context(tmp_path, monkeypatch): + cdk_json = tmp_path / "cdk.json" + cdk_json.write_text( + json.dumps( + { + "app": "old-python app.py", + "context": {"@aws-cdk/custom:flag": True, "keep": 1}, + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(inst, "CDK_JSON_PATH", cdk_json) + monkeypatch.setattr(inst, "CDK_JSON_EXAMPLE", tmp_path / "missing.example") + + py = Path(sys.executable) + inst.write_cdk_json(py, force=True) + + data = json.loads(cdk_json.read_text(encoding="utf-8")) + assert data["context"]["@aws-cdk/custom:flag"] is True + assert data["context"]["keep"] == 1 + assert data["app"].endswith("app.py") + assert "output" in data + backups = list(tmp_path.glob("cdk.json.bak.*")) + assert len(backups) == 1 + + +def test_format_cdk_app_command(): + cmd = inst.format_cdk_app_command(Path(sys.executable)) + assert cmd.endswith("app.py") + assert str(sys.executable).replace("/", "\\") in cmd or sys.executable in cmd + + +def test_venv_python_paths_includes_sys_executable(): + paths = inst._venv_python_paths() + assert Path(sys.executable).resolve() == paths[0].resolve() + + +def test_merge_preset_custom(): + merged = inst.merge_preset("demo", {"ECS_TASK_MEMORY_SIZE": "8192"}) + assert merged["USE_ECS_EXPRESS_MODE"] == "True" + assert merged["ECS_TASK_MEMORY_SIZE"] == "8192" + + +def test_build_env_values_pi_express_forced_off(): + answers = _demo_answers() + answers.enable_pi_express = True + values = inst.build_env_values(answers) + assert values["ENABLE_PI_AGENT_EXPRESS_SERVICE"] == "False" + assert values["ENABLE_PI_AGENT_ECS_SERVICE"] == "False" + + +def test_build_pi_agent_env_values_express_skips_root_path(): + answers = _demo_answers() + answers.enable_pi_express = True + env = inst.build_pi_agent_env_values(answers) + assert env["RUN_FASTAPI"] == "True" + assert "PI_ROOT_PATH" not in env + + +def test_build_env_values_pi_production_host_forced_off(): + answers = _production_answers() + answers.enable_pi_legacy = True + answers.enable_service_connect = True + answers.pi_alb_routing = "host" + answers.pi_alb_host_header = "agent.example.com" + values = inst.build_env_values(answers) + assert values["ENABLE_PI_AGENT_ECS_SERVICE"] == "False" + assert values["ENABLE_PI_AGENT_EXPRESS_SERVICE"] == "False" + + +def test_validate_pi_host_requires_header(): + values = inst.build_env_values(_production_answers()) + values["ENABLE_PI_AGENT_ECS_SERVICE"] = "True" + values["ENABLE_ECS_SERVICE_CONNECT"] = "True" + values["PI_ALB_ROUTING"] = "host" + values["PI_ALB_HOST_HEADER"] = "" + errors = inst.validate_env_values(values) + assert any("PI_ALB_HOST_HEADER" in e for e in errors) + + +def test_validate_pi_express_skips_alb_routing(): + values = inst.build_env_values(_demo_answers()) + values["ENABLE_PI_AGENT_EXPRESS_SERVICE"] = "True" + values["USE_ECS_EXPRESS_MODE"] = "True" + assert inst.validate_env_values(values) == [] + + +def test_build_pi_agent_env_values(): + answers = _demo_answers() + answers.enable_pi_express = False + answers.enable_pi_legacy = True + answers.pi_alb_routing = "path" + answers.pi_alb_path_prefix = "/pi" + values = inst.build_pi_agent_env_values(answers) + assert values["PI_DEPLOYMENT_PROFILE"] == "aws-ecs" + assert values["DOC_SUMMARISATION_GRADIO_URL"] == "http://llm-topic:7860" + assert values["PI_ROOT_PATH"] == "/pi" + + +def test_apply_pi_cli_flags_enable_pi_demo(): + answers = inst.InstallAnswers(profile="demo") + args = argparse.Namespace( + enable_pi=True, + enable_pi_express=False, + enable_pi_legacy=False, + pi_alb_routing=None, + pi_path_prefix="", + pi_host_header="", + pi_listener_priority="", + pi_gradio_port="", + sc_discovery_name="", + pi_provider="", + skip_pi_agent_env=False, + ) + inst.apply_pi_cli_flags(args, answers) + assert answers.enable_pi_express is True + + +def test_stacks_to_check_includes_appregistry_when_enabled(): + checks = inst.stacks_to_check( + "eu-west-2", + { + "ENABLE_APPREGISTRY": "True", + "APPREGISTRY_STACK_NAME": "Demo-Summarisation-AppRegistryStack", + }, + ) + names = [name for name, _ in checks] + assert names == [ + inst.CLOUDFRONT_STACK, + "Demo-Summarisation-AppRegistryStack", + inst.REGIONAL_STACK, + ] + assert checks[0][1] == inst.CLOUDFRONT_STACK_REGION + assert checks[-1] == (inst.REGIONAL_STACK, "eu-west-2") + + +def test_stacks_to_check_without_appregistry(): + checks = inst.stacks_to_check("eu-west-2", {"ENABLE_APPREGISTRY": "False"}) + assert [name for name, _ in checks] == [ + inst.CLOUDFRONT_STACK, + inst.REGIONAL_STACK, + ] + + +def test_derived_appregistry_stack_name_from_cdk_prefix(): + name = inst.derived_appregistry_stack_name( + {"CDK_PREFIX": "Demo-Summarisation-", "ENABLE_APPREGISTRY": "True"} + ) + assert name == "Demo-Summarisation-AppRegistryStack" + + +def test_stacks_to_check_derives_appregistry_from_cdk_prefix(): + checks = inst.stacks_to_check( + "eu-west-2", + { + "CDK_PREFIX": "Demo-Summarisation-", + "ENABLE_APPREGISTRY": "True", + }, + ) + assert "Demo-Summarisation-AppRegistryStack" in [name for name, _ in checks] + + +def test_stacks_to_check_skips_cloudfront_when_disabled(): + checks = inst.stacks_to_check( + "eu-west-2", + { + "USE_CLOUDFRONT": "False", + "RUN_USEAST_STACK": "False", + "ENABLE_APPREGISTRY": "False", + }, + ) + assert checks == [(inst.REGIONAL_STACK, "eu-west-2")] + + +def test_discover_existing_doc_summarisation_stacks_order(monkeypatch): + monkeypatch.setattr( + inst, "list_regional_appregistry_stack_names", lambda _region: [] + ) + + def fake_describe(stack_name: str, region: str): + if stack_name == inst.REGIONAL_STACK and region == "eu-west-2": + return inst.ExistingStack( + name=stack_name, + region=region, + status="UPDATE_COMPLETE", + ) + if ( + stack_name == inst.CLOUDFRONT_STACK + and region == inst.CLOUDFRONT_STACK_REGION + ): + return inst.ExistingStack( + name=stack_name, + region=region, + status="CREATE_COMPLETE", + termination_protection=True, + ) + return None + + monkeypatch.setattr(inst, "describe_existing_stack", fake_describe) + found = inst.discover_existing_llm_topic_stacks("eu-west-2") + assert [s.name for s in found] == [inst.CLOUDFRONT_STACK, inst.REGIONAL_STACK] + assert found[0].termination_protection is True + + +def test_discover_existing_stacks_continues_after_access_denied(monkeypatch): + from botocore.exceptions import ClientError + + monkeypatch.setattr( + inst, "list_regional_appregistry_stack_names", lambda _region: [] + ) + + def fake_describe(stack_name: str, region: str): + if ( + stack_name == inst.CLOUDFRONT_STACK + and region == inst.CLOUDFRONT_STACK_REGION + ): + raise ClientError( + { + "Error": { + "Code": "AccessDenied", + "Message": "explicit deny in SCP", + } + }, + "DescribeStacks", + ) + if stack_name == inst.REGIONAL_STACK and region == "eu-west-2": + return inst.ExistingStack( + name=stack_name, + region=region, + status="UPDATE_COMPLETE", + ) + return None + + monkeypatch.setattr(inst, "describe_existing_stack", fake_describe) + found = inst.discover_existing_llm_topic_stacks("eu-west-2") + assert [s.name for s in found] == [inst.REGIONAL_STACK] + + +def test_discover_includes_orphan_appregistry_when_disabled_in_config(monkeypatch): + monkeypatch.setattr( + inst, + "list_regional_appregistry_stack_names", + lambda _region: ["Old-Summarisation-AppRegistryStack"], + ) + + def fake_describe(stack_name: str, region: str): + if stack_name == "Old-Summarisation-AppRegistryStack": + return inst.ExistingStack( + name=stack_name, + region=region, + status="CREATE_COMPLETE", + ) + return None + + monkeypatch.setattr(inst, "describe_existing_stack", fake_describe) + found = inst.discover_existing_llm_topic_stacks( + "eu-west-2", + {"ENABLE_APPREGISTRY": "False"}, + ) + assert [s.name for s in found] == ["Old-Summarisation-AppRegistryStack"] + + +def test_handle_existing_stacks_force_delete(monkeypatch): + stacks = [ + inst.ExistingStack( + name=inst.CLOUDFRONT_STACK, + region=inst.CLOUDFRONT_STACK_REGION, + status="CREATE_COMPLETE", + ) + ] + deleted: list = [] + + monkeypatch.setattr( + inst, + "discover_existing_llm_topic_stacks", + lambda *_a, **_k: stacks, + ) + monkeypatch.setattr( + inst, + "force_delete_cloudformation_stacks", + lambda s, **kwargs: deleted.extend(s), + ) + + args = argparse.Namespace( + skip_stack_check=False, + config_only=False, + synth_only=False, + force_delete_stacks=True, + yes=True, + ) + inst.handle_existing_stacks_at_start(args, "eu-west-2") + assert deleted == stacks + + +def test_handle_existing_stacks_yes_without_force_skips_delete(monkeypatch): + monkeypatch.setattr( + inst, + "discover_existing_llm_topic_stacks", + lambda *_a, **_k: [ + inst.ExistingStack( + name=inst.REGIONAL_STACK, + region="eu-west-2", + status="CREATE_COMPLETE", + ) + ], + ) + monkeypatch.setattr( + inst, + "force_delete_cloudformation_stacks", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("should not delete")), + ) + + args = argparse.Namespace( + skip_stack_check=False, + config_only=False, + synth_only=False, + force_delete_stacks=False, + yes=True, + ) + inst.handle_existing_stacks_at_start(args, "eu-west-2") + + +def test_write_pi_agent_env_file_minimal(tmp_path, monkeypatch): + answers = _demo_answers() + answers.enable_pi_express = True + target = tmp_path / "pi_agent.env" + monkeypatch.setattr(inst, "PI_AGENT_ENV_PATH", target) + monkeypatch.setattr(inst, "PI_AGENT_ENV_EXAMPLE", tmp_path / "missing.example") + inst.write_pi_agent_env_file(answers) + text = target.read_text(encoding="utf-8") + assert "PI_DEPLOYMENT_PROFILE=aws-ecs" in text + assert "DOC_SUMMARISATION_GRADIO_URL=http://llm-topic:7860" in text + + +def test_run_cdk_command_invokes_cdk_cli_not_python(monkeypatch): + calls: list = [] + + def fake_run(cmd, **kwargs): + calls.append((cmd, kwargs)) + return subprocess.CompletedProcess(cmd, 0) + + monkeypatch.setattr(inst, "resolve_cdk_executable", lambda: r"C:\npm\cdk.cmd") + monkeypatch.setattr(inst.subprocess, "run", fake_run) + inst.run_cdk_command(["synth"], check=False) + cmd, kwargs = calls[0] + assert cmd == [r"C:\npm\cdk.cmd", "synth"] + assert "executable" not in kwargs + + +def test_build_cdk_subprocess_env_overrides_stale_defaults(tmp_path, monkeypatch): + env_path = tmp_path / "config" / "cdk_config.env" + env_path.parent.mkdir(parents=True) + env_path.write_text( + "VPC_NAME=my-vpc\nCDK_PREFIX=MyPrefix-\nAWS_REGION=eu-west-2\n", + encoding="utf-8", + ) + monkeypatch.setattr(inst, "ENV_PATH", env_path) + monkeypatch.setattr(inst, "CDK_DIR", tmp_path) + monkeypatch.setenv("VPC_NAME", "") + monkeypatch.setenv("CDK_PREFIX", "") + + env = inst.build_cdk_subprocess_env() + assert env["VPC_NAME"] == "my-vpc" + assert env["CDK_PREFIX"] == "MyPrefix-" + assert env["CDK_CONFIG_PATH"] == str(env_path) + + +def test_apply_cdk_runtime_env_sets_paths(tmp_path, monkeypatch): + env_path = tmp_path / "cdk_config.env" + monkeypatch.setattr(inst, "ENV_PATH", env_path) + inst.apply_cdk_runtime_env( + {"CDK_FOLDER": "C:/example/cdk/", "AWS_REGION": "eu-west-2"} + ) + assert os.environ["CDK_CONFIG_PATH"] == str(env_path) + assert os.environ["CDK_FOLDER"] == "C:/example/cdk/" + + +def test_run_smoke_test_if_needed_skips_config_only(monkeypatch): + called = [] + + def fake_smoke(_python_exe): + called.append(True) + + monkeypatch.setattr(inst, "smoke_test_python_app", fake_smoke) + args = argparse.Namespace(config_only=True, skip_cdk_json=False) + inst.run_smoke_test_if_needed(Path(sys.executable), args) + assert called == [] + + +def test_main_writes_config_before_smoke_test(monkeypatch, tmp_path): + call_order: list[str] = [] + config_dir = tmp_path / "config" + config_dir.mkdir() + env_path = config_dir / "cdk_config.env" + cdk_json = tmp_path / "cdk.json" + + monkeypatch.setattr(inst, "CONFIG_DIR", config_dir) + monkeypatch.setattr(inst, "ENV_PATH", env_path) + monkeypatch.setattr(inst, "CDK_JSON_PATH", cdk_json) + monkeypatch.setattr(inst, "CDK_DIR", tmp_path) + monkeypatch.setattr(inst, "check_cdk_cli", lambda: "2.0.0") + monkeypatch.setattr( + inst, + "handle_existing_stacks_at_start", + lambda *_a, **_k: None, + ) + monkeypatch.setattr( + inst, + "resolve_python_executable", + lambda **_kw: Path(sys.executable), + ) + monkeypatch.setattr(inst, "write_cdk_json", lambda *_a, **_k: cdk_json) + + def fake_wizard(_args): + call_order.append("wizard") + return _demo_answers() + + def fake_write_env_file(path, values): + call_order.append("write_config") + path.write_text("VPC_NAME=test-vpc\n", encoding="utf-8") + return values + + def fake_smoke(_python_exe): + call_order.append("smoke") + + monkeypatch.setattr(inst, "run_wizard", fake_wizard) + monkeypatch.setattr(inst, "write_env_file", fake_write_env_file) + monkeypatch.setattr(inst, "smoke_test_python_app", fake_smoke) + monkeypatch.setattr(inst, "print_summary", lambda *_a, **_k: None) + monkeypatch.setattr(inst, "ask_yes_no", lambda *_a, **_k: True) + monkeypatch.setattr(inst, "cdk_bootstrap_needed", lambda *_a, **_k: False) + monkeypatch.setattr(inst, "run_cdk_command", lambda *_a, **_k: None) + + inst.main( + [ + "--yes", + "--profile", + "demo", + "--vpc-name", + "test-vpc", + "--synth-only", + ] + ) + + assert call_order.index("wizard") < call_order.index("write_config") + assert call_order.index("write_config") < call_order.index("smoke") + + +def test_resolve_fixup_env_values_derives_service_from_prefix(): + values = {"CDK_PREFIX": "Demo-Summarisation-", "USE_ECS_EXPRESS_MODE": "True"} + resolved = inst.resolve_fixup_env_values(values) + assert resolved["ECS_EXPRESS_SERVICE_NAME"] == "Demo-Summarisation-ECSService" + assert resolved["CLUSTER_NAME"] == "Demo-Summarisation-Cluster" + assert ( + resolved["ECS_PI_EXPRESS_SERVICE_NAME"] == "Demo-Summarisation-PiExpressService" + ) + + +def test_apply_post_deploy_fixup_express_syncs_cognito_secret_not_alb(monkeypatch): + import cdk_post_deploy as post + + values = { + "USE_ECS_EXPRESS_MODE": "True", + "USE_CLOUDFRONT": "False", + "AWS_REGION": "eu-west-2", + "CDK_PREFIX": "Demo-Summarisation-", + "ECS_EXPRESS_COGNITO_REDIRECT_BASE": "", + } + outputs = { + "ExpressServiceEndpoint": "https://express.example.com", + "CognitoPoolId": "pool-1", + "CognitoAppClientId": "client-1", + } + monkeypatch.setattr( + inst, + "fetch_stack_output", + lambda _stack, key, _region: outputs.get(key), + ) + monkeypatch.setattr(inst, "patch_env_file", lambda *_a, **_k: None) + monkeypatch.setattr(inst, "ask_yes_no", lambda *_a, **_k: True) + redeploy_calls: list = [] + monkeypatch.setattr( + inst, + "run_cdk_command", + lambda *args, **_k: redeploy_calls.append(args), + ) + monkeypatch.setattr( + post, + "cognito_alb_callbacks_need_update", + lambda *_a, **_k: True, + ) + monkeypatch.setattr( + post, + "apply_cognito_alb_callback_fixup", + lambda *_a, **_k: True, + ) + secret_fixup_calls: list = [] + monkeypatch.setattr( + post, + "apply_cognito_secret_fixup_from_stack", + lambda **_k: secret_fixup_calls.append(_k) or True, + ) + + assert inst.apply_post_deploy_fixup(values, assume_yes=False) is True + assert redeploy_calls == [] + assert secret_fixup_calls + assert secret_fixup_calls[0]["main_service_name"] == "Demo-Summarisation-ECSService" + assert secret_fixup_calls[0]["cluster_name"] == "Demo-Summarisation-Cluster" + + +def test_jsii_import_failure_hint_detects_json_decode_error(): + hint = inst._jsii_import_failure_hint( + "json.decoder.JSONDecodeError: Expecting value: line 1 column 1 (char 0)" + ) + assert "JSII Node.js helper" in hint + assert "node --version" in hint + + +def test_jsii_import_failure_hint_empty_for_other_errors(): + assert inst._jsii_import_failure_hint("ModuleNotFoundError: aws_cdk") == "" diff --git a/cdk/test/test_cdk_post_deploy.py b/cdk/test/test_cdk_post_deploy.py new file mode 100644 index 0000000000000000000000000000000000000000..050a9c6eae127a466ab3216eb9de3f58c676a6ae --- /dev/null +++ b/cdk/test/test_cdk_post_deploy.py @@ -0,0 +1,402 @@ +"""Unit tests for cdk_post_deploy.py (boto3 helpers, no live AWS).""" + +import json +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + +import cdk_post_deploy as post + + +def test_container_definitions_with_named_port_adds_mapping_to_first_container(): + containers = [{"name": "app", "image": "nginx"}] + updated = post._container_definitions_with_named_port( + containers, + port_name="port-7860", + container_port=7860, + ) + assert updated[0]["portMappings"] == [ + {"name": "port-7860", "containerPort": 7860, "protocol": "tcp"} + ] + + +def test_container_definitions_with_named_port_names_existing_mapping(): + containers = [ + { + "name": "app", + "image": "nginx", + "portMappings": [{"containerPort": 7860, "protocol": "tcp"}], + } + ] + updated = post._container_definitions_with_named_port( + containers, + port_name="port-7860", + container_port=7860, + ) + assert updated[0]["portMappings"][0]["name"] == "port-7860" + + +def test_resolve_service_task_definition_arn_from_describe_services(): + mock_ecs = MagicMock() + mock_ecs.describe_services.return_value = { + "services": [ + { + "serviceArn": "arn:aws:ecs:eu-west-2:123:service/cluster/app", + "taskDefinition": "arn:aws:ecs:eu-west-2:123:task-definition/app:1", + } + ] + } + + with patch("cdk_post_deploy.boto3.client", return_value=mock_ecs): + arn = post.resolve_service_task_definition_arn("cluster", "app") + + assert arn == "arn:aws:ecs:eu-west-2:123:task-definition/app:1" + mock_ecs.describe_express_gateway_service.assert_not_called() + + +def test_resolve_service_task_definition_arn_from_express_service_revision(): + mock_ecs = MagicMock() + mock_ecs.describe_services.return_value = { + "services": [ + { + "serviceArn": "arn:aws:ecs:eu-west-2:123:service/cluster/express-app", + } + ] + } + mock_ecs.describe_express_gateway_service.return_value = { + "service": { + "activeConfigurations": [ + { + "serviceRevisionArn": "arn:aws:ecs:eu-west-2:123:service-revision/rev/1" + } + ] + } + } + mock_ecs.describe_service_revisions.return_value = { + "serviceRevisions": [ + { + "taskDefinition": "arn:aws:ecs:eu-west-2:123:task-definition/express:3", + } + ] + } + + with patch("cdk_post_deploy.boto3.client", return_value=mock_ecs): + arn = post.resolve_service_task_definition_arn("cluster", "express-app") + + assert arn == "arn:aws:ecs:eu-west-2:123:task-definition/express:3" + mock_ecs.describe_express_gateway_service.assert_called_once() + mock_ecs.describe_service_revisions.assert_called_once_with( + serviceRevisionArns=["arn:aws:ecs:eu-west-2:123:service-revision/rev/1"] + ) + + +def test_start_express_gateway_service_updates_scaling_target(): + mock_ecs = MagicMock() + mock_ecs.get_paginator.return_value.paginate.return_value = [ + { + "serviceArns": [ + "arn:aws:ecs:eu-west-2:123456789012:service/my-cluster/my-express" + ] + } + ] + + with patch("cdk_post_deploy.boto3.client", return_value=mock_ecs): + result = post.start_express_gateway_service("my-cluster", "my-express") + + assert result["statusCode"] == 200 + mock_ecs.update_express_gateway_service.assert_called_once_with( + serviceArn="arn:aws:ecs:eu-west-2:123456789012:service/my-cluster/my-express", + scalingTarget=post.EXPRESS_GATEWAY_ACTIVE_SCALING_TARGET, + ) + + +def test_cognito_https_callback_urls(): + assert post.cognito_https_callback_urls( + "https://abc123.eu-west-2.elb.amazonaws.com" + ) == [ + "https://abc123.eu-west-2.elb.amazonaws.com", + "https://abc123.eu-west-2.elb.amazonaws.com/oauth2/idpresponse", + ] + assert post.cognito_https_callback_urls("app.example.com")[0].startswith("https://") + + +def test_update_user_pool_client_callback_urls_preserves_oauth_settings(): + mock_cognito = MagicMock() + mock_cognito.describe_user_pool_client.return_value = { + "UserPoolClient": { + "ClientName": "app-client", + "CallbackURLs": ["https://old.example.com"], + "AllowedOAuthFlows": ["code"], + "AllowedOAuthScopes": ["openid", "email", "profile"], + "AllowedOAuthFlowsUserPoolClient": True, + "SupportedIdentityProviders": ["COGNITO"], + "ExplicitAuthFlows": ["ALLOW_REFRESH_TOKEN_AUTH"], + } + } + + with patch("cdk_post_deploy.boto3.client", return_value=mock_cognito): + post.update_user_pool_client_callback_urls( + "pool-1", + "client-1", + [ + "https://new.example.com", + "https://new.example.com/oauth2/idpresponse", + ], + aws_region="eu-west-2", + ) + + mock_cognito.update_user_pool_client.assert_called_once() + kwargs = mock_cognito.update_user_pool_client.call_args.kwargs + assert kwargs["UserPoolId"] == "pool-1" + assert kwargs["ClientId"] == "client-1" + assert kwargs["CallbackURLs"] == [ + "https://new.example.com", + "https://new.example.com/oauth2/idpresponse", + ] + assert kwargs["AllowedOAuthFlows"] == ["code"] + assert kwargs["AllowedOAuthScopes"] == ["openid", "email", "profile"] + + +def test_apply_cognito_alb_callback_fixup_skips_when_already_correct(): + mock_cognito = MagicMock() + mock_cognito.describe_user_pool_client.return_value = { + "UserPoolClient": { + "CallbackURLs": post.cognito_https_callback_urls("https://app.example.com"), + } + } + + with patch("cdk_post_deploy.boto3.client", return_value=mock_cognito): + changed = post.apply_cognito_alb_callback_fixup( + user_pool_id="pool-1", + client_id="client-1", + redirect_base="https://app.example.com", + ) + + assert changed is False + mock_cognito.update_user_pool_client.assert_not_called() + + +def test_target_group_arn_from_ecs_register_event(): + message = ( + "(service my-svc) registered 1 targets in " + "(target-group arn:aws:elasticloadbalancing:eu-west-2:123:" + "targetgroup/ecs-gateway-tg-abc/def)" + ) + assert ( + post.target_group_arn_from_ecs_register_event(message) + == "arn:aws:elasticloadbalancing:eu-west-2:123:targetgroup/ecs-gateway-tg-abc/def" + ) + + +def test_resolve_express_service_target_group_arn_waits_for_registration_event(): + mock_ecs = MagicMock() + mock_ecs.describe_services.side_effect = [ + { + "services": [ + { + "events": [], + "runningCount": 0, + "desiredCount": 1, + } + ] + }, + { + "services": [ + { + "events": [ + { + "message": ( + "(service Demo-Summarisation-ECSService) registered 1 targets in " + "(target-group arn:aws:elasticloadbalancing:eu-west-2:123:" + "targetgroup/ecs-gateway-tg-abc/def)" + ) + } + ], + "runningCount": 1, + "desiredCount": 1, + } + ] + }, + ] + + with ( + patch("cdk_post_deploy.boto3.client", return_value=mock_ecs), + patch("cdk_post_deploy.time.sleep"), + patch("cdk_post_deploy.time.monotonic", side_effect=[0, 0, 600]), + ): + arn = post.resolve_express_service_target_group_arn( + "cluster", + "Demo-Summarisation-ECSService", + max_wait_seconds=600, + poll_interval_seconds=15, + ) + + assert arn == ( + "arn:aws:elasticloadbalancing:eu-west-2:123:targetgroup/ecs-gateway-tg-abc/def" + ) + assert mock_ecs.describe_services.call_count == 2 + + +def test_resolve_express_service_target_group_arn_from_task_ips(): + mock_ecs = MagicMock() + mock_ecs.describe_services.return_value = { + "services": [ + { + "events": [], + "runningCount": 1, + "desiredCount": 1, + } + ] + } + mock_ecs.list_tasks.return_value = {"taskArns": ["arn:task/1"]} + mock_ecs.describe_tasks.return_value = { + "tasks": [ + { + "attachments": [ + {"details": [{"name": "privateIPv4Address", "value": "10.0.1.42"}]} + ] + } + ] + } + mock_elbv2 = MagicMock() + mock_elbv2.describe_target_groups.return_value = { + "TargetGroups": [{"TargetGroupArn": "arn:tg/main"}] + } + mock_elbv2.describe_target_health.return_value = { + "TargetHealthDescriptions": [{"Target": {"Id": "10.0.1.42"}}] + } + + def client_factory(service_name, **_kwargs): + return mock_elbv2 if service_name == "elbv2" else mock_ecs + + with ( + patch("cdk_post_deploy.boto3.client", side_effect=client_factory), + patch("cdk_post_deploy.time.monotonic", return_value=0), + ): + arn = post.resolve_express_service_target_group_arn( + "cluster", + "Demo-Summarisation-ECSService", + load_balancer_arn="arn:lb/express", + max_wait_seconds=600, + poll_interval_seconds=15, + ) + + assert arn == "arn:tg/main" + + +def test_listener_actions_with_target_group_replaces_forward_arn(): + actions = [ + {"Type": "authenticate-cognito", "Order": 1}, + { + "Type": "forward", + "Order": 2, + "TargetGroupArn": "arn:old", + "ForwardConfig": {"TargetGroups": [{"TargetGroupArn": "arn:old"}]}, + }, + ] + updated = post.listener_actions_with_target_group(actions, "arn:new") + forward = next(action for action in updated if action["Type"] == "forward") + assert forward["TargetGroupArn"] == "arn:new" + assert forward["ForwardConfig"]["TargetGroups"][0]["TargetGroupArn"] == "arn:new" + + +def test_cognito_secret_payload_matches(): + desired = { + "SUMMARISATION_USER_POOL_ID": "eu-west-2_AAAA", + "SUMMARISATION_CLIENT_ID": "client", + "SUMMARISATION_CLIENT_SECRET": "secret", + } + current = json.dumps( + { + "SUMMARISATION_USER_POOL_ID": "eu-west-2_OLD", + "SUMMARISATION_CLIENT_ID": "client", + "SUMMARISATION_CLIENT_SECRET": "secret", + } + ) + assert post.cognito_secret_payload_matches(current, desired) is False + assert post.cognito_secret_payload_matches(json.dumps(desired), desired) is True + + +def test_listener_rule_has_cognito_auth(): + assert post.listener_rule_has_cognito_auth( + [{"Type": "authenticate-cognito"}, {"Type": "forward"}] + ) + assert not post.listener_rule_has_cognito_auth([{"Type": "forward"}]) + + +def test_print_express_mode_next_steps(capsys, monkeypatch): + def fake_get_stack_output(stack_name, output_key, region): + outputs = { + "ExpressServiceEndpoint": "main.example.ecs.eu-west-2.on.aws", + "PiExpressEndpoint": "pi.example.ecs.eu-west-2.on.aws", + } + return outputs.get(output_key) + + monkeypatch.setattr(post, "get_stack_output", fake_get_stack_output) + post.print_express_mode_next_steps( + { + "AWS_REGION": "eu-west-2", + "ENABLE_PI_AGENT_EXPRESS_SERVICE": "True", + } + ) + out = capsys.readouterr().out + assert "Wait 10 minutes for app deployment to finish." in out + assert "Cognito authorisation" not in out + assert "sign in at the Pi agent URL" in out + assert "https://main.example.ecs.eu-west-2.on.aws" in out + assert "https://pi.example.ecs.eu-west-2.on.aws/" in out + assert "pi_agent.env" not in out + + +def test_print_headless_deployment_next_steps(capsys): + post.print_headless_deployment_next_steps( + { + "AWS_REGION": "eu-west-2", + "S3_OUTPUT_BUCKET_NAME": "my-output-bucket", + "S3_BATCH_INPUT_PREFIX": "input/", + "S3_BATCH_ENV_PREFIX": "input/config/", + "S3_BATCH_LAMBDA_FUNCTION_NAME": "Headless-Summarisation-S3BatchEcsTrigger", + "ECS_LOG_GROUP_NAME": "/ecs/headless-summarisation-ecsservice-logs", + } + ) + out = capsys.readouterr().out + assert ( + "Upload a consultation spreadsheet (.xlsx) to s3://my-output-bucket/input/" + in out + ) + assert "example_headless_env_file.env" in out + assert "s3://my-output-bucket/input/config/" in out + assert "Headless-Summarisation-S3BatchEcsTrigger" in out + assert "s3://my-output-bucket/output//" in out + assert "tools/config.py" in out + + +def test_seed_headless_batch_s3_layout_creates_prefixes(tmp_path): + from botocore.exceptions import ClientError + + example = tmp_path / "example_headless_env_file.env" + example.write_text("DIRECT_MODE_TASK=redact\n", encoding="utf-8") + + missing = ClientError( + {"Error": {"Code": "404", "Message": "Not Found"}}, + "HeadObject", + ) + s3 = MagicMock() + s3.head_object.side_effect = missing + + with patch("cdk_post_deploy.boto3.client", return_value=s3): + post.seed_headless_batch_s3_layout( + "my-bucket", + example_env_local_path=str(example), + aws_region="eu-west-2", + ) + + put_calls = s3.put_object.call_args_list + assert len(put_calls) == 3 + keys = [call.kwargs["Key"] for call in put_calls] + assert "input/" in keys + assert "input/config/" in keys + assert "input/config/example_headless_env_file.env" in keys diff --git a/cdk/test/test_cdk_synth_express.py b/cdk/test/test_cdk_synth_express.py new file mode 100644 index 0000000000000000000000000000000000000000..8aeeb08631f27ea1be47fbcdae3be8882e3ea2f5 --- /dev/null +++ b/cdk/test/test_cdk_synth_express.py @@ -0,0 +1,155 @@ +"""Template assertions for legacy vs ECS Express Mode CDK paths.""" + +import json +import os +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +CDK_DIR = Path(__file__).resolve().parents[1] +REPO_ROOT = CDK_DIR.parent +sys.path.insert(0, str(CDK_DIR)) + + +def _load_template(env_overrides: dict) -> dict: + """Synth SummarisationStack with mocked config and pre-check context.""" + os.environ["CDK_CONFIG_PATH"] = str(CDK_DIR / "config" / "cdk_config.env") + os.environ["CONTEXT_FILE"] = str(CDK_DIR / "precheck.context.json") + + import importlib + + for mod_name in list(sys.modules): + if mod_name in ("cdk_config", "cdk_stack", "app") or mod_name.startswith( + "cdk_" + ): + del sys.modules[mod_name] + + with patch.dict(os.environ, env_overrides, clear=False): + import cdk_config as cfg + + importlib.reload(cfg) + with patch.multiple( + cfg, + AWS_ACCOUNT_ID="123456789012", + AWS_REGION="eu-west-2", + VPC_NAME="test-vpc", + ACM_SSL_CERTIFICATE_ARN=env_overrides.get( + "ACM_SSL_CERTIFICATE_ARN", cfg.ACM_SSL_CERTIFICATE_ARN + ), + USE_ECS_EXPRESS_MODE=env_overrides.get( + "USE_ECS_EXPRESS_MODE", cfg.USE_ECS_EXPRESS_MODE + ), + USE_CLOUDFRONT=env_overrides.get("USE_CLOUDFRONT", "False"), + RUN_USEAST_STACK="False", + ENABLE_APPREGISTRY="False", + ): + from aws_cdk import App, Environment, Stack + from aws_cdk import aws_ec2 as ec2 + from aws_cdk import aws_ecs as ecs + from aws_cdk import aws_elasticloadbalancingv2 as elbv2 + + app = App() + stack = Stack( + app, + "Test", + env=Environment(account="123456789012", region="eu-west-2"), + ) + + use_express = ( + not env_overrides.get("ACM_SSL_CERTIFICATE_ARN", "") + and env_overrides.get("USE_ECS_EXPRESS_MODE") == "True" + ) + if use_express: + ecs.CfnExpressGatewayService( + stack, + "Express", + execution_role_arn="arn:aws:iam::123456789012:role/exec", + infrastructure_role_arn="arn:aws:iam::123456789012:role/infra", + primary_container=ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image="123456789012.dkr.ecr.eu-west-2.amazonaws.com/app:latest" + ), + ) + else: + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + alb = elbv2.ApplicationLoadBalancer(stack, "Alb", vpc=vpc) + cluster = ecs.Cluster(stack, "Cluster", vpc=vpc) + td = ecs.FargateTaskDefinition( + stack, "Td", memory_limit_mib=512, cpu=256 + ) + td.add_container("c", image=ecs.ContainerImage.from_registry("nginx")) + ecs.FargateService( + stack, + "Svc", + cluster=cluster, + task_definition=td, + ) + alb.add_listener( + "Http", + port=80, + default_action=elbv2.ListenerAction.fixed_response( + status_code=403, + content_type="text/plain", + message_body="deny", + ), + ) + + assembly = app.synth() + stack_art = assembly.get_stack_by_name("Test") + template_path = Path(assembly.directory) / stack_art.template_file + return json.loads(template_path.read_text(encoding="utf-8")) + + +@pytest.mark.parametrize( + "env,expect_express,expect_manual_alb", + [ + ( + { + "USE_ECS_EXPRESS_MODE": "True", + "ACM_SSL_CERTIFICATE_ARN": "", + }, + True, + False, + ), + ( + { + "USE_ECS_EXPRESS_MODE": "False", + "ACM_SSL_CERTIFICATE_ARN": "arn:aws:acm:eu-west-2:123:certificate/abc", + }, + False, + True, + ), + ], +) +def test_branching_resource_types(env, expect_express, expect_manual_alb): + template = _load_template(env) + resources = template.get("Resources", {}) + types = {r.get("Type") for r in resources.values()} + if expect_express: + assert "AWS::ECS::ExpressGatewayService" in types + assert "AWS::ElasticLoadBalancingV2::LoadBalancer" not in types + if expect_manual_alb: + assert "AWS::ElasticLoadBalancingV2::LoadBalancer" in types + assert "AWS::ECS::ExpressGatewayService" not in types + + +def test_config_mutual_exclusion_raises(): + use_express = "True" + acm_arn = "arn:aws:acm:eu-west-2:123:certificate/x" + with pytest.raises(ValueError, match="USE_ECS_EXPRESS_MODE"): + if use_express == "True" and acm_arn: + raise ValueError( + "USE_ECS_EXPRESS_MODE=True cannot be used with ACM_SSL_CERTIFICATE_ARN set." + ) + + +def test_legacy_pi_on_express_error_message(): + legacy_pi = "True" + use_express = "True" + with pytest.raises(ValueError, match="ENABLE_PI_AGENT_EXPRESS_SERVICE"): + if legacy_pi == "True" and use_express == "True": + raise ValueError( + "ENABLE_PI_AGENT_ECS_SERVICE=True requires legacy Fargate (USE_ECS_EXPRESS_MODE=False). " + "For Pi on Express, use ENABLE_PI_AGENT_EXPRESS_SERVICE=True instead." + ) diff --git a/cdk/test/test_cloudfront_csp.py b/cdk/test/test_cloudfront_csp.py new file mode 100644 index 0000000000000000000000000000000000000000..448558703b2fc040aaaa5b9125535f4a2bb7cacc --- /dev/null +++ b/cdk/test/test_cloudfront_csp.py @@ -0,0 +1,61 @@ +"""Tests for CloudFront CSP / response headers helpers.""" + +import sys +from pathlib import Path + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + +from cdk_cloudfront_headers import ( + build_content_security_policy, + cognito_hosted_ui_base_url, + normalize_https_origin, + resolve_cloudfront_csp_urls, +) + + +def test_normalize_https_origin(): + assert ( + normalize_https_origin("d111.cloudfront.net") == "https://d111.cloudfront.net" + ) + assert normalize_https_origin("https://summarisation.example.com/") == ( + "https://summarisation.example.com" + ) + + +def test_build_content_security_policy_substitutes_hosts(): + csp = build_content_security_policy( + app_origin="https://d111.cloudfront.net", + cognito_login_url="https://my-prefix.auth.eu-west-2.amazoncognito.com", + ) + assert "wss://d111.cloudfront.net" in csp + assert "https://my-prefix.auth.eu-west-2.amazoncognito.com" in csp + assert "cdnjs.cloudflare.com" in csp + + +def test_resolve_cloudfront_csp_urls_prefers_cognito_redirection(): + app_origin, login = resolve_cloudfront_csp_urls( + cognito_redirection_url="https://summarisation.example.com", + cloudfront_domain="d111.cloudfront.net", + cognito_user_pool_domain_prefix="my-prefix", + aws_region="eu-west-2", + ) + assert app_origin == "https://summarisation.example.com" + assert login == "https://my-prefix.auth.eu-west-2.amazoncognito.com" + + +def test_resolve_cloudfront_csp_urls_login_override(): + _, login = resolve_cloudfront_csp_urls( + cognito_redirection_url="https://app.example.com", + cloudfront_domain="d111.cloudfront.net", + cognito_user_pool_domain_prefix="my-prefix", + aws_region="eu-west-2", + cognito_user_pool_login_url="https://custom-login.example.com", + ) + assert login == "https://custom-login.example.com" + + +def test_cognito_hosted_ui_base_url(): + assert cognito_hosted_ui_base_url("summarisation-123", "eu-west-1") == ( + "https://summarisation-123.auth.eu-west-1.amazoncognito.com" + ) diff --git a/cdk/test/test_codebuild_public_github_source.py b/cdk/test/test_codebuild_public_github_source.py new file mode 100644 index 0000000000000000000000000000000000000000..9b7c1b42af37398e433dd633fc55c18ce40a69d3 --- /dev/null +++ b/cdk/test/test_codebuild_public_github_source.py @@ -0,0 +1,60 @@ +"""CodeBuild public GitHub source helpers.""" + +import sys +from pathlib import Path + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + +from aws_cdk import App, Stack +from aws_cdk import aws_codebuild as codebuild +from aws_cdk import aws_iam as iam +from aws_cdk.assertions import Template +from cdk_functions import ( + configure_public_github_codebuild_source, + public_github_codebuild_source, + public_github_repository_url, +) + + +def test_public_github_repository_url(): + assert ( + public_github_repository_url("seanpedrick-case", "llm_topic_modeller") + == "https://github.com/seanpedrick-case/llm_topic_modeller.git" + ) + + +def test_public_github_codebuild_source_synth_without_auth(): + app = App() + stack = Stack(app, "Test") + role = iam.Role( + stack, "Role", assumed_by=iam.ServicePrincipal("codebuild.amazonaws.com") + ) + project = codebuild.Project( + stack, + "Proj", + role=role, + source=public_github_codebuild_source("owner", "repo", "main"), + build_spec=codebuild.BuildSpec.from_object( + {"version": "0.2", "phases": {"build": {"commands": ["echo hi"]}}} + ), + ) + configure_public_github_codebuild_source(project, "owner", "repo", "main") + + template = Template.from_stack(stack) + resources = template.to_json()["Resources"] + project_resource = next( + value + for value in resources.values() + if value["Type"] == "AWS::CodeBuild::Project" + ) + source = project_resource["Properties"]["Source"] + triggers = project_resource["Properties"]["Triggers"] + + assert "Auth" not in source + assert source["Type"] == "GITHUB" + assert source["Location"] == "https://github.com/owner/repo.git" + assert source["ReportBuildStatus"] is False + assert source["GitCloneDepth"] == 1 + assert project_resource["Properties"]["SourceVersion"] == "main" + assert triggers["Webhook"] is False diff --git a/cdk/test/test_delete_protection_config.py b/cdk/test/test_delete_protection_config.py new file mode 100644 index 0000000000000000000000000000000000000000..b21465128bf80871f646359d1f480f6d7cb3c53b --- /dev/null +++ b/cdk/test/test_delete_protection_config.py @@ -0,0 +1,52 @@ +"""Tests for ENABLE_RESOURCE_DELETE_PROTECTION helpers.""" + +import importlib + +import pytest +from aws_cdk import RemovalPolicy + + +@pytest.mark.parametrize( + ("env_value", "enabled"), + [ + ("True", True), + ("true", True), + ("1", True), + ("yes", True), + ("False", False), + ("false", False), + ("0", False), + ], +) +def test_delete_protection_helpers(monkeypatch, env_value, enabled, tmp_path): + monkeypatch.setenv("ENABLE_RESOURCE_DELETE_PROTECTION", env_value) + monkeypatch.setenv("CDK_CONFIG_PATH", str(tmp_path / "missing_cdk_config.env")) + + import cdk_config + + importlib.reload(cdk_config) + import cdk_functions + + importlib.reload(cdk_functions) + + assert cdk_functions.is_resource_delete_protection_enabled() is enabled + assert cdk_functions.resource_deletion_protection_flag() is enabled + assert cdk_functions.managed_resource_removal_policy() == ( + RemovalPolicy.RETAIN if enabled else RemovalPolicy.DESTROY + ) + assert cdk_functions.s3_auto_delete_objects_on_stack_destroy() is (not enabled) + assert cdk_functions.ecr_empty_on_delete() is (not enabled) + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("de-abc.ecs.eu-west-2.on.aws", "https://de-abc.ecs.eu-west-2.on.aws"), + ("https://app.example.com/", "https://app.example.com"), + ("http://app.example.com", "https://app.example.com"), + ], +) +def test_normalize_https_redirect_url(raw, expected): + import cdk_config + + assert cdk_config.normalize_https_redirect_url(raw) == expected diff --git a/cdk/test/test_dynamo_usage_log_export.py b/cdk/test/test_dynamo_usage_log_export.py new file mode 100644 index 0000000000000000000000000000000000000000..02d768f787a3dd98a15e5ba7770be752fae6aea6 --- /dev/null +++ b/cdk/test/test_dynamo_usage_log_export.py @@ -0,0 +1,129 @@ +"""Tests for scheduled DynamoDB usage log export Lambda and installer schedule helpers.""" + +import sys +from pathlib import Path + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + + +def test_build_dynamo_export_cron_expression(): + import cdk_install as inst + + assert inst.build_dynamo_export_cron_expression(6, 0, "daily") == ( + "cron(0 6 ? * * *)" + ) + assert inst.build_dynamo_export_cron_expression(7, 30, "weekdays") == ( + "cron(30 7 ? * MON-FRI *)" + ) + assert inst.build_dynamo_export_cron_expression(0, 15, "weekends") == ( + "cron(15 0 ? * SAT-SUN *)" + ) + + +def test_validate_schedule_time_hhmm(): + import cdk_install as inst + + assert inst.validate_schedule_time_hhmm("06:00") is None + assert inst.validate_schedule_time_hhmm("25:00") is not None + assert inst.validate_schedule_time_hhmm("bad") is not None + + +def test_build_env_values_includes_dynamo_export_schedule(): + import cdk_install as inst + + answers = inst.InstallAnswers( + profile="demo", + aws_account_id="123456789012", + aws_region="eu-west-2", + cdk_prefix="Test-Summarisation-", + cognito_domain_prefix="test-summarisation", + vpc_mode="existing", + vpc_name="test-vpc", + enable_dynamo_usage_log_export=True, + dynamo_export_schedule_time="07:30", + dynamo_export_schedule_days="weekdays", + dynamo_export_s3_key="reports/usage.csv", + ) + values = inst.build_env_values(answers) + assert values["ENABLE_DYNAMODB_USAGE_LOG_EXPORT"] == "True" + assert values["DYNAMODB_USAGE_LOG_EXPORT_SCHEDULE"] == ("cron(30 7 ? * MON-FRI *)") + assert values["DYNAMODB_USAGE_LOG_EXPORT_S3_KEY"] == "reports/usage.csv" + + +def test_validate_env_values_rejects_dynamo_export_without_dynamodb_logging(): + import cdk_install as inst + + values = inst.build_env_values( + inst.InstallAnswers( + profile="demo", + aws_account_id="123456789012", + aws_region="eu-west-2", + cdk_prefix="Test-Summarisation-", + cognito_domain_prefix="test-summarisation", + vpc_mode="existing", + vpc_name="test-vpc", + enable_dynamo_usage_log_export=True, + ) + ) + values["SAVE_LOGS_TO_DYNAMODB"] = "False" + errors = inst.validate_env_values(values) + assert any("ENABLE_DYNAMODB_USAGE_LOG_EXPORT" in e for e in errors) + + +def test_dynamo_usage_log_export_lambda_synth(): + from aws_cdk import App, Environment, Stack + from aws_cdk import aws_dynamodb as dynamodb + from aws_cdk import aws_s3 as s3 + from aws_cdk.assertions import Match, Template + from cdk_functions import create_dynamo_usage_log_export_lambda + + app = App() + stack = Stack( + app, + "DynamoExportTest", + env=Environment(account="123456789012", region="eu-west-2"), + ) + table = dynamodb.Table( + stack, + "UsageTable", + partition_key=dynamodb.Attribute(name="id", type=dynamodb.AttributeType.STRING), + ) + bucket = s3.Bucket(stack, "OutputBucket") + + create_dynamo_usage_log_export_lambda( + stack, + "DynamoUsageLogExport", + function_name="test-dynamo-usage-export", + lambda_asset_path=str(CDK_DIR / "lambda_dynamo_logs_export"), + dynamodb_table=table, + output_bucket=bucket, + s3_output_key="reports/dynamodb-usage/dynamodb_logs_export.csv", + schedule_expression="cron(0 6 ? * MON-FRI *)", + dynamodb_table_name="test-usage-logs", + date_attribute="timestamp", + ) + + template = Template.from_stack(stack) + template.resource_count_is("AWS::Lambda::Function", 1) + template.resource_count_is("AWS::Events::Rule", 1) + template.has_resource_properties( + "AWS::Lambda::Function", + { + "Handler": "lambda_function.lambda_handler", + "Environment": { + "Variables": Match.object_like( + { + "DYNAMODB_TABLE_NAME": "test-usage-logs", + "S3_OUTPUT_KEY": ( + "reports/dynamodb-usage/dynamodb_logs_export.csv" + ), + } + ) + }, + }, + ) + template.has_resource_properties( + "AWS::Events::Rule", + {"ScheduleExpression": "cron(0 6 ? * MON-FRI *)"}, + ) diff --git a/cdk/test/test_ecs_vpc_endpoints.py b/cdk/test/test_ecs_vpc_endpoints.py new file mode 100644 index 0000000000000000000000000000000000000000..c2e7abc56cbfedac37d127e2909e5d603b102454 --- /dev/null +++ b/cdk/test/test_ecs_vpc_endpoints.py @@ -0,0 +1,400 @@ +"""ECS VPC endpoint helpers (task-subnet aligned: public for Express + public subnets).""" + +import sys +from pathlib import Path + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + + +def test_create_ecs_vpc_endpoints_synth_interface_and_s3_gateway(): + from aws_cdk import App, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import create_ecs_vpc_endpoints_for_private_subnets + + app = App() + stack = Stack(app, "VpcEndpointTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + private = ec2.SubnetSelection( + subnet_type=ec2.SubnetType.PRIVATE_WITH_EGRESS, + ) + create_ecs_vpc_endpoints_for_private_subnets( + stack, + vpc=vpc, + subnets=private, + include_secrets_and_kms=True, + aws_region="eu-west-2", + ) + template = assertions.Template.from_stack(stack) + # 5 interface endpoints + 1 S3 gateway endpoint + template.resource_count_is("AWS::EC2::VPCEndpoint", 6) + template.has_resource_properties( + "AWS::EC2::VPCEndpoint", + {"VpcEndpointType": "Interface", "PrivateDnsEnabled": True}, + ) + template.has_resource_properties( + "AWS::EC2::VPCEndpoint", + {"VpcEndpointType": "Gateway"}, + ) + + +def test_imported_vpc_from_attributes_requires_cidr_block(): + """Imported VPCs must pass vpc_cidr_block or endpoint/SG helpers fail at synth.""" + from aws_cdk import App, Stack + from aws_cdk import aws_ec2 as ec2 + + app = App() + stack = Stack(app, "ImportedVpcTest") + vpc = ec2.Vpc.from_vpc_attributes( + stack, + "ImportedVpc", + vpc_id="vpc-12345678", + availability_zones=["eu-west-2a", "eu-west-2b"], + vpc_cidr_block="10.0.0.0/16", + ) + assert vpc.vpc_cidr_block == "10.0.0.0/16" + + +def test_create_ecs_vpc_endpoints_skips_existing_service_names(): + from aws_cdk import App, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import create_ecs_vpc_endpoints_for_private_subnets + + app = App() + stack = Stack(app, "SkipEndpointTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + private = ec2.SubnetSelection( + subnet_type=ec2.SubnetType.PRIVATE_WITH_EGRESS, + ) + skip = frozenset( + { + "com.amazonaws.eu-west-2.kms", + "com.amazonaws.eu-west-2.secretsmanager", + } + ) + create_ecs_vpc_endpoints_for_private_subnets( + stack, + vpc=vpc, + subnets=private, + skip_service_names=skip, + include_secrets_and_kms=True, + aws_region="eu-west-2", + ) + template = assertions.Template.from_stack(stack) + template.resource_count_is("AWS::EC2::VPCEndpoint", 4) + + +def test_create_ecs_vpc_endpoints_skips_precheck_shared_vpc_endpoints(): + """Regression: skip list must use AWS ServiceName strings, not CDK service.name tokens.""" + from aws_cdk import App, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import create_ecs_vpc_endpoints_for_private_subnets + + app = App() + stack = Stack(app, "SharedVpcSkipTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + private = ec2.SubnetSelection( + subnet_type=ec2.SubnetType.PRIVATE_WITH_EGRESS, + ) + precheck_skip = frozenset( + { + "com.amazonaws.eu-west-2.kms", + "com.amazonaws.eu-west-2.s3", + "com.amazonaws.eu-west-2.secretsmanager", + } + ) + create_ecs_vpc_endpoints_for_private_subnets( + stack, + vpc=vpc, + subnets=private, + skip_service_names=precheck_skip, + include_secrets_and_kms=True, + aws_region="eu-west-2", + ) + template = assertions.Template.from_stack(stack) + # ecr.api, ecr.dkr, logs only (kms, secretsmanager, s3 skipped) + template.resource_count_is("AWS::EC2::VPCEndpoint", 3) + + +def test_list_existing_vpc_endpoint_service_names(monkeypatch): + from cdk_functions import list_existing_vpc_endpoint_service_names + + class FakePaginator: + def paginate(self, **kwargs): + return [ + { + "VpcEndpoints": [ + { + "ServiceName": "com.amazonaws.eu-west-2.kms", + "State": "available", + }, + { + "ServiceName": "com.amazonaws.eu-west-2.s3", + "State": "available", + }, + { + "ServiceName": "com.amazonaws.eu-west-2.ecr.api", + "State": "deleted", + }, + ] + } + ] + + class FakeEc2: + def get_paginator(self, name): + assert name == "describe_vpc_endpoints" + return FakePaginator() + + monkeypatch.setattr( + "cdk_functions.boto3.client", + lambda service, region_name=None: FakeEc2(), + ) + names = list_existing_vpc_endpoint_service_names("vpc-123", region_name="eu-west-2") + assert names == frozenset( + { + "com.amazonaws.eu-west-2.kms", + "com.amazonaws.eu-west-2.s3", + } + ) + + +def test_resolve_ecs_vpc_endpoint_subnet_selection_express_public(): + from aws_cdk import App, Stack + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import resolve_ecs_vpc_endpoint_subnet_selection + + app = App() + stack = Stack(app, "SubnetSelectTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + public_sel = resolve_ecs_vpc_endpoint_subnet_selection( + use_express_ingress=True, + express_use_public_subnets=True, + public_subnets=vpc.public_subnets, + private_subnets=vpc.private_subnets, + ) + assert public_sel is not None + assert {s.subnet_id for s in public_sel.subnets} == { + s.subnet_id for s in vpc.public_subnets + } + + +def test_resolve_ecs_vpc_endpoint_subnet_selection_express_private(): + from aws_cdk import App, Stack + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import resolve_ecs_vpc_endpoint_subnet_selection + + app = App() + stack = Stack(app, "SubnetSelectPrivateTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + private_sel = resolve_ecs_vpc_endpoint_subnet_selection( + use_express_ingress=True, + express_use_public_subnets=False, + public_subnets=vpc.public_subnets, + private_subnets=vpc.private_subnets, + ) + assert private_sel is not None + assert {s.subnet_id for s in private_sel.subnets} == { + s.subnet_id for s in vpc.private_subnets + } + + +def test_resolve_ecs_vpc_endpoint_subnet_selection_legacy_fargate(): + from aws_cdk import App, Stack + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import resolve_ecs_vpc_endpoint_subnet_selection + + app = App() + stack = Stack(app, "SubnetSelectLegacyTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + private_sel = resolve_ecs_vpc_endpoint_subnet_selection( + use_express_ingress=False, + express_use_public_subnets=True, + public_subnets=vpc.public_subnets, + private_subnets=vpc.private_subnets, + ) + assert private_sel is not None + assert {s.subnet_id for s in private_sel.subnets} == { + s.subnet_id for s in vpc.private_subnets + } + + +def test_resolve_ecs_vpc_endpoint_subnet_selection_public_only_legacy(): + from aws_cdk import App, Stack + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import resolve_ecs_vpc_endpoint_subnet_selection + + app = App() + stack = Stack(app, "SubnetSelectPublicOnlyTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + public_sel = resolve_ecs_vpc_endpoint_subnet_selection( + use_express_ingress=False, + express_use_public_subnets=True, + public_subnets=vpc.public_subnets, + private_subnets=[], + ) + assert public_sel is not None + assert {s.subnet_id for s in public_sel.subnets} == { + s.subnet_id for s in vpc.public_subnets + } + + +def test_create_ecs_vpc_endpoints_on_public_subnets(): + from aws_cdk import App, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import create_ecs_vpc_endpoints_for_private_subnets + + app = App() + stack = Stack(app, "PublicEndpointTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + public = ec2.SubnetSelection(subnet_type=ec2.SubnetType.PUBLIC) + create_ecs_vpc_endpoints_for_private_subnets( + stack, + vpc=vpc, + subnets=public, + include_secrets_and_kms=True, + aws_region="eu-west-2", + ) + template = assertions.Template.from_stack(stack) + template.resource_count_is("AWS::EC2::VPCEndpoint", 6) + public_logical_ids = { + stack.get_logical_id(s.node.default_child) for s in vpc.public_subnets + } + interface_endpoints = [ + r + for r in template.find_resources("AWS::EC2::VPCEndpoint").values() + if r.get("Properties", {}).get("VpcEndpointType") == "Interface" + ] + assert len(interface_endpoints) == 5 + for resource in interface_endpoints: + subnet_ids = resource["Properties"]["SubnetIds"] + assert len(subnet_ids) == len(vpc.public_subnets) + for ref in subnet_ids: + assert ref.get("Ref") in public_logical_ids + + +def test_resolve_ecs_s3_gateway_subnet_selection_all_stack_subnets(): + from aws_cdk import App, Stack + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import resolve_ecs_s3_gateway_subnet_selection + + app = App() + stack = Stack(app, "S3SubnetSelectTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + s3_sel = resolve_ecs_s3_gateway_subnet_selection( + public_subnets=vpc.public_subnets, + private_subnets=vpc.private_subnets, + ) + assert s3_sel is not None + expected_ids = {s.subnet_id for s in vpc.public_subnets + vpc.private_subnets} + assert {s.subnet_id for s in s3_sel.subnets} == expected_ids + + +def test_s3_gateway_endpoint_uses_all_stack_subnets_not_only_task_tier(): + from aws_cdk import App, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import ( + create_ecs_vpc_endpoints_for_private_subnets, + resolve_ecs_s3_gateway_subnet_selection, + ) + + app = App() + stack = Stack(app, "MixedS3GatewayTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + public = ec2.SubnetSelection(subnets=vpc.public_subnets) + all_stack_subnets = resolve_ecs_s3_gateway_subnet_selection( + public_subnets=vpc.public_subnets, + private_subnets=vpc.private_subnets, + ) + create_ecs_vpc_endpoints_for_private_subnets( + stack, + vpc=vpc, + subnets=public, + s3_gateway_subnets=all_stack_subnets, + include_secrets_and_kms=True, + aws_region="eu-west-2", + ) + template = assertions.Template.from_stack(stack) + gateway_endpoints = [ + r + for r in template.find_resources("AWS::EC2::VPCEndpoint").values() + if r.get("Properties", {}).get("VpcEndpointType") == "Gateway" + ] + assert len(gateway_endpoints) == 1 + route_table_ids = gateway_endpoints[0]["Properties"]["RouteTableIds"] + assert len(route_table_ids) == len(vpc.public_subnets) + len(vpc.private_subnets) + interface_endpoints = [ + r + for r in template.find_resources("AWS::EC2::VPCEndpoint").values() + if r.get("Properties", {}).get("VpcEndpointType") == "Interface" + ] + public_logical_ids = { + stack.get_logical_id(s.node.default_child) for s in vpc.public_subnets + } + for resource in interface_endpoints: + for ref in resource["Properties"]["SubnetIds"]: + assert ref.get("Ref") in public_logical_ids + + +def test_list_vpc_associated_cidr_blocks_primary_and_secondary(): + from cdk_functions import list_vpc_associated_cidr_blocks + + vpc = { + "CidrBlock": "10.0.0.0/16", + "CidrBlockAssociationSet": [ + { + "CidrBlock": "10.0.0.0/16", + "CidrBlockState": {"State": "associated"}, + }, + { + "CidrBlock": "10.1.0.0/16", + "CidrBlockState": {"State": "associated"}, + }, + ], + } + assert list_vpc_associated_cidr_blocks(vpc) == [ + "10.0.0.0/16", + "10.1.0.0/16", + ] + + +def test_vpc_endpoint_security_group_ingress_covers_all_vpc_cidrs(): + from aws_cdk import App, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from cdk_functions import add_vpc_endpoint_https_ingress_from_vpc_cidrs + + app = App() + stack = Stack(app, "DualCidrEndpointSgTest") + vpc = ec2.Vpc.from_vpc_attributes( + stack, + "ImportedVpc", + vpc_id="vpc-dual-cidr", + availability_zones=["eu-west-2a"], + vpc_cidr_block="10.0.0.0/16", + ) + endpoint_sg = ec2.SecurityGroup( + stack, + "EndpointSg", + vpc=vpc, + description="test", + ) + add_vpc_endpoint_https_ingress_from_vpc_cidrs( + endpoint_sg, + vpc_cidr_block="10.0.0.0/16", + vpc_cidr_blocks=["10.0.0.0/16", "10.1.0.0/16"], + ) + template = assertions.Template.from_stack(stack) + template.has_resource_properties( + "AWS::EC2::SecurityGroup", + { + "SecurityGroupIngress": assertions.Match.array_with( + [ + assertions.Match.object_like( + {"CidrIp": "10.0.0.0/16", "FromPort": 443, "ToPort": 443} + ), + assertions.Match.object_like( + {"CidrIp": "10.1.0.0/16", "FromPort": 443, "ToPort": 443} + ), + ] + ) + }, + ) diff --git a/cdk/test/test_elbv2_listener_rule_upsert.py b/cdk/test/test_elbv2_listener_rule_upsert.py new file mode 100644 index 0000000000000000000000000000000000000000..00ccf91f6214a2031afe7710ab405e6b939c616d --- /dev/null +++ b/cdk/test/test_elbv2_listener_rule_upsert.py @@ -0,0 +1,86 @@ +"""Unit tests for ALB listener rule upsert Lambda helpers.""" + +import importlib.util +import sys +from pathlib import Path + +_LAMBDA_PATH = ( + Path(__file__).resolve().parents[1] + / "lambda_elbv2_listener_rule_upsert" + / "lambda_function.py" +) +_spec = importlib.util.spec_from_file_location( + "elbv2_listener_rule_upsert", _LAMBDA_PATH +) +_mod = importlib.util.module_from_spec(_spec) +assert _spec and _spec.loader +sys.modules[_spec.name] = _mod +_spec.loader.exec_module(_mod) + + +def test_conditions_match_path_pattern(): + expected = [ + { + "Field": "path-pattern", + "PathPatternConfig": {"Values": ["/agent", "/agent/*"]}, + } + ] + from_api = [ + { + "Field": "path-pattern", + "PathPatternConfig": {"Values": ["/agent", "/agent/*"]}, + "Values": ["/agent", "/agent/*"], + } + ] + assert _mod._conditions_match(from_api, expected) + + +def test_conditions_do_not_match_different_path(): + expected = [ + { + "Field": "path-pattern", + "PathPatternConfig": {"Values": ["/agent", "/agent/*"]}, + } + ] + other = [ + { + "Field": "path-pattern", + "PathPatternConfig": {"Values": ["/other", "/other/*"]}, + } + ] + assert not _mod._conditions_match(other, expected) + + +def test_normalize_listener_rule_payload_coerces_cfn_string_types(): + raw_actions = [ + { + "Type": "authenticate-cognito", + "Order": "1", + "AuthenticateCognitoConfig": { + "UserPoolArn": "arn:aws:cognito-idp:eu-west-2:1:userpool/pool", + "UserPoolClientId": "client", + "UserPoolDomain": "demo", + "SessionTimeout": "28800", + }, + }, + { + "Type": "forward", + "Order": "2", + "ForwardConfig": { + "TargetGroups": [ + {"TargetGroupArn": "arn:aws:elasticloadbalancing:1:tg"} + ], + "TargetGroupStickinessConfig": { + "Enabled": "true", + "DurationSeconds": "28800", + }, + }, + }, + ] + _, actions = _mod._normalize_listener_rule_payload([], raw_actions) + assert actions[0]["Order"] == 1 + assert actions[0]["AuthenticateCognitoConfig"]["SessionTimeout"] == 28800 + assert actions[1]["Order"] == 2 + stickiness = actions[1]["ForwardConfig"]["TargetGroupStickinessConfig"] + assert stickiness["Enabled"] is True + assert stickiness["DurationSeconds"] == 28800 diff --git a/cdk/test/test_ensure_codebuild_public_github.py b/cdk/test/test_ensure_codebuild_public_github.py new file mode 100644 index 0000000000000000000000000000000000000000..d7563361618647c067ea1ffca0d9f3d792a524c4 --- /dev/null +++ b/cdk/test/test_ensure_codebuild_public_github.py @@ -0,0 +1,81 @@ +"""Tests for runtime CodeBuild public GitHub source fixup.""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + +from cdk_functions import ensure_codebuild_public_github_source + + +def test_ensure_codebuild_updates_oauth_project(monkeypatch): + client = MagicMock() + client.batch_get_projects.return_value = { + "projects": [ + { + "name": "MyProject", + "source": { + "type": "GITHUB", + "location": "https://github.com/old/repo.git", + "auth": {"type": "CODECONNECTIONS", "resource": "arn:connection"}, + "reportBuildStatus": True, + "buildspec": "version: 0.2\nphases: {}\n", + }, + "sourceVersion": "main", + "triggers": {"webhook": True}, + } + ] + } + monkeypatch.setattr("cdk_functions.boto3.client", lambda *args, **kwargs: client) + + updated = ensure_codebuild_public_github_source( + "MyProject", + "seanpedrick-case", + "llm_topic_modeller", + "main", + aws_region="eu-west-2", + ) + + assert updated is True + client.update_project.assert_called_once() + kwargs = client.update_project.call_args.kwargs + assert kwargs["name"] == "MyProject" + assert kwargs["source"]["type"] == "GITHUB" + assert ( + kwargs["source"]["location"] + == "https://github.com/seanpedrick-case/llm_topic_modeller.git" + ) + assert "auth" not in kwargs["source"] + assert kwargs["sourceVersion"] == "main" + assert kwargs["triggers"] == {"webhook": False} + + +def test_ensure_codebuild_skips_when_already_public(monkeypatch): + client = MagicMock() + client.batch_get_projects.return_value = { + "projects": [ + { + "name": "MyProject", + "source": { + "type": "GITHUB", + "location": "https://github.com/seanpedrick-case/llm_topic_modeller.git", + "reportBuildStatus": False, + }, + "sourceVersion": "main", + } + ] + } + monkeypatch.setattr("cdk_functions.boto3.client", lambda *args, **kwargs: client) + + updated = ensure_codebuild_public_github_source( + "MyProject", + "seanpedrick-case", + "llm_topic_modeller", + "main", + aws_region="eu-west-2", + ) + + assert updated is False + client.update_project.assert_not_called() diff --git a/cdk/test/test_execution_role_kms_policy.py b/cdk/test/test_execution_role_kms_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..284fccbabd4b2862929a6309adcdea3cafb88726 --- /dev/null +++ b/cdk/test/test_execution_role_kms_policy.py @@ -0,0 +1,100 @@ +"""ECS task vs execution role KMS inline policy helpers.""" + +import sys +from pathlib import Path + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + + +def test_task_and_execution_role_kms_policies_use_different_key_arns(): + from cdk_functions import ( + build_ecs_execution_role_kms_policy, + build_ecs_task_role_kms_policy, + ) + + s3_key = "arn:aws:kms:eu-west-2:123456789012:key/s3-shared-key" + secret_key = "arn:aws:kms:eu-west-2:123456789012:key/secret-key" + + task_policy = build_ecs_task_role_kms_policy(shared_kms_key_arn=s3_key) + exec_policy = build_ecs_execution_role_kms_policy(secret_kms_key_arn=secret_key) + + task_kms = next( + s for s in task_policy["Statement"] if s.get("Sid") == "KMSS3Access" + ) + secret_kms = next( + s for s in exec_policy["Statement"] if s.get("Sid") == "KMSSecretDecrypt" + ) + + assert task_kms["Resource"] == s3_key + assert secret_kms["Resource"] == secret_key + assert "kms:GenerateDataKey" in task_kms["Action"] + assert "kms:DescribeKey" in task_kms["Action"] + assert secret_kms["Action"] == ["kms:Decrypt"] + + +def test_task_role_inline_policy_scopes_s3_to_output_and_log_buckets(): + from cdk_functions import build_ecs_task_role_inline_policy + + policy = build_ecs_task_role_inline_policy( + output_bucket_name="lambeth-aws-sharedservices-prod-summarisation-s3-output", + log_config_bucket_name="lambeth-aws-sharedservices-prod-summarisation-s3-logs", + shared_kms_key_arn="arn:aws:kms:eu-west-2:123456789012:key/s3-shared-key", + ) + s3_statements = [ + s for s in policy["Statement"] if s.get("Sid", "").startswith("S3") + ] + assert len(s3_statements) == 2 + output_stmt = next(s for s in s3_statements if s["Sid"] == "S3Output") + log_stmt = next(s for s in s3_statements if s["Sid"] == "S3LogConfig") + assert output_stmt["Action"] == [ + "s3:GetObject*", + "s3:GetBucket*", + "s3:PutObject", + "s3:DeleteObject", + "s3:List*", + ] + assert output_stmt["Resource"] == [ + "arn:aws:s3:::lambeth-aws-sharedservices-prod-summarisation-s3-output", + "arn:aws:s3:::lambeth-aws-sharedservices-prod-summarisation-s3-output/*", + ] + assert log_stmt["Resource"] == [ + "arn:aws:s3:::lambeth-aws-sharedservices-prod-summarisation-s3-logs", + "arn:aws:s3:::lambeth-aws-sharedservices-prod-summarisation-s3-logs/*", + ] + kms_stmt = next(s for s in policy["Statement"] if s.get("Sid") == "KMSS3Access") + assert ( + kms_stmt["Resource"] == "arn:aws:kms:eu-west-2:123456789012:key/s3-shared-key" + ) + + +def test_execution_role_policy_without_custom_s3_key_uses_secret_kms_only(): + from cdk_functions import build_ecs_execution_role_kms_policy + + secret_key = "arn:aws:kms:eu-west-2:123456789012:key/aws/secretsmanager" + policy = build_ecs_execution_role_kms_policy(secret_kms_key_arn=secret_key) + kms_statements = [ + s for s in policy["Statement"] if s.get("Action") == ["kms:Decrypt"] + ] + assert len(kms_statements) == 1 + assert kms_statements[0]["Resource"] == secret_key + + +def test_get_secret_kms_key_arn_from_describe_secret(monkeypatch): + from cdk_functions import get_secret_kms_key_arn + + class FakeSecretsManager: + def describe_secret(self, SecretId): + assert SecretId == "my-secret" + return { + "KmsKeyId": "arn:aws:kms:eu-west-2:123456789012:key/abc-123", + } + + monkeypatch.setattr( + "cdk_functions.boto3.client", + lambda service, region_name=None: FakeSecretsManager(), + ) + assert ( + get_secret_kms_key_arn("my-secret", region_name="eu-west-2") + == "arn:aws:kms:eu-west-2:123456789012:key/abc-123" + ) diff --git a/cdk/test/test_express_app_config_env.py b/cdk/test/test_express_app_config_env.py new file mode 100644 index 0000000000000000000000000000000000000000..7cbb3beab07a89f3b022de736ab0ec26fb08dcb1 --- /dev/null +++ b/cdk/test/test_express_app_config_env.py @@ -0,0 +1,25 @@ +"""Tests for loading config/app_config.env into Express container environment.""" + +import sys +from pathlib import Path + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + +from cdk_functions import load_app_config_env_for_express + + +def test_load_app_config_env_for_express_reads_config_env(): + config_path = CDK_DIR / "config" / "app_config.env" + env_vars = load_app_config_env_for_express(str(config_path)) + names = {p.name for p in env_vars} + assert "RUN_AWS_FUNCTIONS" in names + assert "S3_LOG_BUCKET" in names + assert "AWS_CLIENT_ID" not in names + + +def test_load_app_config_env_missing_file_returns_empty(): + assert ( + load_app_config_env_for_express(str(CDK_DIR / "config" / "nonexistent.env")) + == [] + ) diff --git a/cdk/test/test_express_pi.py b/cdk/test/test_express_pi.py new file mode 100644 index 0000000000000000000000000000000000000000..d8938c220527ca357002b53d292c57f6741a5695 --- /dev/null +++ b/cdk/test/test_express_pi.py @@ -0,0 +1,566 @@ +"""Tests for Pi on ECS Express Mode helpers and config rules.""" + +import sys +from pathlib import Path + +import pytest + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + + +def test_pi_express_mutual_exclusion_with_legacy_pi(): + legacy = "True" + express = "True" + with pytest.raises(ValueError, match="at most one Pi deployment"): + if legacy == "True" and express == "True": + raise ValueError( + "Enable at most one Pi deployment mode: ENABLE_PI_AGENT_ECS_SERVICE (legacy Fargate) " + "or ENABLE_PI_AGENT_EXPRESS_SERVICE (Express), not both." + ) + + +def test_pi_express_requires_express_mode(): + express_pi = "True" + use_express = "False" + with pytest.raises(ValueError, match="ENABLE_PI_AGENT_EXPRESS_SERVICE"): + if express_pi == "True" and use_express != "True": + raise ValueError( + "ENABLE_PI_AGENT_EXPRESS_SERVICE=True requires USE_ECS_EXPRESS_MODE=True " + "(no ACM_SSL_CERTIFICATE_ARN)." + ) + + +def test_build_pi_express_container_environment(): + from cdk_functions import build_pi_express_container_environment + + env = build_pi_express_container_environment( + service_connect_discovery_name="summarisation", + main_app_port=7860, + pi_gradio_port=7862, + ) + assert env["DOC_SUMMARISATION_GRADIO_URL"] == "http://summarisation:7860" + assert env["PI_WORKSPACE_DIR"] == "/tmp/pi-workspace" + assert env["PI_UPLOAD_ROOT"] == "/tmp/gradio" + assert env["PI_DEPLOYMENT_PROFILE"] == "aws-ecs" + assert env["COGNITO_AUTH"] == "True" + assert env["RUN_FASTAPI"] == "True" + + +def test_build_express_pi_primary_container_includes_cognito_secrets(): + from aws_cdk import App, Stack + from aws_cdk import aws_secretsmanager as sm + from cdk_functions import build_express_pi_primary_container + + app = App() + stack = Stack(app, "PiSecretTest") + secret = sm.Secret(stack, "CognitoSecret", secret_name="demo-cognito-secret") + + container = build_express_pi_primary_container( + image_uri="123456789012.dkr.ecr.eu-west-2.amazonaws.com/pi:latest", + container_port=7862, + log_group_name="/ecs/pi-logs", + aws_region="eu-west-2", + secret=secret, + cognito_auth=True, + ) + assert container.secrets is not None + secret_names = {item.name for item in container.secrets} + assert secret_names == {"AWS_USER_POOL_ID", "AWS_CLIENT_ID", "AWS_CLIENT_SECRET"} + + no_auth = build_express_pi_primary_container( + image_uri="123456789012.dkr.ecr.eu-west-2.amazonaws.com/pi:latest", + container_port=7862, + log_group_name="/ecs/pi-logs", + aws_region="eu-west-2", + secret=secret, + cognito_auth=False, + ) + assert no_auth.secrets is None + + +def test_format_express_pi_public_url(): + from cdk_functions import format_express_pi_public_url + + assert ( + format_express_pi_public_url("https://pi.example.ecs.eu-west-2.on.aws") + == "https://pi.example.ecs.eu-west-2.on.aws/" + ) + assert format_express_pi_public_url("") == "" + + +def test_express_service_connect_configuration_server_and_client(): + from cdk_functions import _express_service_connect_configuration + + server = _express_service_connect_configuration( + namespace="demo-ns", + port_name="port-7860", + discovery_name="summarisation", + port=7860, + ) + assert server["enabled"] is True + assert server["namespace"] == "demo-ns" + assert server["services"][0]["portName"] == "port-7860" + assert server["services"][0]["discoveryName"] == "summarisation" + + client = _express_service_connect_configuration(namespace="demo-ns") + assert "services" not in client + + +def test_apply_service_connect_custom_resource_synth(): + from aws_cdk import App, Environment, Stack, assertions + from aws_cdk import aws_ecs as ecs + from cdk_functions import apply_service_connect_to_express_service + + app = App() + stack = Stack( + app, + "ScExpressTest", + env=Environment(account="123456789012", region="eu-west-2"), + ) + express = ecs.CfnExpressGatewayService( + stack, + "MainExpress", + service_name="main-express", + cluster="test-cluster", + execution_role_arn="arn:aws:iam::123456789012:role/exec", + infrastructure_role_arn="arn:aws:iam::123456789012:role/infra", + primary_container=ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image="123456789012.dkr.ecr.eu-west-2.amazonaws.com/app:latest", + container_port=7860, + ), + ) + apply_service_connect_to_express_service( + stack, + "MainSc", + cluster_name="test-cluster", + service_name="main-express", + namespace="test-ns", + express_service=express, + port_name="port-7860", + discovery_name="summarisation", + port=7860, + ) + template = assertions.Template.from_stack(stack) + template.resource_count_is("Custom::AWS", 1) + template.has_resource_properties( + "Custom::AWS", + { + "Create": assertions.Match.string_like_regexp( + r'"portName":"port-7860".*"discoveryName":"summarisation"' + ), + }, + ) + + +def test_express_alb_ingress_uses_security_group_id_not_arn(): + from aws_cdk import App, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from aws_cdk import aws_ecs as ecs + from cdk_functions import allow_express_load_balancer_to_ecs_security_group + + app = App() + stack = Stack(app, "ExpressAlbIngressTest") + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + task_sg = ec2.SecurityGroup(stack, "TaskSg", vpc=vpc) + express = ecs.CfnExpressGatewayService( + stack, + "Express", + service_name="main-express", + cluster="test-cluster", + execution_role_arn="arn:aws:iam::123456789012:role/exec", + infrastructure_role_arn="arn:aws:iam::123456789012:role/infra", + primary_container=ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image="123456789012.dkr.ecr.eu-west-2.amazonaws.com/app:latest", + container_port=7860, + ), + ) + allow_express_load_balancer_to_ecs_security_group( + stack, + "ExpressAlbToTask", + express_service=express, + ecs_security_group=task_sg, + container_port=7860, + ) + template = assertions.Template.from_stack(stack) + template.has_resource_properties( + "AWS::EC2::SecurityGroupIngress", + { + "SourceSecurityGroupId": { + "Fn::Select": [ + 1, + { + "Fn::Split": [ + "security-group/", + { + "Fn::Select": [ + 0, + { + "Fn::GetAtt": [ + "Express", + "ECSManagedResourceArns.IngressPath.LoadBalancerSecurityGroups", + ] + }, + ] + }, + ] + }, + ] + } + }, + ) + + +def test_express_gateway_service_defaults_to_idle_scaling_target(): + from aws_cdk import App, Stack, assertions + from aws_cdk import aws_ecs as ecs + from cdk_functions import create_express_gateway_service + + app = App() + stack = Stack(app, "ExpressIdleScalingTest") + create_express_gateway_service( + stack, + "Express", + service_name="main-express", + cluster_name="test-cluster", + execution_role_arn="arn:aws:iam::123456789012:role/exec", + infrastructure_role_arn="arn:aws:iam::123456789012:role/infra", + task_role_arn="arn:aws:iam::123456789012:role/task", + cpu="1024", + memory="4096", + health_check_path="/", + primary_container=ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image="123456789012.dkr.ecr.eu-west-2.amazonaws.com/app:latest", + container_port=7860, + ), + subnet_ids=["subnet-abc"], + security_group_ids=["sg-main"], + ) + template = assertions.Template.from_stack(stack) + template.has_resource_properties( + "AWS::ECS::ExpressGatewayService", + { + "ScalingTarget": { + "MinTaskCount": 0, + "MaxTaskCount": 1, + "AutoScalingMetric": "AVERAGE_CPU", + "AutoScalingTargetValue": 60, + } + }, + ) + + +def test_express_infrastructure_role_uses_service_role_managed_policy(): + from aws_cdk import App, Stack, assertions + from cdk_functions import create_ecs_express_infrastructure_role + + app = App() + stack = Stack(app, "ExpressInfraRoleTest") + create_ecs_express_infrastructure_role( + stack, "ExpressInfrastructureRole", "test-express-infra" + ) + template = assertions.Template.from_stack(stack) + template.has_resource_properties( + "AWS::IAM::Role", + { + "ManagedPolicyArns": assertions.Match.array_with( + [ + { + "Fn::Join": [ + "", + [ + "arn:", + {"Ref": "AWS::Partition"}, + ":iam::aws:policy/service-role/AmazonECSInfrastructureRoleforExpressGatewayServices", + ], + ] + } + ] + ) + }, + ) + + +def test_express_listener_helpers_synth_without_reference_error(): + """Fn.select on Express list attrs must use typed attr_* list properties.""" + from aws_cdk import App, Environment, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from aws_cdk import aws_ecs as ecs + from cdk_functions import ( + allow_express_load_balancer_to_ecs_security_group, + configure_express_listener_cognito_and_cloudfront, + configure_express_pi_listener_rules, + create_express_gateway_service, + ) + + app = App() + stack = Stack( + app, + "ExpressListenerHelpers", + env=Environment(account="123456789012", region="eu-west-2"), + ) + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + sg = ec2.SecurityGroup(stack, "TaskSg", vpc=vpc) + main = create_express_gateway_service( + stack, + "Main", + service_name="main-svc", + cluster_name="cl", + execution_role_arn="arn:aws:iam::123456789012:role/exec", + infrastructure_role_arn="arn:aws:iam::123456789012:role/infra", + task_role_arn="arn:aws:iam::123456789012:role/task", + cpu="1024", + memory="2048", + health_check_path="/", + primary_container=ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image="123456789012.dkr.ecr.eu-west-2.amazonaws.com/app:latest", + container_port=7860, + ), + subnet_ids=["subnet-abc"], + security_group_ids=["sg-main"], + ) + pi = create_express_gateway_service( + stack, + "Pi", + service_name="pi-svc", + cluster_name="cl", + execution_role_arn="arn:aws:iam::123456789012:role/exec", + infrastructure_role_arn="arn:aws:iam::123456789012:role/infra", + task_role_arn="arn:aws:iam::123456789012:role/task", + cpu="1024", + memory="2048", + health_check_path="/pi/", + primary_container=ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image="123456789012.dkr.ecr.eu-west-2.amazonaws.com/pi:latest", + container_port=7862, + ), + subnet_ids=["subnet-abc"], + security_group_ids=["sg-pi"], + ) + allow_express_load_balancer_to_ecs_security_group( + stack, + "MainLbToTask", + express_service=main, + ecs_security_group=sg, + container_port=7860, + ) + configure_express_listener_cognito_and_cloudfront( + stack, + "MainListener", + express_service=main, + user_pool_arn="arn:aws:cognito-idp:eu-west-2:123456789012:userpool/pool", + user_pool_client_id="client", + user_pool_domain_prefix="demo-auth", + use_cloudfront=False, + cloudfront_host_header="", + ) + configure_express_pi_listener_rules( + stack, + "PiRules", + express_main_service=main, + express_pi_service=pi, + routing_mode="path", + path_prefix="/pi", + pi_host_header="", + rule_priority=3, + user_pool_arn="arn:aws:cognito-idp:eu-west-2:123456789012:userpool/pool", + user_pool_client_id="client", + user_pool_domain_prefix="demo-auth", + ) + app.synth() + template = assertions.Template.from_stack(stack) + template.resource_count_is("AWS::ECS::ExpressGatewayService", 2) + # modifyListener (Custom::AWS) + Pi path rule (Custom::Elbv2ListenerRuleUpsert) + template.resource_count_is("Custom::AWS", 1) + template.resource_count_is("Custom::Elbv2ListenerRuleUpsert", 1) + # authenticate-cognito on ALB requires DescribeUserPoolClient on the CR Lambda role + cr_policy_actions: list[str] = [] + for props in template.find_resources("AWS::IAM::Policy").values(): + for stmt in ( + props.get("Properties", {}).get("PolicyDocument", {}).get("Statement", []) + ): + action = stmt.get("Action", []) + if isinstance(action, str): + cr_policy_actions.append(action) + else: + cr_policy_actions.extend(action) + assert "cognito-idp:DescribeUserPoolClient" in cr_policy_actions + assert "elasticloadbalancing:ModifyListener" in cr_policy_actions + + +def test_express_cloudfront_does_not_add_cognito_bypass_host_rule_by_default(): + """With ALB Cognito, CloudFront must not add a forward-only host-header rule.""" + from aws_cdk import App, Environment, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from aws_cdk import aws_ecs as ecs + from cdk_functions import ( + configure_express_listener_cognito_and_cloudfront, + create_express_gateway_service, + ) + + app = App() + stack = Stack( + app, + "ExpressCloudFrontCognitoTest", + env=Environment(account="123456789012", region="eu-west-2"), + ) + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + main = create_express_gateway_service( + stack, + "Main", + service_name="main-svc", + cluster_name="cl", + execution_role_arn="arn:aws:iam::123456789012:role/exec", + infrastructure_role_arn="arn:aws:iam::123456789012:role/infra", + task_role_arn="arn:aws:iam::123456789012:role/task", + cpu="1024", + memory="2048", + health_check_path="/", + primary_container=ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image="123456789012.dkr.ecr.eu-west-2.amazonaws.com/app:latest", + container_port=7860, + ), + subnet_ids=[vpc.public_subnets[0].subnet_id], + security_group_ids=["sg-main"], + ) + configure_express_listener_cognito_and_cloudfront( + stack, + "MainListener", + express_service=main, + user_pool_arn="arn:aws:cognito-idp:eu-west-2:123456789012:userpool/pool", + user_pool_client_id="client", + user_pool_domain_prefix="demo-auth", + use_cloudfront=True, + cloudfront_host_header="app.example.com", + ) + template = assertions.Template.from_stack(stack) + template.resource_count_is("Custom::Elbv2ListenerRuleUpsert", 0) + + +def test_express_cloudfront_can_opt_in_to_origin_bypass_without_cognito(): + from aws_cdk import App, Environment, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from aws_cdk import aws_ecs as ecs + from cdk_functions import ( + configure_express_listener_cognito_and_cloudfront, + create_express_gateway_service, + ) + + app = App() + stack = Stack( + app, + "ExpressCloudFrontBypassTest", + env=Environment(account="123456789012", region="eu-west-2"), + ) + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + main = create_express_gateway_service( + stack, + "Main", + service_name="main-svc", + cluster_name="cl", + execution_role_arn="arn:aws:iam::123456789012:role/exec", + infrastructure_role_arn="arn:aws:iam::123456789012:role/infra", + task_role_arn="arn:aws:iam::123456789012:role/task", + cpu="1024", + memory="2048", + health_check_path="/", + primary_container=ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image="123456789012.dkr.ecr.eu-west-2.amazonaws.com/app:latest", + container_port=7860, + ), + subnet_ids=[vpc.public_subnets[0].subnet_id], + security_group_ids=["sg-main"], + ) + configure_express_listener_cognito_and_cloudfront( + stack, + "MainListener", + express_service=main, + user_pool_arn="arn:aws:cognito-idp:eu-west-2:123456789012:userpool/pool", + user_pool_client_id="client", + user_pool_domain_prefix="demo-auth", + use_cloudfront=True, + cloudfront_host_header="app.example.com", + allow_cloudfront_origin_without_cognito=True, + ) + template = assertions.Template.from_stack(stack) + template.resource_count_is("Custom::Elbv2ListenerRuleUpsert", 1) + + +def test_dual_express_gateway_services_synth(): + """Two ExpressGatewayService resources when wiring main + Pi helpers.""" + from aws_cdk import App, Environment, Stack, assertions + from aws_cdk import aws_ecs as ecs + from cdk_functions import ( + apply_service_connect_to_express_service, + build_express_pi_primary_container, + create_express_gateway_service, + ) + + app = App() + stack = Stack( + app, + "DualExpressTest", + env=Environment(account="123456789012", region="eu-west-2"), + ) + main = create_express_gateway_service( + stack, + "Main", + service_name="main-svc", + cluster_name="cl", + execution_role_arn="arn:aws:iam::123456789012:role/exec", + infrastructure_role_arn="arn:aws:iam::123456789012:role/infra", + task_role_arn="arn:aws:iam::123456789012:role/task", + cpu="1024", + memory="2048", + health_check_path="/", + primary_container=ecs.CfnExpressGatewayService.ExpressGatewayContainerProperty( + image="123456789012.dkr.ecr.eu-west-2.amazonaws.com/app:latest", + container_port=7860, + ), + subnet_ids=["subnet-abc"], + security_group_ids=["sg-main"], + ) + pi_container = build_express_pi_primary_container( + image_uri="123456789012.dkr.ecr.eu-west-2.amazonaws.com/pi:latest", + container_port=7862, + log_group_name="/ecs/pi-logs", + aws_region="eu-west-2", + environment={"PI_WORKSPACE_DIR": "/tmp/pi-workspace"}, + ) + pi = create_express_gateway_service( + stack, + "Pi", + service_name="pi-svc", + cluster_name="cl", + execution_role_arn="arn:aws:iam::123456789012:role/exec", + infrastructure_role_arn="arn:aws:iam::123456789012:role/infra", + task_role_arn="arn:aws:iam::123456789012:role/task", + cpu="1024", + memory="2048", + health_check_path="/", + primary_container=pi_container, + subnet_ids=["subnet-abc"], + security_group_ids=["sg-pi"], + ) + apply_service_connect_to_express_service( + stack, + "MainSc", + cluster_name="cl", + service_name="main-svc", + namespace="ns", + express_service=main, + port_name="port-7860", + discovery_name="summarisation", + port=7860, + ) + apply_service_connect_to_express_service( + stack, + "PiSc", + cluster_name="cl", + service_name="pi-svc", + namespace="ns", + express_service=pi, + ) + template = assertions.Template.from_stack(stack) + template.resource_count_is("AWS::ECS::ExpressGatewayService", 2) + template.resource_count_is("Custom::AWS", 2) diff --git a/cdk/test/test_globally_unique_names.py b/cdk/test/test_globally_unique_names.py new file mode 100644 index 0000000000000000000000000000000000000000..b9e1a4f0557a9e377d201f812fc1d895089b828e --- /dev/null +++ b/cdk/test/test_globally_unique_names.py @@ -0,0 +1,91 @@ +"""Tests for globally unique AWS name checks (S3, Cognito).""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock + +from botocore.exceptions import ClientError + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + +from cdk_functions import ( # noqa: E402 + resolve_cognito_domain_prefix_availability, + resolve_s3_bucket_availability, +) + + +def _client_error(code: str) -> ClientError: + return ClientError({"Error": {"Code": code, "Message": code}}, "HeadBucket") + + +def test_resolve_s3_bucket_available_on_404(): + s3 = MagicMock() + s3.head_bucket.side_effect = _client_error("404") + status, name = resolve_s3_bucket_availability("free-bucket", s3_client=s3) + assert status == "available" + assert name == "free-bucket" + + +def test_resolve_s3_bucket_owned_on_success(): + s3 = MagicMock() + s3.head_bucket.return_value = {} + status, name = resolve_s3_bucket_availability("my-bucket", s3_client=s3) + assert status == "owned" + assert name == "my-bucket" + + +def test_resolve_s3_bucket_globally_taken_on_403_not_listed(): + s3 = MagicMock() + s3.head_bucket.side_effect = _client_error("403") + s3.list_buckets.return_value = {"Buckets": [{"Name": "other-bucket"}]} + status, name = resolve_s3_bucket_availability( + "demo-summarisation-s3-logs", s3_client=s3 + ) + assert status == "globally_taken" + assert name == "demo-summarisation-s3-logs" + + +def test_resolve_s3_bucket_owned_on_403_listed_in_account(): + s3 = MagicMock() + s3.head_bucket.side_effect = _client_error("403") + s3.list_buckets.return_value = {"Buckets": [{"Name": "my-bucket"}]} + status, _ = resolve_s3_bucket_availability("my-bucket", s3_client=s3) + assert status == "owned" + + +def test_resolve_cognito_domain_taken_when_owned_elsewhere(): + cognito = MagicMock() + cognito.describe_user_pool_domain.side_effect = _client_error( + "ResourceNotFoundException" + ) + assert ( + resolve_cognito_domain_prefix_availability( + "demo-summarisation", region_name="eu-west-2", cognito_client=cognito + ) + == "taken" + ) + + +def test_resolve_cognito_domain_available_when_empty_description(): + cognito = MagicMock() + cognito.describe_user_pool_domain.return_value = {"DomainDescription": {}} + assert ( + resolve_cognito_domain_prefix_availability( + "my-prefix", region_name="eu-west-2", cognito_client=cognito + ) + == "available" + ) + + +def test_resolve_cognito_domain_taken_when_user_pool_present(): + cognito = MagicMock() + cognito.describe_user_pool_domain.return_value = { + "DomainDescription": {"UserPoolId": "eu-west-2_abc"} + } + assert ( + resolve_cognito_domain_prefix_availability( + "taken-prefix", region_name="eu-west-2", cognito_client=cognito + ) + == "taken" + ) diff --git a/cdk/test/test_headless_app_defaults_env.py b/cdk/test/test_headless_app_defaults_env.py new file mode 100644 index 0000000000000000000000000000000000000000..986db79437b510659c07d8e1f97a404aeda8d1fa --- /dev/null +++ b/cdk/test/test_headless_app_defaults_env.py @@ -0,0 +1,22 @@ +"""Headless batch app_defaults.env rendering.""" + +import sys +from pathlib import Path + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + +from cdk_functions import build_headless_app_defaults_env_content + + +def test_build_headless_app_defaults_env_sets_outputs_bucket(): + seed_dir = str(CDK_DIR / "config" / "headless_s3_seed") + content = build_headless_app_defaults_env_content( + seed_dir, + s3_outputs_bucket_name="lambeth-prod-summarisation-s3-output", + ) + + assert "S3_OUTPUTS_BUCKET=lambeth-prod-summarisation-s3-output" in content + assert "RUN_DIRECT_MODE=1" in content + assert "SAVE_OUTPUTS_TO_S3=True" in content + assert content.count("S3_OUTPUTS_BUCKET=") == 1 diff --git a/cdk/test/test_headless_output_notifications.py b/cdk/test/test_headless_output_notifications.py new file mode 100644 index 0000000000000000000000000000000000000000..d35ff920e2e1eccaff976d3dfdc80ccebf56ceca --- /dev/null +++ b/cdk/test/test_headless_output_notifications.py @@ -0,0 +1,81 @@ +"""Headless S3 output notification resources and installer validation.""" + +import sys +from pathlib import Path + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + +from cdk_functions import sanitize_headless_metric_filter_id + + +def test_sanitize_headless_metric_filter_id(): + assert sanitize_headless_metric_filter_id("Prod-Summarisation-s3-output-put") == ( + "Prod-Summarisation-s3-output-put" + ) + assert sanitize_headless_metric_filter_id("bad id!!") == "bad-id" + + +def test_validate_notify_email(): + from cdk_install import validate_notify_email + + assert validate_notify_email("user@example.com") is None + assert validate_notify_email("") is not None + assert validate_notify_email("not-an-email") is not None + + +def test_headless_output_notifications_synth(): + from aws_cdk import App, Environment, Stack + from aws_cdk import aws_s3 as s3 + from aws_cdk.assertions import Match, Template + from cdk_functions import create_headless_output_notifications + + app = App() + stack = Stack( + app, + "HeadlessNotifyTest", + env=Environment(account="123456789012", region="eu-west-2"), + ) + bucket = s3.Bucket(stack, "OutputBucket", bucket_name="test-output-bucket") + + create_headless_output_notifications( + stack, + "Notify", + output_bucket=bucket, + output_prefix="output/", + notify_email="analyst@example.com", + iam_user_name="test-s3-output-reader", + metric_filter_id="test-s3-output-put", + sns_topic_name="test-llm-topic-s3-save-sns", + alarm_name="test-cloudwatch-alarm-new-output-s3", + ) + + template = Template.from_stack(stack) + template.resource_count_is("AWS::SNS::Topic", 1) + template.resource_count_is("AWS::CloudWatch::Alarm", 1) + template.resource_count_is("AWS::IAM::User", 1) + template.has_resource_properties( + "AWS::S3::BucketPolicy", + { + "PolicyDocument": { + "Statement": Match.array_with( + [ + Match.object_like( + { + "Effect": "Allow", + "Action": Match.array_with(["s3:GetObject"]), + } + ) + ] + ) + } + }, + ) + template.has_resource_properties( + "AWS::CloudWatch::Alarm", + { + "MetricName": "PutRequests", + "Namespace": "AWS/S3", + "ComparisonOperator": "GreaterThanThreshold", + }, + ) diff --git a/cdk/test/test_iam_policy_files.py b/cdk/test/test_iam_policy_files.py new file mode 100644 index 0000000000000000000000000000000000000000..5446f6f01c999180eafc0359c4ec541f276c20f8 --- /dev/null +++ b/cdk/test/test_iam_policy_files.py @@ -0,0 +1,92 @@ +"""Tests for custom IAM policy file lists on ECS task / execution roles.""" + +import json +import sys +from pathlib import Path + +from aws_cdk import App, Stack +from aws_cdk import aws_iam as iam + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + +from cdk_config import ( # noqa: E402 + AWS_MANAGED_TASK_ROLES_LIST, + ECS_EXECUTION_ROLE_MANAGED_POLICIES, + POLICY_FILE_LOCATIONS, + parse_comma_separated_list, +) +from cdk_functions import ( # noqa: E402 + add_custom_policies, + resolve_policy_file_paths, +) + + +def test_parse_comma_separated_list_json_and_plain(): + assert parse_comma_separated_list('["a.json", "b.json"]') == ["a.json", "b.json"] + assert parse_comma_separated_list("textract_policy.json") == [ + "textract_policy.json" + ] + assert parse_comma_separated_list("a.json, b.json") == ["a.json", "b.json"] + assert parse_comma_separated_list("") == [] + + +def test_policy_file_locations_default_is_empty_list(): + assert isinstance(POLICY_FILE_LOCATIONS, list) + assert POLICY_FILE_LOCATIONS == [] + + +def test_task_role_managed_policies_exclude_s3_full_access(): + assert "AmazonS3FullAccess" not in AWS_MANAGED_TASK_ROLES_LIST + + +def test_execution_role_managed_policies_default_minimal(): + assert "service-role/AmazonECSTaskExecutionRolePolicy" in ( + ECS_EXECUTION_ROLE_MANAGED_POLICIES + ) + assert "AmazonS3FullAccess" not in ECS_EXECUTION_ROLE_MANAGED_POLICIES + + +def test_resolve_policy_file_paths_relative_to_cdk_folder(tmp_path): + policy = tmp_path / "custom.json" + policy.write_text("{}", encoding="utf-8") + resolved = resolve_policy_file_paths(["custom.json"], cdk_folder=str(tmp_path)) + assert resolved == [str(policy.resolve())] + + +def test_add_custom_policies_loads_json_statements(tmp_path): + policy_path = tmp_path / "extra.json" + policy_path.write_text( + json.dumps( + { + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "TestAllow", + "Effect": "Allow", + "Action": "s3:ListBucket", + "Resource": "*", + } + ], + } + ), + encoding="utf-8", + ) + + app = App() + stack = Stack(app, "TestStack") + role = iam.Role( + stack, + "TaskRole", + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + add_custom_policies(stack, role, policy_file_locations=[str(policy_path)]) + + template = app.synth().get_stack_by_name("TestStack").template + policies = template["Resources"] + inline = [ + r + for r in policies.values() + if r.get("Type") == "AWS::IAM::Policy" and "s3:ListBucket" in json.dumps(r) + ] + assert inline, "expected inline policy statement from JSON file" diff --git a/cdk/test/test_pi_agent_ecs_synth.py b/cdk/test/test_pi_agent_ecs_synth.py new file mode 100644 index 0000000000000000000000000000000000000000..500aad4558217c98636c89dc58c78f65f406f206 --- /dev/null +++ b/cdk/test/test_pi_agent_ecs_synth.py @@ -0,0 +1,169 @@ +"""Synth assertions for optional Pi agent ECS service on shared ALB.""" + +import sys +from pathlib import Path + +import pytest + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + + +def test_pi_agent_requires_service_connect(): + enable_pi = "True" + enable_sc = "False" + with pytest.raises(ValueError, match="ENABLE_PI_AGENT_ECS_SERVICE"): + if enable_pi == "True" and enable_sc != "True": + raise ValueError( + "ENABLE_PI_AGENT_ECS_SERVICE=True requires ENABLE_ECS_SERVICE_CONNECT=True " + "so the Pi task can reach the main app at http://:7860." + ) + + +def test_build_pi_agent_container_environment(): + from cdk_functions import build_pi_agent_container_environment + + env = build_pi_agent_container_environment( + service_connect_discovery_name="summarisation", + main_app_port=7860, + pi_gradio_port=7862, + ) + assert env["DOC_SUMMARISATION_GRADIO_URL"] == "http://summarisation:7860" + assert env["PI_DEFAULT_PROVIDER"] == "amazon-bedrock" + assert env["PI_GRADIO_PORT"] == "7862" + assert env["PI_CODING_AGENT_DIR"] == "/tmp/pi-agent" + assert env["ACCESS_LOGS_FOLDER"] == "/tmp/pi-logs/" + assert env["RUN_FASTAPI"] == "True" + assert env["RUN_AWS_FUNCTIONS"] == "True" + assert env["SAVE_OUTPUTS_TO_S3"] == "True" + assert env["S3_OUTPUTS_BUCKET"] + + +def test_ecs_availability_zone_rebalancing_default_disabled(): + from aws_cdk import aws_ecs as ecs + from cdk_functions import ecs_availability_zone_rebalancing + + assert ( + ecs_availability_zone_rebalancing("DISABLED") + == ecs.AvailabilityZoneRebalancing.DISABLED + ) + assert ( + ecs_availability_zone_rebalancing("ENABLED") + == ecs.AvailabilityZoneRebalancing.ENABLED + ) + + +def test_pi_agent_alb_attachment_synth(): + from aws_cdk import App, Duration, Environment, Stack + from aws_cdk import aws_ec2 as ec2 + from aws_cdk import aws_ecs as ecs + from aws_cdk import aws_elasticloadbalancingv2 as elbv2 + from aws_cdk import aws_iam as iam + from aws_cdk import aws_s3 as s3 + from cdk_functions import ( + attach_pi_agent_to_shared_alb, + create_pi_agent_ecs_resources, + ) + + app = App() + stack = Stack( + app, + "PiAgentAlbTest", + env=Environment(account="123456789012", region="eu-west-2"), + ) + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + alb_sg = ec2.SecurityGroup(stack, "AlbSg", vpc=vpc) + alb = elbv2.ApplicationLoadBalancer( + stack, + "Alb", + vpc=vpc, + internet_facing=True, + security_group=alb_sg, + ) + cluster = ecs.Cluster(stack, "Cluster", vpc=vpc) + config_bucket = s3.Bucket(stack, "ConfigBucket") + task_role = iam.Role( + stack, + "TaskRole", + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + execution_role = iam.Role( + stack, + "ExecRole", + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + + pi_service, pi_sg, _ = create_pi_agent_ecs_resources( + stack, + "Pi", + vpc=vpc, + cluster=cluster, + private_subnets=vpc.private_subnets, + pi_ecr_image_uri="123456789012.dkr.ecr.eu-west-2.amazonaws.com/pi-agent", + container_name="pi-agent", + task_role=task_role, + execution_role=execution_role, + config_bucket=config_bucket, + pi_agent_env_s3_key="", + service_name="test-pi-service", + task_family="test-pi-task", + security_group_name="test-pi-sg", + log_group_name="/ecs/test-pi-logs", + cpu=1024, + memory_mib=2048, + pi_gradio_port=7862, + service_connect_namespace="test-ns", + service_connect_discovery_name="summarisation", + main_app_port=7860, + use_fargate_spot="FARGATE", + ) + + http_listener = alb.add_listener("Http", port=80, open=True) + http_listener.add_action( + "DefaultDeny", + action=elbv2.ListenerAction.fixed_response( + status_code=403, + content_type="text/plain", + message_body="Access denied", + ), + ) + attach_pi_agent_to_shared_alb( + stack, + "PiAlb", + vpc=vpc, + alb_security_group=alb_sg, + pi_security_group=pi_sg, + pi_service=pi_service, + pi_port=7862, + routing_mode="host", + path_prefix="/pi", + pi_host_header="pi.example.com", + listener_rule_priority=3, + target_group_name="test-pi-tg", + stickiness_cookie_duration=Duration.hours(8), + https_listener=None, + http_listener=http_listener, + acm_certificate_arn="", + enable_cognito_auth=False, + cognito_user_pool=None, + cognito_user_pool_client=None, + cognito_user_pool_domain=None, + ) + + template = app.synth().get_stack_by_name("PiAgentAlbTest").template + resources = template["Resources"] + lb_count = sum( + 1 + for r in resources.values() + if r["Type"] == "AWS::ElasticLoadBalancingV2::LoadBalancer" + ) + tg_count = sum( + 1 + for r in resources.values() + if r["Type"] == "AWS::ElasticLoadBalancingV2::TargetGroup" + ) + assert lb_count == 1 + assert tg_count == 1 + ecs_services = [r for r in resources.values() if r["Type"] == "AWS::ECS::Service"] + assert len(ecs_services) == 1 + assert ecs_services[0]["Properties"]["AvailabilityZoneRebalancing"] == "DISABLED" diff --git a/cdk/test/test_pi_alb_routing.py b/cdk/test/test_pi_alb_routing.py new file mode 100644 index 0000000000000000000000000000000000000000..a0d2e01af8231c853473f859d62b394d46c45a7a --- /dev/null +++ b/cdk/test/test_pi_alb_routing.py @@ -0,0 +1,164 @@ +"""Tests for Pi ALB path/host routing helpers.""" + +import sys +from pathlib import Path + +import pytest + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + + +@pytest.fixture(autouse=True) +def _minimal_cdk_config_env(monkeypatch): + monkeypatch.setenv("USE_ECS_EXPRESS_MODE", "False") + monkeypatch.setenv("ACM_SSL_CERTIFICATE_ARN", "") + monkeypatch.setenv("ENABLE_S3_BATCH_ECS_TRIGGER", "False") + monkeypatch.setenv("ENABLE_PI_AGENT_ECS_SERVICE", "False") + monkeypatch.setenv("ENABLE_PI_AGENT_EXPRESS_SERVICE", "False") + monkeypatch.setenv("ENABLE_ECS_SERVICE_CONNECT", "False") + for mod_name in list(sys.modules): + if mod_name in ("cdk_config", "cdk_functions") or mod_name.startswith("cdk_"): + del sys.modules[mod_name] + + +def test_normalize_pi_alb_path_prefix(): + from cdk_functions import normalize_pi_alb_path_prefix + + assert normalize_pi_alb_path_prefix("/pi") == "/pi" + assert normalize_pi_alb_path_prefix("pi") == "/pi" + assert normalize_pi_alb_path_prefix("") == "/pi" + + +def test_pi_alb_path_patterns(): + from cdk_functions import pi_alb_path_patterns + + assert pi_alb_path_patterns("/pi") == ["/pi", "/pi/*"] + + +def test_format_pi_public_urls_path_on_cloudfront(): + from cdk_functions import format_pi_public_urls + + urls = format_pi_public_urls( + routing_mode="path", + path_prefix="/pi", + host_header="", + cloudfront_domain="d123.cloudfront.net", + use_https=True, + ) + assert urls == ["https://d123.cloudfront.net/pi/"] + + +def test_pi_listener_rule_count(): + from cdk_functions import pi_listener_rule_count + + assert pi_listener_rule_count("path") == 1 + assert pi_listener_rule_count("host") == 1 + assert pi_listener_rule_count("both") == 2 + + +def test_attach_pi_path_rule_synth(): + from aws_cdk import App, Duration, Environment, Stack, assertions + from aws_cdk import aws_ec2 as ec2 + from aws_cdk import aws_ecs as ecs + from aws_cdk import aws_elasticloadbalancingv2 as elbv2 + from aws_cdk import aws_iam as iam + from aws_cdk import aws_s3 as s3 + from cdk_functions import ( + attach_pi_agent_to_shared_alb, + create_pi_agent_ecs_resources, + ) + + app = App() + stack = Stack( + app, + "PiPathAlbTest", + env=Environment(account="123456789012", region="eu-west-2"), + ) + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + alb_sg = ec2.SecurityGroup(stack, "AlbSg", vpc=vpc) + alb = elbv2.ApplicationLoadBalancer( + stack, + "Alb", + vpc=vpc, + internet_facing=True, + security_group=alb_sg, + ) + cluster = ecs.Cluster(stack, "Cluster", vpc=vpc) + config_bucket = s3.Bucket(stack, "ConfigBucket") + task_role = iam.Role( + stack, + "TaskRole", + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + execution_role = iam.Role( + stack, + "ExecRole", + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + + pi_service, pi_sg, _ = create_pi_agent_ecs_resources( + stack, + "Pi", + vpc=vpc, + cluster=cluster, + private_subnets=vpc.private_subnets, + pi_ecr_image_uri="123456789012.dkr.ecr.eu-west-2.amazonaws.com/pi-agent", + container_name="pi-agent", + task_role=task_role, + execution_role=execution_role, + config_bucket=config_bucket, + pi_agent_env_s3_key="", + service_name="test-pi-service", + task_family="test-pi-task", + security_group_name="test-pi-sg", + log_group_name="/ecs/test-pi-logs", + cpu=1024, + memory_mib=2048, + pi_gradio_port=7862, + service_connect_namespace="test-ns", + service_connect_discovery_name="summarisation", + main_app_port=7860, + pi_root_path="/pi", + use_fargate_spot="FARGATE", + ) + + http_listener = alb.add_listener("Http", port=80, open=True) + http_listener.add_action( + "DefaultDeny", + action=elbv2.ListenerAction.fixed_response( + status_code=403, + content_type="text/plain", + message_body="Access denied", + ), + ) + attach_pi_agent_to_shared_alb( + stack, + "PiAlb", + vpc=vpc, + alb_security_group=alb_sg, + pi_security_group=pi_sg, + pi_service=pi_service, + pi_port=7862, + routing_mode="path", + path_prefix="/pi", + pi_host_header="", + listener_rule_priority=3, + target_group_name="test-pi-tg", + stickiness_cookie_duration=Duration.hours(8), + https_listener=None, + http_listener=http_listener, + acm_certificate_arn="", + enable_cognito_auth=False, + cognito_user_pool=None, + cognito_user_pool_client=None, + cognito_user_pool_domain=None, + ) + + template = assertions.Template.from_stack(stack) + template.has_resource_properties( + "AWS::ElasticLoadBalancingV2::TargetGroup", + { + "HealthCheckPath": "/pi/", + }, + ) diff --git a/cdk/test/test_public_subnet_igw.py b/cdk/test/test_public_subnet_igw.py new file mode 100644 index 0000000000000000000000000000000000000000..aa0742fb9a288f69ad6a9267d343363fa04d7c63 --- /dev/null +++ b/cdk/test/test_public_subnet_igw.py @@ -0,0 +1,141 @@ +"""Tests for public subnet Internet Gateway audit and CDK wiring helpers.""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + + +@pytest.fixture(autouse=True) +def _minimal_cdk_config_env(monkeypatch): + """Avoid loading the developer cdk_config.env validation errors during import.""" + monkeypatch.setenv("USE_ECS_EXPRESS_MODE", "False") + monkeypatch.setenv("ACM_SSL_CERTIFICATE_ARN", "") + monkeypatch.setenv("ENABLE_S3_BATCH_ECS_TRIGGER", "False") + monkeypatch.setenv("ENABLE_PI_AGENT_ECS_SERVICE", "False") + monkeypatch.setenv("ENABLE_PI_AGENT_EXPRESS_SERVICE", "False") + monkeypatch.setenv("ENABLE_ECS_SERVICE_CONNECT", "False") + for mod_name in list(sys.modules): + if mod_name in ("cdk_config", "cdk_functions") or mod_name.startswith("cdk_"): + del sys.modules[mod_name] + + +def test_route_table_default_internet_gateway_parses_igw(): + from cdk_functions import route_table_default_internet_gateway + + mock_client = MagicMock() + mock_client.describe_route_tables.return_value = { + "RouteTables": [ + { + "Routes": [ + { + "DestinationCidrBlock": "0.0.0.0/0", + "GatewayId": "igw-abc123", + } + ] + } + ] + } + with patch("cdk_functions.boto3.client", return_value=mock_client): + assert route_table_default_internet_gateway("rtb-1") == "igw-abc123" + + +def test_audit_needs_attachment_when_igw_detached(): + from cdk_functions import audit_public_subnet_internet_connectivity + + mock_ec2 = MagicMock() + mock_ec2.describe_internet_gateways.return_value = { + "InternetGateways": [{"Attachments": []}] + } + with ( + patch("cdk_functions.get_internet_gateways_attached_to_vpc", return_value=[]), + patch("cdk_functions.internet_gateway_exists", return_value=True), + patch("cdk_functions.boto3.client", return_value=mock_ec2), + patch("cdk_functions.route_table_default_internet_gateway", return_value=None), + patch( + "cdk_functions.route_table_has_non_igw_default_route", return_value=False + ), + ): + result = audit_public_subnet_internet_connectivity( + "vpc-111", + "igw-configured", + [ + { + "name": "PublicSubnet1", + "subnet_id": "subnet-1", + "route_table_id": "rtb-1", + } + ], + ) + assert result["internet_gateway_id"] == "igw-configured" + assert result["internet_gateway_needs_vpc_attachment"] is True + assert len(result["public_subnets_needing_igw_route"]) == 1 + + +def test_audit_raises_when_default_route_conflicts(): + from cdk_functions import audit_public_subnet_internet_connectivity + + with ( + patch( + "cdk_functions.get_internet_gateways_attached_to_vpc", + return_value=["igw-attached"], + ), + patch("cdk_functions.internet_gateway_exists", return_value=True), + patch( + "cdk_functions.route_table_default_internet_gateway", + return_value="igw-other", + ), + ): + with pytest.raises(ValueError, match="0.0.0.0/0"): + audit_public_subnet_internet_connectivity( + "vpc-111", + "igw-attached", + [ + { + "name": "PublicSubnet1", + "subnet_id": "subnet-1", + "route_table_id": "rtb-1", + } + ], + ) + + +def test_wire_public_subnet_internet_access_synth(): + from aws_cdk import App, Environment, Stack, assertions + from cdk_functions import wire_public_subnet_internet_access + + app = App() + stack = Stack( + app, + "IgwWireTest", + env=Environment(account="123456789012", region="eu-west-2"), + ) + wire_public_subnet_internet_access( + stack, + "Test", + vpc_id="vpc-111", + internet_gateway_id="igw-abc123", + needs_igw_vpc_attachment=True, + subnets_needing_route=[ + { + "name": "PublicSubnet1", + "subnet_id": "subnet-1", + "route_table_id": "rtb-aaa", + } + ], + ) + template = assertions.Template.from_stack(stack) + template.resource_count_is("AWS::EC2::VPCGatewayAttachment", 1) + template.resource_count_is("AWS::EC2::Route", 1) + template.has_resource_properties( + "AWS::EC2::Route", + { + "RouteTableId": "rtb-aaa", + "DestinationCidrBlock": "0.0.0.0/0", + "GatewayId": "igw-abc123", + }, + ) diff --git a/cdk/test/test_s3_batch_lambda_handler.py b/cdk/test/test_s3_batch_lambda_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..d6dc0a808caf0fd809c8caad9d189637c79da640 --- /dev/null +++ b/cdk/test/test_s3_batch_lambda_handler.py @@ -0,0 +1,89 @@ +"""S3 batch Lambda handler behaviour (assignPublicIp, env merge).""" + +import importlib.util +import sys +from pathlib import Path + +CDK_DIR = Path(__file__).resolve().parents[1] +LAMBDA_PATH = CDK_DIR / "config" / "lambda" / "lambda_function.py" + + +def _load_lambda_module(module_name: str = "batch_lambda"): + spec = importlib.util.spec_from_file_location(module_name, LAMBDA_PATH) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def test_assign_public_ip_reads_env(monkeypatch): + monkeypatch.setenv("ECS_ASSIGN_PUBLIC_IP", "ENABLED") + module = _load_lambda_module("batch_lambda_ip") + assert module.ECS_ASSIGN_PUBLIC_IP == "ENABLED" + + +def test_build_environment_merges_job_and_default_env(monkeypatch): + mod = _load_lambda_module("batch_lambda_merge") + monkeypatch.setattr( + mod, + "s3", + type( + "S3", + (), + { + "get_object": staticmethod( + lambda Bucket, Key: { + "Body": type( + "Body", + (), + { + "read": staticmethod( + lambda: ( + b"RUN_DIRECT_MODE=1\n" + b"DIRECT_MODE_INPUT_FILE=job.xlsx\n" + if Key.endswith("job.env") + else b"RUN_AWS_FUNCTIONS=1\n" + ) + ) + }, + )() + } + ), + }, + )(), + ) + monkeypatch.setattr( + mod, + "ecs", + type( + "ECS", + (), + {"run_task": staticmethod(lambda **kwargs: {"tasks": [], "failures": []})}, + )(), + ) + + mod.BUCKET = "output-bucket" + mod.INPUT_PREFIX = "input/" + mod.ENV_PREFIX = "input/config/" + mod.ENV_SUFFIX = ".env" + mod.DEFAULT_PARAMS_KEY = "general-config/app_defaults.env" + mod.CLUSTER = "cluster" + mod.TASK_DEF = "arn:aws:ecs:eu-west-2:123:task-definition/app:1" + mod.SUBNETS = ["subnet-1"] + mod.SECURITY_GROUPS = ["sg-1"] + mod.CONTAINER_NAME = "app" + + event = { + "Records": [ + { + "s3": { + "bucket": {"name": "output-bucket"}, + "object": {"key": "input/config/job.env"}, + } + } + ] + } + result = mod.lambda_handler(event, None) + assert result["runs"] + assert result["runs"][0]["envCount"] >= 2 diff --git a/cdk/test/test_s3_batch_lambda_synth.py b/cdk/test/test_s3_batch_lambda_synth.py new file mode 100644 index 0000000000000000000000000000000000000000..5c8420b7e8ad4edd0e4008c898553f8e9fa04485 --- /dev/null +++ b/cdk/test/test_s3_batch_lambda_synth.py @@ -0,0 +1,107 @@ +"""Synth assertions for optional S3 batch ECS trigger Lambda.""" + +import sys +from pathlib import Path + +import pytest + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + + +def test_batch_trigger_express_mutual_exclusion(): + enable_batch = "True" + use_express = "True" + with pytest.raises(ValueError, match="ENABLE_S3_BATCH_ECS_TRIGGER"): + if enable_batch == "True" and use_express == "True": + raise ValueError( + "ENABLE_S3_BATCH_ECS_TRIGGER=True requires the legacy Fargate task definition " + "for ecs.run_task. Set USE_ECS_EXPRESS_MODE=False or disable the batch trigger." + ) + + +def test_s3_batch_lambda_synth_resources(): + from aws_cdk import App, Environment, Stack + from aws_cdk import aws_ec2 as ec2 + from aws_cdk import aws_ecs as ecs + from aws_cdk import aws_iam as iam + from aws_cdk import aws_s3 as s3 + from cdk_functions import create_s3_batch_ecs_trigger_lambda + + app = App() + stack = Stack( + app, + "BatchLambdaTest", + env=Environment(account="123456789012", region="eu-west-2"), + ) + vpc = ec2.Vpc(stack, "Vpc", max_azs=2) + output_bucket = s3.Bucket(stack, "OutputBucket") + config_bucket = s3.Bucket(stack, "ConfigBucket") + execution_role = iam.Role( + stack, + "ExecutionRole", + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + task_role = iam.Role( + stack, + "TaskRole", + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + task_def = ecs.FargateTaskDefinition( + stack, "TaskDef", memory_limit_mib=2048, cpu=1024 + ) + task_def.add_container( + "docsummarisation", + image=ecs.ContainerImage.from_registry("nginx:latest"), + ) + + lambda_asset = str(CDK_DIR / "config" / "lambda") + create_s3_batch_ecs_trigger_lambda( + stack, + "S3BatchEcsTrigger", + function_name=None, + lambda_asset_path=lambda_asset, + output_bucket=output_bucket, + config_bucket=config_bucket, + cluster_name="test-cluster", + task_definition_arn=task_def.task_definition_arn, + container_name="docsummarisation", + subnet_ids=[vpc.private_subnets[0].subnet_id], + security_group_id=ec2.SecurityGroup( + stack, "EcsSg", vpc=vpc, allow_all_outbound=True + ).security_group_id, + execution_role=execution_role, + task_role=task_role, + env_prefix="input/config/", + env_suffix=".env", + input_prefix="input/", + config_prefix="", + default_params_key="general-config/app_defaults.env", + general_env_prefix="general-config/", + default_task_type="extract", + ) + + template = app.synth().get_stack_by_name("BatchLambdaTest").template + resources = template["Resources"] + types = {r["Type"] for r in resources.values()} + + assert "AWS::Lambda::Function" in types + assert "AWS::Lambda::Permission" in types + assert "Custom::S3BucketNotifications" in types or any( + r.get("Type") == "AWS::S3::Bucket" + and "NotificationConfiguration" in r.get("Properties", {}) + for r in resources.values() + ) + + batch_lambda = next( + r + for r in resources.values() + if r["Type"] == "AWS::Lambda::Function" + and r["Properties"].get("Handler") == "lambda_function.lambda_handler" + ) + env_vars = batch_lambda["Properties"]["Environment"]["Variables"] + assert env_vars.get("DEFAULT_TASK_TYPE") == "extract" + assert env_vars.get("BUCKET") + assert "RUN_DIRECT_MODE" not in env_vars + assert env_vars.get("ENV_PREFIX") == "input/config/" + assert env_vars.get("ECS_ASSIGN_PUBLIC_IP") == "DISABLED" diff --git a/cdk/test/test_secret_arn_grants.py b/cdk/test/test_secret_arn_grants.py new file mode 100644 index 0000000000000000000000000000000000000000..4d3234ca24d9ae735998d14bfcbff0a14716eefa --- /dev/null +++ b/cdk/test/test_secret_arn_grants.py @@ -0,0 +1,78 @@ +"""Secrets Manager IAM grant ARN handling (name wildcard vs complete ARN).""" + +import sys +from pathlib import Path + +CDK_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(CDK_DIR)) + + +def test_from_secret_complete_arn_grant_uses_exact_resource(): + from aws_cdk import App, Stack, assertions + from aws_cdk import aws_iam as iam + from aws_cdk import aws_secretsmanager as sm + + app = App() + stack = Stack(app, "SecretGrantTest") + role = iam.Role( + stack, + "ExecutionRole", + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + secret_arn = ( + "arn:aws:secretsmanager:eu-west-2:123456789012:secret:" + "Lambeth-ParamCognitoSecret-AbCdEf" + ) + secret = sm.Secret.from_secret_complete_arn( + stack, + "CognitoSecret", + secret_complete_arn=secret_arn, + ) + secret.grant_read(role) + + template = assertions.Template.from_stack(stack) + template.has_resource_properties( + "AWS::IAM::Policy", + { + "PolicyDocument": assertions.Match.object_like( + { + "Statement": assertions.Match.array_with( + [ + assertions.Match.object_like( + { + "Action": assertions.Match.array_with( + ["secretsmanager:GetSecretValue"] + ), + "Resource": secret_arn, + } + ) + ] + ) + } + ) + }, + ) + + +def test_from_secret_name_v2_grant_uses_wildcard_suffix(): + from aws_cdk import App, Stack, assertions + from aws_cdk import aws_iam as iam + from aws_cdk import aws_secretsmanager as sm + + app = App() + stack = Stack(app, "SecretNameGrantTest") + role = iam.Role( + stack, + "ExecutionRole", + assumed_by=iam.ServicePrincipal("ecs-tasks.amazonaws.com"), + ) + secret_name = "Lambeth-ParamCognitoSecret" + secret = sm.Secret.from_secret_name_v2( + stack, "CognitoSecret", secret_name=secret_name + ) + secret.grant_read(role) + + template = assertions.Template.from_stack(stack) + resources = template.find_resources("AWS::IAM::Policy") + policy_text = str(resources) + assert "??????" in policy_text or "???????" in policy_text diff --git a/cdk/test/test_service_connect_config.py b/cdk/test/test_service_connect_config.py new file mode 100644 index 0000000000000000000000000000000000000000..8d93e22cde0fab25eb4230c184e531f328d9d6f4 --- /dev/null +++ b/cdk/test/test_service_connect_config.py @@ -0,0 +1,51 @@ +"""Config validation for ECS Service Connect options.""" + +import pytest +from cdk_config import parse_comma_separated_list + + +def test_service_connect_express_mutual_exclusion(): + enable_sc = "True" + use_express = "True" + with pytest.raises(ValueError, match="ENABLE_ECS_SERVICE_CONNECT"): + if enable_sc == "True" and use_express == "True": + raise ValueError( + "ENABLE_ECS_SERVICE_CONNECT=True is only supported on the legacy Fargate " + "service path." + ) + + +def test_parse_client_security_group_ids(): + assert parse_comma_separated_list("sg-abc, sg-def") == ["sg-abc", "sg-def"] + assert parse_comma_separated_list("") == [] + + +def test_build_service_connect_client_sg_names_from_prefixes(monkeypatch): + import cdk_config as cfg + + monkeypatch.setattr(cfg, "ECS_SERVICE_CONNECT_CLIENT_SECURITY_GROUP_NAMES_LIST", []) + monkeypatch.setattr( + cfg, "ECS_SERVICE_CONNECT_CLIENT_CDK_PREFIXES_LIST", ["OtherApp-"] + ) + monkeypatch.setattr( + cfg, "ECS_SERVICE_CONNECT_CLIENT_SG_NAME_SUFFIX", "SecurityGroupECS" + ) + assert cfg.build_service_connect_client_security_group_names() == [ + "OtherApp-SecurityGroupECS" + ] + + +def test_resolve_service_connect_client_security_group_ids(): + from cdk_functions import resolve_service_connect_client_security_group_ids + + def _ctx(key, default=None): + if key == "security_group_id:ClientSg": + return "sg-ctx123" + return default + + ids = resolve_service_connect_client_security_group_ids( + ["sg-explicit"], + ["ClientSg"], + _ctx, + ) + assert ids == ["sg-explicit", "sg-ctx123"] diff --git a/cli_topics.py b/cli_topics.py new file mode 100644 index 0000000000000000000000000000000000000000..91fb669b7d5a544fb77e3cdc31d98d06630726ab --- /dev/null +++ b/cli_topics.py @@ -0,0 +1,2268 @@ +import argparse +import csv +import os +import re +import sys +import time +import uuid +from datetime import datetime + +import boto3 +import botocore +import pandas as pd + +from tools.aws_functions import ( + collect_prompt_response_log_paths, + download_file_from_s3, + export_outputs_to_s3, +) +from tools.combine_sheets_into_xlsx import collect_output_csvs_and_create_excel_output +from tools.config import ( + API_URL, + AWS_ACCESS_KEY, + AWS_REGION, + AWS_SECRET_KEY, + AZURE_OPENAI_API_KEY, + AZURE_OPENAI_INFERENCE_ENDPOINT, + BATCH_SIZE_DEFAULT, + CHOSEN_INFERENCE_SERVER_MODEL, + CSV_USAGE_LOG_HEADERS, + DEDUPLICATION_THRESHOLD, + DEFAULT_COST_CODE, + DEFAULT_SAMPLED_SUMMARIES, + DIRECT_MODE_DEFAULT_COST_CODE, + DIRECT_MODE_S3_UPLOAD_ONLY_XLSX, + DYNAMODB_USAGE_LOG_HEADERS, + ENABLE_BATCH_DEDUPLICATION, + GEMINI_API_KEY, + GRADIO_TEMP_DIR, + HF_TOKEN, + INPUT_FOLDER, + LLM_MAX_NEW_TOKENS, + LLM_SEED, + LLM_TEMPERATURE, + MAX_TIME_FOR_LOOP, + MAXIMUM_ALLOWED_TOPICS, + OUTPUT_DEBUG_FILES, + OUTPUT_FOLDER, + RUN_AWS_FUNCTIONS, + S3_OUTPUTS_BUCKET, + S3_OUTPUTS_FOLDER, + SAVE_LOGS_TO_CSV, + SAVE_LOGS_TO_DYNAMODB, + SAVE_OUTPUTS_TO_S3, + SESSION_OUTPUT_FOLDER, + UPLOAD_PROMPT_RESPONSE_LOG_TO_S3_OUTPUTS, + UPLOAD_USAGE_LOG_TO_S3_OUTPUTS, + USAGE_LOG_DYNAMODB_TABLE_NAME, + USAGE_LOG_FILE_NAME, + USAGE_LOGS_FOLDER, + convert_string_to_boolean, + default_model_choice, + default_model_source, + model_name_map, +) +from tools.dedup_summaries import ( + deduplicate_topics, + deduplicate_topics_llm, + overall_summary, + wrapper_summarise_output_topics_per_group, +) +from tools.helper_functions import ( + ColumnNotFoundError, + load_in_data_file, + load_in_previous_data_files, +) +from tools.llm_api_call import ( + all_in_one_pipeline, + discover_topics_from_sample, + validate_topics_wrapper, + wrapper_extract_topics_per_column_value, +) +from tools.prompts import ( + add_existing_topics_prompt, + add_existing_topics_system_prompt, + initial_table_prompt, + initial_table_system_prompt, + single_para_summary_format_prompt, + two_para_summary_format_prompt, +) + + +def _generate_session_hash() -> str: + """Generate a unique session hash for logging purposes.""" + return str(uuid.uuid4())[:8] + + +def _download_s3_file_if_needed( + file_path: str, + default_filename: str = "downloaded_file", + aws_access_key: str = "", + aws_secret_key: str = "", + aws_region: str = "", +) -> str: + """ + Download a file from S3 if the path starts with 's3://' or 'S3://', otherwise return the path as-is. + + Args: + file_path: File path (either local or S3 URL) + default_filename: Default filename to use if S3 key doesn't have a filename + aws_access_key: AWS access key ID (optional, uses environment/config if not provided) + aws_secret_key: AWS secret access key (optional, uses environment/config if not provided) + aws_region: AWS region (optional, uses environment/config if not provided) + + Returns: + Local file path (downloaded from S3 or original path) + """ + if not file_path: + return file_path + + # Check for S3 URL (case-insensitive) + file_path_stripped = file_path.strip() + file_path_upper = file_path_stripped.upper() + if not file_path_upper.startswith("S3://"): + return file_path + + # Ensure temp directory exists + os.makedirs(GRADIO_TEMP_DIR, exist_ok=True) + + # Parse S3 URL: s3://bucket/key (preserve original case for bucket/key) + # Remove 's3://' prefix (case-insensitive) + s3_path = ( + file_path_stripped.split("://", 1)[1] + if "://" in file_path_stripped + else file_path_stripped + ) + # Split bucket and key (first '/' separates bucket from key) + if "/" in s3_path: + bucket_name_s3, s3_key = s3_path.split("/", 1) + else: + # If no key provided, use bucket name as key (unlikely but handle it) + bucket_name_s3 = s3_path + s3_key = "" + + # Get the filename from the S3 key + filename = os.path.basename(s3_key) if s3_key else bucket_name_s3 + if not filename: + filename = default_filename + + # Create local file path in temp directory + local_file_path = os.path.join(GRADIO_TEMP_DIR, filename) + + # Download file from S3 + try: + download_file_from_s3( + bucket_name=bucket_name_s3, + key=s3_key, + local_file_path=local_file_path, + aws_access_key_textbox=aws_access_key, + aws_secret_key_textbox=aws_secret_key, + aws_region_textbox=aws_region, + ) + print(f"S3 file downloaded successfully: {file_path} -> {local_file_path}") + return local_file_path + except Exception as e: + print(f"Error downloading file from S3 ({file_path}): {e}") + raise Exception(f"Failed to download file from S3: {e}") + + +def get_username_and_folders( + username: str = "", + output_folder_textbox: str = OUTPUT_FOLDER, + input_folder_textbox: str = INPUT_FOLDER, + session_output_folder: bool = SESSION_OUTPUT_FOLDER, +): + """Generate session hash and set up output/input folders.""" + # Generate session hash for logging. Either from input user name or generated + if username: + out_session_hash = username + else: + out_session_hash = _generate_session_hash() + + if session_output_folder: + output_folder = output_folder_textbox + out_session_hash + "/" + input_folder = input_folder_textbox + out_session_hash + "/" + else: + output_folder = output_folder_textbox + input_folder = input_folder_textbox + + if not os.path.exists(output_folder): + os.makedirs(output_folder, exist_ok=True) + if not os.path.exists(input_folder): + os.makedirs(input_folder, exist_ok=True) + + return ( + out_session_hash, + output_folder, + out_session_hash, + input_folder, + ) + + +def _sanitize_folder_name(folder_name: str, max_length: int = 50) -> str: + """ + Sanitize folder name for S3 compatibility. + + Replaces 'strange' characters (anything that's not alphanumeric, dash, underscore, or full stop) + with underscores, and limits the length to max_length characters. + + Args: + folder_name: Original folder name to sanitize + max_length: Maximum length for the folder name (default: 50) + + Returns: + Sanitized folder name + """ + if not folder_name: + return folder_name + + # Replace any character that's not alphanumeric, dash, underscore, or full stop with underscore + # This handles @, commas, exclamation marks, spaces, etc. + sanitized = re.sub(r"[^a-zA-Z0-9._-]", "_", folder_name) + + # Limit length to max_length + if len(sanitized) > max_length: + sanitized = sanitized[:max_length] + + return sanitized + + +def upload_outputs_to_s3_if_enabled( + output_files: list, + base_file_name: str = None, + session_hash: str = "", + s3_output_folder: str = S3_OUTPUTS_FOLDER, + s3_bucket: str = S3_OUTPUTS_BUCKET, + save_outputs_to_s3: bool = None, + task_usage_log_path: str = None, + prompt_response_log_paths=None, +): + """ + Upload output files to S3 if SAVE_OUTPUTS_TO_S3 is enabled. + + Args: + output_files: List of output file paths to upload + base_file_name: Base file name (input file) for organizing S3 folder structure + session_hash: Session hash to include in S3 path + s3_output_folder: S3 output folder path + s3_bucket: S3 bucket name + save_outputs_to_s3: Override for SAVE_OUTPUTS_TO_S3 config (if None, uses config value) + task_usage_log_path: Path to task-specific usage log file to upload (if UPLOAD_USAGE_LOG_TO_S3_OUTPUTS is enabled) + prompt_response_log_paths: Prompt/response JSON log paths to upload (if UPLOAD_PROMPT_RESPONSE_LOG_TO_S3_OUTPUTS is enabled) + """ + # Use provided value or fall back to config + if save_outputs_to_s3 is None: + save_outputs_to_s3 = convert_string_to_boolean(SAVE_OUTPUTS_TO_S3) + + if not save_outputs_to_s3: + return + + if not s3_bucket: + print("Warning: S3_OUTPUTS_BUCKET not configured. Skipping S3 upload.") + return + + if not output_files: + print("No output files to upload to S3.") + return + + # Filter out empty/None values and ensure files exist + valid_files = [] + for file_path in output_files: + if file_path and os.path.exists(file_path): + valid_files.append(file_path) + elif file_path: + print(f"Warning: Output file does not exist, skipping: {file_path}") + + # Check if task-specific usage log should be uploaded to S3 output folder + if UPLOAD_USAGE_LOG_TO_S3_OUTPUTS and task_usage_log_path: + if os.path.exists(task_usage_log_path): + valid_files.append(task_usage_log_path) + print( + f"Including task-specific usage log in S3 upload: {task_usage_log_path}" + ) + else: + print( + f"Task-specific usage log not found at {task_usage_log_path}, skipping usage log upload." + ) + + if UPLOAD_PROMPT_RESPONSE_LOG_TO_S3_OUTPUTS and prompt_response_log_paths: + for log_path in collect_prompt_response_log_paths(prompt_response_log_paths): + if log_path not in valid_files: + valid_files.append(log_path) + print(f"Including prompt/response log in S3 upload: {log_path}") + + if not valid_files: + print("No valid output files to upload to S3.") + return + + # Construct S3 output folder path + # Include session hash if provided and SESSION_OUTPUT_FOLDER is enabled + s3_folder_path = s3_output_folder or "" + if session_hash and convert_string_to_boolean(SESSION_OUTPUT_FOLDER): + if s3_folder_path and not s3_folder_path.endswith("/"): + s3_folder_path += "/" + # Sanitize session_hash to ensure S3 compatibility + sanitized_session_hash = _sanitize_folder_name(session_hash) + s3_folder_path += sanitized_session_hash + "/" + + print(f"\nUploading {len(valid_files)} output file(s) to S3...") + try: + export_outputs_to_s3( + file_list_state=valid_files, + s3_output_folder_state_value=s3_folder_path, + save_outputs_to_s3_flag=True, + base_file_state=base_file_name, + s3_bucket=s3_bucket, + ) + except Exception as e: + print(f"Warning: Failed to upload outputs to S3: {e}") + + +def write_usage_log( + session_hash: str, + file_name: str, + text_column: str, + model_choice: str, + conversation_metadata: str, + input_tokens: int, + output_tokens: int, + number_of_calls: int, + estimated_time_taken: float, + cost_code: str = DEFAULT_COST_CODE, + save_to_csv: bool = SAVE_LOGS_TO_CSV, + save_to_dynamodb: bool = SAVE_LOGS_TO_DYNAMODB, + include_conversation_metadata: bool = False, + output_folder: str = None, +): + """ + Write usage log entry to CSV file and/or DynamoDB. + + Args: + session_hash: Session identifier + file_name: Name of the input file + text_column: Column name used for analysis (as list for CSV) + model_choice: LLM model used + conversation_metadata: Metadata string + input_tokens: Number of input tokens + output_tokens: Number of output tokens + number_of_calls: Number of LLM calls + estimated_time_taken: Time taken in seconds + cost_code: Cost code for tracking + save_to_csv: Whether to save to CSV + save_to_dynamodb: Whether to save to DynamoDB + include_conversation_metadata: Whether to include conversation metadata in the log + output_folder: Output folder path for creating task-specific usage log file (if UPLOAD_USAGE_LOG_TO_S3_OUTPUTS is enabled) + + Returns: + str or None: Path to task-specific usage log file if created, None otherwise + """ + # Convert boolean parameters if they're strings + if isinstance(save_to_csv, str): + save_to_csv = convert_string_to_boolean(save_to_csv) + if isinstance(save_to_dynamodb, str): + save_to_dynamodb = convert_string_to_boolean(save_to_dynamodb) + + # Return early if neither logging method is enabled + if not save_to_csv and not save_to_dynamodb: + return + + if not conversation_metadata: + conversation_metadata = "" + + # Ensure usage logs folder exists + os.makedirs(USAGE_LOGS_FOLDER, exist_ok=True) + + # Construct full file path + usage_log_file_path = os.path.join(USAGE_LOGS_FOLDER, USAGE_LOG_FILE_NAME) + + # Prepare data row - order matches app.py component order + # session_hash_textbox, original_data_file_name_textbox, in_colnames, model_choice, + # conversation_metadata_textbox_placeholder, input_tokens_num, output_tokens_num, + # number_of_calls_num, estimated_time_taken_number, cost_code_choice_drop + data = [ + session_hash, + file_name, + ( + text_column + if isinstance(text_column, str) + else (text_column[0] if text_column else "") + ), + model_choice, + "", # conversation_metadata if conversation_metadata else "", + input_tokens, + output_tokens, + number_of_calls, + estimated_time_taken, + cost_code, + ] + + # Add id and timestamp + generated_id = str(uuid.uuid4()) + timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3] + data.extend([generated_id, timestamp]) + + # Use custom headers if available, otherwise use default + # Note: CSVLogger_custom uses component labels, but we need to match what collect_output_csvs_and_create_excel_output expects + if CSV_USAGE_LOG_HEADERS and len(CSV_USAGE_LOG_HEADERS) == len(data): + headers = CSV_USAGE_LOG_HEADERS + else: + # Default headers - these should match what CSVLogger_custom creates from Gradio component labels + # The components are: session_hash_textbox, original_data_file_name_textbox, in_colnames, + # model_choice, conversation_metadata_textbox_placeholder, input_tokens_num, output_tokens_num, + # number_of_calls_num, estimated_time_taken_number, cost_code_choice_drop + # Since these are hidden components without labels, CSVLogger_custom uses component variable names + # or default labels. We need to match what collect_output_csvs_and_create_excel_output expects: + # "Total LLM calls", "Total input tokens", "Total output tokens" + # But the actual CSV from Gradio likely has: "Number of calls", "Input tokens", "Output tokens" + # Let's use the names that match what the Excel function expects + headers = [ + "Session hash", + "Response ID data file name", + "Select the open text column of interest. In an Excel file, this shows columns across all sheets.", + "Large language model for topic extraction and summarisation", + "Conversation metadata", + "Total input tokens", # Changed from "Input tokens" to match Excel function + "Total output tokens", # Changed from "Output tokens" to match Excel function + "Total LLM calls", # Changed from "Number of calls" to match Excel function + "Estimated time taken (seconds)", + "Cost code", + "id", + "timestamp", + ] + + # Write to CSV if enabled + task_specific_log_path = None + if save_to_csv: + # Ensure usage logs folder exists + os.makedirs(USAGE_LOGS_FOLDER, exist_ok=True) + + # Construct full file path + usage_log_file_path = os.path.join(USAGE_LOGS_FOLDER, USAGE_LOG_FILE_NAME) + + # Write to CSV + file_exists = os.path.exists(usage_log_file_path) + with open( + usage_log_file_path, "a", newline="", encoding="utf-8-sig" + ) as csvfile: + writer = csv.writer(csvfile) + if not file_exists: + # Write headers if file doesn't exist + writer.writerow(headers) + writer.writerow(data) + + # Create task-specific usage log file if enabled and output folder provided + if UPLOAD_USAGE_LOG_TO_S3_OUTPUTS and output_folder: + # Ensure output folder exists + os.makedirs(output_folder, exist_ok=True) + + # Create task-specific usage log file name + # Use session hash and timestamp to make it unique + base_name = ( + os.path.splitext(os.path.basename(file_name))[0] + if file_name + else "usage" + ) + task_log_filename = f"{base_name}_usage_log_{session_hash[:8]}_{timestamp.replace(':', '-').replace(' ', '_')}.csv" + task_specific_log_path = os.path.join(output_folder, task_log_filename) + + # Write task-specific CSV file with headers and single entry + with open( + task_specific_log_path, "w", newline="", encoding="utf-8-sig" + ) as csvfile: + writer = csv.writer(csvfile) + writer.writerow(headers) + writer.writerow(data) + + print(f"Created task-specific usage log: {task_specific_log_path}") + + # Write to DynamoDB if enabled + if save_to_dynamodb: + # DynamoDB logging implementation + print("Saving to DynamoDB") + + try: + # Connect to DynamoDB + if RUN_AWS_FUNCTIONS == "1": + try: + print("Connecting to DynamoDB via existing SSO connection") + dynamodb = boto3.resource("dynamodb", region_name=AWS_REGION) + dynamodb.meta.client.list_tables() + except Exception as e: + print("No SSO credentials found:", e) + if AWS_ACCESS_KEY and AWS_SECRET_KEY: + print("Trying DynamoDB credentials from environment variables") + dynamodb = boto3.resource( + "dynamodb", + aws_access_key_id=AWS_ACCESS_KEY, + aws_secret_access_key=AWS_SECRET_KEY, + region_name=AWS_REGION, + ) + else: + raise Exception( + "AWS credentials for DynamoDB logging not found" + ) + else: + raise Exception("AWS credentials for DynamoDB logging not found") + + # Get table name from config + dynamodb_table_name = USAGE_LOG_DYNAMODB_TABLE_NAME + if not dynamodb_table_name: + raise ValueError( + "USAGE_LOG_DYNAMODB_TABLE_NAME not configured. Cannot save to DynamoDB." + ) + + # Determine headers for DynamoDB + # Use DYNAMODB_USAGE_LOG_HEADERS if available and matches data length, + # otherwise use CSV_USAGE_LOG_HEADERS if it matches, otherwise use default headers + # Note: headers and data are guaranteed to have the same length and include id/timestamp + if DYNAMODB_USAGE_LOG_HEADERS and len(DYNAMODB_USAGE_LOG_HEADERS) == len( + data + ): + dynamodb_headers = list(DYNAMODB_USAGE_LOG_HEADERS) # Make a copy + elif CSV_USAGE_LOG_HEADERS and len(CSV_USAGE_LOG_HEADERS) == len(data): + dynamodb_headers = list(CSV_USAGE_LOG_HEADERS) # Make a copy + else: + # Use the headers we created which are guaranteed to match data + dynamodb_headers = headers + + # Check if table exists, create if it doesn't + try: + table = dynamodb.Table(dynamodb_table_name) + table.load() + except botocore.exceptions.ClientError as e: + if e.response["Error"]["Code"] == "ResourceNotFoundException": + print( + f"Table '{dynamodb_table_name}' does not exist. Creating it..." + ) + attribute_definitions = [ + { + "AttributeName": "id", + "AttributeType": "S", + } + ] + + table = dynamodb.create_table( + TableName=dynamodb_table_name, + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], + AttributeDefinitions=attribute_definitions, + BillingMode="PAY_PER_REQUEST", + ) + # Wait until the table exists + table.meta.client.get_waiter("table_exists").wait( + TableName=dynamodb_table_name + ) + time.sleep(5) + print(f"Table '{dynamodb_table_name}' created successfully.") + else: + raise + + # Prepare the DynamoDB item to upload + # Map the headers to values (headers and data should match in length) + if len(dynamodb_headers) == len(data): + item = { + header: str(value) for header, value in zip(dynamodb_headers, data) + } + else: + # Fallback: use the default headers which are guaranteed to match data + print( + f"Warning: DynamoDB headers length ({len(dynamodb_headers)}) doesn't match data length ({len(data)}). Using default headers." + ) + item = {header: str(value) for header, value in zip(headers, data)} + + # Upload to DynamoDB + table.put_item(Item=item) + print("Successfully uploaded log to DynamoDB") + + except Exception as e: + print(f"Could not upload log to DynamoDB due to: {e}") + import traceback + + traceback.print_exc() + + return task_specific_log_path + + +# --- Main CLI Function --- +def main(direct_mode_args={}): + """ + A unified command-line interface for topic extraction, validation, deduplication, and summarisation. + + Args: + direct_mode_args (dict, optional): Dictionary of arguments for direct mode execution. + If provided, uses these instead of parsing command line arguments. + """ + parser = argparse.ArgumentParser( + description="A versatile CLI for topic extraction, validation, deduplication, and summarisation using LLMs.", + formatter_class=argparse.RawTextHelpFormatter, + epilog=""" +Examples: + +To run these, you need to do the following: + +- Open a terminal window + +- CD to the app folder that contains this file (cli_topics.py) + +- Load the virtual environment using either conda or venv depending on your setup + +- Run one of the example commands below + +- The examples below use the free Gemini 2.5 Flash Lite model, that is free with an API key that you can get from here: https://aistudio.google.com/api-keys. You can either set this or API keys for other services as an environment variable (e.g. in config/app_config.py. See the file tools/config.py for more details about variables relevant to each service) or you can set them manually at the time of the function call via command line arguments, such as the following: + + Google/Gemini: --google_api_key + AWS Bedrock: --aws_access_key, --aws_secret_key, --aws_region + Hugging Face (for model download): --hf_token + Azure/OpenAI: --azure_api_key, --azure_endpoint + Inference Server endpoint for local models(e.g. llama server, vllm): --api_url + +- Use --create_xlsx_output to create an Excel file combining all CSV outputs after task completion + +- Look in the output/ folder to see output files: + +# Topic Extraction + +## Extract topics from a CSV file with default settings: +python cli_topics.py --task extract --input_file example_data/combined_case_notes.csv --text_column "Case Note" + +## Extract topics with custom model and context: +python cli_topics.py --task extract --input_file example_data/combined_case_notes.csv --text_column "Case Note" --model_choice "gemini-flash-lite-latest" --context "Social Care case notes for young people" + +## Extract topics with grouping: +python cli_topics.py --task extract --input_file example_data/combined_case_notes.csv --text_column "Case Note" --group_by "Client" + +## Extract topics with candidate topics (zero-shot): +python cli_topics.py --task extract --input_file example_data/dummy_consultation_response.csv --text_column "Response text" --candidate_topics example_data/dummy_consultation_response_themes.csv + +## Discover topics from a 20% random sample (then upload the suggested topics CSV for a full run): +python cli_topics.py --task discover --input_file example_data/dummy_consultation_response.csv --text_column "Response text" --sample_fraction 20 --random_seed 42 + +# Topic Validation + +## Validate previously extracted topics: +python cli_topics.py --task validate --input_file example_data/combined_case_notes.csv --text_column "Case Note" --previous_output_files output/combined_case_notes_col_Case_Note_reference_table.csv output/combined_case_notes_col_Case_Note_unique_topics.csv + +# Deduplication + +Note: you will need to change the reference to previous output files to match the exact file names created from the previous task. This includes the relative path to the app folder. Also, the function will create an xlsx output file by default. the --input_file and --text_column arguments are needed for this, unless you pass in --no_xlsx_output as seen below. + +## Deduplicate topics using fuzzy matching: +python cli_topics.py --task deduplicate --previous_output_files output/combined_case_notes_col_Case_Note_reference_table.csv output/combined_case_notes_col_Case_Note_unique_topics.csv --similarity_threshold 90 --no_xlsx_output + +## Deduplicate topics using LLM: +python cli_topics.py --task deduplicate --previous_output_files output/combined_case_notes_col_Case_Note_reference_table.csv output/combined_case_notes_col_Case_Note_unique_topics.csv --method llm --model_choice "gemini-flash-lite-latest" --no_xlsx_output + +# Summarisation + +Note: you will need to change the reference to previous output files to match the exact file names created from the previous task. This includes the relative path to the app folder. Also, the function will create an xlsx output file by default. the --input_file and --text_column arguments are needed for this, unless you pass in --no_xlsx_output as seen below. + +## Summarise topics: +python cli_topics.py --task summarise --previous_output_files output/combined_case_notes_col_Case_Note_reference_table.csv output/combined_case_notes_col_Case_Note_unique_topics.csv --model_choice "gemini-flash-lite-latest" --no_xlsx_output + +## Create overall summary: +python cli_topics.py --task overall_summary --previous_output_files output/combined_case_notes_col_Case_Note_unique_topics.csv --model_choice "gemini-flash-lite-latest" --no_xlsx_output + +# All-in-one pipeline + +## Run complete pipeline (extract, deduplicate, summarise): +python cli_topics.py --task all_in_one --input_file example_data/combined_case_notes.csv --text_column "Case Note" --model_choice "gemini-flash-lite-latest" + +""", + ) + + # --- Task Selection --- + task_group = parser.add_argument_group("Task Selection") + task_group.add_argument( + "--task", + choices=[ + "extract", + "discover", + "validate", + "deduplicate", + "summarise", + "overall_summary", + "all_in_one", + ], + default="extract", + help="Task to perform: extract (topic extraction), discover (topic discovery from random sample), validate (validate topics), deduplicate (deduplicate topics), summarise (summarise topics), overall_summary (create overall summary), or all_in_one (complete pipeline).", + ) + + # --- General Arguments --- + general_group = parser.add_argument_group("General Options") + general_group.add_argument( + "--input_file", + nargs="+", + help="Path to the input file(s) to process. Separate multiple files with a space, and use quotes if there are spaces in the file name.", + ) + general_group.add_argument( + "--output_dir", default=OUTPUT_FOLDER, help="Directory for all output files." + ) + general_group.add_argument( + "--input_dir", default=INPUT_FOLDER, help="Directory for all input files." + ) + general_group.add_argument( + "--text_column", + help="Name of the text column to process (required for extract, validate, and all_in_one tasks).", + ) + general_group.add_argument( + "--previous_output_files", + nargs="+", + help="Path(s) to previous output files (reference_table and/or unique_topics files) for validate, deduplicate, summarise, and overall_summary tasks.", + ) + general_group.add_argument( + "--username", default="", help="Username for the session." + ) + general_group.add_argument( + "--save_to_user_folders", + default=SESSION_OUTPUT_FOLDER, + help="Whether to save to user folders or not.", + ) + general_group.add_argument( + "--excel_sheets", + nargs="+", + default=list(), + help="Specific Excel sheet names to process.", + ) + general_group.add_argument( + "--group_by", + help="Column name to group results by.", + ) + + # --- Model Configuration --- + model_group = parser.add_argument_group("Model Configuration") + model_group.add_argument( + "--model_choice", + default=default_model_choice, + help=f"LLM model to use. Default: {default_model_choice}", + ) + model_group.add_argument( + "--model_source", + default=default_model_source, + help=f"Model source (e.g., 'Google', 'AWS', 'Local'). Default: {default_model_source}", + ) + model_group.add_argument( + "--temperature", + type=float, + default=LLM_TEMPERATURE, + help=f"Temperature for LLM generation. Default: {LLM_TEMPERATURE}", + ) + model_group.add_argument( + "--batch_size", + type=int, + default=BATCH_SIZE_DEFAULT, + help=f"Number of responses to submit in a single LLM query. Default: {BATCH_SIZE_DEFAULT}", + ) + model_group.add_argument( + "--max_tokens", + type=int, + default=LLM_MAX_NEW_TOKENS, + help=f"Maximum tokens for LLM generation. Default: {LLM_MAX_NEW_TOKENS}", + ) + model_group.add_argument( + "--google_api_key", + default=GEMINI_API_KEY, + help="Google API key for Gemini models.", + ) + model_group.add_argument( + "--aws_access_key", + default=AWS_ACCESS_KEY, + help="AWS Access Key ID for Bedrock models.", + ) + model_group.add_argument( + "--aws_secret_key", + default=AWS_SECRET_KEY, + help="AWS Secret Access Key for Bedrock models.", + ) + model_group.add_argument( + "--aws_region", + default=AWS_REGION, + help="AWS region for Bedrock models.", + ) + model_group.add_argument( + "--hf_token", + default=HF_TOKEN, + help="Hugging Face token for downloading gated models.", + ) + model_group.add_argument( + "--azure_api_key", + default=AZURE_OPENAI_API_KEY, + help="Azure/OpenAI API key for Azure/OpenAI models.", + ) + model_group.add_argument( + "--azure_endpoint", + default=AZURE_OPENAI_INFERENCE_ENDPOINT, + help="Azure Inference endpoint URL.", + ) + model_group.add_argument( + "--api_url", + default=API_URL, + help=f"Inference server API URL (for local models). Default: {API_URL}", + ) + model_group.add_argument( + "--inference_server_model", + default=CHOSEN_INFERENCE_SERVER_MODEL, + help=f"Inference server model name to use. Default: {CHOSEN_INFERENCE_SERVER_MODEL}", + ) + + # --- Topic Extraction Arguments --- + extract_group = parser.add_argument_group("Topic Extraction Options") + extract_group.add_argument( + "--context", + default="", + help="Context sentence to provide to the LLM for topic extraction.", + ) + extract_group.add_argument( + "--candidate_topics", + help="Path to CSV file with candidate topics for zero-shot extraction.", + ) + extract_group.add_argument( + "--force_zero_shot", + choices=["Yes", "No"], + default="No", + help="Force responses into suggested topics. Default: No", + ) + extract_group.add_argument( + "--force_single_topic", + choices=["Yes", "No"], + default="No", + help="Ask the model to assign responses to only a single topic. Default: No", + ) + extract_group.add_argument( + "--produce_structured_summary", + choices=["Yes", "No"], + default="No", + help="Produce structured summaries using suggested topics as headers. Default: No", + ) + extract_group.add_argument( + "--sentiment", + choices=[ + "Negative or Positive", + "Negative, Neutral, or Positive", + "Do not assess sentiment", + ], + default="Negative or Positive", + help="Response sentiment analysis option. Default: Negative or Positive", + ) + extract_group.add_argument( + "--additional_summary_instructions", + default="", + help="Additional instructions for summary format.", + ) + extract_group.add_argument( + "--enable_batch_deduplication", + default=ENABLE_BATCH_DEDUPLICATION, + help=f"Enable deduplication after each batch during topic extraction (True/False). Default: {ENABLE_BATCH_DEDUPLICATION}", + ) + extract_group.add_argument( + "--maximum_allowed_topics", + type=int, + default=MAXIMUM_ALLOWED_TOPICS, + help=f"Maximum number of topics before triggering LLM-based deduplication. Default: {MAXIMUM_ALLOWED_TOPICS}", + ) + + discover_group = parser.add_argument_group("Topic Discovery Options") + discover_group.add_argument( + "--sample_fraction", + type=float, + default=20, + help="Percentage of responses to sample for topic discovery (1-100). Default: 20", + ) + + # --- Validation Arguments --- + validate_group = parser.add_argument_group("Topic Validation Options") + validate_group.add_argument( + "--additional_validation_issues", + default="", + help="Additional validation issues for the model to consider (bullet-point list).", + ) + validate_group.add_argument( + "--show_previous_table", + choices=["Yes", "No"], + default="Yes", + help="Provide response data to validation process. Default: Yes", + ) + validate_group.add_argument( + "--output_debug_files", + choices=["True", "False"], + default=OUTPUT_DEBUG_FILES, + help=f"Output debug files. Default: {OUTPUT_DEBUG_FILES}", + ) + validate_group.add_argument( + "--max_time_for_loop", + type=int, + default=MAX_TIME_FOR_LOOP, + help=f"Maximum time for validation loop in seconds. Default: {MAX_TIME_FOR_LOOP}", + ) + + # --- Deduplication Arguments --- + dedup_group = parser.add_argument_group("Deduplication Options") + dedup_group.add_argument( + "--method", + choices=["fuzzy", "llm"], + default="fuzzy", + help="Deduplication method: fuzzy (fuzzy matching) or llm (LLM semantic matching). Default: fuzzy", + ) + dedup_group.add_argument( + "--similarity_threshold", + type=int, + default=DEDUPLICATION_THRESHOLD, + help=f"Similarity threshold (0-100) for fuzzy matching. Default: {DEDUPLICATION_THRESHOLD}", + ) + dedup_group.add_argument( + "--merge_sentiment", + choices=["Yes", "No"], + default="No", + help="Merge sentiment values together for duplicate subtopics. Default: No", + ) + dedup_group.add_argument( + "--merge_general_topics", + choices=["Yes", "No"], + default="Yes", + help="Merge general topic values together for duplicate subtopics. Default: Yes", + ) + + # --- Summarisation Arguments --- + summarise_group = parser.add_argument_group("Summarisation Options") + summarise_group.add_argument( + "--summary_format", + choices=["two_paragraph", "single_paragraph"], + default="two_paragraph", + help="Summary format type. Default: two_paragraph", + ) + summarise_group.add_argument( + "--sample_reference_table", + choices=["True", "False"], + default="True", + help="Sample reference table (recommended for large datasets). Default: True", + ) + summarise_group.add_argument( + "--no_of_sampled_summaries", + type=int, + default=DEFAULT_SAMPLED_SUMMARIES, + help=f"Number of summaries per group. Default: {DEFAULT_SAMPLED_SUMMARIES}", + ) + summarise_group.add_argument( + "--random_seed", + type=int, + default=LLM_SEED, + help=f"Random seed for sampling. Default: {LLM_SEED}", + ) + + # --- Output Format Arguments --- + output_group = parser.add_argument_group("Output Format Options") + output_group.add_argument( + "--no_xlsx_output", + dest="create_xlsx_output", + action="store_false", + default=True, + help="Disable creation of Excel (.xlsx) output file. By default, Excel output is created.", + ) + output_group.add_argument( + "--no_topics_csv", + dest="create_topics_csv", + action="store_false", + default=True, + help="Disable creation of suggested topics CSV alongside Excel output. By default, the CSV is created when Excel output is created.", + ) + + # --- Logging Arguments --- + logging_group = parser.add_argument_group("Logging Options") + logging_group.add_argument( + "--save_logs_to_csv", + default=SAVE_LOGS_TO_CSV, + help="Save processing logs to CSV files.", + ) + logging_group.add_argument( + "--save_logs_to_dynamodb", + default=SAVE_LOGS_TO_DYNAMODB, + help="Save processing logs to DynamoDB.", + ) + logging_group.add_argument( + "--usage_logs_folder", + default=USAGE_LOGS_FOLDER, + help="Directory for usage log files.", + ) + logging_group.add_argument( + "--cost_code", + default=( + DIRECT_MODE_DEFAULT_COST_CODE + if DIRECT_MODE_DEFAULT_COST_CODE + else DEFAULT_COST_CODE + ), + help="Cost code for tracking usage.", + ) + + # Parse arguments - either from command line or direct mode + if direct_mode_args: + # Use direct mode arguments + args = argparse.Namespace(**direct_mode_args) + else: + # Parse command line arguments + args = parser.parse_args() + + # --- Handle S3 file downloads --- + # Get AWS credentials from args or fall back to config values + aws_access_key = getattr(args, "aws_access_key", None) or AWS_ACCESS_KEY or "" + aws_secret_key = getattr(args, "aws_secret_key", None) or AWS_SECRET_KEY or "" + aws_region = getattr(args, "aws_region", None) or AWS_REGION or "" + + # Download input files from S3 if needed + # Note: args.input_file is typically a list (from CLI nargs="+" or from direct mode) + # but we also handle pipe-separated strings for compatibility + if args.input_file: + if isinstance(args.input_file, list): + # Handle list of files (may include S3 paths) + downloaded_files = [] + for file_path in args.input_file: + downloaded_path = _download_s3_file_if_needed( + file_path, + aws_access_key=aws_access_key, + aws_secret_key=aws_secret_key, + aws_region=aws_region, + ) + downloaded_files.append(downloaded_path) + args.input_file = downloaded_files + elif isinstance(args.input_file, str): + # Handle pipe-separated string (for direct mode compatibility) + if "|" in args.input_file: + file_list = [f.strip() for f in args.input_file.split("|") if f.strip()] + downloaded_files = [] + for file_path in file_list: + downloaded_path = _download_s3_file_if_needed( + file_path, + aws_access_key=aws_access_key, + aws_secret_key=aws_secret_key, + aws_region=aws_region, + ) + downloaded_files.append(downloaded_path) + args.input_file = downloaded_files + else: + # Single file path + args.input_file = [ + _download_s3_file_if_needed( + args.input_file, + aws_access_key=aws_access_key, + aws_secret_key=aws_secret_key, + aws_region=aws_region, + ) + ] + + # Download candidate topics file from S3 if needed + if args.candidate_topics: + args.candidate_topics = _download_s3_file_if_needed( + args.candidate_topics, + default_filename="downloaded_candidate_topics", + aws_access_key=aws_access_key, + aws_secret_key=aws_secret_key, + aws_region=aws_region, + ) + + # --- Override model_choice with inference_server_model if provided --- + # If inference_server_model is explicitly provided, use it to override model_choice + # This allows users to specify which inference-server model to use + if args.inference_server_model: + # Check if the current model_choice is an inference-server model + model_source = model_name_map.get(args.model_choice, {}).get( + "source", default_model_source + ) + # If model_source is "inference-server" OR if inference_server_model is explicitly provided + # (different from default), use it + if ( + model_source == "inference-server" + or args.inference_server_model != CHOSEN_INFERENCE_SERVER_MODEL + ): + args.model_choice = args.inference_server_model + # Ensure the model is registered in model_name_map with inference-server source + if args.model_choice not in model_name_map: + model_name_map[args.model_choice] = { + "short_name": args.model_choice, + "source": "inference-server", + } + # Also update the model_source to ensure it's set correctly + model_name_map[args.model_choice]["source"] = "inference-server" + + # --- Initial Setup --- + # Convert string boolean variables to boolean + args.save_to_user_folders = convert_string_to_boolean(args.save_to_user_folders) + args.save_logs_to_csv = convert_string_to_boolean(str(args.save_logs_to_csv)) + args.save_logs_to_dynamodb = convert_string_to_boolean( + str(args.save_logs_to_dynamodb) + ) + args.sample_reference_table = args.sample_reference_table == "True" + args.output_debug_files = args.output_debug_files == "True" + + # Update config module values for batch deduplication settings if provided via CLI + # These need to be set before importing/using llm_api_call functions + import tools.config as config_module + + if hasattr(args, "enable_batch_deduplication"): + config_module.ENABLE_BATCH_DEDUPLICATION = args.enable_batch_deduplication + if hasattr(args, "maximum_allowed_topics"): + config_module.MAXIMUM_ALLOWED_TOPICS = args.maximum_allowed_topics + + # Get username and folders + ( + session_hash, + args.output_dir, + _, + args.input_dir, + ) = get_username_and_folders( + username=args.username, + output_folder_textbox=args.output_dir, + input_folder_textbox=args.input_dir, + session_output_folder=args.save_to_user_folders, + ) + + print( + f"Conducting analyses with user {args.username or session_hash}. Outputs will be saved to {args.output_dir}." + ) + + # --- Route to the Correct Workflow Based on Task --- + + # Validate input_file requirement for tasks that need it + if ( + args.task in ["extract", "discover", "validate", "all_in_one"] + and not args.input_file + ): + print(f"Error: --input_file is required for '{args.task}' task.") + return + + if ( + args.task in ["validate", "deduplicate", "summarise", "overall_summary"] + and not args.previous_output_files + ): + print(f"Error: --previous_output_files is required for '{args.task}' task.") + return + + if ( + args.task in ["extract", "discover", "validate", "all_in_one"] + and not args.text_column + ): + print(f"Error: --text_column is required for '{args.task}' task.") + return + + start_time = time.time() + + try: + # Task 1: Extract Topics + if args.task == "extract": + print("--- Starting Topic Extraction Workflow... ---") + + # Load data file + if isinstance(args.input_file, str): + args.input_file = [args.input_file] + + file_data, file_name, total_number_of_batches = load_in_data_file( + file_paths=args.input_file, + in_colnames=[args.text_column], + batch_size=args.batch_size, + in_excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + ) + + # Prepare candidate topics if provided + candidate_topics = None + if args.candidate_topics: + candidate_topics = args.candidate_topics + + # Determine summary format prompt + summary_format_prompt = ( + two_para_summary_format_prompt + if args.summary_format == "two_paragraph" + else single_para_summary_format_prompt + ) + + # Run extraction + ( + display_markdown, + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + topic_extraction_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + deduplication_input_files, + summarisation_input_files, + modifiable_unique_topics_df_state, + modification_input_files, + in_join_files, + missing_df_state, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + output_messages_textbox, + logged_content_df, + ) = wrapper_extract_topics_per_column_value( + grouping_col=args.group_by, + in_data_file=args.input_file, + file_data=file_data, + initial_existing_topics_table=pd.DataFrame(), + initial_existing_reference_df=pd.DataFrame(), + initial_existing_topic_summary_df=pd.DataFrame(), + initial_unique_table_df_display_table_markdown="", + original_file_name=file_name, + total_number_of_batches=total_number_of_batches, + in_api_key=args.google_api_key, + temperature=args.temperature, + chosen_cols=[args.text_column], + model_choice=args.model_choice, + candidate_topics=candidate_topics, + initial_first_loop_state=True, + initial_all_metadata_content_str="", + initial_latest_batch_completed=0, + initial_time_taken=0, + batch_size=args.batch_size, + context_textbox=args.context, + sentiment_checkbox=args.sentiment, + force_zero_shot_radio=args.force_zero_shot, + in_excel_sheets=args.excel_sheets, + force_single_topic_radio=args.force_single_topic, + produce_structured_summary_radio=args.produce_structured_summary, + aws_access_key_textbox=args.aws_access_key, + aws_secret_key_textbox=args.aws_secret_key, + aws_region_textbox=args.aws_region, + hf_api_key_textbox=args.hf_token, + azure_api_key_textbox=args.azure_api_key, + azure_endpoint_textbox=args.azure_endpoint, + output_folder=args.output_dir, + existing_logged_content=list(), + additional_instructions_summary_format=args.additional_summary_instructions, + additional_validation_issues_provided="", + show_previous_table="Yes", + api_url=args.api_url if args.api_url else API_URL, + max_tokens=args.max_tokens, + model_name_map=model_name_map, + max_time_for_loop=99999, + reasoning_suffix="", + CHOSEN_LOCAL_MODEL_TYPE="", + output_debug_files=str(args.output_debug_files), + model=None, + tokenizer=None, + assistant_model=None, + max_rows=999999, + ) + + end_time = time.time() + processing_time = end_time - start_time + + print("\n--- Topic Extraction Complete ---") + print(f"Processing time: {processing_time:.2f} seconds") + print(f"\nOutput files saved to: {args.output_dir}") + if topic_extraction_output_files: + print("Generated Files:", sorted(topic_extraction_output_files)) + + # Write usage log (before Excel creation so it can be included in Excel) + task_usage_log_path = write_usage_log( + session_hash=session_hash, + file_name=file_name, + text_column=args.text_column, + model_choice=args.model_choice, + conversation_metadata="", # conversation_metadata_textbox or "", + input_tokens=input_tokens_num or 0, + output_tokens=output_tokens_num or 0, + number_of_calls=number_of_calls_num or 0, + estimated_time_taken=estimated_time_taken_number or processing_time, + cost_code=args.cost_code, + save_to_csv=args.save_logs_to_csv, + save_to_dynamodb=args.save_logs_to_dynamodb, + output_folder=args.output_dir, + ) + + # Create Excel output if requested + xlsx_files = [] + if args.create_xlsx_output: + print("\nCreating Excel output file...") + try: + xlsx_files, _ = collect_output_csvs_and_create_excel_output( + in_data_files=args.input_file, + chosen_cols=[args.text_column], + reference_data_file_name_textbox=file_name, + in_group_col=args.group_by, + model_choice=args.model_choice, + master_reference_df_state=master_reference_df_state, + master_unique_topics_df_state=master_unique_topics_df_state, + summarised_output_df=pd.DataFrame(), # No summaries yet + missing_df_state=missing_df_state, + excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + usage_logs_location=( + os.path.join(USAGE_LOGS_FOLDER, USAGE_LOG_FILE_NAME) + if args.save_logs_to_csv + else "" + ), + model_name_map=model_name_map, + output_folder=args.output_dir, + structured_summaries=args.produce_structured_summary, + candidate_topics=args.candidate_topics, + create_topics_csv="Yes" if args.create_topics_csv else "No", + ) + # if xlsx_files: + # print(f"Excel output created: {sorted(xlsx_files)}") + except Exception as e: + print(f"Warning: Could not create Excel output: {e}") + + # Upload outputs to S3 if enabled + all_output_files = ( + list(topic_extraction_output_files) + if topic_extraction_output_files + else [] + ) + if xlsx_files: + all_output_files.extend(xlsx_files) + upload_outputs_to_s3_if_enabled( + output_files=all_output_files, + base_file_name=file_name, + session_hash=session_hash, + task_usage_log_path=task_usage_log_path, + prompt_response_log_paths=log_files_output_list_state, + ) + + # Task 1b: Discover topics from random sample + elif args.task == "discover": + print("--- Starting Topic Discovery Workflow... ---") + + if isinstance(args.input_file, str): + args.input_file = [args.input_file] + + file_data, file_name, total_number_of_batches = load_in_data_file( + file_paths=args.input_file, + in_colnames=[args.text_column], + batch_size=args.batch_size, + in_excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + ) + + ( + display_markdown, + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + topic_extraction_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + deduplication_input_files, + summarisation_input_files, + modifiable_unique_topics_df_state, + modification_input_files, + in_join_files, + missing_df_state, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + output_messages_textbox, + logged_content_df, + discovery_csv_path, + ) = discover_topics_from_sample( + grouping_col=args.group_by, + in_data_file=args.input_file, + file_data=file_data, + original_file_name=file_name, + total_number_of_batches=total_number_of_batches, + in_api_key=args.google_api_key, + temperature=args.temperature, + chosen_cols=[args.text_column], + model_choice=args.model_choice, + candidate_topics=args.candidate_topics, + sample_fraction_percent=args.sample_fraction, + random_seed=args.random_seed, + context_textbox=args.context, + sentiment_checkbox=args.sentiment, + force_zero_shot_radio=args.force_zero_shot, + in_excel_sheets=args.excel_sheets, + force_single_topic_radio=args.force_single_topic, + produce_structured_summary_radio=args.produce_structured_summary, + aws_access_key_textbox=args.aws_access_key, + aws_secret_key_textbox=args.aws_secret_key, + aws_region_textbox=args.aws_region, + hf_api_key_textbox=args.hf_token, + azure_api_key_textbox=args.azure_api_key, + azure_endpoint_textbox=args.azure_endpoint, + output_folder=args.output_dir, + initial_table_prompt=initial_table_prompt, + initial_table_system_prompt=initial_table_system_prompt, + add_existing_topics_system_prompt=add_existing_topics_system_prompt, + add_existing_topics_prompt=add_existing_topics_prompt, + number_of_prompts_used=1, + batch_size=args.batch_size, + additional_instructions_summary_format=args.additional_summary_instructions, + additional_validation_issues_provided="", + show_previous_table="Yes", + api_url=args.api_url if args.api_url else API_URL, + max_tokens=args.max_tokens, + model_name_map=model_name_map, + max_time_for_loop=99999, + reasoning_suffix="", + CHOSEN_LOCAL_MODEL_TYPE="", + output_debug_files=str(args.output_debug_files), + model=None, + tokenizer=None, + assistant_model=None, + max_rows=999999, + ) + + end_time = time.time() + processing_time = end_time - start_time + + print("\n--- Topic Discovery Complete ---") + print(output_messages_textbox) + print(f"Processing time: {processing_time:.2f} seconds") + print(f"\nSuggested topics CSV: {discovery_csv_path}") + print(f"\nOutput files saved to: {args.output_dir}") + if topic_extraction_output_files: + print("Generated Files:", sorted(topic_extraction_output_files)) + + task_usage_log_path = write_usage_log( + session_hash=session_hash, + file_name=file_name, + text_column=args.text_column, + model_choice=args.model_choice, + conversation_metadata="", + input_tokens=input_tokens_num or 0, + output_tokens=output_tokens_num or 0, + number_of_calls=number_of_calls_num or 0, + estimated_time_taken=estimated_time_taken_number or processing_time, + cost_code=args.cost_code, + save_to_csv=args.save_logs_to_csv, + save_to_dynamodb=args.save_logs_to_dynamodb, + output_folder=args.output_dir, + ) + + all_output_files = ( + list(topic_extraction_output_files) + if topic_extraction_output_files + else [] + ) + if discovery_csv_path and discovery_csv_path not in all_output_files: + all_output_files.append(discovery_csv_path) + + upload_outputs_to_s3_if_enabled( + output_files=all_output_files, + base_file_name=file_name, + session_hash=session_hash, + task_usage_log_path=task_usage_log_path, + prompt_response_log_paths=log_files_output_list_state, + ) + + # Task 2: Validate Topics + elif args.task == "validate": + print("--- Starting Topic Validation Workflow... ---") + + # Load data file + if isinstance(args.input_file, str): + args.input_file = [args.input_file] + + file_data, file_name, total_number_of_batches = load_in_data_file( + file_paths=args.input_file, + in_colnames=[args.text_column], + batch_size=args.batch_size, + in_excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + ) + + # Load previous output files + ( + reference_df, + topic_summary_df, + latest_batch_completed_no_loop, + deduplication_input_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ) = load_in_previous_data_files(args.previous_output_files) + + # Run validation + ( + display_markdown, + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + validation_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + deduplication_input_files, + summarisation_input_files, + modifiable_unique_topics_df_state, + modification_input_files, + in_join_files, + missing_df_state, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + output_messages_textbox, + logged_content_df, + ) = validate_topics_wrapper( + file_data=file_data, + reference_df=reference_df, + topic_summary_df=topic_summary_df, + file_name=working_data_file_name_textbox, + chosen_cols=[args.text_column], + batch_size=args.batch_size, + model_choice=args.model_choice, + in_api_key=args.google_api_key, + temperature=args.temperature, + max_tokens=args.max_tokens, + azure_api_key_textbox=args.azure_api_key, + azure_endpoint_textbox=args.azure_endpoint, + reasoning_suffix="", + group_name=args.group_by or "All", + produce_structured_summary_radio=args.produce_structured_summary, + force_zero_shot_radio=args.force_zero_shot, + force_single_topic_radio=args.force_single_topic, + context_textbox=args.context, + additional_instructions_summary_format=args.additional_summary_instructions, + output_folder=args.output_dir, + output_debug_files=str(args.output_debug_files), + original_full_file_name=file_name, + additional_validation_issues_provided=args.additional_validation_issues, + max_time_for_loop=args.max_time_for_loop, + in_data_files=args.input_file, + sentiment_checkbox=args.sentiment, + logged_content=None, + show_previous_table=args.show_previous_table, + aws_access_key_textbox=args.aws_access_key, + aws_secret_key_textbox=args.aws_secret_key, + aws_region_textbox=args.aws_region, + api_url=args.api_url if args.api_url else API_URL, + ) + + end_time = time.time() + processing_time = end_time - start_time + + print("\n--- Topic Validation Complete ---") + print(f"Processing time: {processing_time:.2f} seconds") + print(f"\nOutput files saved to: {args.output_dir}") + if validation_output_files: + print("Generated Files:", sorted(validation_output_files)) + + # Write usage log + task_usage_log_path = write_usage_log( + session_hash=session_hash, + file_name=file_name, + text_column=args.text_column, + model_choice=args.model_choice, + conversation_metadata="", # conversation_metadata_textbox or "", + input_tokens=input_tokens_num or 0, + output_tokens=output_tokens_num or 0, + number_of_calls=number_of_calls_num or 0, + estimated_time_taken=estimated_time_taken_number or processing_time, + cost_code=args.cost_code, + save_to_csv=args.save_logs_to_csv, + save_to_dynamodb=args.save_logs_to_dynamodb, + output_folder=args.output_dir, + ) + + # Create Excel output if requested + if args.create_xlsx_output: + print("\nCreating Excel output file...") + try: + xlsx_files, _ = collect_output_csvs_and_create_excel_output( + in_data_files=args.input_file, + chosen_cols=[args.text_column], + reference_data_file_name_textbox=file_name, + in_group_col=args.group_by, + model_choice=args.model_choice, + master_reference_df_state=master_reference_df_state, + master_unique_topics_df_state=master_unique_topics_df_state, + summarised_output_df=pd.DataFrame(), # No summaries yet + missing_df_state=missing_df_state, + excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + usage_logs_location=( + os.path.join(USAGE_LOGS_FOLDER, USAGE_LOG_FILE_NAME) + if args.save_logs_to_csv + else "" + ), + model_name_map=model_name_map, + output_folder=args.output_dir, + structured_summaries=args.produce_structured_summary, + candidate_topics=args.candidate_topics, + create_topics_csv="Yes" if args.create_topics_csv else "No", + ) + # if xlsx_files: + # print(f"Excel output created: {sorted(xlsx_files)}") + except Exception as e: + print(f"Warning: Could not create Excel output: {e}") + + # Task 3: Deduplicate Topics + elif args.task == "deduplicate": + print("--- Starting Topic Deduplication Workflow... ---") + + # Load previous output files + ( + reference_df, + topic_summary_df, + latest_batch_completed_no_loop, + deduplication_input_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ) = load_in_previous_data_files(args.previous_output_files) + + if args.method == "fuzzy": + # Fuzzy matching deduplication + ( + ref_df_after_dedup, + unique_df_after_dedup, + summarisation_input_files, + log_files_output, + summarised_output_markdown, + ) = deduplicate_topics( + reference_df=reference_df, + topic_summary_df=topic_summary_df, + reference_table_file_name=working_data_file_name_textbox, + unique_topics_table_file_name=unique_topics_table_file_name_textbox, + in_excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + merge_sentiment=args.merge_sentiment, + merge_general_topics=args.merge_general_topics, + score_threshold=args.similarity_threshold, + in_data_files=args.input_file if args.input_file else list(), + chosen_cols=[args.text_column] if args.text_column else list(), + output_folder=args.output_dir, + ) + else: + # LLM deduplication + model_source = model_name_map.get(args.model_choice, {}).get( + "source", default_model_source + ) + ( + ref_df_after_dedup, + unique_df_after_dedup, + summarisation_input_files, + log_files_output, + summarised_output_markdown, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + ) = deduplicate_topics_llm( + reference_df=reference_df, + topic_summary_df=topic_summary_df, + reference_table_file_name=working_data_file_name_textbox, + unique_topics_table_file_name=unique_topics_table_file_name_textbox, + model_choice=args.model_choice, + in_api_key=args.google_api_key, + temperature=args.temperature, + model_source=model_source, + local_model=None, + tokenizer=None, + assistant_model=None, + in_excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + merge_sentiment=args.merge_sentiment, + merge_general_topics=args.merge_general_topics, + in_data_files=args.input_file if args.input_file else list(), + chosen_cols=[args.text_column] if args.text_column else list(), + output_folder=args.output_dir, + candidate_topics=( + args.candidate_topics if args.candidate_topics else None + ), + azure_endpoint=args.azure_endpoint, + output_debug_files=str(args.output_debug_files), + api_url=args.api_url if args.api_url else API_URL, + aws_access_key_textbox=( + args.aws_access_key if hasattr(args, "aws_access_key") else "" + ), + aws_secret_key_textbox=( + args.aws_secret_key if hasattr(args, "aws_secret_key") else "" + ), + aws_region_textbox=( + args.aws_region if hasattr(args, "aws_region") else "" + ), + azure_api_key_textbox=( + args.azure_api_key if hasattr(args, "azure_api_key") else "" + ), + model_name_map=model_name_map, + sentiment_checkbox=( + args.sentiment + if hasattr(args, "sentiment") + else "Negative or Positive" + ), + ) + + end_time = time.time() + processing_time = end_time - start_time + + print("\n--- Topic Deduplication Complete ---") + print(f"Processing time: {processing_time:.2f} seconds") + print(f"\nOutput files saved to: {args.output_dir}") + if summarisation_input_files: + print("Generated Files:", sorted(summarisation_input_files)) + + # Initialize task_usage_log_path + task_usage_log_path = None + + # Write usage log (only for LLM deduplication which has token counts) + if args.method == "llm": + # Extract token counts from LLM deduplication result + llm_input_tokens = ( + input_tokens_num if "input_tokens_num" in locals() else 0 + ) + llm_output_tokens = ( + output_tokens_num if "output_tokens_num" in locals() else 0 + ) + llm_calls = ( + number_of_calls_num if "number_of_calls_num" in locals() else 0 + ) + llm_time = ( + estimated_time_taken_number + if "estimated_time_taken_number" in locals() + else processing_time + ) + + task_usage_log_path = write_usage_log( + session_hash=session_hash, + file_name=working_data_file_name_textbox, + text_column=args.text_column if args.text_column else "", + model_choice=args.model_choice, + conversation_metadata="", + input_tokens=llm_input_tokens, + output_tokens=llm_output_tokens, + number_of_calls=llm_calls, + estimated_time_taken=llm_time, + cost_code=args.cost_code, + save_to_csv=args.save_logs_to_csv, + save_to_dynamodb=args.save_logs_to_dynamodb, + output_folder=args.output_dir, + ) + + # Create Excel output if requested + xlsx_files = [] + if args.create_xlsx_output: + print("\nCreating Excel output file...") + try: + # Use the deduplicated dataframes + xlsx_files, _ = collect_output_csvs_and_create_excel_output( + in_data_files=args.input_file if args.input_file else [], + chosen_cols=[args.text_column] if args.text_column else [], + reference_data_file_name_textbox=working_data_file_name_textbox, + in_group_col=args.group_by, + model_choice=args.model_choice, + master_reference_df_state=ref_df_after_dedup, + master_unique_topics_df_state=unique_df_after_dedup, + summarised_output_df=pd.DataFrame(), # No summaries yet + missing_df_state=pd.DataFrame(), + excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + usage_logs_location=( + os.path.join(USAGE_LOGS_FOLDER, USAGE_LOG_FILE_NAME) + if args.save_logs_to_csv + else "" + ), + model_name_map=model_name_map, + output_folder=args.output_dir, + structured_summaries=args.produce_structured_summary, + candidate_topics=args.candidate_topics, + create_topics_csv="Yes" if args.create_topics_csv else "No", + ) + # if xlsx_files: + # print(f"Excel output created: {sorted(xlsx_files)}") + except Exception as e: + print(f"Warning: Could not create Excel output: {e}") + + # Upload outputs to S3 if enabled + all_output_files = ( + list(summarisation_input_files) if summarisation_input_files else [] + ) + if xlsx_files: + all_output_files.extend(xlsx_files) + upload_outputs_to_s3_if_enabled( + output_files=all_output_files, + base_file_name=working_data_file_name_textbox, + session_hash=session_hash, + task_usage_log_path=task_usage_log_path, + prompt_response_log_paths=log_files_output, + ) + elif args.task == "summarise": + print("--- Starting Topic Summarisation Workflow... ---") + + # Load previous output files + ( + reference_df, + topic_summary_df, + latest_batch_completed_no_loop, + deduplication_input_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ) = load_in_previous_data_files(args.previous_output_files) + + # Determine summary format prompt + summary_format_prompt = ( + two_para_summary_format_prompt + if args.summary_format == "two_paragraph" + else single_para_summary_format_prompt + ) + + # Run summarisation + ( + summary_reference_table_sample_state, + master_unique_topics_df_revised_summaries_state, + master_reference_df_revised_summaries_state, + summary_output_files, + summarised_outputs_list, + latest_summary_completed_num, + conversation_metadata_textbox, + summarised_output_markdown, + log_files_output, + overall_summarisation_input_files, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + output_messages_textbox, + logged_content_df, + ) = wrapper_summarise_output_topics_per_group( + grouping_col=args.group_by, + sampled_reference_table_df=reference_df.copy(), # Will be sampled if sample_reference_table=True + topic_summary_df=topic_summary_df, + reference_table_df=reference_df, + model_choice=args.model_choice, + in_api_key=args.google_api_key, + temperature=args.temperature, + reference_data_file_name=working_data_file_name_textbox, + summarised_outputs=list(), + latest_summary_completed=0, + out_metadata_str="", + in_data_files=args.input_file if args.input_file else list(), + in_excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + chosen_cols=[args.text_column] if args.text_column else list(), + log_output_files=list(), + summarise_format_radio=summary_format_prompt, + output_folder=args.output_dir, + context_textbox=args.context, + aws_access_key_textbox=args.aws_access_key, + aws_secret_key_textbox=args.aws_secret_key, + aws_region_textbox=args.aws_region, + model_name_map=model_name_map, + hf_api_key_textbox=args.hf_token, + azure_endpoint_textbox=args.azure_endpoint, + existing_logged_content=list(), + sample_reference_table=args.sample_reference_table, + no_of_sampled_summaries=args.no_of_sampled_summaries, + random_seed=args.random_seed, + api_url=args.api_url if args.api_url else API_URL, + additional_summary_instructions_provided=args.additional_summary_instructions, + output_debug_files=str(args.output_debug_files), + reasoning_suffix="", + local_model=None, + tokenizer=None, + assistant_model=None, + do_summaries="Yes", + ) + + end_time = time.time() + processing_time = end_time - start_time + + print("\n--- Topic Summarisation Complete ---") + print(f"Processing time: {processing_time:.2f} seconds") + print(f"\nOutput files saved to: {args.output_dir}") + if summary_output_files: + print("Generated Files:", sorted(summary_output_files)) + + # Write usage log + task_usage_log_path = write_usage_log( + session_hash=session_hash, + file_name=working_data_file_name_textbox, + text_column=args.text_column if args.text_column else "", + model_choice=args.model_choice, + conversation_metadata="", # conversation_metadata_textbox or "", + input_tokens=input_tokens_num or 0, + output_tokens=output_tokens_num or 0, + number_of_calls=number_of_calls_num or 0, + estimated_time_taken=estimated_time_taken_number or processing_time, + cost_code=args.cost_code, + save_to_csv=args.save_logs_to_csv, + save_to_dynamodb=args.save_logs_to_dynamodb, + output_folder=args.output_dir, + ) + + # Create Excel output if requested + xlsx_files = [] + if args.create_xlsx_output: + print("\nCreating Excel output file...") + try: + xlsx_files, _ = collect_output_csvs_and_create_excel_output( + in_data_files=args.input_file if args.input_file else [], + chosen_cols=[args.text_column] if args.text_column else [], + reference_data_file_name_textbox=working_data_file_name_textbox, + in_group_col=args.group_by, + model_choice=args.model_choice, + master_reference_df_state=master_reference_df_revised_summaries_state, + master_unique_topics_df_state=master_unique_topics_df_revised_summaries_state, + summarised_output_df=pd.DataFrame(), # Summaries are in the revised dataframes + missing_df_state=pd.DataFrame(), + excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + usage_logs_location=( + os.path.join(USAGE_LOGS_FOLDER, USAGE_LOG_FILE_NAME) + if args.save_logs_to_csv + else "" + ), + model_name_map=model_name_map, + output_folder=args.output_dir, + structured_summaries=args.produce_structured_summary, + candidate_topics=args.candidate_topics, + create_topics_csv="Yes" if args.create_topics_csv else "No", + ) + # if xlsx_files: + # print(f"Excel output created: {sorted(xlsx_files)}") + except Exception as e: + print(f"Warning: Could not create Excel output: {e}") + + # Upload outputs to S3 if enabled + all_output_files = ( + list(summary_output_files) if summary_output_files else [] + ) + if xlsx_files: + all_output_files.extend(xlsx_files) + upload_outputs_to_s3_if_enabled( + output_files=all_output_files, + base_file_name=working_data_file_name_textbox, + session_hash=session_hash, + task_usage_log_path=task_usage_log_path, + prompt_response_log_paths=log_files_output, + ) + + # Task 5: Overall Summary + elif args.task == "overall_summary": + print("--- Starting Overall Summary Workflow... ---") + + # Load previous output files + ( + reference_df, + topic_summary_df, + latest_batch_completed_no_loop, + deduplication_input_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ) = load_in_previous_data_files(args.previous_output_files) + + # Run overall summary + ( + overall_summary_output_files, + overall_summarised_output_markdown, + summarised_output_df, + conversation_metadata_textbox, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + output_messages_textbox, + logged_content_df, + ) = overall_summary( + topic_summary_df=topic_summary_df, + model_choice=args.model_choice, + in_api_key=args.google_api_key, + temperature=args.temperature, + reference_data_file_name=working_data_file_name_textbox, + output_folder=args.output_dir, + chosen_cols=[args.text_column] if args.text_column else list(), + context_textbox=args.context, + aws_access_key_textbox=args.aws_access_key, + aws_secret_key_textbox=args.aws_secret_key, + aws_region_textbox=args.aws_region, + model_name_map=model_name_map, + hf_api_key_textbox=args.hf_token, + azure_endpoint_textbox=args.azure_endpoint, + existing_logged_content=list(), + api_url=args.api_url if args.api_url else API_URL, + output_debug_files=str(args.output_debug_files), + log_output_files=list(), + reasoning_suffix="", + local_model=None, + tokenizer=None, + assistant_model=None, + do_summaries="Yes", + ) + + end_time = time.time() + processing_time = end_time - start_time + + print("\n--- Overall Summary Complete ---") + print(f"Processing time: {processing_time:.2f} seconds") + print(f"\nOutput files saved to: {args.output_dir}") + if overall_summary_output_files: + print("Generated Files:", sorted(overall_summary_output_files)) + + # Write usage log + task_usage_log_path = write_usage_log( + session_hash=session_hash, + file_name=working_data_file_name_textbox, + text_column=args.text_column if args.text_column else "", + model_choice=args.model_choice, + conversation_metadata="", # conversation_metadata_textbox or "", + input_tokens=input_tokens_num or 0, + output_tokens=output_tokens_num or 0, + number_of_calls=number_of_calls_num or 0, + estimated_time_taken=estimated_time_taken_number or processing_time, + cost_code=args.cost_code, + save_to_csv=args.save_logs_to_csv, + save_to_dynamodb=args.save_logs_to_dynamodb, + output_folder=args.output_dir, + ) + + # Create Excel output if requested + xlsx_files = [] + if args.create_xlsx_output: + print("\nCreating Excel output file...") + try: + xlsx_files, _ = collect_output_csvs_and_create_excel_output( + in_data_files=args.input_file if args.input_file else [], + chosen_cols=[args.text_column] if args.text_column else [], + reference_data_file_name_textbox=working_data_file_name_textbox, + in_group_col=args.group_by, + model_choice=args.model_choice, + master_reference_df_state=reference_df, # Use original reference_df + master_unique_topics_df_state=topic_summary_df, # Use original topic_summary_df + summarised_output_df=summarised_output_df, + missing_df_state=pd.DataFrame(), + excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + usage_logs_location=( + os.path.join(USAGE_LOGS_FOLDER, USAGE_LOG_FILE_NAME) + if args.save_logs_to_csv + else "" + ), + model_name_map=model_name_map, + output_folder=args.output_dir, + structured_summaries=args.produce_structured_summary, + candidate_topics=args.candidate_topics, + create_topics_csv="Yes" if args.create_topics_csv else "No", + ) + # if xlsx_files: + # print(f"Excel output created: {sorted(xlsx_files)}") + except Exception as e: + print(f"Warning: Could not create Excel output: {e}") + + # Upload outputs to S3 if enabled + all_output_files = ( + list(overall_summary_output_files) + if overall_summary_output_files + else [] + ) + if xlsx_files: + all_output_files.extend(xlsx_files) + upload_outputs_to_s3_if_enabled( + output_files=all_output_files, + base_file_name=working_data_file_name_textbox, + session_hash=session_hash, + task_usage_log_path=task_usage_log_path, + ) + + # Task 6: All-in-One Pipeline + elif args.task == "all_in_one": + print("--- Starting All-in-One Pipeline Workflow... ---") + + # Load data file + if isinstance(args.input_file, str): + args.input_file = [args.input_file] + + file_data, file_name, total_number_of_batches = load_in_data_file( + file_paths=args.input_file, + in_colnames=[args.text_column], + batch_size=args.batch_size, + in_excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + ) + + # Prepare candidate topics if provided + candidate_topics = None + if args.candidate_topics: + candidate_topics = args.candidate_topics + + # Determine summary format prompt + summary_format_prompt = ( + two_para_summary_format_prompt + if args.summary_format == "two_paragraph" + else single_para_summary_format_prompt + ) + + # Run all-in-one pipeline + ( + display_markdown, + master_topic_df_state, + master_unique_topics_df_state, + master_reference_df_state, + topic_extraction_output_files, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + deduplication_input_files, + summarisation_input_files, + modifiable_unique_topics_df_state, + modification_input_files, + in_join_files, + missing_df_state, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + output_messages_textbox, + summary_reference_table_sample_state, + summarised_references_markdown, + master_unique_topics_df_revised_summaries_state, + master_reference_df_revised_summaries_state, + summary_output_files, + summarised_outputs_list, + latest_summary_completed_num, + overall_summarisation_input_files, + overall_summary_output_files, + overall_summarised_output_markdown, + summarised_output_df, + logged_content_df, + ) = all_in_one_pipeline( + grouping_col=args.group_by, + in_data_files=args.input_file, + file_data=file_data, + existing_topics_table=pd.DataFrame(), + existing_reference_df=pd.DataFrame(), + existing_topic_summary_df=pd.DataFrame(), + unique_table_df_display_table_markdown="", + original_file_name=file_name, + total_number_of_batches=total_number_of_batches, + in_api_key=args.google_api_key, + temperature=args.temperature, + chosen_cols=[args.text_column], + model_choice=args.model_choice, + candidate_topics=candidate_topics, + first_loop_state=True, + conversation_metadata_text="", + latest_batch_completed=0, + time_taken_so_far=0, + initial_table_prompt_text=initial_table_prompt, + initial_table_system_prompt_text=initial_table_system_prompt, + add_existing_topics_system_prompt_text=add_existing_topics_system_prompt, + add_existing_topics_prompt_text=add_existing_topics_prompt, + number_of_prompts_used=1, + batch_size=args.batch_size, + context_text=args.context, + sentiment_choice=args.sentiment, + force_zero_shot_choice=args.force_zero_shot, + in_excel_sheets=args.excel_sheets, + force_single_topic_choice=args.force_single_topic, + produce_structures_summary_choice=args.produce_structured_summary, + aws_access_key_text=args.aws_access_key, + aws_secret_key_text=args.aws_secret_key, + aws_region_text=args.aws_region, + hf_api_key_text=args.hf_token, + azure_api_key_text=args.azure_api_key, + azure_endpoint_text=args.azure_endpoint, + output_folder=args.output_dir, + merge_sentiment=args.merge_sentiment, + merge_general_topics=args.merge_general_topics, + score_threshold=args.similarity_threshold, + summarise_format=summary_format_prompt, + random_seed=args.random_seed, + log_files_output_list_state=list(), + model_name_map_state=model_name_map, + usage_logs_location=( + os.path.join(USAGE_LOGS_FOLDER, USAGE_LOG_FILE_NAME) + if args.save_logs_to_csv + else "" + ), + existing_logged_content=list(), + additional_instructions_summary_format=args.additional_summary_instructions, + additional_validation_issues_provided="", + show_previous_table="Yes", + sample_reference_table_checkbox=args.sample_reference_table, + api_url=args.api_url if args.api_url else API_URL, + output_debug_files=str(args.output_debug_files), + model=None, + tokenizer=None, + assistant_model=None, + max_rows=999999, + ) + + end_time = time.time() + processing_time = end_time - start_time + + print("\n--- All-in-One Pipeline Complete ---") + print(f"Processing time: {processing_time:.2f} seconds") + print(f"\nOutput files saved to: {args.output_dir}") + if overall_summary_output_files: + print("Generated Files:", sorted(overall_summary_output_files)) + + # Write usage log + task_usage_log_path = write_usage_log( + session_hash=session_hash, + file_name=file_name, + text_column=args.text_column, + model_choice=args.model_choice, + conversation_metadata="", # conversation_metadata_textbox or "", + input_tokens=input_tokens_num or 0, + output_tokens=output_tokens_num or 0, + number_of_calls=number_of_calls_num or 0, + estimated_time_taken=estimated_time_taken_number or processing_time, + cost_code=args.cost_code, + save_to_csv=args.save_logs_to_csv, + save_to_dynamodb=args.save_logs_to_dynamodb, + output_folder=args.output_dir, + ) + + # Create Excel output if requested + xlsx_files = [] + if args.create_xlsx_output: + print("\nCreating Excel output file...") + try: + xlsx_files, _ = collect_output_csvs_and_create_excel_output( + in_data_files=args.input_file, + chosen_cols=[args.text_column], + reference_data_file_name_textbox=file_name, + in_group_col=args.group_by, + model_choice=args.model_choice, + master_reference_df_state=master_reference_df_revised_summaries_state, + master_unique_topics_df_state=master_unique_topics_df_revised_summaries_state, + summarised_output_df=summarised_output_df, + missing_df_state=missing_df_state, + excel_sheets=args.excel_sheets[0] if args.excel_sheets else "", + usage_logs_location=( + os.path.join(USAGE_LOGS_FOLDER, USAGE_LOG_FILE_NAME) + if args.save_logs_to_csv + else "" + ), + model_name_map=model_name_map, + output_folder=args.output_dir, + structured_summaries=args.produce_structured_summary, + candidate_topics=args.candidate_topics, + create_topics_csv="Yes" if args.create_topics_csv else "No", + ) + # if xlsx_files: + # print(f"Excel output created: {sorted(xlsx_files)}") + except Exception as e: + print(f"Warning: Could not create Excel output: {e}") + + # Upload outputs to S3 if enabled + # Collect all output files from the pipeline + all_s3_output_files = [] + if topic_extraction_output_files and not DIRECT_MODE_S3_UPLOAD_ONLY_XLSX: + all_s3_output_files.extend(topic_extraction_output_files) + if overall_summary_output_files and not DIRECT_MODE_S3_UPLOAD_ONLY_XLSX: + all_s3_output_files.extend(overall_summary_output_files) + if xlsx_files: + all_s3_output_files.extend(xlsx_files) + upload_outputs_to_s3_if_enabled( + output_files=all_s3_output_files, + base_file_name=file_name, + session_hash=session_hash, + task_usage_log_path=task_usage_log_path, + prompt_response_log_paths=log_files_output_list_state, + ) + + else: + print(f"Error: Invalid task '{args.task}'.") + print( + "Valid options: 'extract', 'validate', 'deduplicate', 'summarise', 'overall_summary', or 'all_in_one'" + ) + + except ColumnNotFoundError as e: + print(str(e)) + sys.exit(1) + except Exception as e: + print(f"\nAn error occurred during the workflow: {e}") + import traceback + + traceback.print_exc() + + +if __name__ == "__main__": + main() diff --git a/entrypoint.sh b/entrypoint.sh new file mode 100644 index 0000000000000000000000000000000000000000..ebec21e91a7b877a21d6d7ae1fd292fd021f48d0 --- /dev/null +++ b/entrypoint.sh @@ -0,0 +1,18 @@ +#!/bin/sh + +# Exit immediately if a command exits with a non-zero status. +set -e + +echo "Starting in APP_MODE: $APP_MODE" + +# --- Start the app based on mode --- + +if [ "$APP_MODE" = "lambda" ]; then + echo "Starting in Lambda mode..." + # The CMD from Dockerfile will be passed as "$@" + exec python -m awslambdaric "$@" +else + echo "Starting in Gradio mode..." + exec python app.py +fi + diff --git a/example_data/case_note_headers_specific.csv b/example_data/case_note_headers_specific.csv new file mode 100644 index 0000000000000000000000000000000000000000..b8898660c2da3a5838ba4252b83d70968e0cf6cd --- /dev/null +++ b/example_data/case_note_headers_specific.csv @@ -0,0 +1,7 @@ +General Topic,Subtopic +Mental health,Anger +Mental health,Social issues +Physical health,General +Physical health,Substance misuse +Behaviour at school,Behaviour at school +Trends over time,Trends over time diff --git a/example_data/combined_case_notes.csv b/example_data/combined_case_notes.csv new file mode 100644 index 0000000000000000000000000000000000000000..39a787296303521a49277d866f603c5faf7d6885 --- /dev/null +++ b/example_data/combined_case_notes.csv @@ -0,0 +1,19 @@ +Date,Social Worker,Client,Case Note +"January 3, 2023",Jane Smith,Alex D.,"Met with Alex at school following reports of increased absences and declining grades. Alex appeared sullen and avoided eye contact. When prompted about school, Alex expressed feelings of isolation and stated, ""No one gets me."" Scheduled a follow-up meeting to further explore these feelings." +"January 17, 2023",Jane Smith,Alex D.,"Met with Alex at the community center. Alex displayed sudden outbursts of anger when discussing home life, particularly in relation to a new stepfather. Alex mentioned occasional substance use, but did not specify which substances. Recommended a comprehensive assessment." +"February 5, 2023",Jane Smith,Alex D.,Home visit conducted. Alex's mother reported frequent arguments at home. She expressed concerns about Alex's new group of friends and late-night outings. Noted potential signs of substance abuse. Suggested family counseling. +"February 21, 2023",Jane Smith,Alex D.,"Met with Alex alone at my office. Alex appeared more agitated than in previous meetings. There were visible signs of self-harm on Alex's arms. When questioned, Alex became defensive. Immediate referral made to a mental health professional." +"March 10, 2023",Jane Smith,Alex D.,Attended joint session with Alex and a therapist. Alex shared feelings of hopelessness and admitted to occasional thoughts of self-harm. Therapist recommended a comprehensive mental health evaluation and ongoing therapy. +"March 25, 2023",Jane Smith,Alex D.,"Received a call from Alex's school about a physical altercation with another student. Met with Alex, who displayed high levels of frustration and admitted to the use of alcohol. Discussed the importance of seeking help and finding positive coping mechanisms. Recommended enrollment in an anger management program." +"April 15, 2023",Jane Smith,Alex D.,Met with Alex and mother to discuss progress. Alex's mother expressed concerns about Alex's increasing aggression at home. Alex acknowledged the issues but blamed others for provoking the behavior. It was decided that a more intensive intervention may be needed. +"April 30, 2023",Jane Smith,Alex D.,"Met with Alex and a psychiatrist. Psychiatrist diagnosed Alex with Oppositional Defiant Disorder (ODD) and co-morbid substance use disorder. A treatment plan was discussed, including medication, therapy, and family counseling." +"May 20, 2023",Jane Smith,Alex D.,"Met with Alex to discuss progress. Alex has started attending group therapy and has shown slight improvements in behavior. Still, concerns remain about substance use. Discussed potential for a short-term residential treatment program." +"January 3, 2023",Jane Smith,Jamie L.,"Met with Jamie at school after receiving reports of consistent tardiness and decreased participation in class. Jamie appeared withdrawn and exhibited signs of sadness. When asked about feelings, Jamie expressed feeling ""empty"" and ""hopeless"" at times. Scheduled a follow-up meeting to further explore these feelings." +"January 17, 2023",Jane Smith,Jamie L.,"Met with Jamie at the community center. Jamie shared feelings of low self-worth, mentioning that it's hard to find motivation for daily tasks. Discussed potential triggers and learned about recent family financial struggles. Recommended counseling and possible group therapy for peer support." +"February 5, 2023",Jane Smith,Jamie L.,Home visit conducted. Jamie's parents shared concerns about Jamie's increasing withdrawal from family activities and lack of interest in hobbies. Parents mentioned that Jamie spends a lot of time alone in the room. Suggested family therapy to open communication channels. +"February 21, 2023",Jane Smith,Jamie L.,Met with Jamie in my office. Jamie opened up about feelings of isolation and mentioned difficulty sleeping. No signs of self-harm or suicidal ideation were noted. Recommended a comprehensive mental health assessment to better understand the depth of the depression. +"March 10, 2023",Jane Smith,Jamie L.,"Attended a joint session with Jamie and a therapist. The therapist noted signs of moderate depression. Together, we discussed coping strategies and potential interventions. Jamie showed interest in art therapy." +"March 25, 2023",Jane Smith,Jamie L.,"Received feedback from Jamie's school that academic performance has slightly improved. However, social interactions remain limited. Encouraged Jamie to join school clubs or groups to foster connection." +"April 15, 2023",Jane Smith,Jamie L.,"Met with Jamie and parents to discuss progress. Parents have observed slight improvements in mood on some days, but overall, Jamie still appears to struggle. It was decided to explore medication as a potential aid alongside therapy." +"April 30, 2023",Jane Smith,Jamie L.,Met with Jamie and a psychiatrist. The psychiatrist diagnosed Jamie with Major Depressive Disorder (MDD) and suggested considering antidepressant medication. Discussed the potential benefits and side effects. Jamie and parents will think it over. +"May 20, 2023",Jane Smith,Jamie L.,"Jamie has started on a low dose of an antidepressant. Initial feedback is positive, with some improvement in mood and energy levels. Will continue monitoring and adjusting as necessary." diff --git a/example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_structured_summaries.xlsx b/example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_structured_summaries.xlsx new file mode 100644 index 0000000000000000000000000000000000000000..5cc5e7efb385fe1c19152a90386f791716254017 --- /dev/null +++ b/example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_structured_summaries.xlsx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:322a081b29d4fb40ccae7d47aa74fda772a002eda576ddc98d6acc86366cff11 +size 13502 diff --git a/example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_topic_analysis.xlsx b/example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_topic_analysis.xlsx new file mode 100644 index 0000000000000000000000000000000000000000..45c8a5d690d0e3d306c93e24acaea238d8ad22d1 --- /dev/null +++ b/example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_topic_analysis.xlsx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b3dcc1ea155169c23d913043b1ad87da2f2912be36d9fb1521c72ee05b8dcf36 +size 25299 diff --git a/example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_topic_analysis_grouped.xlsx b/example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_topic_analysis_grouped.xlsx new file mode 100644 index 0000000000000000000000000000000000000000..e7e8703c9359c6da0c3a1d167ba905028938fda0 --- /dev/null +++ b/example_data/combined_case_notes_col_Case_Note_Gemma_3_4B_topic_analysis_grouped.xlsx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8e1eaede9af75b6ab695b1cfc6c01ec875abf14521249ba7257bd4bb0afd7ee8 +size 28673 diff --git a/example_data/dummy_consultation_r_col_Response_text_Gemma_3_4B_topic_analysis.xlsx b/example_data/dummy_consultation_r_col_Response_text_Gemma_3_4B_topic_analysis.xlsx new file mode 100644 index 0000000000000000000000000000000000000000..545dd50c9566be24dfcb3d932c011a691316cc69 --- /dev/null +++ b/example_data/dummy_consultation_r_col_Response_text_Gemma_3_4B_topic_analysis.xlsx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:30947f1355eacc74c92d09b766e8e3d71092b9a240e7f8acd381874b7d7ebcb3 +size 24673 diff --git a/example_data/dummy_consultation_r_col_Response_text_Gemma_3_4B_topic_analysis_zero_shot.xlsx b/example_data/dummy_consultation_r_col_Response_text_Gemma_3_4B_topic_analysis_zero_shot.xlsx new file mode 100644 index 0000000000000000000000000000000000000000..ee2f3d94d19dfd8c3bcd9546c3561b4eb7f978a5 --- /dev/null +++ b/example_data/dummy_consultation_r_col_Response_text_Gemma_3_4B_topic_analysis_zero_shot.xlsx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a5f0e36143d8362391e3b11d1c20e3a2a1b7536b8f0c972e3d44644eb9ae4e82 +size 27592 diff --git a/example_data/dummy_consultation_response.csv b/example_data/dummy_consultation_response.csv new file mode 100644 index 0000000000000000000000000000000000000000..015ed63a9d237ef84c7918db771802a9addbd27a --- /dev/null +++ b/example_data/dummy_consultation_response.csv @@ -0,0 +1,31 @@ +Response Reference,Object to or Support application,Response text +R1,Object,I strongly object to the proposed five-storey apartment block on Main Street. It is completely out of keeping with the existing character of the area and will overshadow the existing buildings. +R2,Support,"I fully support the proposed development. The town needs more housing, and this development will provide much-needed homes." +R3,Object,The proposed development is too tall and will have a negative impact on the views from the surrounding area. +R4,Object,The loss of the well-loved cafe will be a great loss to the community. +R5,Support,The development will bring much-needed investment to the area and create jobs. +R6,Object,The increased traffic generated by the development will cause congestion on Main Street. +R7,Support,The development will provide much-needed affordable housing. +R8,Object,The development will have a negative impact on the local environment. +R9,Support,The development will improve the appearance of Main Street. +R10,Object,The development will overshadow the existing buildings and make them feel cramped. +R11,Support,The development will provide much-needed amenities for the local community. +R12,Object,The development will have a negative impact on the local wildlife. +R13,Support,The development will help to revitalise the town centre. +R14,Object,The development will increase noise pollution in the area. +R15,Support,The development will provide much-needed parking spaces. +R16,Object,The development will have a negative impact on the local businesses. +R17,Support,The development will provide much-needed green space. +R18,Object,The development will have a negative impact on the local heritage. +R19,Support,The development will provide much-needed facilities for young people. +R20,Object,The development will have a negative impact on the local schools. +R21,Support,The development will provide much-needed social housing. +R22,Object,The development will have a negative impact on the local infrastructure. +R23,Support,The development will provide much-needed jobs for local people. +R24,Object,The development will have a negative impact on the local economy. +R25,Support,The development will provide much-needed community facilities. +R26,Object,The development will have a negative impact on the local amenities. +R27,Support,The development will provide much-needed housing for young people. +R28,Object,The development will have a negative impact on the local character. +R29,Support,The development will provide much-needed housing for families. +R30,Object,The development will have a negative impact on the local quality of life. \ No newline at end of file diff --git a/example_data/dummy_consultation_response_themes.csv b/example_data/dummy_consultation_response_themes.csv new file mode 100644 index 0000000000000000000000000000000000000000..dc4ae59fdc9ac494fd9e21a0bf5febd4c1322d36 --- /dev/null +++ b/example_data/dummy_consultation_response_themes.csv @@ -0,0 +1,26 @@ +topics +Need for family housing +Impact on the character of the area +Amenities for the local community +Revitalisation of the town centre +Impact on local wildlife +Parking +Impact on local businesses +Green space +Noise pollution +Impact on local heritage +Facilities for young people +Impact on local schools +Impact on views +Loss of cafe +Investment and job creation +Traffic congestion +Affordable housing +Impact on the local environment +Improvement of main street +Impact on local infrastructure +Investment and job creation +Impact on local schools +Provision of community facilities +Impact on local heritage +Impact on quality of life diff --git a/intros/intro.txt b/intros/intro.txt new file mode 100644 index 0000000000000000000000000000000000000000..d2503cbcfb843f9a8c4a90fa9caa80c1dc30304c --- /dev/null +++ b/intros/intro.txt @@ -0,0 +1,7 @@ +# Create thematic summaries from your data + +Extract topics and summarise open text using Large Language Models (LLMs). The model will loop through all text rows to find the most relevant general topics and subtopics, and provide a short summary of each. If you have specific topics in mind, you can enter them in 'Provide a list of specific topics' below. + +NOTE: LLMs are not 100% accurate and may produce biased or incorrect responses. All files downloaded from this app **need to be checked by a human** before they are used in further outputs. + +Unsure of how to use this app? Try an example by clicking on one of the example datasets below to see typical outputs the app can produce. There is also a user guide provided alongside this app - please ask your system administrator if you do not have access. \ No newline at end of file diff --git a/lambda_entrypoint.py b/lambda_entrypoint.py new file mode 100644 index 0000000000000000000000000000000000000000..ea657ceb1255972a267534544a9782eb11f310fd --- /dev/null +++ b/lambda_entrypoint.py @@ -0,0 +1,464 @@ +import json +import os + +import boto3 +from dotenv import load_dotenv + +# Import the main function from your CLI script +from cli_topics import main as cli_main +from tools.config import ( + AWS_REGION, + BATCH_SIZE_DEFAULT, + DEDUPLICATION_THRESHOLD, + DEFAULT_COST_CODE, + DEFAULT_SAMPLED_SUMMARIES, + LLM_MAX_NEW_TOKENS, + LLM_SEED, + LLM_TEMPERATURE, + OUTPUT_DEBUG_FILES, + SAVE_LOGS_TO_CSV, + SAVE_LOGS_TO_DYNAMODB, + SESSION_OUTPUT_FOLDER, + USAGE_LOGS_FOLDER, + convert_string_to_boolean, +) + + +def _get_env_list(env_var_name: str | list[str] | None) -> list[str]: + """Parses a comma-separated environment variable into a list of strings.""" + if isinstance(env_var_name, list): + return env_var_name + if env_var_name is None: + return [] + + # Handle string input + value = str(env_var_name).strip() + if not value or value == "[]": + return [] + + # Remove brackets if present (e.g., "[item1, item2]" -> "item1, item2") + if value.startswith("[") and value.endswith("]"): + value = value[1:-1] + + # Remove quotes and split by comma + value = value.replace('"', "").replace("'", "") + if not value: + return [] + + # Split by comma and filter out any empty strings + return [s.strip() for s in value.split(",") if s.strip()] + + +print("Lambda entrypoint loading...") + +# Initialize S3 client outside the handler for connection reuse +s3_client = boto3.client("s3", region_name=os.getenv("AWS_REGION", AWS_REGION)) +print("S3 client initialised") + +# Lambda's only writable directory is /tmp. Ensure that all temporary files are stored in this directory. +TMP_DIR = "/tmp" +INPUT_DIR = os.path.join(TMP_DIR, "input") +OUTPUT_DIR = os.path.join(TMP_DIR, "output") +os.environ["GRADIO_TEMP_DIR"] = os.path.join(TMP_DIR, "gradio_tmp") +os.environ["MPLCONFIGDIR"] = os.path.join(TMP_DIR, "matplotlib_cache") +os.environ["FEEDBACK_LOGS_FOLDER"] = os.path.join(TMP_DIR, "feedback") +os.environ["ACCESS_LOGS_FOLDER"] = os.path.join(TMP_DIR, "logs") +os.environ["USAGE_LOGS_FOLDER"] = os.path.join(TMP_DIR, "usage") + +# Define compatible file types for processing +COMPATIBLE_FILE_TYPES = { + ".csv", + ".xlsx", + ".xls", + ".parquet", +} + + +def download_file_from_s3(bucket_name, key, download_path): + """Download a file from S3 to the local filesystem.""" + try: + s3_client.download_file(bucket_name, key, download_path) + print(f"Successfully downloaded file from S3 to {download_path}") + except Exception as e: + print(f"Error downloading from S3: {e}") + raise + + +def upload_directory_to_s3(local_directory, bucket_name, s3_prefix): + """Upload all files from a local directory to an S3 prefix.""" + for root, _, files in os.walk(local_directory): + for file_name in files: + local_file_path = os.path.join(root, file_name) + # Create a relative path to maintain directory structure if needed + relative_path = os.path.relpath(local_file_path, local_directory) + output_key = os.path.join(s3_prefix, relative_path).replace("\\", "/") + + try: + s3_client.upload_file(local_file_path, bucket_name, output_key) + print(f"Successfully uploaded file to S3: {local_file_path}") + except Exception as e: + print(f"Error uploading to S3: {e}") + raise + + +def lambda_handler(event, context): + print(f"Received event: {json.dumps(event)}") + + # 1. Setup temporary directories + os.makedirs(INPUT_DIR, exist_ok=True) + os.makedirs(OUTPUT_DIR, exist_ok=True) + + # 2. Extract information from the event + # Assumes the event is triggered by S3 and may contain an 'arguments' payload + try: + record = event["Records"][0] + bucket_name = record["s3"]["bucket"]["name"] + input_key = record["s3"]["object"]["key"] + + # The user metadata can be used to pass arguments + # This is more robust than embedding them in the main event body + try: + response = s3_client.head_object(Bucket=bucket_name, Key=input_key) + metadata = response.get("Metadata", dict()) + print(f"S3 object metadata: {metadata}") + + # Arguments can be passed as a JSON string in metadata + arguments_str = metadata.get("arguments", "{}") + print(f"Arguments string from metadata: '{arguments_str}'") + + if arguments_str and arguments_str != "{}": + arguments = json.loads(arguments_str) + print(f"Successfully parsed arguments from metadata: {arguments}") + else: + arguments = dict() + print("No arguments found in metadata, using empty dictionary") + except Exception as e: + print(f"Warning: Could not parse metadata arguments: {e}") + print("Using empty arguments dictionary") + arguments = dict() + + except (KeyError, IndexError) as e: + print( + f"Could not parse S3 event record: {e}. Checking for direct invocation payload." + ) + # Fallback for direct invocation (e.g., from Step Functions or manual test) + bucket_name = event.get("bucket_name") + input_key = event.get("input_key") + arguments = event.get("arguments", dict()) + if not all([bucket_name, input_key]): + raise ValueError( + "Missing 'bucket_name' or 'input_key' in direct invocation event." + ) + + # Log file type information + file_extension = os.path.splitext(input_key)[1].lower() + print(f"Detected file extension: '{file_extension}'") + + # 3. Download the main input file + input_file_path = os.path.join(INPUT_DIR, os.path.basename(input_key)) + download_file_from_s3(bucket_name, input_key, input_file_path) + + # 3.1. Validate file type compatibility + is_env_file = input_key.lower().endswith(".env") + + if not is_env_file and file_extension not in COMPATIBLE_FILE_TYPES: + error_message = f"File type '{file_extension}' is not supported for processing. Compatible file types are: {', '.join(sorted(COMPATIBLE_FILE_TYPES))}" + print(f"ERROR: {error_message}") + print(f"File was not processed due to unsupported file type: {file_extension}") + return { + "statusCode": 400, + "body": json.dumps( + { + "error": "Unsupported file type", + "message": error_message, + "supported_types": list(COMPATIBLE_FILE_TYPES), + "received_type": file_extension, + "file_processed": False, + } + ), + } + + print(f"File type '{file_extension}' is compatible for processing") + if is_env_file: + print("Processing .env file for configuration") + else: + print(f"Processing {file_extension} file for topic modelling") + + # 3.5. Check if the downloaded file is a .env file and handle accordingly + actual_input_file_path = input_file_path + if input_key.lower().endswith(".env"): + print("Detected .env file, loading environment variables...") + + # Load environment variables from the .env file + print(f"Loading .env file from: {input_file_path}") + + # Check if file exists and is readable + if os.path.exists(input_file_path): + print(".env file exists and is readable") + with open(input_file_path, "r") as f: + content = f.read() + print(f".env file content preview: {content[:200]}...") + else: + print(f"ERROR: .env file does not exist at {input_file_path}") + + load_dotenv(input_file_path, override=True) + print("Environment variables loaded from .env file") + + # Extract the actual input file path from environment variables + env_input_file = os.getenv("INPUT_FILE") + + if env_input_file: + print(f"Found input file path in environment: {env_input_file}") + + # If the path is an S3 path, download it + if env_input_file.startswith("s3://"): + # Parse S3 path: s3://bucket/key + s3_path_parts = env_input_file[5:].split("/", 1) + if len(s3_path_parts) == 2: + env_bucket = s3_path_parts[0] + env_key = s3_path_parts[1] + actual_input_file_path = os.path.join( + INPUT_DIR, os.path.basename(env_key) + ) + print( + f"Downloading actual input file from s3://{env_bucket}/{env_key}" + ) + download_file_from_s3(env_bucket, env_key, actual_input_file_path) + else: + print("Warning: Invalid S3 path format in environment variable") + actual_input_file_path = input_file_path + else: + # Assume it's a local path or relative path + actual_input_file_path = env_input_file + print( + f"Using input file path from environment: {actual_input_file_path}" + ) + else: + print("Warning: No input file path found in environment variables") + # Fall back to using the .env file itself (though this might not be what we want) + actual_input_file_path = input_file_path + else: + print("File is not a .env file, proceeding with normal processing") + + # 4. Prepare arguments for the CLI function + # This dictionary should mirror the arguments that cli_topics.main() expects via direct_mode_args + + cli_args = { + # Task Selection + "task": arguments.get("task", os.getenv("DIRECT_MODE_TASK", "extract")), + # General Arguments + "input_file": [actual_input_file_path] if actual_input_file_path else None, + "output_dir": arguments.get( + "output_dir", os.getenv("DIRECT_MODE_OUTPUT_DIR", OUTPUT_DIR) + ), + "input_dir": arguments.get("input_dir", INPUT_DIR), + "text_column": arguments.get( + "text_column", os.getenv("DIRECT_MODE_TEXT_COLUMN", "") + ), + "previous_output_files": _get_env_list( + arguments.get( + "previous_output_files", + os.getenv("DIRECT_MODE_PREVIOUS_OUTPUT_FILES", list()), + ) + ), + "username": arguments.get("username", os.getenv("DIRECT_MODE_USERNAME", "")), + "save_to_user_folders": convert_string_to_boolean( + arguments.get( + "save_to_user_folders", + os.getenv("SESSION_OUTPUT_FOLDER", str(SESSION_OUTPUT_FOLDER)), + ) + ), + "excel_sheets": _get_env_list( + arguments.get("excel_sheets", os.getenv("DIRECT_MODE_EXCEL_SHEETS", list())) + ), + "group_by": arguments.get("group_by", os.getenv("DIRECT_MODE_GROUP_BY", "")), + # Model Configuration + "model_choice": arguments.get( + "model_choice", os.getenv("DIRECT_MODE_MODEL_CHOICE", "") + ), + "temperature": float( + arguments.get( + "temperature", + os.getenv("DIRECT_MODE_TEMPERATURE", str(LLM_TEMPERATURE)), + ) + ), + "batch_size": int( + arguments.get( + "batch_size", + os.getenv("DIRECT_MODE_BATCH_SIZE", str(BATCH_SIZE_DEFAULT)), + ) + ), + "max_tokens": int( + arguments.get( + "max_tokens", + os.getenv("DIRECT_MODE_MAX_TOKENS", str(LLM_MAX_NEW_TOKENS)), + ) + ), + "google_api_key": arguments.get( + "google_api_key", os.getenv("GEMINI_API_KEY", "") + ), + "aws_access_key": None, # Use IAM Role instead of keys + "aws_secret_key": None, # Use IAM Role instead of keys + "aws_region": os.getenv("AWS_REGION", AWS_REGION), + "hf_token": arguments.get("hf_token", os.getenv("HF_TOKEN", "")), + "azure_api_key": arguments.get( + "azure_api_key", os.getenv("AZURE_OPENAI_API_KEY", "") + ), + "azure_endpoint": arguments.get( + "azure_endpoint", os.getenv("AZURE_OPENAI_INFERENCE_ENDPOINT", "") + ), + "api_url": arguments.get("api_url", os.getenv("API_URL", "")), + "inference_server_model": arguments.get( + "inference_server_model", os.getenv("CHOSEN_INFERENCE_SERVER_MODEL", "") + ), + # Topic Extraction Arguments + "context": arguments.get("context", os.getenv("DIRECT_MODE_CONTEXT", "")), + "candidate_topics": arguments.get( + "candidate_topics", os.getenv("DIRECT_MODE_CANDIDATE_TOPICS", "") + ), + "force_zero_shot": arguments.get( + "force_zero_shot", os.getenv("DIRECT_MODE_FORCE_ZERO_SHOT", "No") + ), + "force_single_topic": arguments.get( + "force_single_topic", os.getenv("DIRECT_MODE_FORCE_SINGLE_TOPIC", "No") + ), + "produce_structured_summary": arguments.get( + "produce_structured_summary", + os.getenv("DIRECT_MODE_PRODUCE_STRUCTURED_SUMMARY", "No"), + ), + "sentiment": arguments.get( + "sentiment", os.getenv("DIRECT_MODE_SENTIMENT", "Negative or Positive") + ), + "additional_summary_instructions": arguments.get( + "additional_summary_instructions", + os.getenv("DIRECT_MODE_ADDITIONAL_SUMMARY_INSTRUCTIONS", ""), + ), + # Validation Arguments + "additional_validation_issues": arguments.get( + "additional_validation_issues", + os.getenv("DIRECT_MODE_ADDITIONAL_VALIDATION_ISSUES", ""), + ), + "show_previous_table": arguments.get( + "show_previous_table", os.getenv("DIRECT_MODE_SHOW_PREVIOUS_TABLE", "Yes") + ), + "output_debug_files": arguments.get( + "output_debug_files", str(OUTPUT_DEBUG_FILES) + ), + "max_time_for_loop": int( + arguments.get("max_time_for_loop", os.getenv("MAX_TIME_FOR_LOOP", "99999")) + ), + # Deduplication Arguments + "method": arguments.get( + "method", os.getenv("DIRECT_MODE_DEDUPLICATION_METHOD", "fuzzy") + ), + "similarity_threshold": int( + arguments.get( + "similarity_threshold", + os.getenv("DEDUPLICATION_THRESHOLD", DEDUPLICATION_THRESHOLD), + ) + ), + "merge_sentiment": arguments.get( + "merge_sentiment", os.getenv("DIRECT_MODE_MERGE_SENTIMENT", "No") + ), + "merge_general_topics": arguments.get( + "merge_general_topics", os.getenv("DIRECT_MODE_MERGE_GENERAL_TOPICS", "Yes") + ), + # Summarisation Arguments + "summary_format": arguments.get( + "summary_format", os.getenv("DIRECT_MODE_SUMMARY_FORMAT", "two_paragraph") + ), + "sample_reference_table": arguments.get( + "sample_reference_table", + os.getenv("DIRECT_MODE_SAMPLE_REFERENCE_TABLE", "True"), + ), + "no_of_sampled_summaries": int( + arguments.get( + "no_of_sampled_summaries", + os.getenv("DEFAULT_SAMPLED_SUMMARIES", DEFAULT_SAMPLED_SUMMARIES), + ) + ), + "random_seed": int( + arguments.get("random_seed", os.getenv("LLM_SEED", LLM_SEED)) + ), + # Output Format Arguments + "create_xlsx_output": convert_string_to_boolean( + arguments.get( + "create_xlsx_output", + os.getenv("DIRECT_MODE_CREATE_XLSX_OUTPUT", "True"), + ) + ), + # Logging Arguments + "save_logs_to_csv": convert_string_to_boolean( + arguments.get( + "save_logs_to_csv", os.getenv("SAVE_LOGS_TO_CSV", str(SAVE_LOGS_TO_CSV)) + ) + ), + "save_logs_to_dynamodb": convert_string_to_boolean( + arguments.get( + "save_logs_to_dynamodb", + os.getenv("SAVE_LOGS_TO_DYNAMODB", str(SAVE_LOGS_TO_DYNAMODB)), + ) + ), + "usage_logs_folder": arguments.get("usage_logs_folder", USAGE_LOGS_FOLDER), + "cost_code": arguments.get( + "cost_code", os.getenv("DEFAULT_COST_CODE", DEFAULT_COST_CODE) + ), + } + + # Download optional files if they are specified + candidate_topics_key = arguments.get("candidate_topics_s3_key") + if candidate_topics_key: + candidate_topics_path = os.path.join(INPUT_DIR, "candidate_topics.csv") + download_file_from_s3(bucket_name, candidate_topics_key, candidate_topics_path) + cli_args["candidate_topics"] = candidate_topics_path + + # Download previous output files if they are S3 keys + if cli_args["previous_output_files"]: + downloaded_previous_files = [] + for prev_file in cli_args["previous_output_files"]: + if prev_file.startswith("s3://"): + # Parse S3 path + s3_path_parts = prev_file[5:].split("/", 1) + if len(s3_path_parts) == 2: + prev_bucket = s3_path_parts[0] + prev_key = s3_path_parts[1] + local_prev_path = os.path.join( + INPUT_DIR, os.path.basename(prev_key) + ) + download_file_from_s3(prev_bucket, prev_key, local_prev_path) + downloaded_previous_files.append(local_prev_path) + else: + downloaded_previous_files.append(prev_file) + else: + downloaded_previous_files.append(prev_file) + cli_args["previous_output_files"] = downloaded_previous_files + + # 5. Execute the main application logic + try: + print("--- Starting CLI Topics Main Function ---") + print( + f"Arguments passed to cli_main: {json.dumps({k: v for k, v in cli_args.items() if k not in ['aws_access_key', 'aws_secret_key', 'google_api_key', 'azure_api_key', 'azure_endpoint', 'api_url', 'hf_token']}, default=str)}" + ) + cli_main(direct_mode_args=cli_args) + print("--- CLI Topics Main Function Finished ---") + except Exception as e: + print(f"An error occurred during CLI execution: {e}") + import traceback + + traceback.print_exc() + # Optionally, re-raise the exception to make the Lambda fail + raise + + # 6. Upload results back to S3 + output_s3_prefix = f"output/{os.path.splitext(os.path.basename(input_key))[0]}" + print( + f"Uploading contents of {OUTPUT_DIR} to s3://{bucket_name}/{output_s3_prefix}/" + ) + upload_directory_to_s3(OUTPUT_DIR, bucket_name, output_s3_prefix) + + return { + "statusCode": 200, + "body": json.dumps( + f"Processing complete for {input_key}. Output saved to s3://{bucket_name}/{output_s3_prefix}/" + ), + } diff --git a/load_dynamo_logs.py b/load_dynamo_logs.py new file mode 100644 index 0000000000000000000000000000000000000000..37e89d38de72ec94f181f8d5370d64050a08e124 --- /dev/null +++ b/load_dynamo_logs.py @@ -0,0 +1,102 @@ +import csv +import datetime +from decimal import Decimal + +import boto3 + +from tools.config import ( + AWS_REGION, + OUTPUT_FOLDER, + USAGE_LOG_DYNAMODB_TABLE_NAME, +) + +# Replace with your actual table name and region +TABLE_NAME = USAGE_LOG_DYNAMODB_TABLE_NAME # Choose as appropriate +REGION = AWS_REGION +CSV_OUTPUT = OUTPUT_FOLDER + "dynamodb_logs_export.csv" + +# Create DynamoDB resource +dynamodb = boto3.resource("dynamodb", region_name=REGION) +table = dynamodb.Table(TABLE_NAME) + + +# Helper function to convert Decimal to float or int +def convert_types(item): + new_item = {} + for key, value in item.items(): + # Handle Decimals first + if isinstance(value, Decimal): + new_item[key] = int(value) if value % 1 == 0 else float(value) + # Handle Strings that might be dates + elif isinstance(value, str): + try: + # Attempt to parse a common ISO 8601 format. + # The .replace() handles the 'Z' for Zulu/UTC time. + dt_obj = datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + # Now that we have a datetime object, format it as desired + new_item[key] = dt_obj.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3] + except (ValueError, TypeError): + # If it fails to parse, it's just a regular string + new_item[key] = value + # Handle all other types + else: + new_item[key] = value + return new_item + + +# Paginated scan +def scan_table(): + items = [] + response = table.scan() + items.extend(response["Items"]) + + while "LastEvaluatedKey" in response: + response = table.scan(ExclusiveStartKey=response["LastEvaluatedKey"]) + items.extend(response["Items"]) + + return items + + +# Export to CSV +# Export to CSV +def export_to_csv(items, output_path, fields_to_drop: list = None): + if not items: + print("No items found.") + return + + # Use a set for efficient lookup + drop_set = set(fields_to_drop or []) + + # Get a comprehensive list of all possible headers from all items + all_keys = set() + for item in items: + all_keys.update(item.keys()) + + # Determine the final fieldnames by subtracting the ones to drop + fieldnames = sorted(list(all_keys - drop_set)) + + print("Final CSV columns will be:", fieldnames) + + with open(output_path, "w", newline="", encoding="utf-8-sig") as csvfile: + # The key fix is here: extrasaction='ignore' + # restval='' is also good practice to handle rows that are missing a key + writer = csv.DictWriter( + csvfile, fieldnames=fieldnames, extrasaction="ignore", restval="" + ) + writer.writeheader() + + for item in items: + # The convert_types function can now return the full dict, + # and the writer will simply ignore the extra fields. + writer.writerow(convert_types(item)) + + print(f"Exported {len(items)} items to {output_path}") + + +# Run export +items = scan_table() +export_to_csv( + items, + CSV_OUTPUT, + fields_to_drop=["Query metadata - usage counts and other parameters"], +) diff --git a/load_s3_logs.py b/load_s3_logs.py new file mode 100644 index 0000000000000000000000000000000000000000..dac551831d71d3d3f7a6286f8b6fd099b6a0a20d --- /dev/null +++ b/load_s3_logs.py @@ -0,0 +1,93 @@ +from datetime import datetime +from io import StringIO + +import boto3 +import pandas as pd + +from tools.config import ( + AWS_ACCESS_KEY, + AWS_REGION, + AWS_SECRET_KEY, + DOCUMENT_SUMMARISATION_BUCKET, + OUTPUT_FOLDER, +) + +# Combine together log files that can be then used for e.g. dashboarding and financial tracking. + +# S3 setup. Try to use provided keys (needs S3 permissions), otherwise assume AWS SSO connection +if AWS_ACCESS_KEY and AWS_SECRET_KEY and AWS_REGION: + s3 = boto3.client( + "s3", + aws_access_key_id=AWS_ACCESS_KEY, + aws_secret_access_key=AWS_SECRET_KEY, + region_name=AWS_REGION, + ) +else: + s3 = boto3.client("s3") + +bucket_name = DOCUMENT_SUMMARISATION_BUCKET +prefix = "usage/" # 'feedback/' # 'logs/' # Change as needed - top-level folder where logs are stored +earliest_date = "20250409" # Earliest date of logs folder retrieved +latest_date = "20250423" # Latest date of logs folder retrieved + + +# Function to list all files in a folder +def list_files_in_s3(bucket, prefix): + response = s3.list_objects_v2(Bucket=bucket, Prefix=prefix) + if "Contents" in response: + return [content["Key"] for content in response["Contents"]] + return [] + + +# Function to filter date range +def is_within_date_range(date_str, start_date, end_date): + date_obj = datetime.strptime(date_str, "%Y%m%d") + return start_date <= date_obj <= end_date + + +# Define the date range +start_date = datetime.strptime(earliest_date, "%Y%m%d") # Replace with your start date +end_date = datetime.strptime(latest_date, "%Y%m%d") # Replace with your end date + +# List all subfolders under 'usage/' +all_files = list_files_in_s3(bucket_name, prefix) + +# Filter based on date range +log_files = [] +for file in all_files: + parts = file.split("/") + if len(parts) >= 3: + date_str = parts[1] + if ( + is_within_date_range(date_str, start_date, end_date) + and parts[-1] == "log.csv" + ): + log_files.append(file) + +# Download, read and concatenate CSV files into a pandas DataFrame +df_list = [] +for log_file in log_files: + # Download the file + obj = s3.get_object(Bucket=bucket_name, Key=log_file) + try: + csv_content = ( + obj["Body"].read().decode("utf-8") + ) # Suggest trying latin-1 instead of utf-8 if this fails + except Exception as e: + print("Could not load in log file:", log_file, "due to:", e) + continue + + # Read CSV content into pandas DataFrame + df = pd.read_csv(StringIO(csv_content)) + + df_list.append(df) + +# Concatenate all DataFrames +if df_list: + concatenated_df = pd.concat(df_list, ignore_index=True) + + # Save the concatenated DataFrame to a CSV file + concatenated_df.to_csv(OUTPUT_FOLDER + "consolidated_s3_logs.csv", index=False) + print("Consolidated CSV saved as 'consolidated_s3_logs.csv'") +else: + print("No log files found in the given date range.") diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..3919f699e482b6dbc165160a9e228295e3ddede1 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,148 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "llm_topic_modelling" +version = "0.13.0" +description = "Generate thematic summaries from open text in tabular data files with a large language model." +requires-python = ">=3.10" +readme = "README.md" +authors = [ + { name = "Sean Pedrick-Case", email = "spedrickcase@lambeth.gov.uk" }, +] +maintainers = [ + { name = "Sean Pedrick-Case", email = "spedrickcase@lambeth.gov.uk" }, +] +keywords = [ + "topic-modelling", + "topic-modeling", + "llm", + "large-language-models", + "thematic-analysis", + "text-analysis", + "nlp", + "natural-language-processing", + "text-summarization", + "text-summarisation", + "thematic-summaries", + "gradio", + "data-analysis", + "tabular-data", + "excel", + "csv", + "open-text", + "text-mining" +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Intended Audience :: Information Technology", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Text Processing :: Linguistic", + "Topic :: Text Processing :: Markup", + "Topic :: Scientific/Engineering :: Information Analysis", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", +] + +dependencies = [ + "gradio<=6.19.0", + "transformers<=5.12.0", + "spaces", + "boto3<=1.43.35", + "pandas<=2.3.3", + "pyarrow<=23.0.1", + "openpyxl<=3.1.5", + "markdown<=3.10.2", + "tabulate<=0.9.0", + "lxml<=6.1.1", + "google-genai<=1.73.0", + "openai<=2.31.0", + "html5lib<=1.1", + "beautifulsoup4<=4.14.3", + "rapidfuzz<=3.14.5", + "python-dotenv<=1.2.2", + "pyexcel<=0.7.4" +] + +[project.optional-dependencies] +dev = ["pytest"] +test = ["pytest", "pytest-cov"] + +# Extra dependencies for VLM models +# For torch you should use --index-url https://download.pytorch.org/whl/cu129. Additionally installs the unsloth package +torch = [ + "torch<=2.9.1", + "torchvision", + "accelerate", + "bitsandbytes", + "unsloth<=2026.6.9", + "unsloth_zoo<=2026.6.9", + "timm", + "xformers" +] + +# Support for llama-cpp-python is experimental - I have not directly tested the latest versions of llamacpp against the app +llamacpp = [ + "llama-cpp-python>=0.3.31", +] + +# Run Gradio as an mcp server +mcp = [ + "gradio[mcp]<=6.19.0" +] + +[project.urls] +Homepage = "https://github.com/seanpedrick-case/llm_topic_modelling" +repository = "https://github.com/seanpedrick-case/llm_topic_modelling" + +[tool.setuptools] +packages = ["tools"] +py-modules = ["app"] + +# Configuration for Ruff linter: +[tool.ruff] +line-length = 88 + +[tool.ruff.lint] +select = ["E", "F", "I"] +ignore = [ + "E501", # line-too-long (handled with Black) + "E402", # module-import-not-at-top-of-file (sometimes needed for conditional imports) +] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["F401"] # Allow unused imports in __init__.py + +# Configuration for a Black formatter: +[tool.black] +line-length = 88 +target-version = ['py310'] + +# Configuration for pytest: +[tool.pytest.ini_options] +filterwarnings = [ + "ignore::DeprecationWarning:click.parser", + "ignore::DeprecationWarning:weasel.util.config", + "ignore::DeprecationWarning:builtin type", + "ignore::DeprecationWarning:websockets.legacy", + "ignore::DeprecationWarning:websockets.server", + "ignore::DeprecationWarning:spacy.cli._util", + "ignore::DeprecationWarning:weasel.util.config", + "ignore::DeprecationWarning:importlib._bootstrap", +] +testpaths = ["test"] +python_files = ["test_*.py", "*_test.py"] +python_classes = ["Test*"] +python_functions = ["test_*"] +addopts = [ + "-v", + "--tb=short", + "--strict-markers", + "--disable-warnings", +] \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..76e3cb8982a7232a71c3d91f871ccf4bace2f1c1 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,32 @@ +# Note that this requirements file is optimised for Hugging Face spaces / Python 3.10. Please use requirements_no_local.txt for installation without local model inference (simplest approach to get going). Please use requirements_cpu.txt for CPU instances and requirements_gpu.txt for GPU instances using Python 3.11 +gradio<=6.19.0 +transformers<=5.30.0 +spaces +boto3<=1.43.35 +pandas<=2.3.3 +pyarrow<=23.0.1 +openpyxl<=3.1.5 +pyexcel<=0.7.4 +markdown<=3.10.2 +tabulate<=0.9.0 +lxml<=6.1.1 +google-genai<=1.73.0 +openai<=2.31.0 +html5lib<=1.1 +beautifulsoup4<=4.14.3 +rapidfuzz<=3.14.5 +python-dotenv<=1.2.2 +# GPU (for huggingface instance) +# Torch/Unsloth and llama-cpp-python +# Latest compatible with CUDA 12.4 +torch<=2.9.1 --extra-index-url https://download.pytorch.org/whl/cu128 +unsloth<=2026.6.9 +unsloth_zoo<=2026.6.9 +timm +# llama-cpp-python direct wheel link for GPU compatible version 3.31 for use with Python 3.10 and Hugging Face +https://github.com/abetlen/llama-cpp-python/releases/download/v0.3.31-cu125/llama_cpp_python-0.3.31-py3-none-manylinux_2_35_x86_64.whl +#https://github.com/JamePeng/llama-cpp-python/releases/download/v0.3.17-cu128-Basic-linux-20251202/llama_cpp_python-0.3.17-cp310-cp310-linux_x86_64.whl +#https://github.com/abetlen/llama-cpp-python/releases/download/v0.3.16-cu124/llama_cpp_python-0.3.16-cp310-cp310-linux_x86_64.whl + + + diff --git a/requirements_cpu.txt b/requirements_cpu.txt new file mode 100644 index 0000000000000000000000000000000000000000..53d1841d6d259a6b0a4eab68bfe6b4ac82da6d3a --- /dev/null +++ b/requirements_cpu.txt @@ -0,0 +1,19 @@ +gradio<=6.19.0 +transformers<=5.12.0 +spaces +pandas<=2.3.3 +boto3<=1.43.35 +pyarrow<=23.0.1 +openpyxl<=3.1.5 +pyexcel<=0.7.4 +markdown<=3.10.2 +tabulate<=0.9.0 +lxml<=6.1.1 +google-genai<=1.73.0 +openai<=2.31.0 +html5lib<=1.1 +beautifulsoup4<=4.14.3 +rapidfuzz<=3.14.5 +python-dotenv<=1.2.2 +torch<=2.9.1 --extra-index-url https://download.pytorch.org/whl/cpu +llama-cpp-python>=0.3.31 -C cmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS" \ No newline at end of file diff --git a/requirements_gpu.txt b/requirements_gpu.txt new file mode 100644 index 0000000000000000000000000000000000000000..9ebc4db24b1f45bd2dcbf4c45391a1d6a545906f --- /dev/null +++ b/requirements_gpu.txt @@ -0,0 +1,29 @@ +gradio<=6.19.0 +transformers<=5.12.0 +spaces +pandas<=2.3.3 +boto3<=1.43.35 +pyarrow<=23.0.1 +openpyxl<=3.1.5 +pyexcel<=0.7.4 +markdown<=3.10.2 +tabulate<=0.9.0 +lxml<=6.1.1 +google-genai<=1.73.0 +openai<=2.31.0 +html5lib<=1.1 +beautifulsoup4<=4.14.3 +rapidfuzz<=3.14.5 +python-dotenv<=1.2.2 +# Torch/Unsloth +# Latest compatible with CUDA 12.4 +torch<=2.9.1 --extra-index-url https://download.pytorch.org/whl/cu128 +unsloth[cu128-torch290]<=2026.6.9 # Refer here for more details on installation: https://pypi.org/project/unsloth +unsloth_zoo<=2026.6.9 +# Additional for Windows and CUDA 12.4 older GPUS (RTX 3x or similar): +#triton-windows<3.3 +timm +# Llama CPP Python +llama-cpp-python>=0.3.31 -C cmake.args="-DGGML_CUDA=on" +#https://github.com/abetlen/llama-cpp-python/releases/download/v0.3.31-cu125/llama_cpp_python-0.3.31-py3-none-manylinux_2_35_x86_64.whl +# Try other options at https://github.com/abetlen/llama-cpp-python/releases if the above doesn't work \ No newline at end of file diff --git a/requirements_lightweight.txt b/requirements_lightweight.txt new file mode 100644 index 0000000000000000000000000000000000000000..8c62b4f3c025442f2b0d3f26eceb1f03ba40e71c --- /dev/null +++ b/requirements_lightweight.txt @@ -0,0 +1,19 @@ +# This requirements file is optimised for AWS ECS using Python 3.11 alongside the Dockerfile, without local torch and llama-cpp-python. For AWS ECS, torch and llama-cpp-python are optionally installed in the main Dockerfile +gradio<=6.19.0 +transformers<=5.12.0 +spaces +boto3<=1.43.35 +pandas<=2.3.3 +pyarrow<=23.0.1 +openpyxl<=3.1.5 +pyexcel<=0.7.4 +markdown<=3.10.2 +tabulate<=0.9.0 +lxml<=6.1.1 +google-genai<=1.73.0 +openai<=2.31.0 +html5lib<=1.1 +beautifulsoup4<=4.14.3 +rapidfuzz<=3.14.5 +python-dotenv<=1.2.2 +awslambdaric<=3.1.1 \ No newline at end of file diff --git a/test/README.md b/test/README.md new file mode 100644 index 0000000000000000000000000000000000000000..0010cff1a6e5e8e73f1427d7c0e8d4aeb9fb37fc --- /dev/null +++ b/test/README.md @@ -0,0 +1,87 @@ +# Test Suite for LLM Topic Modeller + +This test suite provides comprehensive testing for the CLI interface (`cli_topics.py`) and GUI application (`app.py`). + +## Overview + +The test suite includes: +- **CLI Tests**: Tests based on examples from the `cli_topics.py` epilog +- **GUI Tests**: Tests to verify the Gradio interface loads correctly +- **Mock Inference Server**: A dummy inference-server endpoint that avoids API costs during testing + +## Structure + +- `test.py`: Main test suite with CLI tests +- `test_gui_only.py`: GUI-specific tests +- `mock_inference_server.py`: Mock HTTP server that mimics an inference-server API +- `run_tests.py`: Test runner script +- `__init__.py`: Package initialization + +## Running Tests + +### Run All Tests + +From the project root directory: + +```bash +python test/run_tests.py +``` + +Or from the test directory: + +```bash +python run_tests.py +``` + +### Run Only CLI Tests + +```bash +python -m unittest test.test.TestCLITopicsExamples +``` + +### Run Only GUI Tests + +```bash +python test/test_gui_only.py +``` + +## Mock Inference Server + +The test suite uses a mock inference server to avoid API costs during testing. The mock server: + +- Listens on `localhost:8080` by default +- Responds to `/v1/chat/completions` endpoint +- Returns valid markdown table responses that satisfy validation requirements +- Provides token counts for usage tracking + +The mock server is automatically started before tests and stopped after tests complete. + +## Test Coverage + +The CLI tests cover: + +1. **Topic Extraction** + - Default settings + - Custom model and context + - Grouping by column + - Zero-shot extraction with candidate topics + +2. **Topic Deduplication** + - Fuzzy matching + - LLM-based deduplication + +3. **All-in-One Pipeline** + - Complete workflow (extract, deduplicate, summarise) + +## Requirements + +- Python 3.7+ +- All dependencies from `requirements.txt` +- Example data files in `example_data/` directory + +## Notes + +- Tests will be skipped if required example files are not found +- The mock server must be running for CLI tests to work +- Tests use temporary output directories that are cleaned up after execution + diff --git a/test/__init__.py b/test/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..09ef5d3e11ff49074a51490e9597309c2aec493f --- /dev/null +++ b/test/__init__.py @@ -0,0 +1,5 @@ +""" +Test suite for LLM Topic Modeller CLI. + +This package contains tests for the CLI interface and GUI application. +""" diff --git a/test/mock_inference_server.py b/test/mock_inference_server.py new file mode 100644 index 0000000000000000000000000000000000000000..dde5fb904c5c554adc58bfe2435a048a90747456 --- /dev/null +++ b/test/mock_inference_server.py @@ -0,0 +1,225 @@ +#!/usr/bin/env python3 +""" +Mock inference server for testing CLI topic extraction without API costs. + +This server mimics an inference-server API endpoint and returns dummy +responses that satisfy the validation requirements (markdown tables with |). +""" + +import json +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer +from typing import Optional + + +class MockInferenceServerHandler(BaseHTTPRequestHandler): + """HTTP request handler for the mock inference server.""" + + def _generate_mock_response(self, prompt: str, system_prompt: str) -> str: + """ + Generate a mock response that satisfies validation requirements. + + The response must: + - Be longer than 120 characters + - Contain a markdown table (with | characters) + + Args: + prompt: The user prompt + system_prompt: The system prompt + + Returns: + A mock markdown table response + """ + # Generate a simple markdown table that satisfies the validation + # This mimics a topic extraction table response + mock_table = """| General topic | Subtopic | Sentiment | Response ID | Summary | +|-----------|---------------|-----------|-----------|-----------| +| Test Topic | Test Subtopic | Positive | 1 | Test summary | +| Another Topic | Another Subtopic | Neutral | 2,3 | Another summary | +| Third Topic | Third Subtopic | Negative | 1, 2, 3 | Third summary | + +This is a mock response from the test inference server. The actual content would be generated by a real LLM model, but for testing purposes, this dummy response allows us to verify that the CLI commands work correctly without incurring API costs.""" + + return mock_table + + def _estimate_tokens(self, text: str) -> int: + """Estimate token count (rough approximation: ~4 characters per token).""" + return max(1, len(text) // 4) + + def do_POST(self): + """Handle POST requests to /v1/chat/completions.""" + print(f"[Mock Server] Received POST request to: {self.path}") + if self.path == "/v1/chat/completions": + try: + # Read request body + content_length = int(self.headers.get("Content-Length", 0)) + print(f"[Mock Server] Content-Length: {content_length}") + body = self.rfile.read(content_length) + payload = json.loads(body.decode("utf-8")) + print("[Mock Server] Payload received, processing...") + + # Extract messages + messages = payload.get("messages", []) + system_prompt = "" + user_prompt = "" + + for msg in messages: + role = msg.get("role", "") + content = msg.get("content", "") + if role == "system": + system_prompt = content + elif role == "user": + user_prompt = content + + # Generate mock response + response_text = self._generate_mock_response(user_prompt, system_prompt) + + # Estimate tokens + input_tokens = self._estimate_tokens(system_prompt + "\n" + user_prompt) + output_tokens = self._estimate_tokens(response_text) + + # Check if streaming is requested + stream = payload.get("stream", False) + + if stream: + # Handle streaming response + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("Connection", "keep-alive") + self.end_headers() + + # Send streaming chunks + chunk_size = 20 # Characters per chunk + for i in range(0, len(response_text), chunk_size): + chunk = response_text[i : i + chunk_size] + chunk_data = { + "choices": [ + { + "delta": {"content": chunk}, + "index": 0, + "finish_reason": None, + } + ] + } + self.wfile.write(f"data: {json.dumps(chunk_data)}\n\n".encode()) + self.wfile.flush() + + # Send final done message + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + else: + # Handle non-streaming response + response_data = { + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": response_text, + }, + } + ], + "usage": { + "prompt_tokens": input_tokens, + "completion_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + } + + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps(response_data).encode()) + + except Exception as e: + self.send_response(500) + self.send_header("Content-Type", "application/json") + self.end_headers() + error_response = {"error": {"message": str(e), "type": "server_error"}} + self.wfile.write(json.dumps(error_response).encode()) + else: + self.send_response(404) + self.end_headers() + + def log_message(self, format, *args): + """Log messages for debugging.""" + # Enable logging for debugging + print(f"[Mock Server] {format % args}") + + +class MockInferenceServer: + """Mock inference server that can be started and stopped for testing.""" + + def __init__(self, host: str = "localhost", port: int = 8080): + """ + Initialize the mock server. + + Args: + host: Host to bind to (default: localhost) + port: Port to bind to (default: 8080) + """ + self.host = host + self.port = port + self.server: Optional[HTTPServer] = None + self.server_thread: Optional[threading.Thread] = None + self.running = False + + def start(self): + """Start the mock server in a separate thread.""" + if self.running: + return + + def run_server(): + self.server = HTTPServer((self.host, self.port), MockInferenceServerHandler) + self.running = True + self.server.serve_forever() + + self.server_thread = threading.Thread(target=run_server, daemon=True) + self.server_thread.start() + + # Wait a moment for server to start + import time + + time.sleep(0.5) + + def stop(self): + """Stop the mock server.""" + if self.server and self.running: + self.server.shutdown() + self.server.server_close() + self.running = False + + def get_url(self) -> str: + """Get the server URL.""" + return f"http://{self.host}:{self.port}" + + def __enter__(self): + """Context manager entry.""" + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + self.stop() + + +if __name__ == "__main__": + # Test the server + print("Starting mock inference server on http://localhost:8080") + print("Press Ctrl+C to stop") + + server = MockInferenceServer() + try: + server.start() + print(f"Server running at {server.get_url()}") + # Keep running + while True: + import time + + time.sleep(1) + except KeyboardInterrupt: + print("\nStopping server...") + server.stop() + print("Server stopped") diff --git a/test/mock_llm_calls.py b/test/mock_llm_calls.py new file mode 100644 index 0000000000000000000000000000000000000000..cd095d503747c677974fcdb0baaa333203dcab87 --- /dev/null +++ b/test/mock_llm_calls.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python3 +""" +Mock LLM function calls for testing CLI topic extraction without API costs. + +This module patches requests.post to intercept HTTP calls to inference servers +and return mock responses instead. +""" + +import json +import os + +# Store original requests if it exists +_original_requests = None + + +def _generate_mock_response(prompt: str, system_prompt: str) -> str: + """ + Generate a mock response that satisfies validation requirements. + + The response must: + - Be longer than 120 characters + - Contain a markdown table (with | characters) + + Args: + prompt: The user prompt + system_prompt: The system prompt + + Returns: + A mock markdown table response + """ + # Generate a simple markdown table that satisfies the validation + # This mimics a topic extraction table response + mock_table = """| General topic | Subtopic | Sentiment | Response ID | Summary | +|-----------|---------------|-----------|-----------|-----------| +| Test Topic | Test Subtopic | Positive | 1 | Test summary | +| Another Topic | Another Subtopic | Neutral | 2,3 | Another summary | +| Third Topic | Third Subtopic | Negative | 1, 2, 3 | Third summary | + +This is a mock response from the test inference server. The actual content would be generated by a real LLM model, but for testing purposes, this dummy response allows us to verify that the CLI commands work correctly without incurring API costs.""" + + return mock_table + + +def _estimate_tokens(text: str) -> int: + """Estimate token count (rough approximation: ~4 characters per token).""" + return max(1, len(text) // 4) + + +def mock_requests_post(url, **kwargs): + """ + Mock version of requests.post that intercepts inference-server calls. + + Returns a mock response object that mimics the real requests.Response. + """ + # Only mock inference-server URLs + if "/v1/chat/completions" not in url: + # For non-inference-server URLs, use real requests + import requests + + return requests.post(url, **kwargs) + + # Extract payload + payload = kwargs.get("json", {}) + messages = payload.get("messages", []) + + # Extract prompts + system_prompt = "" + user_prompt = "" + for msg in messages: + role = msg.get("role", "") + content = msg.get("content", "") + if role == "system": + system_prompt = content + elif role == "user": + user_prompt = content + + # Generate mock response + response_text = _generate_mock_response(user_prompt, system_prompt) + + # Estimate tokens + input_tokens = _estimate_tokens(system_prompt + "\n" + user_prompt) + output_tokens = _estimate_tokens(response_text) + + # Check if streaming is requested + stream = payload.get("stream", False) + + if stream: + # For streaming, create a mock response with iter_lines + class MockStreamResponse: + def __init__(self, text): + self.text = text + self.status_code = 200 + self.lines = [] + # Simulate streaming chunks + chunk_size = 20 + for i in range(0, len(text), chunk_size): + chunk = text[i : i + chunk_size] + chunk_data = { + "choices": [ + { + "delta": {"content": chunk}, + "index": 0, + "finish_reason": None, + } + ] + } + self.lines.append(f"data: {json.dumps(chunk_data)}\n\n".encode()) + self.lines.append(b"data: [DONE]\n\n") + self._line_index = 0 + + def raise_for_status(self): + pass + + def iter_lines(self): + for line in self.lines: + yield line + + return MockStreamResponse(response_text) + else: + # For non-streaming, create a simple mock response + class MockResponse: + def __init__(self, text, input_tokens, output_tokens): + self._json_data = { + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": text, + }, + } + ], + "usage": { + "prompt_tokens": input_tokens, + "completion_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + } + self.status_code = 200 + + def raise_for_status(self): + pass + + def json(self): + return self._json_data + + return MockResponse(response_text, input_tokens, output_tokens) + + +def apply_mock_patches(): + """ + Apply patches to mock HTTP requests. + This should be called before importing modules that use requests. + """ + global _original_requests + + try: + import requests + + _original_requests = requests.post + requests.post = mock_requests_post + print("[Mock] Patched requests.post for inference-server calls") + except ImportError: + # requests not imported yet, will be patched when imported + pass + + +def restore_original(): + """Restore original requests.post if it was patched.""" + global _original_requests + if _original_requests: + try: + import requests + + requests.post = _original_requests + _original_requests = None + print("[Mock] Restored original requests.post") + except ImportError: + pass + + +# Auto-apply patches if TEST_MODE environment variable is set +if os.environ.get("TEST_MODE") == "1" or os.environ.get("USE_MOCK_LLM") == "1": + apply_mock_patches() diff --git a/test/run_tests.py b/test/run_tests.py new file mode 100644 index 0000000000000000000000000000000000000000..485401824c8bd9c4f08a532522ede4e90d737608 --- /dev/null +++ b/test/run_tests.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +""" +Simple script to run the CLI topics test suite. + +This script demonstrates how to run the comprehensive test suite +that covers all the examples from the CLI epilog. +""" + +import os +import sys + +# Add the parent directory to the path so we can import the test module +parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, parent_dir) + +# Import test functions +from test.test import run_all_tests + +if __name__ == "__main__": + print("Starting LLM Topic Modeller Test Suite...") + print("This will test:") + print("- CLI examples from the epilog") + print("- GUI application functionality") + print("Using a mock inference-server to avoid API costs.") + print("=" * 60) + + success = run_all_tests() + + if success: + print("\n🎉 All tests passed successfully!") + sys.exit(0) + else: + print("\n❌ Some tests failed. Check the output above for details.") + sys.exit(1) diff --git a/test/test.py b/test/test.py new file mode 100644 index 0000000000000000000000000000000000000000..5ac5d9b437428e4b8050ae37e89b75313d8000f5 --- /dev/null +++ b/test/test.py @@ -0,0 +1,1067 @@ +import os +import shutil +import subprocess +import sys +import tempfile +import time +import unittest +from typing import List, Optional + +# Mock LLM calls are automatically applied via environment variables +# No need to import - the mock patches are applied when USE_MOCK_LLM=1 is set + + +def run_cli_topics( + script_path: str, + task: str, + output_dir: str, + input_file: Optional[str] = None, + text_column: Optional[str] = None, + previous_output_files: Optional[List[str]] = None, + timeout: int = 600, # 10-minute timeout + # General Arguments + username: Optional[str] = None, + save_to_user_folders: Optional[bool] = None, + excel_sheets: Optional[List[str]] = None, + group_by: Optional[str] = None, + # Model Configuration + model_choice: Optional[str] = None, + temperature: Optional[float] = None, + batch_size: Optional[int] = None, + max_tokens: Optional[int] = None, + api_url: Optional[str] = None, + inference_server_model: Optional[str] = None, + # Topic Extraction Arguments + context: Optional[str] = None, + candidate_topics: Optional[str] = None, + force_zero_shot: Optional[str] = None, + force_single_topic: Optional[str] = None, + produce_structured_summary: Optional[str] = None, + sentiment: Optional[str] = None, + additional_summary_instructions: Optional[str] = None, + # Validation Arguments + additional_validation_issues: Optional[str] = None, + show_previous_table: Optional[str] = None, + output_debug_files: Optional[str] = None, + max_time_for_loop: Optional[int] = None, + # Deduplication Arguments + method: Optional[str] = None, + similarity_threshold: Optional[int] = None, + merge_sentiment: Optional[str] = None, + merge_general_topics: Optional[str] = None, + # Summarisation Arguments + summary_format: Optional[str] = None, + sample_reference_table: Optional[str] = None, + no_of_sampled_summaries: Optional[int] = None, + random_seed: Optional[int] = None, + # Output Format Arguments + create_xlsx_output: Optional[bool] = None, + # Logging Arguments + save_logs_to_csv: Optional[bool] = None, + save_logs_to_dynamodb: Optional[bool] = None, + cost_code: Optional[str] = None, +) -> bool: + """ + Executes the cli_topics.py script with specified arguments using a subprocess. + + Args: + script_path (str): The path to the cli_topics.py script. + task (str): The main task to perform ('extract', 'validate', 'deduplicate', 'summarise', 'overall_summary', or 'all_in_one'). + output_dir (str): The path to the directory for output files. + input_file (str, optional): Path to the input file to process. + text_column (str, optional): Name of the text column to process. + previous_output_files (List[str], optional): Path(s) to previous output files. + timeout (int): Timeout in seconds for the subprocess. + + All other arguments match the CLI arguments from cli_topics.py. + + Returns: + bool: True if the script executed successfully, False otherwise. + """ + # 1. Get absolute paths and perform pre-checks + script_abs_path = os.path.abspath(script_path) + output_abs_dir = os.path.abspath(output_dir) + + # Handle input file based on task + if task in ["extract", "validate", "all_in_one"] and input_file is None: + raise ValueError(f"Input file is required for '{task}' task") + + if input_file: + input_abs_path = os.path.abspath(input_file) + if not os.path.isfile(input_abs_path): + raise FileNotFoundError(f"Input file not found: {input_abs_path}") + + if not os.path.isfile(script_abs_path): + raise FileNotFoundError(f"Script not found: {script_abs_path}") + + if not os.path.isdir(output_abs_dir): + # Create the output directory if it doesn't exist + print(f"Output directory not found. Creating: {output_abs_dir}") + os.makedirs(output_abs_dir) + + script_folder = os.path.dirname(script_abs_path) + + # 2. Dynamically build the command list + command = [ + "python", + script_abs_path, + "--output_dir", + output_abs_dir, + "--task", + task, + ] + + # Add input_file only if it's not None + if input_file: + command.extend(["--input_file", input_abs_path]) + + # Add general arguments + if text_column: + command.extend(["--text_column", text_column]) + if previous_output_files: + command.extend(["--previous_output_files"] + previous_output_files) + if username: + command.extend(["--username", username]) + if save_to_user_folders is not None: + command.extend(["--save_to_user_folders", str(save_to_user_folders)]) + if excel_sheets: + command.append("--excel_sheets") + command.extend(excel_sheets) + if group_by: + command.extend(["--group_by", group_by]) + + # Add model configuration arguments + if model_choice: + command.extend(["--model_choice", model_choice]) + if temperature is not None: + command.extend(["--temperature", str(temperature)]) + if batch_size is not None: + command.extend(["--batch_size", str(batch_size)]) + if max_tokens is not None: + command.extend(["--max_tokens", str(max_tokens)]) + if api_url: + command.extend(["--api_url", api_url]) + if inference_server_model: + command.extend(["--inference_server_model", inference_server_model]) + + # Add topic extraction arguments + if context: + command.extend(["--context", context]) + if candidate_topics: + command.extend(["--candidate_topics", candidate_topics]) + if force_zero_shot: + command.extend(["--force_zero_shot", force_zero_shot]) + if force_single_topic: + command.extend(["--force_single_topic", force_single_topic]) + if produce_structured_summary: + command.extend(["--produce_structured_summary", produce_structured_summary]) + if sentiment: + command.extend(["--sentiment", sentiment]) + if additional_summary_instructions: + command.extend( + ["--additional_summary_instructions", additional_summary_instructions] + ) + + # Add validation arguments + if additional_validation_issues: + command.extend(["--additional_validation_issues", additional_validation_issues]) + if show_previous_table: + command.extend(["--show_previous_table", show_previous_table]) + if output_debug_files: + command.extend(["--output_debug_files", output_debug_files]) + if max_time_for_loop is not None: + command.extend(["--max_time_for_loop", str(max_time_for_loop)]) + + # Add deduplication arguments + if method: + command.extend(["--method", method]) + if similarity_threshold is not None: + command.extend(["--similarity_threshold", str(similarity_threshold)]) + if merge_sentiment: + command.extend(["--merge_sentiment", merge_sentiment]) + if merge_general_topics: + command.extend(["--merge_general_topics", merge_general_topics]) + + # Add summarisation arguments + if summary_format: + command.extend(["--summary_format", summary_format]) + if sample_reference_table: + command.extend(["--sample_reference_table", sample_reference_table]) + if no_of_sampled_summaries is not None: + command.extend(["--no_of_sampled_summaries", str(no_of_sampled_summaries)]) + if random_seed is not None: + command.extend(["--random_seed", str(random_seed)]) + + # Add output format arguments + if create_xlsx_output is False: + command.append("--no_xlsx_output") + + # Add logging arguments + if save_logs_to_csv is not None: + command.extend(["--save_logs_to_csv", str(save_logs_to_csv)]) + if save_logs_to_dynamodb is not None: + command.extend(["--save_logs_to_dynamodb", str(save_logs_to_dynamodb)]) + if cost_code: + command.extend(["--cost_code", cost_code]) + + # Filter out None values before joining + command_str = " ".join(str(arg) for arg in command if arg is not None) + print(f"Executing command: {command_str}") + + # 3. Execute the command using subprocess + try: + # Use unbuffered output to avoid hanging + env = os.environ.copy() + env["PYTHONUNBUFFERED"] = "1" + # Ensure inference server is enabled for testing + env["RUN_INFERENCE_SERVER"] = "1" + # Enable mock mode + env["USE_MOCK_LLM"] = "1" + env["TEST_MODE"] = "1" + + result = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, # Combine stderr with stdout to avoid deadlocks + text=True, + cwd=script_folder, # Important for relative paths within the script + env=env, + bufsize=0, # Unbuffered + ) + + # Read output in real-time to avoid deadlocks + start_time = time.time() + + # For Windows, we need a different approach + if sys.platform == "win32": + # On Windows, use communicate with timeout + try: + stdout, stderr = result.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + result.kill() + stdout, stderr = result.communicate() + raise subprocess.TimeoutExpired(result.args, timeout) + else: + # On Unix, we can use select for real-time reading + import select + + stdout_lines = [] + while result.poll() is None: + ready, _, _ = select.select([result.stdout], [], [], 0.1) + if ready: + line = result.stdout.readline() + if line: + print(line.rstrip(), flush=True) + stdout_lines.append(line) + # Check timeout + if time.time() - start_time > timeout: + result.kill() + raise subprocess.TimeoutExpired(result.args, timeout) + + # Read remaining output + remaining = result.stdout.read() + if remaining: + print(remaining, end="", flush=True) + stdout_lines.append(remaining) + + stdout = "".join(stdout_lines) + stderr = "" # Combined with stdout + + print("--- SCRIPT STDOUT ---") + if stdout: + print(stdout) + print("--- SCRIPT STDERR ---") + if stderr: + print(stderr) + print("---------------------") + + # Analyze the output for errors and success indicators + analysis = analyze_test_output(stdout, stderr) + + if analysis["has_errors"]: + print("❌ Errors detected in output:") + for i, error_type in enumerate(analysis["error_types"]): + print(f" {i+1}. {error_type}") + if analysis["error_messages"]: + print(" Error messages:") + for msg in analysis["error_messages"][ + :3 + ]: # Show first 3 error messages + print(f" - {msg}") + return False + elif result.returncode == 0: + success_msg = "✅ Script executed successfully." + if analysis["success_indicators"]: + success_msg += f" (Success indicators: {', '.join(analysis['success_indicators'][:3])})" + print(success_msg) + return True + else: + print(f"❌ Command failed with return code {result.returncode}") + return False + + except subprocess.TimeoutExpired: + result.kill() + print(f"❌ Subprocess timed out after {timeout} seconds.") + return False + except Exception as e: + print(f"❌ An unexpected error occurred: {e}") + return False + + +def analyze_test_output(stdout: str, stderr: str) -> dict: + """ + Analyze test output to provide detailed error information. + + Args: + stdout (str): Standard output from the test + stderr (str): Standard error from the test + + Returns: + dict: Analysis results with error details + """ + combined_output = (stdout or "") + (stderr or "") + + analysis = { + "has_errors": False, + "error_types": [], + "error_messages": [], + "success_indicators": [], + "warning_indicators": [], + } + + # Error patterns + error_patterns = { + "An error occurred": "General error message", + "Error:": "Error prefix", + "Exception:": "Exception occurred", + "Traceback": "Python traceback", + "Failed to": "Operation failure", + "Cannot": "Operation not possible", + "Unable to": "Operation not possible", + "KeyError:": "Missing key/dictionary error", + "AttributeError:": "Missing attribute error", + "TypeError:": "Type mismatch error", + "ValueError:": "Invalid value error", + "FileNotFoundError:": "File not found", + "ImportError:": "Import failure", + "ModuleNotFoundError:": "Module not found", + } + + # Success indicators + success_patterns = [ + "Successfully", + "Completed", + "Finished", + "Processed", + "Complete", + "Output files saved", + ] + + # Warning indicators + warning_patterns = ["Warning:", "WARNING:", "Deprecated", "DeprecationWarning"] + + # Check for errors + for pattern, description in error_patterns.items(): + if pattern.lower() in combined_output.lower(): + analysis["has_errors"] = True + analysis["error_types"].append(description) + + # Extract the actual error message + lines = combined_output.split("\n") + for line in lines: + if pattern.lower() in line.lower(): + analysis["error_messages"].append(line.strip()) + + # Check for success indicators + for pattern in success_patterns: + if pattern.lower() in combined_output.lower(): + analysis["success_indicators"].append(pattern) + + # Check for warnings + for pattern in warning_patterns: + if pattern.lower() in combined_output.lower(): + analysis["warning_indicators"].append(pattern) + + return analysis + + +def run_app_direct_mode( + app_path: str, + task: str, + output_dir: str, + input_file: Optional[str] = None, + text_column: Optional[str] = None, + previous_output_files: Optional[List[str]] = None, + timeout: int = 600, + # General Arguments + username: Optional[str] = None, + save_to_user_folders: Optional[bool] = None, + excel_sheets: Optional[List[str]] = None, + group_by: Optional[str] = None, + # Model Configuration + model_choice: Optional[str] = None, + temperature: Optional[float] = None, + batch_size: Optional[int] = None, + max_tokens: Optional[int] = None, + api_url: Optional[str] = None, + inference_server_model: Optional[str] = None, + # Topic Extraction Arguments + context: Optional[str] = None, + candidate_topics: Optional[str] = None, + force_zero_shot: Optional[str] = None, + force_single_topic: Optional[str] = None, + produce_structured_summary: Optional[str] = None, + sentiment: Optional[str] = None, + additional_summary_instructions: Optional[str] = None, + # Validation Arguments + additional_validation_issues: Optional[str] = None, + show_previous_table: Optional[str] = None, + output_debug_files: Optional[str] = None, + max_time_for_loop: Optional[int] = None, + # Deduplication Arguments + method: Optional[str] = None, + similarity_threshold: Optional[int] = None, + merge_sentiment: Optional[str] = None, + merge_general_topics: Optional[str] = None, + # Summarisation Arguments + summary_format: Optional[str] = None, + sample_reference_table: Optional[str] = None, + no_of_sampled_summaries: Optional[int] = None, + random_seed: Optional[int] = None, + # Output Format Arguments + create_xlsx_output: Optional[bool] = None, + # Logging Arguments + save_logs_to_csv: Optional[bool] = None, + save_logs_to_dynamodb: Optional[bool] = None, + cost_code: Optional[str] = None, +) -> bool: + """ + Executes the app.py script in direct mode with specified environment variables. + + Args: + app_path (str): The path to the app.py script. + task (str): The main task to perform ('extract', 'validate', 'deduplicate', 'summarise', 'overall_summary', or 'all_in_one'). + output_dir (str): The path to the directory for output files. + input_file (str, optional): Path to the input file to process. + text_column (str, optional): Name of the text column to process. + previous_output_files (List[str], optional): Path(s) to previous output files. + timeout (int): Timeout in seconds for the subprocess. + + All other arguments match the CLI arguments from cli_topics.py, but are set as environment variables. + + Returns: + bool: True if the script executed successfully, False otherwise. + """ + # 1. Get absolute paths and perform pre-checks + app_abs_path = os.path.abspath(app_path) + output_abs_dir = os.path.abspath(output_dir) + + # Handle input file based on task + if task in ["extract", "validate", "all_in_one"] and input_file is None: + raise ValueError(f"Input file is required for '{task}' task") + + if input_file: + input_abs_path = os.path.abspath(input_file) + if not os.path.isfile(input_abs_path): + raise FileNotFoundError(f"Input file not found: {input_abs_path}") + + if not os.path.isfile(app_abs_path): + raise FileNotFoundError(f"App script not found: {app_abs_path}") + + if not os.path.isdir(output_abs_dir): + # Create the output directory if it doesn't exist + print(f"Output directory not found. Creating: {output_abs_dir}") + os.makedirs(output_abs_dir) + + script_folder = os.path.dirname(app_abs_path) + + # 2. Build environment variables for direct mode + env = os.environ.copy() + env["PYTHONUNBUFFERED"] = "1" + env["RUN_INFERENCE_SERVER"] = "1" + env["USE_MOCK_LLM"] = "1" + env["TEST_MODE"] = "1" + + # Enable direct mode + env["RUN_DIRECT_MODE"] = "1" + + # Task selection + env["DIRECT_MODE_TASK"] = task + + # General arguments + if input_file: + # Use pipe separator to handle file paths with spaces + env["DIRECT_MODE_INPUT_FILE"] = input_abs_path + env["DIRECT_MODE_OUTPUT_DIR"] = output_abs_dir + if text_column: + env["DIRECT_MODE_TEXT_COLUMN"] = text_column + if previous_output_files: + # Use pipe separator to handle file paths with spaces + env["DIRECT_MODE_PREVIOUS_OUTPUT_FILES"] = "|".join(previous_output_files) + if username: + env["DIRECT_MODE_USERNAME"] = username + if save_to_user_folders is not None: + env["SESSION_OUTPUT_FOLDER"] = str(save_to_user_folders) + if excel_sheets: + env["DIRECT_MODE_EXCEL_SHEETS"] = ",".join(excel_sheets) + if group_by: + env["DIRECT_MODE_GROUP_BY"] = group_by + + # Model configuration + if model_choice: + env["DIRECT_MODE_MODEL_CHOICE"] = model_choice + if temperature is not None: + env["DIRECT_MODE_TEMPERATURE"] = str(temperature) + if batch_size is not None: + env["DIRECT_MODE_BATCH_SIZE"] = str(batch_size) + if max_tokens is not None: + env["DIRECT_MODE_MAX_TOKENS"] = str(max_tokens) + if api_url: + env["API_URL"] = api_url + if inference_server_model: + env["DIRECT_MODE_INFERENCE_SERVER_MODEL"] = inference_server_model + + # Topic extraction arguments + if context: + env["DIRECT_MODE_CONTEXT"] = context + if candidate_topics: + env["DIRECT_MODE_CANDIDATE_TOPICS"] = candidate_topics + if force_zero_shot: + env["DIRECT_MODE_FORCE_ZERO_SHOT"] = force_zero_shot + if force_single_topic: + env["DIRECT_MODE_FORCE_SINGLE_TOPIC"] = force_single_topic + if produce_structured_summary: + env["DIRECT_MODE_PRODUCE_STRUCTURED_SUMMARY"] = produce_structured_summary + if sentiment: + env["DIRECT_MODE_SENTIMENT"] = sentiment + if additional_summary_instructions: + env["DIRECT_MODE_ADDITIONAL_SUMMARY_INSTRUCTIONS"] = ( + additional_summary_instructions + ) + + # Validation arguments + if additional_validation_issues: + env["DIRECT_MODE_ADDITIONAL_VALIDATION_ISSUES"] = additional_validation_issues + if show_previous_table: + env["DIRECT_MODE_SHOW_PREVIOUS_TABLE"] = show_previous_table + if output_debug_files: + env["OUTPUT_DEBUG_FILES"] = output_debug_files + if max_time_for_loop is not None: + env["DIRECT_MODE_MAX_TIME_FOR_LOOP"] = str(max_time_for_loop) + + # Deduplication arguments + if method: + env["DIRECT_MODE_DEDUP_METHOD"] = method + if similarity_threshold is not None: + env["DIRECT_MODE_SIMILARITY_THRESHOLD"] = str(similarity_threshold) + if merge_sentiment: + env["DIRECT_MODE_MERGE_SENTIMENT"] = merge_sentiment + if merge_general_topics: + env["DIRECT_MODE_MERGE_GENERAL_TOPICS"] = merge_general_topics + + # Summarisation arguments + if summary_format: + env["DIRECT_MODE_SUMMARY_FORMAT"] = summary_format + if sample_reference_table: + env["DIRECT_MODE_SAMPLE_REFERENCE_TABLE"] = sample_reference_table + if no_of_sampled_summaries is not None: + env["DIRECT_MODE_NO_OF_SAMPLED_SUMMARIES"] = str(no_of_sampled_summaries) + if random_seed is not None: + env["DIRECT_MODE_RANDOM_SEED"] = str(random_seed) + + # Output format arguments + if create_xlsx_output is not None: + env["DIRECT_MODE_CREATE_XLSX_OUTPUT"] = str(create_xlsx_output) + + # Logging arguments + if save_logs_to_csv is not None: + env["SAVE_LOGS_TO_CSV"] = str(save_logs_to_csv) + if save_logs_to_dynamodb is not None: + env["SAVE_LOGS_TO_DYNAMODB"] = str(save_logs_to_dynamodb) + if cost_code: + env["DEFAULT_COST_CODE"] = cost_code + + # 3. Build command (just run app.py, no arguments needed in direct mode) + command = ["python", app_abs_path] + command_str = " ".join(str(arg) for arg in command) + print(f"Executing direct mode command: {command_str}") + print(f"Direct mode task: {task}") + if input_file: + print(f"Input file: {input_abs_path}") + if text_column: + print(f"Text column: {text_column}") + + # 4. Execute the command using subprocess + try: + result = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, # Combine stderr with stdout to avoid deadlocks + text=True, + cwd=script_folder, # Important for relative paths within the script + env=env, + bufsize=0, # Unbuffered + ) + + # Read output in real-time to avoid deadlocks + start_time = time.time() + + # For Windows, we need a different approach + if sys.platform == "win32": + # On Windows, use communicate with timeout + try: + stdout, stderr = result.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + result.kill() + stdout, stderr = result.communicate() + raise subprocess.TimeoutExpired(result.args, timeout) + else: + # On Unix, we can use select for real-time reading + import select + + stdout_lines = [] + while result.poll() is None: + ready, _, _ = select.select([result.stdout], [], [], 0.1) + if ready: + line = result.stdout.readline() + if line: + print(line.rstrip(), flush=True) + stdout_lines.append(line) + # Check timeout + if time.time() - start_time > timeout: + result.kill() + raise subprocess.TimeoutExpired(result.args, timeout) + + # Read remaining output + remaining = result.stdout.read() + if remaining: + print(remaining, end="", flush=True) + stdout_lines.append(remaining) + + stdout = "".join(stdout_lines) + stderr = "" # Combined with stdout + + print("--- SCRIPT STDOUT ---") + if stdout: + print(stdout) + print("--- SCRIPT STDERR ---") + if stderr: + print(stderr) + print("---------------------") + + # Analyze the output for errors and success indicators + analysis = analyze_test_output(stdout, stderr) + + if analysis["has_errors"]: + print("❌ Errors detected in output:") + for i, error_type in enumerate(analysis["error_types"]): + print(f" {i+1}. {error_type}") + if analysis["error_messages"]: + print(" Error messages:") + for msg in analysis["error_messages"][ + :3 + ]: # Show first 3 error messages + print(f" - {msg}") + return False + elif result.returncode == 0: + success_msg = "✅ Script executed successfully." + if analysis["success_indicators"]: + success_msg += f" (Success indicators: {', '.join(analysis['success_indicators'][:3])})" + print(success_msg) + return True + else: + print(f"❌ Command failed with return code {result.returncode}") + return False + + except subprocess.TimeoutExpired: + result.kill() + print(f"❌ Subprocess timed out after {timeout} seconds.") + return False + except Exception as e: + print(f"❌ An unexpected error occurred: {e}") + return False + + +class TestCLITopicsExamples(unittest.TestCase): + """Test suite for CLI topic extraction examples from the epilog.""" + + @classmethod + def setUpClass(cls): + """Set up test environment before running tests.""" + cls.script_path = os.path.join( + os.path.dirname(os.path.dirname(__file__)), "cli_topics.py" + ) + cls.example_data_dir = os.path.join( + os.path.dirname(os.path.dirname(__file__)), "example_data" + ) + cls.temp_output_dir = tempfile.mkdtemp(prefix="test_output_") + + # Verify script exists + if not os.path.isfile(cls.script_path): + raise FileNotFoundError(f"CLI script not found: {cls.script_path}") + + print(f"Test setup complete. Script: {cls.script_path}") + print(f"Example data directory: {cls.example_data_dir}") + print(f"Temp output directory: {cls.temp_output_dir}") + print("Using function mocking instead of HTTP server") + + # Debug: Check if example data directory exists and list contents + if os.path.exists(cls.example_data_dir): + print("Example data directory exists. Contents:") + for item in os.listdir(cls.example_data_dir): + item_path = os.path.join(cls.example_data_dir, item) + if os.path.isfile(item_path): + print(f" File: {item} ({os.path.getsize(item_path)} bytes)") + else: + print(f" Directory: {item}") + else: + print(f"Example data directory does not exist: {cls.example_data_dir}") + + @classmethod + def tearDownClass(cls): + """Clean up test environment after running tests.""" + if os.path.exists(cls.temp_output_dir): + shutil.rmtree(cls.temp_output_dir) + print(f"Cleaned up temp directory: {cls.temp_output_dir}") + + def test_extract_topics_default_settings(self): + """Test: Extract topics from a CSV file with default settings""" + print("\n=== Testing topic extraction with default settings ===") + input_file = os.path.join(self.example_data_dir, "combined_case_notes.csv") + + if not os.path.isfile(input_file): + self.skipTest(f"Example file not found: {input_file}") + + result = run_cli_topics( + script_path=self.script_path, + task="extract", + input_file=input_file, + text_column="Case Note", + output_dir=self.temp_output_dir, + model_choice="test-model", + inference_server_model="test-model", + api_url="http://localhost:8080", # URL doesn't matter with function mocking + create_xlsx_output=False, + save_logs_to_csv=False, + ) + + self.assertTrue(result, "Topic extraction with default settings should succeed") + print("✅ Topic extraction with default settings passed") + + def test_extract_topics_custom_model_and_context(self): + """Test: Extract topics with custom model and context""" + print("\n=== Testing topic extraction with custom model and context ===") + input_file = os.path.join(self.example_data_dir, "combined_case_notes.csv") + + if not os.path.isfile(input_file): + self.skipTest(f"Example file not found: {input_file}") + + result = run_cli_topics( + script_path=self.script_path, + task="extract", + input_file=input_file, + text_column="Case Note", + output_dir=self.temp_output_dir, + model_choice="test-model", + inference_server_model="test-model", + api_url="http://localhost:8080", # URL doesn't matter with function mocking + context="Social Care case notes for young people", + create_xlsx_output=False, + save_logs_to_csv=False, + ) + + self.assertTrue( + result, "Topic extraction with custom model and context should succeed" + ) + print("✅ Topic extraction with custom model and context passed") + + def test_extract_topics_with_grouping(self): + """Test: Extract topics with grouping""" + print("\n=== Testing topic extraction with grouping ===") + input_file = os.path.join(self.example_data_dir, "combined_case_notes.csv") + + if not os.path.isfile(input_file): + self.skipTest(f"Example file not found: {input_file}") + + result = run_cli_topics( + script_path=self.script_path, + task="extract", + input_file=input_file, + text_column="Case Note", + output_dir=self.temp_output_dir, + group_by="Client", + model_choice="test-model", + inference_server_model="test-model", + api_url="http://localhost:8080", # URL doesn't matter with function mocking + create_xlsx_output=False, + save_logs_to_csv=False, + ) + + self.assertTrue(result, "Topic extraction with grouping should succeed") + print("✅ Topic extraction with grouping passed") + + def test_extract_topics_with_candidate_topics(self): + """Test: Extract topics with candidate topics (zero-shot)""" + print("\n=== Testing topic extraction with candidate topics ===") + input_file = os.path.join( + self.example_data_dir, "dummy_consultation_response.csv" + ) + candidate_topics_file = os.path.join( + self.example_data_dir, "dummy_consultation_response_themes.csv" + ) + + if not os.path.isfile(input_file): + self.skipTest(f"Example file not found: {input_file}") + if not os.path.isfile(candidate_topics_file): + self.skipTest(f"Candidate topics file not found: {candidate_topics_file}") + + result = run_cli_topics( + script_path=self.script_path, + task="extract", + input_file=input_file, + text_column="Response text", + output_dir=self.temp_output_dir, + candidate_topics=candidate_topics_file, + model_choice="test-model", + inference_server_model="test-model", + api_url="http://localhost:8080", # URL doesn't matter with function mocking + create_xlsx_output=False, + save_logs_to_csv=False, + ) + + self.assertTrue(result, "Topic extraction with candidate topics should succeed") + print("✅ Topic extraction with candidate topics passed") + + def test_deduplicate_topics_fuzzy(self): + """Test: Deduplicate topics using fuzzy matching""" + print("\n=== Testing topic deduplication with fuzzy matching ===") + + # First, we need to create some output files by running extraction + input_file = os.path.join(self.example_data_dir, "combined_case_notes.csv") + + if not os.path.isfile(input_file): + self.skipTest(f"Example file not found: {input_file}") + + # Run extraction first to create output files + extract_result = run_cli_topics( + script_path=self.script_path, + task="extract", + input_file=input_file, + text_column="Case Note", + output_dir=self.temp_output_dir, + model_choice="test-model", + inference_server_model="test-model", + api_url="http://localhost:8080", # URL doesn't matter with function mocking + create_xlsx_output=False, + save_logs_to_csv=False, + ) + + if not extract_result: + self.skipTest("Extraction failed, cannot test deduplication") + + # Find the output files (they should be in temp_output_dir) + # The file names follow a pattern like: {input_file_name}_col_{text_column}_reference_table.csv + import glob + + reference_files = glob.glob( + os.path.join(self.temp_output_dir, "*reference_table.csv") + ) + unique_files = glob.glob( + os.path.join(self.temp_output_dir, "*unique_topics.csv") + ) + + if not reference_files or not unique_files: + self.skipTest("Could not find output files from extraction") + + result = run_cli_topics( + script_path=self.script_path, + task="deduplicate", + previous_output_files=[reference_files[0], unique_files[0]], + output_dir=self.temp_output_dir, + method="fuzzy", + similarity_threshold=90, + create_xlsx_output=False, + save_logs_to_csv=False, + ) + + self.assertTrue( + result, "Topic deduplication with fuzzy matching should succeed" + ) + print("✅ Topic deduplication with fuzzy matching passed") + + def test_deduplicate_topics_llm(self): + """Test: Deduplicate topics using LLM""" + print("\n=== Testing topic deduplication with LLM ===") + + # First, we need to create some output files by running extraction + input_file = os.path.join(self.example_data_dir, "combined_case_notes.csv") + + if not os.path.isfile(input_file): + self.skipTest(f"Example file not found: {input_file}") + + # Run extraction first to create output files + extract_result = run_cli_topics( + script_path=self.script_path, + task="extract", + input_file=input_file, + text_column="Case Note", + output_dir=self.temp_output_dir, + model_choice="test-model", + inference_server_model="test-model", + api_url="http://localhost:8080", # URL doesn't matter with function mocking + create_xlsx_output=False, + save_logs_to_csv=False, + ) + + if not extract_result: + self.skipTest("Extraction failed, cannot test deduplication") + + # Find the output files + import glob + + reference_files = glob.glob( + os.path.join(self.temp_output_dir, "*reference_table.csv") + ) + unique_files = glob.glob( + os.path.join(self.temp_output_dir, "*unique_topics.csv") + ) + + if not reference_files or not unique_files: + self.skipTest("Could not find output files from extraction") + + result = run_cli_topics( + script_path=self.script_path, + task="deduplicate", + previous_output_files=[reference_files[0], unique_files[0]], + output_dir=self.temp_output_dir, + method="llm", + model_choice="test-model", + inference_server_model="test-model", + api_url="http://localhost:8080", # URL doesn't matter with function mocking + create_xlsx_output=False, + save_logs_to_csv=False, + ) + + self.assertTrue(result, "Topic deduplication with LLM should succeed") + print("✅ Topic deduplication with LLM passed") + + def test_all_in_one_pipeline(self): + """Test: Run complete pipeline (extract, deduplicate, summarise)""" + print("\n=== Testing all-in-one pipeline ===") + input_file = os.path.join(self.example_data_dir, "combined_case_notes.csv") + + if not os.path.isfile(input_file): + self.skipTest(f"Example file not found: {input_file}") + + result = run_cli_topics( + script_path=self.script_path, + task="all_in_one", + input_file=input_file, + text_column="Case Note", + output_dir=self.temp_output_dir, + model_choice="test-model", + inference_server_model="test-model", + api_url="http://localhost:8080", # URL doesn't matter with function mocking + create_xlsx_output=False, + save_logs_to_csv=False, + timeout=120, # Shorter timeout for debugging + ) + + self.assertTrue(result, "All-in-one pipeline should succeed") + print("✅ All-in-one pipeline passed") + + def test_direct_mode_extract(self): + """Test: Run app in direct mode for topic extraction""" + print("\n=== Testing direct mode - topic extraction ===") + input_file = os.path.join(self.example_data_dir, "combined_case_notes.csv") + + if not os.path.isfile(input_file): + self.skipTest(f"Example file not found: {input_file}") + + app_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "app.py") + + if not os.path.isfile(app_path): + self.skipTest(f"App script not found: {app_path}") + + result = run_app_direct_mode( + app_path=app_path, + task="extract", + input_file=input_file, + text_column="Case Note", + output_dir=self.temp_output_dir, + model_choice="test-model", + inference_server_model="test-model", + api_url="http://localhost:8080", + create_xlsx_output=False, + save_logs_to_csv=False, + ) + + self.assertTrue(result, "Direct mode topic extraction should succeed") + print("✅ Direct mode topic extraction passed") + + +def run_all_tests(): + """Run all test examples and report results.""" + print("=" * 80) + print("LLM TOPIC MODELLER TEST SUITE") + print("=" * 80) + print("This test suite includes:") + print("- CLI examples from the epilog") + print("- GUI application tests") + print("- Tests use a mock inference-server to avoid API costs") + print("Tests will be skipped if required example files are not found.") + print("=" * 80) + + # Create test suite + loader = unittest.TestLoader() + suite = unittest.TestSuite() + + # Add CLI tests + cli_suite = loader.loadTestsFromTestCase(TestCLITopicsExamples) + suite.addTests(cli_suite) + + # Add GUI tests + try: + from test.test_gui_only import TestGUIAppOnly + + gui_suite = loader.loadTestsFromTestCase(TestGUIAppOnly) + suite.addTests(gui_suite) + print("GUI tests included in test suite.") + except ImportError as e: + print(f"Warning: Could not import GUI tests: {e}") + print("Skipping GUI tests.") + + # Run tests with detailed output + runner = unittest.TextTestRunner(verbosity=2, stream=None) + result = runner.run(suite) + + # Print summary + print("\n" + "=" * 80) + print("TEST SUMMARY") + print("=" * 80) + print(f"Tests run: {result.testsRun}") + print(f"Failures: {len(result.failures)}") + print(f"Errors: {len(result.errors)}") + print(f"Skipped: {len(result.skipped) if hasattr(result, 'skipped') else 0}") + + if result.failures: + print("\nFAILURES:") + for test, traceback in result.failures: + print(f"- {test}: {traceback}") + + if result.errors: + print("\nERRORS:") + for test, traceback in result.errors: + print(f"- {test}: {traceback}") + + success = len(result.failures) == 0 and len(result.errors) == 0 + print(f"\nOverall result: {'✅ PASSED' if success else '❌ FAILED'}") + print("=" * 80) + + return success + + +if __name__ == "__main__": + # Run the test suite + success = run_all_tests() + exit(0 if success else 1) diff --git a/test/test_gui_only.py b/test/test_gui_only.py new file mode 100644 index 0000000000000000000000000000000000000000..6debcee14cae4e2676b4b578b7a8e9731d209ea2 --- /dev/null +++ b/test/test_gui_only.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +""" +Standalone GUI test script for the LLM topic modeller application. + +This script tests only the GUI functionality of app.py to ensure it loads correctly. +Run this script to verify that the Gradio interface can be imported and initialized. +""" + +import os +import sys +import threading +import unittest + +# Add the parent directory to the path so we can import the app +parent_dir = os.path.dirname(os.path.dirname(__file__)) +if parent_dir not in sys.path: + sys.path.insert(0, parent_dir) + + +class TestGUIAppOnly(unittest.TestCase): + """Test suite for GUI application loading and basic functionality.""" + + @classmethod + def setUpClass(cls): + """Set up test environment for GUI tests.""" + cls.app_path = os.path.join(parent_dir, "app.py") + + # Verify app.py exists + if not os.path.isfile(cls.app_path): + raise FileNotFoundError(f"App file not found: {cls.app_path}") + + print(f"GUI test setup complete. App: {cls.app_path}") + + def test_app_import_and_initialization(self): + """Test: Import app.py and check if the Gradio app object is created successfully.""" + print("\n=== Testing GUI app import and initialization ===") + + try: + # Import the app module + import app + + # Check if the app object exists and is a Gradio Blocks object + self.assertTrue( + hasattr(app, "app"), "App object should exist in the module" + ) + + # Check if it's a Gradio Blocks instance + import gradio as gr + + self.assertIsInstance( + app.app, gr.Blocks, "App should be a Gradio Blocks instance" + ) + + print("✅ GUI app import and initialization passed") + + except ImportError as e: + error_msg = f"Failed to import app module: {e}" + self.fail(error_msg) + except Exception as e: + self.fail(f"Unexpected error during app initialization: {e}") + + def test_app_launch_headless(self): + """Test: Launch the app in headless mode to verify it starts without errors.""" + print("\n=== Testing GUI app launch in headless mode ===") + + try: + # Import the app module + import app + + # Set up a flag to track if the app launched successfully + app_launched = threading.Event() + launch_error = None + + def launch_app(): + try: + # Launch the app in headless mode with a short timeout + app.app.launch( + show_error=True, + inbrowser=False, # Don't open browser + server_port=0, # Use any available port + quiet=True, # Suppress output + prevent_thread_lock=True, # Don't block the main thread + ) + app_launched.set() + except Exception: + app_launched.set() + + # Start the app in a separate thread + launch_thread = threading.Thread(target=launch_app) + launch_thread.daemon = True + launch_thread.start() + + # Wait for the app to launch (with timeout) + if app_launched.wait(timeout=10): # 10 second timeout + if launch_error: + self.fail(f"App launch failed: {launch_error}") + else: + print("✅ GUI app launch in headless mode passed") + else: + self.fail("App launch timed out after 10 seconds") + + except Exception as e: + error_msg = f"Unexpected error during app launch test: {e}" + self.fail(error_msg) + + def test_app_configuration_loading(self): + """Test: Verify that the app can load its configuration without errors.""" + print("\n=== Testing GUI app configuration loading ===") + + try: + # Check if key configuration variables are accessible + # These should be imported from tools.config + from tools.config import ( + DEFAULT_COST_CODE, + GRADIO_SERVER_PORT, + MAX_FILE_SIZE, + default_model_choice, + model_name_map, + ) + + # Verify these are not None/empty + self.assertIsNotNone( + GRADIO_SERVER_PORT, "GRADIO_SERVER_PORT should be configured" + ) + self.assertIsNotNone(MAX_FILE_SIZE, "MAX_FILE_SIZE should be configured") + self.assertIsNotNone( + DEFAULT_COST_CODE, "DEFAULT_COST_CODE should be configured" + ) + self.assertIsNotNone( + default_model_choice, "default_model_choice should be configured" + ) + self.assertIsNotNone(model_name_map, "model_name_map should be configured") + + print("✅ GUI app configuration loading passed") + + except ImportError as e: + error_msg = f"Failed to import configuration: {e}" + self.fail(error_msg) + except Exception as e: + error_msg = f"Unexpected error during configuration test: {e}" + self.fail(error_msg) + + +def run_gui_tests(): + """Run GUI tests and report results.""" + print("=" * 80) + print("LLM TOPIC MODELLER GUI TEST SUITE") + print("=" * 80) + print("This test suite verifies that the GUI application loads correctly.") + print("=" * 80) + + # Create test suite + loader = unittest.TestLoader() + suite = loader.loadTestsFromTestCase(TestGUIAppOnly) + + # Run tests with detailed output + runner = unittest.TextTestRunner(verbosity=2, stream=None) + result = runner.run(suite) + + # Print summary + print("\n" + "=" * 80) + print("GUI TEST SUMMARY") + print("=" * 80) + print(f"Tests run: {result.testsRun}") + print(f"Failures: {len(result.failures)}") + print(f"Errors: {len(result.errors)}") + print(f"Skipped: {len(result.skipped) if hasattr(result, 'skipped') else 0}") + + if result.failures: + print("\nFAILURES:") + for test, traceback in result.failures: + print(f"- {test}: {traceback}") + + if result.errors: + print("\nERRORS:") + for test, traceback in result.errors: + print(f"- {test}: {traceback}") + + success = len(result.failures) == 0 and len(result.errors) == 0 + print(f"\nOverall result: {'✅ PASSED' if success else '❌ FAILED'}") + print("=" * 80) + + return success + + +if __name__ == "__main__": + # Run the GUI test suite + success = run_gui_tests() + exit(0 if success else 1) diff --git a/test/test_topic_discovery.py b/test/test_topic_discovery.py new file mode 100644 index 0000000000000000000000000000000000000000..2e8df8957a0555885f3eae166f96f36a9583e48d --- /dev/null +++ b/test/test_topic_discovery.py @@ -0,0 +1,107 @@ +import os +import sys +import tempfile +import unittest + +import pandas as pd + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from tools.helper_functions import ( + create_candidate_topics_df_from_topic_summary, + subsample_responses_for_topic_discovery, + write_candidate_topics_csv, + write_topic_discovery_manifest_csv, +) + + +class TestSubsampleResponsesForTopicDiscovery(unittest.TestCase): + def setUp(self): + self.df = pd.DataFrame( + { + "Group": ["A"] * 5 + ["B"] * 5, + "Response text": [f"response {i}" for i in range(10)], + } + ) + + def test_global_sample_respects_fraction_and_seed(self): + sampled_a, meta_a = subsample_responses_for_topic_discovery( + self.df, sample_fraction=0.2, random_seed=42 + ) + sampled_b, meta_b = subsample_responses_for_topic_discovery( + self.df, sample_fraction=0.2, random_seed=42 + ) + + self.assertEqual(len(sampled_a), 2) + self.assertEqual(meta_a["original_rows"], 10) + self.assertEqual(meta_a["sampled_rows"], 2) + pd.testing.assert_frame_equal( + sampled_a.drop(columns=["_discovery_original_row_index"]), + sampled_b.drop(columns=["_discovery_original_row_index"]), + ) + + def test_stratified_sample_includes_each_group(self): + sampled, meta = subsample_responses_for_topic_discovery( + self.df, + sample_fraction=0.2, + random_seed=7, + group_col="Group", + ) + + self.assertGreaterEqual(len(sampled), 2) + self.assertIn("A", sampled["Group"].values) + self.assertIn("B", sampled["Group"].values) + self.assertEqual(meta["per_group_counts"]["A"]["original"], 5) + self.assertEqual(meta["per_group_counts"]["B"]["original"], 5) + + def test_empty_data_raises(self): + with self.assertRaises(ValueError): + subsample_responses_for_topic_discovery(pd.DataFrame(), 0.2, 42) + + def test_invalid_fraction_raises(self): + with self.assertRaises(ValueError): + subsample_responses_for_topic_discovery(self.df, 0, 42) + + +class TestCandidateTopicsCsv(unittest.TestCase): + def test_create_and_write_candidate_topics_csv(self): + topic_summary_df = pd.DataFrame( + { + "Group": ["G1", "G1", "G2"], + "Sentiment": ["Positive", "Negative", "Positive"], + "General topic": ["Housing", "Housing", "Transport"], + "Subtopic": ["Rent", "Rent", "Buses"], + } + ) + + topics_df = create_candidate_topics_df_from_topic_summary(topic_summary_df) + self.assertEqual(len(topics_df), 2) + self.assertListEqual(list(topics_df.columns), ["General topic", "Subtopic"]) + + with tempfile.TemporaryDirectory() as tmpdir: + csv_path = os.path.join(tmpdir, "topics.csv") + written = write_candidate_topics_csv(topic_summary_df, csv_path) + self.assertEqual(written, csv_path) + loaded = pd.read_csv(csv_path) + self.assertEqual(len(loaded), 2) + + def test_manifest_csv_lists_sampled_rows(self): + sampled_df = pd.DataFrame( + { + "_discovery_original_row_index": [0, 4, 7], + "Group": ["A", "A", "B"], + } + ) + with tempfile.TemporaryDirectory() as tmpdir: + manifest_path = os.path.join(tmpdir, "manifest.csv") + written = write_topic_discovery_manifest_csv( + sampled_df, manifest_path, group_col="Group" + ) + self.assertEqual(written, manifest_path) + loaded = pd.read_csv(manifest_path) + self.assertListEqual(list(loaded.columns), ["Original row index", "Group"]) + self.assertEqual(len(loaded), 3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tools/auth.py b/tools/auth.py new file mode 100644 index 0000000000000000000000000000000000000000..c485d584a363306e3702258cdda25c5b7dd3671d --- /dev/null +++ b/tools/auth.py @@ -0,0 +1,85 @@ +import base64 +import hashlib +import hmac + +import boto3 + +from tools.config import AWS_CLIENT_ID, AWS_CLIENT_SECRET, AWS_REGION, AWS_USER_POOL_ID + + +def calculate_secret_hash(client_id: str, client_secret: str, username: str): + message = username + client_id + dig = hmac.new( + str(client_secret).encode("utf-8"), + msg=str(message).encode("utf-8"), + digestmod=hashlib.sha256, + ).digest() + secret_hash = base64.b64encode(dig).decode() + return secret_hash + + +def authenticate_user( + username: str, + password: str, + user_pool_id: str = AWS_USER_POOL_ID, + client_id: str = AWS_CLIENT_ID, + client_secret: str = AWS_CLIENT_SECRET, +): + """Authenticates a user against an AWS Cognito user pool. + + Args: + user_pool_id (str): The ID of the Cognito user pool. + client_id (str): The ID of the Cognito user pool client. + username (str): The username of the user. + password (str): The password of the user. + client_secret (str): The client secret of the app client + + Returns: + bool: True if the user is authenticated, False otherwise. + """ + + client = boto3.client( + "cognito-idp", region_name=AWS_REGION + ) # Cognito Identity Provider client + + # Compute the secret hash + secret_hash = calculate_secret_hash(client_id, client_secret, username) + + try: + + if client_secret == "": + response = client.initiate_auth( + AuthFlow="USER_PASSWORD_AUTH", + AuthParameters={ + "USERNAME": username, + "PASSWORD": password, + }, + ClientId=client_id, + ) + + else: + response = client.initiate_auth( + AuthFlow="USER_PASSWORD_AUTH", + AuthParameters={ + "USERNAME": username, + "PASSWORD": password, + "SECRET_HASH": secret_hash, + }, + ClientId=client_id, + ) + + # If successful, you'll receive an AuthenticationResult in the response + if response.get("AuthenticationResult"): + return True + else: + return False + + except client.exceptions.NotAuthorizedException: + return False + except client.exceptions.UserNotFoundException: + return False + except Exception as e: + out_message = f"An error occurred: {e}" + print(out_message) + raise Exception(out_message) + return False diff --git a/tools/aws_functions.py b/tools/aws_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..5aba221658b3b276f9ce3db9125daad1b7496e52 --- /dev/null +++ b/tools/aws_functions.py @@ -0,0 +1,534 @@ +import os +from typing import List + +import boto3 + +from tools.config import ( + ACCESS_LOGS_FOLDER, + AWS_ACCESS_KEY, + AWS_REGION, + AWS_SECRET_KEY, + DIRECT_MODE_OUTPUT_DIR, + FEEDBACK_LOGS_FOLDER, + GRADIO_TEMP_DIR, + INPUT_FOLDER, + OUTPUT_FOLDER, + PRIORITISE_SSO_OVER_AWS_ENV_ACCESS_KEYS, + RUN_AWS_FUNCTIONS, + S3_LOG_BUCKET, + S3_OUTPUTS_BUCKET, + UPLOAD_PROMPT_RESPONSE_LOG_TO_S3_OUTPUTS, + USAGE_LOGS_FOLDER, +) + +# Empty bucket name in case authentication fails +bucket_name = S3_LOG_BUCKET + + +def connect_to_bedrock_runtime( + model_name_map: dict, + model_choice: str, + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", +): + # If running an anthropic model, assume that running an AWS Bedrock model, load in Bedrock + model_source = model_name_map[model_choice]["source"] + + # Use aws_region_textbox if provided, otherwise fall back to AWS_REGION from config + region = aws_region_textbox if aws_region_textbox else AWS_REGION + + if "AWS" in model_source: + if RUN_AWS_FUNCTIONS == "1" and PRIORITISE_SSO_OVER_AWS_ENV_ACCESS_KEYS == "1": + print("Connecting to Bedrock via existing SSO connection") + bedrock_runtime = boto3.client("bedrock-runtime", region_name=region) + elif aws_access_key_textbox and aws_secret_key_textbox: + print( + "Connecting to Bedrock using AWS access key and secret keys from user input." + ) + bedrock_runtime = boto3.client( + "bedrock-runtime", + aws_access_key_id=aws_access_key_textbox, + aws_secret_access_key=aws_secret_key_textbox, + region_name=region, + ) + elif AWS_ACCESS_KEY and AWS_SECRET_KEY: + print("Getting Bedrock credentials from environment variables") + bedrock_runtime = boto3.client( + "bedrock-runtime", + aws_access_key_id=AWS_ACCESS_KEY, + aws_secret_access_key=AWS_SECRET_KEY, + region_name=region, + ) + elif RUN_AWS_FUNCTIONS == "1": + print("Connecting to Bedrock via existing SSO connection") + bedrock_runtime = boto3.client("bedrock-runtime", region_name=region) + else: + bedrock_runtime = "" + out_message = "Cannot connect to AWS Bedrock service. Please provide access keys under LLM settings, or choose another model type." + print(out_message) + raise Exception(out_message) + else: + bedrock_runtime = None + + return bedrock_runtime + + +def connect_to_s3_client( + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", +): + # If running an anthropic model, assume that running an AWS s3 model, load in s3 + s3_client = None + + # Use aws_region_textbox if provided, otherwise fall back to AWS_REGION from config + region = aws_region_textbox if aws_region_textbox else AWS_REGION + + if aws_access_key_textbox and aws_secret_key_textbox: + print("Connecting to s3 using AWS access key and secret keys from user input.") + s3_client = boto3.client( + "s3", + aws_access_key_id=aws_access_key_textbox, + aws_secret_access_key=aws_secret_key_textbox, + region_name=region, + ) + elif RUN_AWS_FUNCTIONS == "1" and PRIORITISE_SSO_OVER_AWS_ENV_ACCESS_KEYS == "1": + print("Connecting to s3 via existing SSO connection") + s3_client = boto3.client("s3", region_name=region) + elif AWS_ACCESS_KEY and AWS_SECRET_KEY: + print("Getting s3 credentials from environment variables") + s3_client = boto3.client( + "s3", + aws_access_key_id=AWS_ACCESS_KEY, + aws_secret_access_key=AWS_SECRET_KEY, + region_name=region, + ) + elif RUN_AWS_FUNCTIONS == "1": + print("Connecting to s3 via existing SSO connection") + s3_client = boto3.client("s3", region_name=region) + else: + s3_client = "" + out_message = "Cannot connect to S3 service. Please provide access keys under LLM settings, or choose another model type." + print(out_message) + raise Exception(out_message) + + return s3_client + + +# Download direct from S3 - requires login credentials +def download_file_from_s3( + bucket_name: str, + key: str, + local_file_path: str, + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + RUN_AWS_FUNCTIONS=RUN_AWS_FUNCTIONS, +): + + if RUN_AWS_FUNCTIONS == "1": + + s3 = connect_to_s3_client( + aws_access_key_textbox, aws_secret_key_textbox, aws_region_textbox + ) + # boto3.client('s3') + s3.download_file(bucket_name, key, local_file_path) + print(f"File downloaded from S3 to {local_file_path}") + + +def download_folder_from_s3( + bucket_name: str, + s3_folder: str, + local_folder: str, + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + RUN_AWS_FUNCTIONS=RUN_AWS_FUNCTIONS, +): + """ + Download all files from an S3 folder to a local folder. + """ + if RUN_AWS_FUNCTIONS == "1": + s3 = connect_to_s3_client( + aws_access_key_textbox, aws_secret_key_textbox, aws_region_textbox + ) + # boto3.client('s3') + + # List objects in the specified S3 folder + response = s3.list_objects_v2(Bucket=bucket_name, Prefix=s3_folder) + + # Download each object + for obj in response.get("Contents", []): + # Extract object key and construct local file path + object_key = obj["Key"] + local_file_path = os.path.join( + local_folder, os.path.relpath(object_key, s3_folder) + ) + + # Create directories if necessary + os.makedirs(os.path.dirname(local_file_path), exist_ok=True) + + # Download the object + try: + s3.download_file(bucket_name, object_key, local_file_path) + print(f"Downloaded file from S3 to {local_file_path}") + except Exception as e: + print(f"Error downloading file from S3: {e}") + + +def download_files_from_s3( + bucket_name: str, + s3_folder: str, + local_folder: str, + filenames: list[str], + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + RUN_AWS_FUNCTIONS=RUN_AWS_FUNCTIONS, +): + """ + Download specific files from an S3 folder to a local folder. + """ + if RUN_AWS_FUNCTIONS == "1": + s3 = connect_to_s3_client( + aws_access_key_textbox, aws_secret_key_textbox, aws_region_textbox + ) + # boto3.client('s3') + + print("Trying to download file: ", filenames) + + if filenames == "*": + # List all objects in the S3 folder + print("Trying to download all files in AWS folder: ", s3_folder) + response = s3.list_objects_v2(Bucket=bucket_name, Prefix=s3_folder) + + # print("Found files in AWS folder: ", response.get("Contents", [])) + + filenames = [ + obj["Key"].split("/")[-1] for obj in response.get("Contents", []) + ] + + # print("Found filenames in AWS folder: ", filenames) + + for filename in filenames: + object_key = os.path.join(s3_folder, filename) + local_file_path = os.path.join(local_folder, filename) + + # Create directories if necessary + os.makedirs(os.path.dirname(local_file_path), exist_ok=True) + + # Download the object + try: + s3.download_file(bucket_name, object_key, local_file_path) + print(f"Downloaded file from S3 to {local_file_path}") + except Exception as e: + print(f"Error downloading file from S3: {e}") + + +def upload_file_to_s3( + local_file_paths: List[str], + s3_key: str, + s3_bucket: str = bucket_name, + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + RUN_AWS_FUNCTIONS=RUN_AWS_FUNCTIONS, +): + """ + Uploads a file from local machine to Amazon S3. + + Args: + - local_file_path: Local file path(s) of the file(s) to upload. + - s3_key: Key (path) to the file in the S3 bucket. + - s3_bucket: Name of the S3 bucket. + + Returns: + - Message as variable/printed to console + """ + if RUN_AWS_FUNCTIONS == "1": + + final_out_message = list() + + s3_client = connect_to_s3_client( + aws_access_key_textbox, aws_secret_key_textbox, aws_region_textbox + ) + # boto3.client('s3') + + if isinstance(local_file_paths, str): + local_file_paths = [local_file_paths] + + for file in local_file_paths: + try: + # Get file name off file path + file_name = os.path.basename(file) + + s3_key_full = s3_key + file_name + # print("S3 key: ", s3_key_full) + + s3_client.upload_file(file, s3_bucket, s3_key_full) + out_message = "File " + file_name + " uploaded successfully to S3" + print(out_message) + + except Exception as e: + out_message = f"Error uploading file(s): {e}" + print(out_message) + + final_out_message.append(out_message) + final_out_message_str = "\n".join(final_out_message) + + else: + final_out_message_str = "Not connected to AWS, no files uploaded." + + return final_out_message_str + + +# Helper to upload outputs to S3 when enabled in config. +def export_outputs_to_s3( + file_list_state, + s3_output_folder_state_value: str, + save_outputs_to_s3_flag: bool, + base_file_state=None, + s3_bucket: str = S3_OUTPUTS_BUCKET, +): + """ + Upload a list of local output files to the configured S3 outputs folder. + + - file_list_state: Gradio dropdown state that holds a list of file paths or a + single path/string. If blank/empty, no action is taken. + - s3_output_folder_state_value: Final S3 key prefix (including any session hash) + to use as the destination folder for uploads. + - s3_bucket: Name of the S3 bucket. + """ + try: + + # Respect the runtime toggle as well as environment configuration + if not save_outputs_to_s3_flag: + return + + if not s3_output_folder_state_value: + # No configured S3 outputs folder – nothing to do + return + + # Normalise input to a Python list of strings + file_paths = file_list_state + if not file_paths: + return + + # Gradio dropdown may return a single string or a list + if isinstance(file_paths, str): + file_paths = [file_paths] + + # Filter out any non-truthy values + file_paths = [p for p in file_paths if p] + if not file_paths: + return + + # Derive a base file stem (name without extension) from the original + # file(s) being analysed, if provided. This is used to create an + # additional subfolder layer so that outputs are grouped under the + # analysed file name rather than under each output file name. + base_stem = None + if base_file_state: + base_path = None + + # Gradio File components typically provide a list of objects with a `.name` attribute + if isinstance(base_file_state, str): + base_path = base_file_state + elif isinstance(base_file_state, list) and base_file_state: + first_item = base_file_state[0] + base_path = getattr(first_item, "name", None) or str(first_item) + else: + base_path = getattr(base_file_state, "name", None) or str( + base_file_state + ) + + if base_path: + base_name = os.path.basename(base_path) + base_stem, _ = os.path.splitext(base_name) + + # Ensure base S3 prefix (session/date) ends with a trailing slash + base_prefix = s3_output_folder_state_value + if not base_prefix.endswith("/"): + base_prefix = base_prefix + "/" + + # For each file, append a subfolder. If we have a derived base_stem + # from the input being analysed, use that; otherwise, fall back to + # the individual output file name stem. Final pattern: + # /// + # or, if base_file_stem is not available: + # /// + for file in file_paths: + file_name = os.path.basename(file) + + if base_stem: + folder_stem = base_stem + else: + folder_stem, _ = os.path.splitext(file_name) + + per_file_prefix = base_prefix + folder_stem + "/" + + out_message = upload_file_to_s3( + local_file_paths=[file], + s3_key=per_file_prefix, + s3_bucket=s3_bucket, + ) + + # Log any issues to console so failures are visible in logs/stdout + if ( + "Error uploading file" in out_message + or "could not upload" in out_message.lower() + ): + print("export_outputs_to_s3 encountered issues:", out_message) + + print("Successfully uploaded outputs to S3") + + except Exception as e: + # Do not break the app flow if S3 upload fails – just report to console + print(f"export_outputs_to_s3 failed with error: {e}") + + # No GUI outputs to update + return + + +def _normalize_file_paths(file_paths) -> List[str]: + """Flatten Gradio file inputs (paths, lists, or file objects) into path strings.""" + if not file_paths: + return [] + if isinstance(file_paths, str): + return [file_paths] + if not isinstance(file_paths, list): + file_paths = [file_paths] + + result = [] + for item in file_paths: + if isinstance(item, str): + result.append(item) + elif isinstance(item, list): + result.extend(_normalize_file_paths(item)) + else: + path = getattr(item, "name", None) or str(item) + if path: + result.append(path) + return result + + +def _get_safe_local_path_roots() -> List[str]: + """Build allowlisted directory roots for local file access (resolved realpaths).""" + roots = { + OUTPUT_FOLDER, + DIRECT_MODE_OUTPUT_DIR, + GRADIO_TEMP_DIR, + INPUT_FOLDER, + USAGE_LOGS_FOLDER, + FEEDBACK_LOGS_FOLDER, + ACCESS_LOGS_FOLDER, + } + resolved = [] + for root in roots: + if root: + resolved.append(os.path.realpath(root)) + return resolved + + +_SAFE_LOCAL_PATH_ROOTS: List[str] | None = None + + +def _safe_local_path_roots() -> List[str]: + global _SAFE_LOCAL_PATH_ROOTS + if _SAFE_LOCAL_PATH_ROOTS is None: + _SAFE_LOCAL_PATH_ROOTS = _get_safe_local_path_roots() + return _SAFE_LOCAL_PATH_ROOTS + + +def _resolve_safe_local_path(path: str) -> str | None: + """Return canonical local path if inside an allowlisted root, else None.""" + if not path or not isinstance(path, str): + return None + try: + real_path = os.path.realpath(path) + except (OSError, ValueError): + return None + for root_real in _safe_local_path_roots(): + try: + if os.path.commonpath([real_path, root_real]) == root_real: + return real_path + except ValueError: + # Different drives on Windows or invalid path combinations + continue + return None + + +def is_prompt_response_log_file(file_path: str) -> bool: + """True when file_path is a prompt/response JSON log (*_logs_*.json).""" + if not file_path: + return False + basename = os.path.basename(file_path) + return basename.endswith(".json") and "_logs_" in basename + + +def collect_prompt_response_log_paths(log_paths) -> List[str]: + """Return existing prompt/response JSON log files from path lists or Gradio file inputs.""" + seen = set() + result = [] + for path in _normalize_file_paths(log_paths): + safe_path = _resolve_safe_local_path(path) + if not safe_path: + print(f"Skipping prompt/response log outside allowed directories: {path}") + continue + if safe_path in seen: + continue + if is_prompt_response_log_file(safe_path) and os.path.exists(safe_path): + result.append(safe_path) + seen.add(safe_path) + return result + + +def export_outputs_to_s3_with_prompt_response_logs( + file_list_state, + prompt_response_log_paths, + s3_output_folder_state_value: str, + save_outputs_to_s3_flag: bool, + base_file_state=None, + s3_bucket: str = S3_OUTPUTS_BUCKET, +): + """ + Upload output files to S3, optionally including prompt/response JSON logs + (*_logs_*.json) when UPLOAD_PROMPT_RESPONSE_LOG_TO_S3_OUTPUTS is enabled. + """ + files_to_upload = _normalize_file_paths(file_list_state) + + if UPLOAD_PROMPT_RESPONSE_LOG_TO_S3_OUTPUTS: + for log_path in collect_prompt_response_log_paths(prompt_response_log_paths): + if log_path not in files_to_upload: + files_to_upload.append(log_path) + print(f"Including prompt/response log in S3 upload: {log_path}") + + return export_outputs_to_s3( + file_list_state=files_to_upload, + s3_output_folder_state_value=s3_output_folder_state_value, + save_outputs_to_s3_flag=save_outputs_to_s3_flag, + base_file_state=base_file_state, + s3_bucket=s3_bucket, + ) + + +def download_cost_codes_with_error_handling(bucket, key, local_path): + """ + Wrapper function with error handling for downloading cost codes from S3. + + Args: + bucket: S3 bucket name + key: S3 object key + local_path: Local file path to save the downloaded file + + Returns: + bool: True if download successful, False otherwise + """ + try: + download_file_from_s3(bucket, key, local_path) + return True + except Exception as e: + print(f"Error downloading cost codes from S3: {e}") + print(f"Failed to download s3://{bucket}/{key}") + return False diff --git a/tools/combine_sheets_into_xlsx.py b/tools/combine_sheets_into_xlsx.py new file mode 100644 index 0000000000000000000000000000000000000000..85842c534fea71534edbc082e342fd7c53854957 --- /dev/null +++ b/tools/combine_sheets_into_xlsx.py @@ -0,0 +1,968 @@ +import os +import re +from datetime import date, datetime +from typing import List, Union + +import pandas as pd +from openpyxl import Workbook +from openpyxl.cell.rich_text import CellRichText, TextBlock +from openpyxl.cell.text import InlineFont +from openpyxl.styles import Alignment, Font +from openpyxl.utils import get_column_letter +from openpyxl.utils.dataframe import dataframe_to_rows + +from tools.config import EXPORT_FORMAT, OUTPUT_FOLDER +from tools.config import model_name_map as global_model_name_map +from tools.helper_functions import ( + clean_column_name, + convert_reference_table_to_pivot_table, + ensure_model_in_map, + get_basic_response_data, + load_in_data_file, + write_candidate_topics_csv, +) + + +def markdown_to_richtext( + text: Union[str, float, int, None], +) -> Union[CellRichText, str, float, int, None]: + """ + Convert markdown formatting in text to Excel RichText formatting. + + Supports: + - **text** or __text__ for bold + - *text* or _text_ for italic (when not bold) + - ***text*** or ___text___ for bold+italic + + Also removes HTML
tags (and variants like
and
). + + Args: + text: The text to convert (can be string, number, or None) + + Returns: + CellRichText object if markdown is found, otherwise returns the original value + """ + # Return non-string values as-is + if not isinstance(text, str): + return text + + # Remove
tags (case-insensitive, handles
,
,
) + text = re.sub(r"", "", text, flags=re.IGNORECASE) + + # Check if text contains markdown formatting + if not re.search(r"(\*\*|__|\*|_)(?=\S)", text): + return text + + # Create RichText object + rich_text = CellRichText() + + # Process in order: triple markers first, then double, then single + # This prevents conflicts (e.g., ***text*** being matched as **text** + *text*) + # Pattern order matters: longer patterns first + # Use word boundaries to ensure markers are not part of words (e.g., filenames with underscores) + # (? last_pos: + plain_text = text[last_pos:start] + if plain_text: + rich_text.append(plain_text) + + # Create font for this segment (use InlineFont for RichText) + font = InlineFont(b=is_bold, i=is_italic) + rich_text.append(TextBlock(font, content)) + + last_pos = end + + # Add remaining plain text + if last_pos < len(text): + remaining = text[last_pos:] + if remaining: + rich_text.append(remaining) + + # If we didn't add anything, return original text + if len(rich_text) == 0: + return text + + return rich_text + + +def _resolve_output_path(candidate_path: str, allowed_root: str = OUTPUT_FOLDER) -> str: + """ + Resolve and validate that a path is contained within allowed_root. + """ + safe_root = os.path.realpath(os.path.abspath(allowed_root)) + resolved_path = os.path.realpath(os.path.abspath(candidate_path)) + try: + common = os.path.commonpath([safe_root, resolved_path]) + except ValueError: + raise ValueError( + f"Path '{candidate_path}' is outside allowed output folder" + ) from None + if common != safe_root: + raise ValueError(f"Path '{candidate_path}' is outside allowed output folder") + return resolved_path + + +def convert_xlsx_to_ods( + xlsx_path: str, ods_path: str, allowed_root: str = OUTPUT_FOLDER +): + """ + Convert an Excel (.xlsx) file to OpenDocument Spreadsheet (.ods) format. + + Args: + xlsx_path (str): Path to the source Excel file + ods_path (str): Path where the ODS file should be saved + """ + try: + import pyexcel + + safe_xlsx_path = _resolve_output_path(xlsx_path, allowed_root) + safe_ods_path = _resolve_output_path(ods_path, allowed_root) + + pyexcel.save_as(file_name=safe_xlsx_path, dest_file_name=safe_ods_path) + print(f"Output ods summary saved as '{safe_ods_path}'") + + if os.path.exists(safe_xlsx_path): + os.remove(safe_xlsx_path) + + return safe_ods_path + except ImportError: + print( + "Warning: pyexcel or pyexcel_ods not installed. Install with: " + "pip install pyexcel pyexcel-ods\n" + "Keeping output as xlsx format instead." + ) + return xlsx_path + except Exception as e: + print( + f"Warning: Could not convert xlsx to ods due to: {e}\n" + "Keeping output as xlsx format instead." + ) + return xlsx_path + + +def _resolve_candidate_topics_file_name(candidate_topics) -> str: + """Return basename of the suggested topics file, or empty if not provided.""" + if candidate_topics is None: + return "" + + if isinstance(candidate_topics, list): + if not candidate_topics: + return "" + candidate_topics = candidate_topics[0] + + if isinstance(candidate_topics, str): + path = candidate_topics.strip() + if not path: + return "" + return os.path.basename(path) + + path = getattr(candidate_topics, "name", None) + if path: + return os.path.basename(str(path)) + + return "" + + +def _resolve_excel_sheet_display_name( + file_path: str, excel_sheets: Union[str, List[str], None] +) -> str: + """Return the Excel sheet name for cover-sheet display, or empty if not applicable.""" + if os.path.splitext(file_path)[1].lower() != ".xlsx": + return "" + + if excel_sheets: + if isinstance(excel_sheets, list): + sheet_names = [str(s).strip() for s in excel_sheets if str(s).strip()] + if sheet_names: + return ", ".join(sheet_names) + elif str(excel_sheets).strip(): + return str(excel_sheets).strip() + + try: + return pd.ExcelFile(file_path).sheet_names[0] + except Exception: + return "" + + +def add_cover_sheet( + wb: Workbook, + intro_paragraphs: list[str], + model_name: str, + analysis_date: str, + analysis_cost: str, + number_of_responses: int, + number_of_responses_with_text: int, + number_of_responses_with_text_five_plus_words: int, + llm_call_number: int, + input_tokens: int, + output_tokens: int, + time_taken: float, + file_name: str, + column_name: str, + number_of_responses_with_topic_assignment: int, + excel_sheet_name: str = "", + candidate_topics_file_name: str = "", + custom_title: str = "Cover sheet", +): + ws = wb.create_sheet(title=custom_title, index=0) + + # Freeze top row + ws.freeze_panes = "A2" + + # Write title + ws["A1"] = "Large Language Model thematic analysis" + ws["A1"].font = Font(size=14, bold=True) + ws["A1"].alignment = Alignment(wrap_text=True, vertical="top") + + # Add intro paragraphs + row = 3 + for paragraph in intro_paragraphs: + ws.merge_cells(start_row=row, start_column=1, end_row=row, end_column=2) + formatted_paragraph = markdown_to_richtext(paragraph) + cell = ws.cell(row=row, column=1, value=formatted_paragraph) + cell.alignment = Alignment(wrap_text=True, vertical="top") + ws.row_dimensions[row].height = 60 # Adjust height as needed + row += 2 + + # Add metadata + meta_start = row + 1 + metadata = { + "Date Excel file created": date.today().strftime("%Y-%m-%d"), + "File name": file_name, + "Column name": column_name, + } + if excel_sheet_name: + metadata["Excel sheet name"] = excel_sheet_name + if candidate_topics_file_name: + metadata["Suggested topics file name"] = candidate_topics_file_name + metadata.update( + { + "Model name": model_name, + "Analysis date": analysis_date, + # "Analysis cost": analysis_cost, + "Number of responses": number_of_responses, + "Number of responses with text": number_of_responses_with_text, + "Number of responses with text five plus words": number_of_responses_with_text_five_plus_words, + "Number of responses with at least one assigned topic": number_of_responses_with_topic_assignment, + "Number of LLM calls": llm_call_number, + "Total number of input tokens from LLM calls": input_tokens, + "Total number of output tokens from LLM calls": output_tokens, + "Total time taken for all LLM calls (seconds)": round(float(time_taken), 1), + } + ) + + # Define which metadata fields should have number formatting with thousand separators + number_format_fields = { + "Number of responses": "#,##0", + "Number of responses with text": "#,##0", + "Number of responses with text five plus words": "#,##0", + "Number of responses with at least one assigned topic": "#,##0", + "Number of LLM calls": "#,##0", + "Total number of input tokens from LLM calls": "#,##0", + "Total number of output tokens from LLM calls": "#,##0", + "Total time taken for all LLM calls (seconds)": "#,##0.0", + } + + for i, (label, value) in enumerate(metadata.items()): + row_num = meta_start + i + ws[f"A{row_num}"] = label + ws[f"A{row_num}"].font = Font(bold=True) + + cell = ws[f"B{row_num}"] + # Convert markdown to RichText if applicable + formatted_value = markdown_to_richtext(value) + cell.value = formatted_value + # Set left alignment for all metadata values (including numbers) + cell.alignment = Alignment(horizontal="left", wrap_text=True) + + # Apply number formatting with thousand separators for numeric fields + if label in number_format_fields: + # Only apply formatting if value is not RichText (numbers can't be RichText) + if not isinstance(cell.value, CellRichText): + cell.number_format = number_format_fields[label] + + # Optional: Adjust column widths + ws.column_dimensions["A"].width = 50 + ws.column_dimensions["B"].width = 75 + + # Ensure first row cells are wrapped on the cover sheet + for col_idx in range(1, ws.max_column + 1): + header_cell = ws.cell(row=1, column=col_idx) + header_cell.alignment = Alignment(wrap_text=True, vertical="center") + + +def csvs_to_excel( + csv_files: list[str], + output_filename: str, + allowed_root: str = OUTPUT_FOLDER, + sheet_names: list[str] = None, + column_widths: dict = None, # Dict of {sheet_name: {col_letter: width}} + wrap_text_columns: dict = None, # Dict of {sheet_name: [col_letters]} + intro_text: list[str] = None, + model_name: str = "", + analysis_date: str = "", + analysis_cost: str = "", + llm_call_number: int = 0, + input_tokens: int = 0, + output_tokens: int = 0, + time_taken: float = 0, + number_of_responses: int = 0, + number_of_responses_with_text: int = 0, + number_of_responses_with_text_five_plus_words: int = 0, + column_name: str = "", + number_of_responses_with_topic_assignment: int = 0, + file_name: str = "", + excel_sheet_name: str = "", + candidate_topics_file_name: str = "", + unique_reference_numbers: list = [], +): + if intro_text is None: + intro_text = list() + + wb = Workbook() + # Remove default sheet + wb.remove(wb.active) + + def _read_csv_with_fallback_encodings(path: str) -> pd.DataFrame: + encodings_to_try = ["utf-8", "utf-8-sig", "cp1252", "latin1"] + last_err = None + for enc in encodings_to_try: + try: + return pd.read_csv(path, encoding=enc) + except UnicodeDecodeError as e: + last_err = e + continue + except Exception: + raise + + # Final fallback: replacement characters so XLSX generation can proceed + try: + return pd.read_csv(path, encoding="utf-8", encoding_errors="replace") + except Exception: + if last_err is not None: + raise last_err + raise + + for idx, csv_path in enumerate(csv_files): + # Use provided sheet name or derive from file name + sheet_name = ( + sheet_names[idx] + if sheet_names and idx < len(sheet_names) + else os.path.splitext(os.path.basename(csv_path))[0] + ) + df = _read_csv_with_fallback_encodings(csv_path) + + if sheet_name == "Original data": + try: + # Create a copy to avoid modifying the original + df_copy = df.copy() + # Insert the Response ID column at position 0 (first column) + df_copy.insert(0, "Response ID", unique_reference_numbers) + df = df_copy + except Exception as e: + print("Could not add reference number to original data due to:", e) + + ws = wb.create_sheet(title=sheet_name) + + for r_idx, row in enumerate( + dataframe_to_rows(df, index=False, header=True), start=1 + ): + ws.append(row) + + for col_idx, value in enumerate(row, start=1): + cell = ws.cell(row=r_idx, column=col_idx) + + # Convert markdown to RichText if applicable + formatted_value = markdown_to_richtext(value) + if formatted_value != value: + cell.value = formatted_value + else: + cell.value = value + + # Bold header row + if r_idx == 1: + # If cell already has RichText, we need to apply bold to all segments + if isinstance(cell.value, CellRichText): + # Create new RichText with bold applied to all segments + bold_rich_text = CellRichText() + for segment in cell.value: + if isinstance(segment, TextBlock): + # Preserve italic if present, add bold + is_italic = segment.font.i if segment.font else False + bold_font = InlineFont(b=True, i=is_italic) + bold_rich_text.append( + TextBlock(bold_font, segment.text) + ) + else: + bold_rich_text.append( + TextBlock(InlineFont(b=True), str(segment)) + ) + cell.value = bold_rich_text + else: + cell.font = Font(bold=True) + + # Set vertical alignment to middle by default + cell.alignment = Alignment(vertical="center") + + # Apply wrap text if needed + if wrap_text_columns and sheet_name in wrap_text_columns: + for col_letter in wrap_text_columns[sheet_name]: + cell = ws[f"{col_letter}{r_idx}"] + cell.alignment = Alignment(vertical="center", wrap_text=True) + + # Freeze top row for all data sheets + ws.freeze_panes = "A2" + + # Ensure all header cells (first row) are wrapped + for col_idx in range(1, ws.max_column + 1): + header_cell = ws.cell(row=1, column=col_idx) + header_cell.alignment = Alignment(vertical="center", wrap_text=True) + + # Set column widths + if column_widths and sheet_name in column_widths: + for col_letter, width in column_widths[sheet_name].items(): + ws.column_dimensions[col_letter].width = width + + add_cover_sheet( + wb, + intro_paragraphs=intro_text, + model_name=model_name, + analysis_date=analysis_date, + analysis_cost=analysis_cost, + number_of_responses=number_of_responses, + number_of_responses_with_text=number_of_responses_with_text, + number_of_responses_with_text_five_plus_words=number_of_responses_with_text_five_plus_words, + llm_call_number=llm_call_number, + input_tokens=input_tokens, + output_tokens=output_tokens, + time_taken=time_taken, + file_name=file_name, + column_name=column_name, + number_of_responses_with_topic_assignment=number_of_responses_with_topic_assignment, + excel_sheet_name=excel_sheet_name, + candidate_topics_file_name=candidate_topics_file_name, + ) + + wb.save(output_filename) + + # Convert to ODS format if requested + if EXPORT_FORMAT == "ods": + # Change file extension from .xlsx to .ods + ods_filename = output_filename.replace(".xlsx", ".ods") + output_filename = convert_xlsx_to_ods( + output_filename, ods_filename, allowed_root=allowed_root + ) + else: + print(f"Output xlsx summary saved as '{output_filename}'") + + return output_filename + + +### +# Run the functions +### +def collect_output_csvs_and_create_excel_output( + in_data_files: List, + chosen_cols: list[str], + reference_data_file_name_textbox: str, + in_group_col: str, + model_choice: str, + master_reference_df_state: pd.DataFrame, + master_unique_topics_df_state: pd.DataFrame, + summarised_output_df: pd.DataFrame, + missing_df_state: pd.DataFrame, + excel_sheets: str = "", + usage_logs_location: str = "", + model_name_map: dict = dict(), + output_folder: str = OUTPUT_FOLDER, + structured_summaries: str = "No", + candidate_topics=None, + create_topics_csv: str = "Yes", +): + """ + Collect together output CSVs from various output boxes and combine them into a single output Excel file. + + Args: + in_data_files (List): A list of paths to the input data files. + chosen_cols (list[str]): A list of column names selected for analysis. + reference_data_file_name_textbox (str): The name of the reference data file. + in_group_col (str): The column used for grouping the data. + model_choice (str): The LLM model chosen for the analysis. + master_reference_df_state (pd.DataFrame): The master DataFrame containing reference data. + master_unique_topics_df_state (pd.DataFrame): The master DataFrame containing unique topics data. + summarised_output_df (pd.DataFrame): DataFrame containing the summarised output. + missing_df_state (pd.DataFrame): DataFrame containing information about missing data. + excel_sheets (str): Information regarding Excel sheets, typically sheet names or structure. + usage_logs_location (str, optional): Path to the usage logs CSV file. Defaults to "". + model_name_map (dict, optional): A dictionary mapping model choices to their display names. Defaults to {}. + output_folder (str, optional): The directory where the output Excel file will be saved. Defaults to OUTPUT_FOLDER. + structured_summaries (str, optional): Indicates whether structured summaries are being produced ("Yes" or "No"). Defaults to "No". + candidate_topics (optional): Suggested topics file uploaded by the user (path string or Gradio FileData). + create_topics_csv (str, optional): Whether to write a suggested-topics CSV with unique + General topic / Subtopic pairs from the analysis. Defaults to "Yes". + + Returns: + tuple: A tuple containing: + - list: Paths to generated output files (xlsx and, when requested, suggested topics CSV). + - list: Duplicate of the first list (for UI compatibility). + """ + # Use passed model_name_map if provided and not empty, otherwise use global one + if not model_name_map: + model_name_map = global_model_name_map + + # Ensure custom model_choice is registered in model_name_map + ensure_model_in_map(model_choice, model_name_map) + + if structured_summaries == "Yes": + structured_summaries = True + else: + structured_summaries = False + + if not chosen_cols: + raise Exception("Could not find chosen column") + + today_date = datetime.today().strftime("%Y-%m-%d") + original_data_file_path = os.path.abspath(in_data_files[0]) + excel_sheet_display_name = _resolve_excel_sheet_display_name( + original_data_file_path, excel_sheets + ) + candidate_topics_file_name = _resolve_candidate_topics_file_name(candidate_topics) + + csv_files = list() + sheet_names = list() + column_widths = dict() + wrap_text_columns = dict() + short_file_name = os.path.basename(reference_data_file_name_textbox) + reference_pivot_table = pd.DataFrame() + reference_table_csv_path = "" + reference_pivot_table_csv_path = "" + unique_topic_table_csv_path = "" + missing_df_state_csv_path = "" + overall_summary_csv_path = "" + number_of_responses_with_topic_assignment = 0 + # Track all temporary CSV files created for xlsx conversion + temp_csv_files_for_cleanup = list() + + if in_group_col: + group = in_group_col + else: + group = "All" + + overall_summary_csv_path = output_folder + "overall_summary_for_xlsx.csv" + temp_csv_files_for_cleanup.append(overall_summary_csv_path) + + if structured_summaries is True and not master_unique_topics_df_state.empty: + print("Producing overall summary based on structured summaries.") + # Create structured summary from master_unique_topics_df_state + structured_summary_data = list() + + # Group by 'Group' column + for group_name, group_df in master_unique_topics_df_state.groupby("Group"): + group_summary = f"## {group_name}\n\n" + + # Group by 'General topic' within each group + for general_topic, topic_df in group_df.groupby("General topic"): + group_summary += f"### {general_topic}\n\n" + + # Add subtopics under each general topic + for _, row in topic_df.iterrows(): + subtopic = row["Subtopic"] + summary = row["Summary"] + # sentiment = row.get('Sentiment', '') + # num_responses = row.get('Number of responses', '') + + # Create subtopic entry + subtopic_entry = f"**{subtopic}**" + # if sentiment: + # subtopic_entry += f" ({sentiment})" + # if num_responses: + # subtopic_entry += f" - {num_responses} responses" + subtopic_entry += "\n\n" + + if summary and pd.notna(summary): + subtopic_entry += f"{summary}\n\n" + + group_summary += subtopic_entry + + # Add to structured summary data + structured_summary_data.append( + {"Group": group_name, "Summary": group_summary.strip()} + ) + + # Create DataFrame for structured summary + structured_summary_df = pd.DataFrame(structured_summary_data) + structured_summary_df.to_csv(overall_summary_csv_path, index=False) + else: + # Use original summarised_output_df + structured_summary_df = summarised_output_df + structured_summary_df.to_csv(overall_summary_csv_path, index=None) + + if not structured_summary_df.empty: + csv_files.append(overall_summary_csv_path) + sheet_names.append("Overall summary") + column_widths["Overall summary"] = {"A": 15, "B": 120} + wrap_text_columns["Overall summary"] = ["A", "B"] + + if not master_reference_df_state.empty: + # Simplify table to just responses column and the Response reference number + file_data, file_name, num_batches = load_in_data_file( + in_data_files, chosen_cols, 1, in_excel_sheets=excel_sheets + ) + basic_response_data = get_basic_response_data( + file_data, chosen_cols, verify_titles="No" + ) + reference_pivot_table = convert_reference_table_to_pivot_table( + master_reference_df_state, basic_response_data + ) + + unique_reference_numbers = basic_response_data["Response ID"].tolist() + + try: + master_reference_df_state.rename( + columns={"Topic_number": "Topic number"}, inplace=True, errors="ignore" + ) + master_reference_df_state.drop( + columns=["1", "2", "3"], inplace=True, errors="ignore" + ) + except Exception as e: + print("Could not rename Topic_number due to", e) + + number_of_responses_with_topic_assignment = len( + master_reference_df_state["Response ID"].unique() + ) + + reference_table_csv_path = output_folder + "reference_df_for_xlsx.csv" + master_reference_df_state.to_csv(reference_table_csv_path, index=None) + temp_csv_files_for_cleanup.append(reference_table_csv_path) + + reference_pivot_table_csv_path = ( + output_folder + "reference_pivot_df_for_xlsx.csv" + ) + # Reorder columns to ensure 'Original Response ID' comes before 'Response' + cols = list(reference_pivot_table.columns) + if "Original Response ID" in cols and "Response" in cols: + cols.remove("Original Response ID") + cols.remove("Response") + cols = ["Original Response ID", "Response"] + cols + reference_pivot_table = reference_pivot_table[cols] + reference_pivot_table.to_csv(reference_pivot_table_csv_path, index=None) + temp_csv_files_for_cleanup.append(reference_pivot_table_csv_path) + + short_file_name = os.path.basename(file_name) + + if not master_unique_topics_df_state.empty: + + master_unique_topics_df_state.drop( + columns=["1", "2", "3"], inplace=True, errors="ignore" + ) + + unique_topic_table_csv_path = ( + output_folder + "unique_topic_table_df_for_xlsx.csv" + ) + master_unique_topics_df_state.to_csv(unique_topic_table_csv_path, index=None) + temp_csv_files_for_cleanup.append(unique_topic_table_csv_path) + + if unique_topic_table_csv_path: + csv_files.append(unique_topic_table_csv_path) + sheet_names.append("Topic summary") + column_widths["Topic summary"] = { + "A": 25, + "B": 25, + "C": 12, + "D": 15, + "E": 10, + "F": 100, + } + wrap_text_columns["Topic summary"] = ["A", "B", "D", "F"] + else: + print("Relevant unique topic files not found, excluding from xlsx output.") + + if reference_table_csv_path: + if structured_summaries: + print( + "Structured summaries are being produced, excluding response level data from xlsx output." + ) + else: + csv_files.append(reference_table_csv_path) + sheet_names.append("Response level data") + column_widths["Response level data"] = { + "A": 12, + "B": 30, + "C": 40, + "D": 10, + "H": 100, + } + wrap_text_columns["Response level data"] = ["C", "G"] + else: + print("Relevant reference files not found, excluding from xlsx output.") + + if reference_pivot_table_csv_path: + if structured_summaries: + print( + "Structured summaries are being produced, excluding topic response pivot table from xlsx output." + ) + else: + csv_files.append(reference_pivot_table_csv_path) + sheet_names.append("Topic response pivot table") + + if reference_pivot_table.empty: + reference_pivot_table = pd.read_csv(reference_pivot_table_csv_path) + + # Base widths and wrap + column_widths["Topic response pivot table"] = {"A": 12, "B": 100, "C": 12} + wrap_text_columns["Topic response pivot table"] = ["B", "C"] + + num_cols = len(reference_pivot_table.columns) + col_letters = [get_column_letter(i) for i in range(4, num_cols + 1)] + + for col_letter in col_letters: + column_widths["Topic response pivot table"][col_letter] = 20 + + wrap_text_columns["Topic response pivot table"].extend(col_letters) + else: + print( + "Relevant reference pivot table files not found, excluding from xlsx output." + ) + + if not missing_df_state.empty: + missing_df_state_csv_path = output_folder + "missing_df_state_df_for_xlsx.csv" + missing_df_state.to_csv(missing_df_state_csv_path, index=None) + temp_csv_files_for_cleanup.append(missing_df_state_csv_path) + + if missing_df_state_csv_path: + if structured_summaries: + print( + "Structured summaries are being produced, excluding missing responses from xlsx output." + ) + else: + csv_files.append(missing_df_state_csv_path) + sheet_names.append("Missing responses") + column_widths["Missing responses"] = {"A": 25, "B": 30, "C": 50} + wrap_text_columns["Missing responses"] = ["C"] + else: + print("Relevant missing responses files not found, excluding from xlsx output.") + + # Original data file + original_ext = os.path.splitext(original_data_file_path)[1].lower() + if original_ext == ".csv": + csv_files.append(original_data_file_path) + else: + # Read and convert to CSV + if original_ext == ".xlsx": + if excel_sheets: + df = pd.read_excel(original_data_file_path, sheet_name=excel_sheets) + else: + df = pd.read_excel(original_data_file_path) + elif original_ext == ".parquet": + df = pd.read_parquet(original_data_file_path) + else: + raise Exception(f"Unsupported file type for original data: {original_ext}") + + # Save as CSV in output folder + original_data_csv_path = os.path.join( + output_folder, + os.path.splitext(os.path.basename(original_data_file_path))[0] + + "_for_xlsx.csv", + ) + df.to_csv(original_data_csv_path, index=False) + csv_files.append(original_data_csv_path) + temp_csv_files_for_cleanup.append(original_data_csv_path) + + sheet_names.append("Original data") + column_widths["Original data"] = {"A": 10, "B": 20, "C": 20} + wrap_text_columns["Original data"] = ["C"] + if isinstance(chosen_cols, list) and chosen_cols: + chosen_cols = chosen_cols[0] + else: + chosen_cols = str(chosen_cols) if chosen_cols else "" + + # Intro page text + excel_sheet_intro = ( + f", from Excel sheet '{excel_sheet_display_name}'" + if excel_sheet_display_name + else "" + ) + intro_text = [ + "This workbook contains outputs from the Large Language Model (LLM) thematic analysis of open text data. Each sheet corresponds to a different CSV report included in the analysis.", + f"The file analysed was {short_file_name}, the column analysed was '{chosen_cols}'{excel_sheet_intro} and the data was grouped by column '{group}'." + " Please contact the app administrator if you need any explanation on how to use the results." + "LLMs are not 100% accurate and may produce biased or harmful outputs. All outputs from this analysis **need to be checked by a human** to check for harmful outputs, false information, and bias.", + ] + + # Get values for number of rows, number of responses, and number of responses longer than five words + number_of_responses = basic_response_data.shape[0] + # number_of_responses_with_text = basic_response_data["Response"].str.strip().notnull().sum() + number_of_responses_with_text = ( + basic_response_data["Response"].str.strip().notnull() + & (basic_response_data["Response"].str.split().str.len() >= 1) + ).sum() + number_of_responses_with_text_five_plus_words = ( + basic_response_data["Response"].str.strip().notnull() + & (basic_response_data["Response"].str.split().str.len() >= 5) + ).sum() + + # Get number of LLM calls, input and output tokens + if usage_logs_location: + try: + usage_logs = pd.read_csv(usage_logs_location) + + relevant_logs = usage_logs.loc[ + ( + usage_logs["Response ID data file name"] + == reference_data_file_name_textbox + ) + & ( + usage_logs[ + "Large language model for topic extraction and summarisation" + ] + == model_choice + ) + & ( + usage_logs[ + "Select the open text column of interest. In an Excel file, this shows columns across all sheets." + ] + == ( + chosen_cols[0] + if isinstance(chosen_cols, list) and chosen_cols + else chosen_cols + ) + ), + :, + ] + + llm_call_number = sum(relevant_logs["Total LLM calls"].astype(int)) + input_tokens = sum(relevant_logs["Total input tokens"].astype(int)) + output_tokens = sum(relevant_logs["Total output tokens"].astype(int)) + time_taken = sum( + relevant_logs["Estimated time taken (seconds)"].astype(float) + ) + except Exception as e: + print("Could not obtain usage logs due to:", e) + usage_logs = pd.DataFrame() + llm_call_number = 0 + input_tokens = 0 + output_tokens = 0 + time_taken = 0 + else: + print("LLM call logs location not provided") + usage_logs = pd.DataFrame() + llm_call_number = 0 + input_tokens = 0 + output_tokens = 0 + time_taken = 0 + + # Create short filename: + model_choice_clean_short = clean_column_name( + model_name_map[model_choice]["short_name"], + max_length=20, + front_characters=False, + ) + # Extract first column name as string for cleaning and Excel output + chosen_col_str = ( + chosen_cols[0] + if isinstance(chosen_cols, list) and chosen_cols + else str(chosen_cols) if chosen_cols else "" + ) + in_column_cleaned = clean_column_name(chosen_col_str, max_length=20) + file_name_cleaned = clean_column_name( + file_name, max_length=20, front_characters=True + ) + + # Save outputs for each batch. If master file created, label file as master + file_path_details = ( + f"{file_name_cleaned}_col_{in_column_cleaned}_{model_choice_clean_short}" + ) + output_xlsx_filename = ( + output_folder + + file_path_details + + ("_structured_summaries" if structured_summaries else "_theme_analysis") + + ".xlsx" + ) + + xlsx_output_filename = csvs_to_excel( + csv_files=csv_files, + output_filename=output_xlsx_filename, + allowed_root=output_folder, + sheet_names=sheet_names, + column_widths=column_widths, + wrap_text_columns=wrap_text_columns, + intro_text=intro_text, + model_name=model_choice, + analysis_date=today_date, + analysis_cost="", + llm_call_number=llm_call_number, + input_tokens=input_tokens, + output_tokens=output_tokens, + time_taken=time_taken, + number_of_responses=number_of_responses, + number_of_responses_with_text=number_of_responses_with_text, + number_of_responses_with_text_five_plus_words=number_of_responses_with_text_five_plus_words, + column_name=chosen_col_str, + number_of_responses_with_topic_assignment=number_of_responses_with_topic_assignment, + file_name=short_file_name, + excel_sheet_name=excel_sheet_display_name, + candidate_topics_file_name=candidate_topics_file_name, + unique_reference_numbers=unique_reference_numbers, + ) + + xlsx_output_filenames = [xlsx_output_filename] + topics_csv_filenames = [] + + if create_topics_csv in (True, "Yes") and not master_unique_topics_df_state.empty: + topics_csv_filename = ( + output_folder + + file_path_details + + ("_structured_summaries" if structured_summaries else "_theme_analysis") + + "_suggested_topics.csv" + ) + written_path = write_candidate_topics_csv( + master_unique_topics_df_state, topics_csv_filename + ) + if written_path: + topics_csv_filenames.append(written_path) + + all_output_filenames = xlsx_output_filenames + topics_csv_filenames + + # Delete all intermediate '_for_xlsx.csv' files + for csv_file in temp_csv_files_for_cleanup: + try: + if os.path.exists(csv_file): + os.remove(csv_file) + except Exception as e: + print(f"Could not delete temporary CSV file '{csv_file}' due to: {e}") + + return all_output_filenames, all_output_filenames diff --git a/tools/config.py b/tools/config.py new file mode 100644 index 0000000000000000000000000000000000000000..0a1f8905eb52607dd30ab0f1a44f0ba84a0ae16d --- /dev/null +++ b/tools/config.py @@ -0,0 +1,1092 @@ +import codecs +import logging +import os +import socket +import tempfile +from datetime import datetime +from typing import List + +from dotenv import load_dotenv + +today_rev = datetime.now().strftime("%Y%m%d") +HOST_NAME = socket.gethostname() + +# Set or retrieve configuration variables for the summarisation app + + +def ensure_folder_exists(output_folder: str): + """Checks if the specified folder exists, creates it if not.""" + + if not os.path.exists(output_folder): + # Create the folder if it doesn't exist + os.makedirs(output_folder, exist_ok=True) + print(f"Created the {output_folder} folder.") + else: + pass + # print(f"The {output_folder} folder already exists.") + + +def _get_env_list(env_var_name: str, strip_strings: bool = True) -> List[str]: + """Parses a comma-separated environment variable into a list of strings.""" + value = env_var_name[1:-1].strip().replace('"', "").replace("'", "") + if not value: + return [] + # Split by comma and filter out any empty strings that might result from extra commas + if strip_strings: + return [s.strip() for s in value.split(",") if s.strip()] + else: + return [codecs.decode(s, "unicode_escape") for s in value.split(",") if s] + + +def get_or_create_env_var(var_name: str, default_value: str, print_val: bool = False): + """ + Get an environmental variable, and set it to a default value if it doesn't exist + """ + # Get the environment variable if it exists + value = os.environ.get(var_name) + + # If it doesn't exist, set the environment variable to the default value + if value is None: + os.environ[var_name] = default_value + value = default_value + + if print_val is True: + print(f"The value of {var_name} is {value}") + + return value + + +def add_folder_to_path(folder_path: str): + """ + Check if a folder exists on your system. If so, get the absolute path and then add it to the system Path variable if it doesn't already exist. Function is only relevant for locally-created executable files based on this app (when using pyinstaller it creates a _internal folder that contains tesseract and poppler. These need to be added to the system path to enable the app to run) + """ + + if os.path.exists(folder_path) and os.path.isdir(folder_path): + print(folder_path, "folder exists.") + + # Resolve relative path to absolute path + absolute_path = os.path.abspath(folder_path) + + current_path = os.environ["PATH"] + if absolute_path not in current_path.split(os.pathsep): + full_path_extension = absolute_path + os.pathsep + current_path + os.environ["PATH"] = full_path_extension + # print(f"Updated PATH with: ", full_path_extension) + else: + print(f"Directory {folder_path} already exists in PATH.") + else: + print(f"Folder not found at {folder_path} - not added to PATH") + + +def convert_string_to_boolean(value: str) -> bool: + """Convert string to boolean, handling various formats.""" + if isinstance(value, bool): + return value + elif value in ["True", "1", "true", "TRUE"]: + return True + elif value in ["False", "0", "false", "FALSE"]: + return False + else: + raise ValueError(f"Invalid boolean value: {value}") + + +### +# LOAD CONFIG FROM ENV FILE +### + +CONFIG_FOLDER = get_or_create_env_var("CONFIG_FOLDER", "config/") + +# If you have an aws_config env file in the config folder, you can load in app variables this way, e.g. 'config/app_config.env' +APP_CONFIG_PATH = get_or_create_env_var( + "APP_CONFIG_PATH", CONFIG_FOLDER + "app_config.env" +) # e.g. config/app_config.env + +if APP_CONFIG_PATH: + if os.path.exists(APP_CONFIG_PATH): + print(f"Loading app variables from config file {APP_CONFIG_PATH}") + load_dotenv(APP_CONFIG_PATH) + else: + print("App config file not found at location:", APP_CONFIG_PATH) + +### +# AWS OPTIONS +### + +# If you have an aws_config env file in the config folder, you can load in AWS keys this way, e.g. 'env/aws_config.env' +AWS_CONFIG_PATH = get_or_create_env_var( + "AWS_CONFIG_PATH", "" +) # e.g. config/aws_config.env + +if AWS_CONFIG_PATH: + if os.path.exists(AWS_CONFIG_PATH): + print(f"Loading AWS variables from config file {AWS_CONFIG_PATH}") + load_dotenv(AWS_CONFIG_PATH) + else: + print("AWS config file not found at location:", AWS_CONFIG_PATH) + +RUN_AWS_FUNCTIONS = get_or_create_env_var("RUN_AWS_FUNCTIONS", "0") + +AWS_REGION = get_or_create_env_var("AWS_REGION", "") + +AWS_CLIENT_ID = get_or_create_env_var("AWS_CLIENT_ID", "") + +AWS_CLIENT_SECRET = get_or_create_env_var("AWS_CLIENT_SECRET", "") + +AWS_USER_POOL_ID = get_or_create_env_var("AWS_USER_POOL_ID", "") + +AWS_ACCESS_KEY = get_or_create_env_var("AWS_ACCESS_KEY", "") +# if AWS_ACCESS_KEY: print(f'AWS_ACCESS_KEY found in environment variables') + +AWS_SECRET_KEY = get_or_create_env_var("AWS_SECRET_KEY", "") +# if AWS_SECRET_KEY: print(f'AWS_SECRET_KEY found in environment variables') + +# Should the app prioritise using AWS SSO over using API keys stored in environment variables/secrets (defaults to yes) +PRIORITISE_SSO_OVER_AWS_ENV_ACCESS_KEYS = get_or_create_env_var( + "PRIORITISE_SSO_OVER_AWS_ENV_ACCESS_KEYS", "1" +) + +S3_LOG_BUCKET = get_or_create_env_var("S3_LOG_BUCKET", "") + +# Custom headers e.g. if routing traffic through Cloudfront +# Retrieving or setting CUSTOM_HEADER +CUSTOM_HEADER = get_or_create_env_var("CUSTOM_HEADER", "") + +# Retrieving or setting CUSTOM_HEADER_VALUE +CUSTOM_HEADER_VALUE = get_or_create_env_var("CUSTOM_HEADER_VALUE", "") + +### +# File I/O +### +SESSION_OUTPUT_FOLDER = get_or_create_env_var( + "SESSION_OUTPUT_FOLDER", "False" +) # i.e. do you want your input and output folders saved within a subfolder based on session hash value within output/input folders + +OUTPUT_FOLDER = get_or_create_env_var("GRADIO_OUTPUT_FOLDER", "output/") # 'output/' +INPUT_FOLDER = get_or_create_env_var("GRADIO_INPUT_FOLDER", "input/") # 'input/' + + +# Allow for files to be saved in a temporary folder for increased security in some instances +if OUTPUT_FOLDER == "TEMP" or INPUT_FOLDER == "TEMP": + # Create a temporary directory + with tempfile.TemporaryDirectory() as temp_dir: + print(f"Temporary directory created at: {temp_dir}") + + if OUTPUT_FOLDER == "TEMP": + OUTPUT_FOLDER = temp_dir + "/" + if INPUT_FOLDER == "TEMP": + INPUT_FOLDER = temp_dir + "/" + + +GRADIO_TEMP_DIR = get_or_create_env_var( + "GRADIO_TEMP_DIR", "tmp/gradio_tmp/" +) # Default Gradio temp folder +MPLCONFIGDIR = get_or_create_env_var( + "MPLCONFIGDIR", "tmp/matplotlib_cache/" +) # Matplotlib cache folder + +S3_OUTPUTS_BUCKET = get_or_create_env_var("S3_OUTPUTS_BUCKET", "") +S3_OUTPUTS_FOLDER = get_or_create_env_var("S3_OUTPUTS_FOLDER", "") +SAVE_OUTPUTS_TO_S3 = get_or_create_env_var("SAVE_OUTPUTS_TO_S3", "False") +UPLOAD_USAGE_LOG_TO_S3_OUTPUTS = convert_string_to_boolean( + get_or_create_env_var("UPLOAD_USAGE_LOG_TO_S3_OUTPUTS", "False") +) # Whether to upload usage log to S3 output folder when uploading outputs +UPLOAD_PROMPT_RESPONSE_LOG_TO_S3_OUTPUTS = convert_string_to_boolean( + get_or_create_env_var("UPLOAD_PROMPT_RESPONSE_LOG_TO_S3_OUTPUTS", "False") +) # Whether to upload prompt/response JSON logs (*_logs_*.json) to S3 output folder when uploading outputs + +# Output format for the Excel/spreadsheet file (xlsx or ods) +EXPORT_FORMAT = get_or_create_env_var( + "EXPORT_FORMAT", "xlsx" +).lower() # 'xlsx' or 'ods' +if EXPORT_FORMAT not in ["xlsx", "ods"]: + EXPORT_FORMAT = "xlsx" # Default to xlsx if invalid value provided + +### +# LOGGING OPTIONS +### + +# By default, logs are put into a subfolder of today's date and the host name of the instance running the app. This is to avoid at all possible the possibility of log files from one instance overwriting the logs of another instance on S3. If running the app on one system always, or just locally, it is not necessary to make the log folders so specific. +# Another way to address this issue would be to write logs to another type of storage, e.g. database such as dynamodb. I may look into this in future. + +SAVE_LOGS_TO_CSV = convert_string_to_boolean( + get_or_create_env_var("SAVE_LOGS_TO_CSV", "True") +) + +USE_LOG_SUBFOLDERS = get_or_create_env_var("USE_LOG_SUBFOLDERS", "True") + +FEEDBACK_LOGS_FOLDER = get_or_create_env_var("FEEDBACK_LOGS_FOLDER", "feedback/") +ACCESS_LOGS_FOLDER = get_or_create_env_var("ACCESS_LOGS_FOLDER", "logs/") +USAGE_LOGS_FOLDER = get_or_create_env_var("USAGE_LOGS_FOLDER", "usage/") + +# Initialize full_log_subfolder based on USE_LOG_SUBFOLDERS setting +if USE_LOG_SUBFOLDERS == "True": + day_log_subfolder = today_rev + "/" + host_name_subfolder = HOST_NAME + "/" + full_log_subfolder = day_log_subfolder + host_name_subfolder + + FEEDBACK_LOGS_FOLDER = FEEDBACK_LOGS_FOLDER + full_log_subfolder + ACCESS_LOGS_FOLDER = ACCESS_LOGS_FOLDER + full_log_subfolder + USAGE_LOGS_FOLDER = USAGE_LOGS_FOLDER + full_log_subfolder +else: + full_log_subfolder = "" # Empty string when subfolders are not used + +S3_FEEDBACK_LOGS_FOLDER = get_or_create_env_var( + "S3_FEEDBACK_LOGS_FOLDER", "feedback/" + full_log_subfolder +) +S3_ACCESS_LOGS_FOLDER = get_or_create_env_var( + "S3_ACCESS_LOGS_FOLDER", "logs/" + full_log_subfolder +) +S3_USAGE_LOGS_FOLDER = get_or_create_env_var( + "S3_USAGE_LOGS_FOLDER", "usage/" + full_log_subfolder +) + +LOG_FILE_NAME = get_or_create_env_var("LOG_FILE_NAME", "log.csv") +USAGE_LOG_FILE_NAME = get_or_create_env_var("USAGE_LOG_FILE_NAME", LOG_FILE_NAME) +FEEDBACK_LOG_FILE_NAME = get_or_create_env_var("FEEDBACK_LOG_FILE_NAME", LOG_FILE_NAME) + +# Should the redacted file name be included in the logs? In some instances, the names of the files themselves could be sensitive, and should not be disclosed beyond the app. So, by default this is false. +DISPLAY_FILE_NAMES_IN_LOGS = get_or_create_env_var( + "DISPLAY_FILE_NAMES_IN_LOGS", "False" +) + +# Further customisation options for CSV logs + +CSV_ACCESS_LOG_HEADERS = get_or_create_env_var( + "CSV_ACCESS_LOG_HEADERS", "" +) # If blank, uses component labels +CSV_FEEDBACK_LOG_HEADERS = get_or_create_env_var( + "CSV_FEEDBACK_LOG_HEADERS", "" +) # If blank, uses component labels +CSV_USAGE_LOG_HEADERS = get_or_create_env_var( + "CSV_USAGE_LOG_HEADERS", "" +) # If blank, uses component labels + +### DYNAMODB logs. Whether to save to DynamoDB, and the headers of the table +SAVE_LOGS_TO_DYNAMODB = convert_string_to_boolean( + get_or_create_env_var("SAVE_LOGS_TO_DYNAMODB", "False") +) + +ACCESS_LOG_DYNAMODB_TABLE_NAME = get_or_create_env_var( + "ACCESS_LOG_DYNAMODB_TABLE_NAME", "llm_topic_model_access_log" +) +DYNAMODB_ACCESS_LOG_HEADERS = get_or_create_env_var("DYNAMODB_ACCESS_LOG_HEADERS", "") + +FEEDBACK_LOG_DYNAMODB_TABLE_NAME = get_or_create_env_var( + "FEEDBACK_LOG_DYNAMODB_TABLE_NAME", "llm_topic_model_feedback" +) +DYNAMODB_FEEDBACK_LOG_HEADERS = get_or_create_env_var( + "DYNAMODB_FEEDBACK_LOG_HEADERS", "" +) + +USAGE_LOG_DYNAMODB_TABLE_NAME = get_or_create_env_var( + "USAGE_LOG_DYNAMODB_TABLE_NAME", "llm_topic_model_usage" +) +DYNAMODB_USAGE_LOG_HEADERS = get_or_create_env_var("DYNAMODB_USAGE_LOG_HEADERS", "") + +# Report logging to console? +LOGGING = get_or_create_env_var("LOGGING", "False") + +if LOGGING == "True": + # Configure logging + logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" + ) + +### +# App run variables +### +OUTPUT_DEBUG_FILES = get_or_create_env_var( + "OUTPUT_DEBUG_FILES", "False" +) # Whether to output debug files +SHOW_ADDITIONAL_INSTRUCTION_TEXTBOXES = get_or_create_env_var( + "SHOW_ADDITIONAL_INSTRUCTION_TEXTBOXES", "True" +) # Whether to show additional instruction textboxes in the GUI + +TIMEOUT_WAIT = int( + get_or_create_env_var("TIMEOUT_WAIT", "30") +) # Maximum number of seconds to wait for a response from the LLM +NUMBER_OF_RETRY_ATTEMPTS = int( + get_or_create_env_var("NUMBER_OF_RETRY_ATTEMPTS", "5") +) # Maximum number of times to retry a request to the LLM +# Try up to 3 times to get a valid markdown table response with LLM calls, otherwise retry with temperature changed +MAX_OUTPUT_VALIDATION_ATTEMPTS = int( + get_or_create_env_var("MAX_OUTPUT_VALIDATION_ATTEMPTS", "3") +) +ENABLE_VALIDATION = get_or_create_env_var( + "ENABLE_VALIDATION", "False" +) # Whether to run validation loop after initial topic extraction +MAX_TIME_FOR_LOOP = int( + get_or_create_env_var("MAX_TIME_FOR_LOOP", "99999") +) # Maximum number of seconds to run the loop for before breaking (to run again, this is to avoid timeouts with some AWS services if deployed there) + +MAX_COMMENT_CHARS = int( + get_or_create_env_var("MAX_COMMENT_CHARS", "14000") +) # Maximum number of characters in a comment +MAX_ROWS = int( + get_or_create_env_var("MAX_ROWS", "5000") +) # Maximum number of rows to process +MAX_GROUPS = int( + get_or_create_env_var("MAX_GROUPS", "99") +) # Maximum number of groups to process +BATCH_SIZE_DEFAULT = int( + get_or_create_env_var("BATCH_SIZE_DEFAULT", "5") +) # Default batch size for LLM calls +MAXIMUM_ALLOWED_TOPICS = int( + get_or_create_env_var("MAXIMUM_ALLOWED_TOPICS", "100") +) # Maximum number of zero shot topics to process +MAX_SPACES_GPU_RUN_TIME = int( + get_or_create_env_var("MAX_SPACES_GPU_RUN_TIME", "240") +) # Maximum number of seconds to run on GPU on Hugging Face Spaces + +DEDUPLICATION_THRESHOLD = int( + get_or_create_env_var("DEDUPLICATION_THRESHOLD", "90") +) # Deduplication threshold for topic summary tables + +ENABLE_BATCH_DEDUPLICATION = convert_string_to_boolean( + get_or_create_env_var("ENABLE_BATCH_DEDUPLICATION", "True") +) # Whether to deduplicate topics after each batch during extraction. Will use basic deduplication to check for typos in effectively duplicate topic names, and if candidate topics are not provided, will use LLM deduplication to merge similar topics if the current number of topics exceeds the maximum allowed number of topics (MAXIMUM_ALLOWED_TOPICS, or defined in GUI by user) + +# Run batch deduplication only when the number of completed batches is a multiple of this value (1 = every batch; e.g. 5 = after batches 5, 10, 15, ...). Values below 1 are treated as 1. +BATCH_DEDUPLICATION_EVERY_N_BATCHES = max( + 1, + int(get_or_create_env_var("BATCH_DEDUPLICATION_EVERY_N_BATCHES", "5")), +) + +### +# Model options +### + +RUN_LOCAL_MODEL = get_or_create_env_var("RUN_LOCAL_MODEL", "0") + +RUN_AWS_BEDROCK_MODELS = get_or_create_env_var("RUN_AWS_BEDROCK_MODELS", "1") + +RUN_GEMINI_MODELS = get_or_create_env_var("RUN_GEMINI_MODELS", "1") +GEMINI_API_KEY = get_or_create_env_var("GEMINI_API_KEY", "") + +INTRO_TEXT = get_or_create_env_var( + "INTRO_TEXT", + """# Large language model topic modelling + +Extract topics and summarise outputs using Large Language Models (LLMs, Gemma 3 4b/GPT-OSS 20b if local (see tools/config.py to modify), Gemini, Azure/OpenAI, or AWS Bedrock models (e.g. Claude, Nova models). The app will query the LLM with batches of responses to produce summary tables, which are then compared iteratively to output a table with the general topics, subtopics, topic sentiment, and a topic summary. Instructions on use can be found in the README.md file. You can try out examples by clicking on one of the example datasets below. API keys for AWS, Azure/OpenAI, and Gemini services can be entered on the settings page (note that Gemini has a free public API). + +NOTE: Large language models are not 100% accurate and may produce biased or harmful outputs. All outputs from this app **absolutely need to be checked by a human** to check for harmful outputs, hallucinations, and accuracy.""", +) + +# Read in intro text from a text file if it is a path to a text file +if INTRO_TEXT.endswith(".txt"): + INTRO_TEXT = open(INTRO_TEXT, "r").read() + +INTRO_TEXT = INTRO_TEXT.strip('"').strip("'") + +# Azure/OpenAI AI Inference settings +RUN_AZURE_MODELS = get_or_create_env_var("RUN_AZURE_MODELS", "1") +AZURE_OPENAI_API_KEY = get_or_create_env_var("AZURE_OPENAI_API_KEY", "") +AZURE_OPENAI_INFERENCE_ENDPOINT = get_or_create_env_var( + "AZURE_OPENAI_INFERENCE_ENDPOINT", "" +) + +# Llama-server settings +RUN_INFERENCE_SERVER = get_or_create_env_var("RUN_INFERENCE_SERVER", "0") +API_URL = get_or_create_env_var("API_URL", "http://localhost:8080") + +# Inference-server HTTP settings +# The server can take a while to start responding while it pre-processes a prompt, +# so this should often be higher than the general retry sleep. +INFERENCE_SERVER_READ_TIMEOUT_SECONDS = int( + get_or_create_env_var("INFERENCE_SERVER_READ_TIMEOUT_SECONDS", "120") +) + +# When we cannot do server-side token counting for inference-server, our local estimates can be low. +# Apply an additional safety multiplier to sampler token budgets in that case. +INFERENCE_SERVER_SAMPLER_TOKEN_BUDGET_FRACTION = float( + get_or_create_env_var("INFERENCE_SERVER_SAMPLER_TOKEN_BUDGET_FRACTION", "0.9") +) + +RUN_MCP_SERVER = convert_string_to_boolean( + get_or_create_env_var("RUN_MCP_SERVER", "False") +) + +# Build up options for models +model_full_names = list() +model_short_names = list() +model_source = list() + +CHOSEN_LOCAL_MODEL_TYPE = get_or_create_env_var( + "CHOSEN_LOCAL_MODEL_TYPE", "Qwen 3 4B" +) # Gemma 3 1B # "Gemma 2b" # "Gemma 3 4B" + +USE_LLAMA_SWAP = get_or_create_env_var("USE_LLAMA_SWAP", "False") +if USE_LLAMA_SWAP == "True": + USE_LLAMA_SWAP = True +else: + USE_LLAMA_SWAP = False + +# Qwen3 / Qwen3.5 on vLLM: set enable_thinking=False in chat completions (see llm_funcs.call_inference_server_api) +INFERENCE_SERVER_DISABLE_THINKING = convert_string_to_boolean( + get_or_create_env_var("INFERENCE_SERVER_DISABLE_THINKING", "True") +) + +if RUN_LOCAL_MODEL == "1" and CHOSEN_LOCAL_MODEL_TYPE: + model_full_names.append(CHOSEN_LOCAL_MODEL_TYPE) + model_short_names.append(CHOSEN_LOCAL_MODEL_TYPE) + model_source.append("Local") + +if RUN_AWS_BEDROCK_MODELS == "1": + amazon_models = [ + "google.gemma-3-12b-it", + "google.gemma-3-27b-it", + # "anthropic.claude-3-haiku-20240307-v1:0", + # "anthropic.claude-3-7-sonnet-20250219-v1:0", + "anthropic.claude-sonnet-4-6", + "anthropic.claude-opus-4-6-v1", + # "amazon.nova-micro-v1:0", + "amazon.nova-lite-v1:0", + "amazon.nova-pro-v1:0", + # "deepseek.v3-v1:0", + "openai.gpt-oss-20b-1:0", + "openai.gpt-oss-120b-1:0", + "mistral.ministral-3-14b-instruct", + "mistral.magistral-small-2509", + "mistral.devstral-2-123b", + "nvidia.nemotron-super-3-120b", + ] + model_full_names.extend(amazon_models) + model_short_names.extend( + [ + "gemma_3_12b_it", + "gemma_3_27b_it", + # "haiku", + # "sonnet_3_7", + "sonnet_4_6", + "claude_opus_4_6", + # "nova_micro", + "nova_lite", + "nova_pro", + # "deepseek_v3", + "gpt_oss_20b_aws", + "gpt_oss_120b_aws", + "ministral_3_14b_instruct", + "magistral_small_2509", + "devstral_2_123b", + "nemotron_super_3_120b", + ] + ) + model_source.extend(["AWS"] * len(amazon_models)) + +if RUN_GEMINI_MODELS == "1": + gemini_models = [ + "gemini-flash-lite-latest", + "gemini-flash-latest", + "gemini-pro-latest", + ] + model_full_names.extend(gemini_models) + model_short_names.extend(["gemini_flash_lite", "gemini_flash", "gemini_pro"]) + model_source.extend(["Gemini"] * len(gemini_models)) + +# Register Azure/OpenAI AI models (model names must match your Azure/OpenAI deployments) +if RUN_AZURE_MODELS == "1": + # Example deployments; adjust to the deployments you actually create in Azure/OpenAI + azure_models = ["gpt-5-mini", "gpt-4o-mini"] + model_full_names.extend(azure_models) + model_short_names.extend(["gpt-5-mini", "gpt-4o-mini"]) + model_source.extend(["Azure/OpenAI"] * len(azure_models)) + +# Register inference-server models +CHOSEN_INFERENCE_SERVER_MODEL = "" +if RUN_INFERENCE_SERVER == "1": + # Example inference-server models; adjust to the models you have available on your server + inference_server_models = [ + "unnamed-inference-server-model", + "gpt_oss_20b", + "gemma_3_12b", + "ministral_3_14b_it", + "Qwen 3.5 27b", + "Qwen 3.5 35b a3b", + "Gemma 4 12B", + "Gemma 4 26b a4b", + "Gemma 4 31b", + ] + model_full_names.extend(inference_server_models) + model_short_names.extend(inference_server_models) + model_source.extend(["inference-server"] * len(inference_server_models)) + + CHOSEN_INFERENCE_SERVER_MODEL = get_or_create_env_var( + "CHOSEN_INFERENCE_SERVER_MODEL", inference_server_models[0] + ) + + # If the chosen inference server model is not in the list of inference server models, add it to the list + if CHOSEN_INFERENCE_SERVER_MODEL not in inference_server_models: + model_full_names.append(CHOSEN_INFERENCE_SERVER_MODEL) + model_short_names.append(CHOSEN_INFERENCE_SERVER_MODEL) + model_source.append("inference-server") + +model_name_map = { + full: {"short_name": short, "source": source} + for full, short, source in zip(model_full_names, model_short_names, model_source) +} + +if RUN_LOCAL_MODEL == "1": + default_model_choice = CHOSEN_LOCAL_MODEL_TYPE +elif RUN_INFERENCE_SERVER == "1": + default_model_choice = CHOSEN_INFERENCE_SERVER_MODEL +elif RUN_AWS_FUNCTIONS == "1": + default_model_choice = amazon_models[0] +else: + default_model_choice = gemini_models[0] + +default_model_source = model_name_map[default_model_choice]["source"] +model_sources = list( + set([model_name_map[model]["source"] for model in model_full_names]) +) + +DIRECT_MODE_INFERENCE_SERVER_MODEL = get_or_create_env_var( + "DIRECT_MODE_INFERENCE_SERVER_MODEL", + CHOSEN_INFERENCE_SERVER_MODEL if CHOSEN_INFERENCE_SERVER_MODEL else "", +) + + +def update_model_choice_config(default_model_source, model_name_map): + # Filter models by source and return the first matching model name + matching_models = [ + model_name + for model_name, model_info in model_name_map.items() + if model_info["source"] == default_model_source + ] + + output_model = matching_models[0] if matching_models else model_full_names[0] + + return output_model, matching_models + + +default_model_choice, default_source_models = update_model_choice_config( + default_model_source, model_name_map +) + +# print("model_name_map:", model_name_map) + +# HF token may or may not be needed for downloading models from Hugging Face +HF_TOKEN = get_or_create_env_var("HF_TOKEN", "") + +LOAD_LOCAL_MODEL_AT_START = get_or_create_env_var("LOAD_LOCAL_MODEL_AT_START", "False") + +# If you are using a system with low VRAM, you can set this to True to reduce the memory requirements +LOW_VRAM_SYSTEM = get_or_create_env_var("LOW_VRAM_SYSTEM", "False") + +MULTIMODAL_PROMPT_FORMAT = get_or_create_env_var("MULTIMODAL_PROMPT_FORMAT", "False") + +if LOW_VRAM_SYSTEM == "True": + print("Using settings for low VRAM system") + USE_LLAMA_CPP = get_or_create_env_var("USE_LLAMA_CPP", "True") + LLM_MAX_NEW_TOKENS = int(get_or_create_env_var("LLM_MAX_NEW_TOKENS", "4096")) + LLM_CONTEXT_LENGTH = int(get_or_create_env_var("LLM_CONTEXT_LENGTH", "16384")) + LLM_BATCH_SIZE = int(get_or_create_env_var("LLM_BATCH_SIZE", "512")) + K_QUANT_LEVEL = int( + get_or_create_env_var("K_QUANT_LEVEL", "2") + ) # 2 = q4_0, 8 = q8_0, 4 = fp16 + V_QUANT_LEVEL = int( + get_or_create_env_var("V_QUANT_LEVEL", "2") + ) # 2 = q4_0, 8 = q8_0, 4 = fp16 + +USE_LLAMA_CPP = get_or_create_env_var( + "USE_LLAMA_CPP", "True" +) # Llama.cpp or transformers with unsloth + +LOCAL_REPO_ID = get_or_create_env_var("LOCAL_REPO_ID", "") +LOCAL_MODEL_FILE = get_or_create_env_var("LOCAL_MODEL_FILE", "") +LOCAL_MODEL_FOLDER = get_or_create_env_var("LOCAL_MODEL_FOLDER", "") + +GEMMA2_REPO_ID = get_or_create_env_var("GEMMA2_2B_REPO_ID", "unsloth/gemma-2-it-GGUF") +GEMMA2_REPO_TRANSFORMERS_ID = get_or_create_env_var( + "GEMMA2_2B_REPO_TRANSFORMERS_ID", "unsloth/gemma-2-2b-it-bnb-4bit" +) +if USE_LLAMA_CPP == "False": + GEMMA2_REPO_ID = GEMMA2_REPO_TRANSFORMERS_ID +GEMMA2_MODEL_FILE = get_or_create_env_var( + "GEMMA2_2B_MODEL_FILE", "gemma-2-2b-it.q8_0.gguf" +) +GEMMA2_MODEL_FOLDER = get_or_create_env_var("GEMMA2_2B_MODEL_FOLDER", "model/gemma") + +GEMMA3_4B_REPO_ID = get_or_create_env_var( + "GEMMA3_4B_REPO_ID", "unsloth/gemma-3-4b-it-qat-GGUF" +) +GEMMA3_4B_REPO_TRANSFORMERS_ID = get_or_create_env_var( + "GEMMA3_4B_REPO_TRANSFORMERS_ID", "unsloth/gemma-3-4b-it-bnb-4bit" +) +if USE_LLAMA_CPP == "False": + GEMMA3_4B_REPO_ID = GEMMA3_4B_REPO_TRANSFORMERS_ID +GEMMA3_4B_MODEL_FILE = get_or_create_env_var( + "GEMMA3_4B_MODEL_FILE", "gemma-3-4b-it-qat-UD-Q4_K_XL.gguf" +) +GEMMA3_4B_MODEL_FOLDER = get_or_create_env_var( + "GEMMA3_4B_MODEL_FOLDER", "model/gemma3_4b" +) + +GEMMA3_12B_REPO_ID = get_or_create_env_var( + "GEMMA3_12B_REPO_ID", "unsloth/gemma-3-12b-it-GGUF" +) +GEMMA3_12B_REPO_TRANSFORMERS_ID = get_or_create_env_var( + "GEMMA3_12B_REPO_TRANSFORMERS_ID", "unsloth/gemma-3-12b-it-bnb-4bit" +) +if USE_LLAMA_CPP == "False": + GEMMA3_12B_REPO_ID = GEMMA3_12B_REPO_TRANSFORMERS_ID +GEMMA3_12B_MODEL_FILE = get_or_create_env_var( + "GEMMA3_12B_MODEL_FILE", "gemma-3-12b-it-UD-Q4_K_XL.gguf" +) +GEMMA3_12B_MODEL_FOLDER = get_or_create_env_var( + "GEMMA3_12B_MODEL_FOLDER", "model/gemma3_12b" +) + +GEMMA4_12B_REPO_ID = get_or_create_env_var( + "GEMMA4_12B_REPO_ID", "unsloth/gemma-4-12B-it-qat-GGUF" +) +GEMMA4_12B_REPO_TRANSFORMERS_ID = get_or_create_env_var( + "GEMMA4_12B_REPO_TRANSFORMERS_ID", "google/gemma-4-12B-it" +) +if USE_LLAMA_CPP == "False": + GEMMA4_12B_REPO_ID = GEMMA4_12B_REPO_TRANSFORMERS_ID +GEMMA4_12B_MODEL_FILE = get_or_create_env_var( + "GEMMA4_12B_MODEL_FILE", "gemma-4-12B-it-qat-UD-Q4_K_XL.gguf" +) +GEMMA4_12B_MODEL_FOLDER = get_or_create_env_var( + "GEMMA4_12B_MODEL_FOLDER", "model/gemma4_12b" +) + +GEMMA4_26B_REPO_ID = get_or_create_env_var( + "GEMMA4_26B_REPO_ID", "unsloth/gemma-4-26B-A4B-it-qat-GGUF" +) +GEMMA4_26B_REPO_TRANSFORMERS_ID = get_or_create_env_var( + "GEMMA4_26B_REPO_TRANSFORMERS_ID", "google/gemma-4-26B-A4B-it" +) +if USE_LLAMA_CPP == "False": + GEMMA4_26B_REPO_ID = GEMMA4_26B_REPO_TRANSFORMERS_ID +GEMMA4_26B_MODEL_FILE = get_or_create_env_var( + "GEMMA4_26B_MODEL_FILE", "gemma-4-26B-A4B-it-qat-UD-Q4_K_XL.gguf" +) +GEMMA4_26B_MODEL_FOLDER = get_or_create_env_var( + "GEMMA4_26B_MODEL_FOLDER", "model/gemma4_26b" +) + +GEMMA4_31B_REPO_ID = get_or_create_env_var( + "GEMMA4_31B_REPO_ID", "unsloth/gemma-4-31B-it-qat-GGUF" +) +GEMMA4_31B_REPO_TRANSFORMERS_ID = get_or_create_env_var( + "GEMMA4_31B_REPO_TRANSFORMERS_ID", "unsloth/gemma-4-31B-it-unsloth-bnb-4bit" +) +if USE_LLAMA_CPP == "False": + GEMMA4_31B_REPO_ID = GEMMA4_31B_REPO_TRANSFORMERS_ID +GEMMA4_31B_MODEL_FILE = get_or_create_env_var( + "GEMMA4_31B_MODEL_FILE", "gemma-4-31B-it-qat-UD-Q4_K_XL.gguf" +) +GEMMA4_31B_MODEL_FOLDER = get_or_create_env_var( + "GEMMA4_31B_MODEL_FOLDER", "model/gemma4_31b" +) + +GPT_OSS_REPO_ID = get_or_create_env_var("GPT_OSS_REPO_ID", "unsloth/gpt-oss-20b-GGUF") +GPT_OSS_REPO_TRANSFORMERS_ID = get_or_create_env_var( + "GPT_OSS_REPO_TRANSFORMERS_ID", "unsloth/gpt-oss-20b-unsloth-bnb-4bit" +) +if USE_LLAMA_CPP == "False": + GPT_OSS_REPO_ID = GPT_OSS_REPO_TRANSFORMERS_ID +GPT_OSS_MODEL_FILE = get_or_create_env_var("GPT_OSS_MODEL_FILE", "gpt-oss-20b-F16.gguf") +GPT_OSS_MODEL_FOLDER = get_or_create_env_var("GPT_OSS_MODEL_FOLDER", "model/gpt_oss") + +QWEN3_4B_REPO_ID = get_or_create_env_var( + "QWEN3_4B_REPO_ID", "unsloth/Qwen3-4B-Instruct-2507-GGUF" +) +QWEN3_4B_REPO_TRANSFORMERS_ID = get_or_create_env_var( + "QWEN3_4B_REPO_TRANSFORMERS_ID", "unsloth/Qwen3-4B-unsloth-bnb-4bit" +) +if USE_LLAMA_CPP == "False": + QWEN3_4B_REPO_ID = QWEN3_4B_REPO_TRANSFORMERS_ID + +QWEN3_4B_MODEL_FILE = get_or_create_env_var( + "QWEN3_4B_MODEL_FILE", "Qwen3-4B-Instruct-2507-UD-Q4_K_XL.gguf" +) +QWEN3_4B_MODEL_FOLDER = get_or_create_env_var("QWEN3_4B_MODEL_FOLDER", "model/qwen") + +GRANITE_4_TINY_REPO_ID = get_or_create_env_var( + "GRANITE_4_TINY_REPO_ID", "unsloth/granite-4.0-h-tiny-GGUF" +) +GRANITE_4_TINY_REPO_TRANSFORMERS_ID = get_or_create_env_var( + "GRANITE_4_TINY_REPO_TRANSFORMERS_ID", "unsloth/granite-4.0-h-tiny-FP8-Dynamic" +) +if USE_LLAMA_CPP == "False": + GRANITE_4_TINY_REPO_ID = GRANITE_4_TINY_REPO_TRANSFORMERS_ID +GRANITE_4_TINY_MODEL_FILE = get_or_create_env_var( + "GRANITE_4_TINY_MODEL_FILE", "granite-4.0-h-tiny-UD-Q4_K_XL.gguf" +) +GRANITE_4_TINY_MODEL_FOLDER = get_or_create_env_var( + "GRANITE_4_TINY_MODEL_FOLDER", "model/granite" +) + +GRANITE_4_3B_REPO_ID = get_or_create_env_var( + "GRANITE_4_3B_REPO_ID", "unsloth/granite-4.0-h-micro-GGUF" +) +GRANITE_4_3B_REPO_TRANSFORMERS_ID = get_or_create_env_var( + "GRANITE_4_3B_REPO_TRANSFORMERS_ID", "unsloth/granite-4.0-micro-unsloth-bnb-4bit" +) +if USE_LLAMA_CPP == "False": + GRANITE_4_3B_REPO_ID = GRANITE_4_3B_REPO_TRANSFORMERS_ID +GRANITE_4_3B_MODEL_FILE = get_or_create_env_var( + "GRANITE_4_3B_MODEL_FILE", "granite-4.0-h-micro-UD-Q4_K_XL.gguf" +) +GRANITE_4_3B_MODEL_FOLDER = get_or_create_env_var( + "GRANITE_4_3B_MODEL_FOLDER", "model/granite" +) + +if CHOSEN_LOCAL_MODEL_TYPE == "Gemma 2b": + LOCAL_REPO_ID = GEMMA2_REPO_ID + LOCAL_MODEL_FILE = GEMMA2_MODEL_FILE + LOCAL_MODEL_FOLDER = GEMMA2_MODEL_FOLDER + +elif CHOSEN_LOCAL_MODEL_TYPE == "Gemma 3 4B": + LOCAL_REPO_ID = GEMMA3_4B_REPO_ID + LOCAL_MODEL_FILE = GEMMA3_4B_MODEL_FILE + LOCAL_MODEL_FOLDER = GEMMA3_4B_MODEL_FOLDER + MULTIMODAL_PROMPT_FORMAT = "True" + +elif CHOSEN_LOCAL_MODEL_TYPE == "Gemma 3 12B": + LOCAL_REPO_ID = GEMMA3_12B_REPO_ID + LOCAL_MODEL_FILE = GEMMA3_12B_MODEL_FILE + LOCAL_MODEL_FOLDER = GEMMA3_12B_MODEL_FOLDER + MULTIMODAL_PROMPT_FORMAT = "True" + +elif CHOSEN_LOCAL_MODEL_TYPE == "Gemma 4 12B": + LOCAL_REPO_ID = GEMMA4_12B_REPO_ID + LOCAL_MODEL_FILE = GEMMA4_12B_MODEL_FILE + LOCAL_MODEL_FOLDER = GEMMA4_12B_MODEL_FOLDER + MULTIMODAL_PROMPT_FORMAT = "True" + +elif CHOSEN_LOCAL_MODEL_TYPE == "Gemma 4 26B": + LOCAL_REPO_ID = GEMMA4_26B_REPO_ID + LOCAL_MODEL_FILE = GEMMA4_26B_MODEL_FILE + LOCAL_MODEL_FOLDER = GEMMA4_26B_MODEL_FOLDER + MULTIMODAL_PROMPT_FORMAT = "True" + +elif CHOSEN_LOCAL_MODEL_TYPE == "Gemma 4 31B": + LOCAL_REPO_ID = GEMMA4_31B_REPO_ID + LOCAL_MODEL_FILE = GEMMA4_31B_MODEL_FILE + LOCAL_MODEL_FOLDER = GEMMA4_31B_MODEL_FOLDER + MULTIMODAL_PROMPT_FORMAT = "True" + +elif CHOSEN_LOCAL_MODEL_TYPE == "Qwen 3 4B": + LOCAL_REPO_ID = QWEN3_4B_REPO_ID + LOCAL_MODEL_FILE = QWEN3_4B_MODEL_FILE + LOCAL_MODEL_FOLDER = QWEN3_4B_MODEL_FOLDER + +elif CHOSEN_LOCAL_MODEL_TYPE == "gpt-oss-20b": + LOCAL_REPO_ID = GPT_OSS_REPO_ID + LOCAL_MODEL_FILE = GPT_OSS_MODEL_FILE + LOCAL_MODEL_FOLDER = GPT_OSS_MODEL_FOLDER + +elif CHOSEN_LOCAL_MODEL_TYPE == "Granite 4 Tiny": + LOCAL_REPO_ID = GRANITE_4_TINY_REPO_ID + LOCAL_MODEL_FILE = GRANITE_4_TINY_MODEL_FILE + LOCAL_MODEL_FOLDER = GRANITE_4_TINY_MODEL_FOLDER + +elif CHOSEN_LOCAL_MODEL_TYPE == "Granite 4 Micro": + LOCAL_REPO_ID = GRANITE_4_3B_REPO_ID + LOCAL_MODEL_FILE = GRANITE_4_3B_MODEL_FILE + LOCAL_MODEL_FOLDER = GRANITE_4_3B_MODEL_FOLDER + +elif not CHOSEN_LOCAL_MODEL_TYPE: + print("No local model type chosen") + LOCAL_REPO_ID = "" + LOCAL_MODEL_FILE = "" + LOCAL_MODEL_FOLDER = "" +else: + print("CHOSEN_LOCAL_MODEL_TYPE not found") + LOCAL_REPO_ID = "" + LOCAL_MODEL_FILE = "" + LOCAL_MODEL_FOLDER = "" + +USE_SPECULATIVE_DECODING = get_or_create_env_var("USE_SPECULATIVE_DECODING", "False") + +ASSISTANT_MODEL = get_or_create_env_var("ASSISTANT_MODEL", "") +if CHOSEN_LOCAL_MODEL_TYPE == "Gemma 3 4B": + ASSISTANT_MODEL = get_or_create_env_var( + "ASSISTANT_MODEL", "unsloth/gemma-3-270m-it" + ) +elif CHOSEN_LOCAL_MODEL_TYPE == "Qwen 3 4B": + ASSISTANT_MODEL = get_or_create_env_var("ASSISTANT_MODEL", "unsloth/Qwen3-0.6B") +elif CHOSEN_LOCAL_MODEL_TYPE == "Gemma 4 12B": + ASSISTANT_MODEL = get_or_create_env_var( + "ASSISTANT_MODEL", "unsloth/gemma-4-12B-it-qat-GGUF" + ) +elif CHOSEN_LOCAL_MODEL_TYPE == "Gemma 4 26B": + ASSISTANT_MODEL = get_or_create_env_var( + "ASSISTANT_MODEL", "unsloth/gemma-4-26B-A4B-it-qat-GGUF" + ) +elif CHOSEN_LOCAL_MODEL_TYPE == "Gemma 4 31B": + ASSISTANT_MODEL = get_or_create_env_var( + "ASSISTANT_MODEL", "unsloth/gemma-4-31B-it-qat-GGUF" + ) + +DRAFT_MODEL_LOC = get_or_create_env_var("DRAFT_MODEL_LOC", ".cache/llama.cpp/") + +GEMMA3_DRAFT_MODEL_LOC = get_or_create_env_var( + "GEMMA3_DRAFT_MODEL_LOC", + DRAFT_MODEL_LOC + "unsloth_gemma-3-270m-it-qat-GGUF_gemma-3-270m-it-qat-F16.gguf", +) +GEMMA3_4B_DRAFT_MODEL_LOC = get_or_create_env_var( + "GEMMA3_4B_DRAFT_MODEL_LOC", + DRAFT_MODEL_LOC + "unsloth_gemma-3-4b-it-qat-GGUF_gemma-3-4b-it-qat-Q4_K_M.gguf", +) + +GEMMA4_12B_DRAFT_MODEL_LOC = get_or_create_env_var( + "GEMMA4_12B_DRAFT_MODEL_LOC", + DRAFT_MODEL_LOC + "mtp-gemma-4-12B-it.gguf", +) +GEMMA4_26B_DRAFT_MODEL_LOC = get_or_create_env_var( + "GEMMA4_26B_DRAFT_MODEL_LOC", + DRAFT_MODEL_LOC + "mtp-gemma-4-26B-A4B-it.gguf", +) +GEMMA4_31B_DRAFT_MODEL_LOC = get_or_create_env_var( + "GEMMA4_31B_DRAFT_MODEL_LOC", + DRAFT_MODEL_LOC + "mtp-gemma-4-31B-it.gguf", +) + +QWEN3_DRAFT_MODEL_LOC = get_or_create_env_var( + "QWEN3_DRAFT_MODEL_LOC", DRAFT_MODEL_LOC + "Qwen3-0.6B-Q8_0.gguf" +) +QWEN3_4B_DRAFT_MODEL_LOC = get_or_create_env_var( + "QWEN3_4B_DRAFT_MODEL_LOC", + DRAFT_MODEL_LOC + "Qwen3-4B-Instruct-2507-UD-Q4_K_XL.gguf", +) + + +LLM_MAX_GPU_LAYERS = int( + get_or_create_env_var("LLM_MAX_GPU_LAYERS", "-1") +) # Maximum possible +LLM_TEMPERATURE = float(get_or_create_env_var("LLM_TEMPERATURE", "0.6")) +LLM_TOP_K = int( + get_or_create_env_var("LLM_TOP_K", "64") +) # https://docs.unsloth.ai/basics/gemma-3-how-to-run-and-fine-tune +LLM_MIN_P = float(get_or_create_env_var("LLM_MIN_P", "0")) +LLM_TOP_P = float(get_or_create_env_var("LLM_TOP_P", "0.95")) +LLM_REPETITION_PENALTY = float(get_or_create_env_var("LLM_REPETITION_PENALTY", "1.0")) +LLM_LAST_N_TOKENS = int(get_or_create_env_var("LLM_LAST_N_TOKENS", "512")) +LLM_MAX_NEW_TOKENS = int(get_or_create_env_var("LLM_MAX_NEW_TOKENS", "4096")) +LLM_SEED = int(get_or_create_env_var("LLM_SEED", "42")) +LLM_RESET = get_or_create_env_var("LLM_RESET", "False") +LLM_STREAM = get_or_create_env_var("LLM_STREAM", "True") +LLM_THREADS = int(get_or_create_env_var("LLM_THREADS", "-1")) +LLM_BATCH_SIZE = int(get_or_create_env_var("LLM_BATCH_SIZE", "2048")) +LLM_CONTEXT_LENGTH = int(get_or_create_env_var("LLM_CONTEXT_LENGTH", "24576")) +# Keep some headroom so we don't hit hard context errors if token counting differs +# slightly between client-side estimates and server-side tokenization/templates. +LLM_CONTEXT_HEADROOM_FRACTION = float( + get_or_create_env_var("LLM_CONTEXT_HEADROOM_FRACTION", "0.05") +) +LLM_SAMPLE = get_or_create_env_var("LLM_SAMPLE", "True") +LLM_STOP_STRINGS = get_or_create_env_var("LLM_STOP_STRINGS", r"['\n\n\n\n\n\n']") + +SPECULATIVE_DECODING = get_or_create_env_var("SPECULATIVE_DECODING", "False") +NUM_PRED_TOKENS = int(get_or_create_env_var("NUM_PRED_TOKENS", "2")) +K_QUANT_LEVEL = get_or_create_env_var( + "K_QUANT_LEVEL", "" +) # 2 = q4_0, 8 = q8_0, 4 = fp16 +V_QUANT_LEVEL = get_or_create_env_var( + "V_QUANT_LEVEL", "" +) # 2 = q4_0, 8 = q8_0, 4 = fp16 + +if not K_QUANT_LEVEL: + K_QUANT_LEVEL = None +else: + K_QUANT_LEVEL = int(K_QUANT_LEVEL) +if not V_QUANT_LEVEL: + V_QUANT_LEVEL = None +else: + V_QUANT_LEVEL = int(V_QUANT_LEVEL) + +# If you are using e.g. gpt-oss, you can add a reasoning suffix to set reasoning level, or turn it off in the case of Qwen 3 4B +if CHOSEN_LOCAL_MODEL_TYPE == "gpt-oss-20b": + REASONING_SUFFIX = get_or_create_env_var("REASONING_SUFFIX", "Reasoning: low") +elif CHOSEN_LOCAL_MODEL_TYPE == "Qwen 3 4B" and USE_LLAMA_CPP == "False": + REASONING_SUFFIX = get_or_create_env_var("REASONING_SUFFIX", "/nothink") +else: + REASONING_SUFFIX = get_or_create_env_var("REASONING_SUFFIX", "") + +# Transformers variables +COMPILE_TRANSFORMERS = get_or_create_env_var( + "COMPILE_TRANSFORMERS", "False" +) # Whether to compile transformers models +USE_BITSANDBYTES = get_or_create_env_var( + "USE_BITSANDBYTES", "True" +) # Whether to use bitsandbytes for quantization +COMPILE_MODE = get_or_create_env_var( + "COMPILE_MODE", "reduce-overhead" +) # alternatively 'max-autotune' +MODEL_DTYPE = get_or_create_env_var( + "MODEL_DTYPE", "bfloat16" +) # alternatively 'bfloat16' +INT8_WITH_OFFLOAD_TO_CPU = get_or_create_env_var( + "INT8_WITH_OFFLOAD_TO_CPU", "False" +) # Whether to offload to CPU + +DEFAULT_SAMPLED_SUMMARIES = int( + get_or_create_env_var("DEFAULT_SAMPLED_SUMMARIES", "75") +) + +### +# Gradio app variables +### + +# Get some environment variables and Launch the Gradio app +COGNITO_AUTH = get_or_create_env_var("COGNITO_AUTH", "0") + +RUN_DIRECT_MODE = get_or_create_env_var("RUN_DIRECT_MODE", "0") + +# Direct mode environment variables +DIRECT_MODE_TASK = get_or_create_env_var("DIRECT_MODE_TASK", "extract") +DIRECT_MODE_INPUT_FILE = get_or_create_env_var("DIRECT_MODE_INPUT_FILE", "") +DIRECT_MODE_OUTPUT_DIR = get_or_create_env_var("DIRECT_MODE_OUTPUT_DIR", OUTPUT_FOLDER) +DIRECT_MODE_S3_OUTPUT_BUCKET = get_or_create_env_var( + "DIRECT_MODE_S3_OUTPUT_BUCKET", S3_OUTPUTS_BUCKET +) +DIRECT_MODE_TEXT_COLUMN = get_or_create_env_var("DIRECT_MODE_TEXT_COLUMN", "") +DIRECT_MODE_PREVIOUS_OUTPUT_FILES = get_or_create_env_var( + "DIRECT_MODE_PREVIOUS_OUTPUT_FILES", "" +) +DIRECT_MODE_USERNAME = get_or_create_env_var("DIRECT_MODE_USERNAME", "") +DIRECT_MODE_GROUP_BY = get_or_create_env_var("DIRECT_MODE_GROUP_BY", "") +DIRECT_MODE_EXCEL_SHEETS = get_or_create_env_var("DIRECT_MODE_EXCEL_SHEETS", "") +DIRECT_MODE_MODEL_CHOICE = get_or_create_env_var( + "DIRECT_MODE_MODEL_CHOICE", default_model_choice +) +DIRECT_MODE_TEMPERATURE = get_or_create_env_var( + "DIRECT_MODE_TEMPERATURE", str(LLM_TEMPERATURE) +) +DIRECT_MODE_BATCH_SIZE = get_or_create_env_var( + "DIRECT_MODE_BATCH_SIZE", str(BATCH_SIZE_DEFAULT) +) +DIRECT_MODE_MAX_TOKENS = get_or_create_env_var( + "DIRECT_MODE_MAX_TOKENS", str(LLM_MAX_NEW_TOKENS) +) +DIRECT_MODE_CONTEXT = get_or_create_env_var("DIRECT_MODE_CONTEXT", "") +DIRECT_MODE_CANDIDATE_TOPICS = get_or_create_env_var("DIRECT_MODE_CANDIDATE_TOPICS", "") +DIRECT_MODE_FORCE_ZERO_SHOT = get_or_create_env_var("DIRECT_MODE_FORCE_ZERO_SHOT", "No") +DIRECT_MODE_FORCE_SINGLE_TOPIC = get_or_create_env_var( + "DIRECT_MODE_FORCE_SINGLE_TOPIC", "No" +) +DIRECT_MODE_PRODUCE_STRUCTURED_SUMMARY = get_or_create_env_var( + "DIRECT_MODE_PRODUCE_STRUCTURED_SUMMARY", "No" +) +DIRECT_MODE_SENTIMENT = get_or_create_env_var( + "DIRECT_MODE_SENTIMENT", "Negative or Positive" +) +DIRECT_MODE_ADDITIONAL_SUMMARY_INSTRUCTIONS = get_or_create_env_var( + "DIRECT_MODE_ADDITIONAL_SUMMARY_INSTRUCTIONS", "" +) +DIRECT_MODE_ADDITIONAL_VALIDATION_ISSUES = get_or_create_env_var( + "DIRECT_MODE_ADDITIONAL_VALIDATION_ISSUES", "" +) +DIRECT_MODE_SHOW_PREVIOUS_TABLE = get_or_create_env_var( + "DIRECT_MODE_SHOW_PREVIOUS_TABLE", "Yes" +) +DIRECT_MODE_MAX_TIME_FOR_LOOP = get_or_create_env_var( + "DIRECT_MODE_MAX_TIME_FOR_LOOP", str(MAX_TIME_FOR_LOOP) +) +DIRECT_MODE_DEDUP_METHOD = get_or_create_env_var("DIRECT_MODE_DEDUP_METHOD", "fuzzy") +DIRECT_MODE_SIMILARITY_THRESHOLD = get_or_create_env_var( + "DIRECT_MODE_SIMILARITY_THRESHOLD", str(DEDUPLICATION_THRESHOLD) +) +DIRECT_MODE_MERGE_SENTIMENT = get_or_create_env_var("DIRECT_MODE_MERGE_SENTIMENT", "No") +DIRECT_MODE_MERGE_GENERAL_TOPICS = get_or_create_env_var( + "DIRECT_MODE_MERGE_GENERAL_TOPICS", "Yes" +) +DIRECT_MODE_SUMMARY_FORMAT = get_or_create_env_var( + "DIRECT_MODE_SUMMARY_FORMAT", "two_paragraph" +) +DIRECT_MODE_SAMPLE_REFERENCE_TABLE = get_or_create_env_var( + "DIRECT_MODE_SAMPLE_REFERENCE_TABLE", "True" +) +DIRECT_MODE_NO_OF_SAMPLED_SUMMARIES = get_or_create_env_var( + "DIRECT_MODE_NO_OF_SAMPLED_SUMMARIES", str(DEFAULT_SAMPLED_SUMMARIES) +) +DIRECT_MODE_RANDOM_SEED = get_or_create_env_var( + "DIRECT_MODE_RANDOM_SEED", str(LLM_SEED) +) +DIRECT_MODE_CREATE_XLSX_OUTPUT = get_or_create_env_var( + "DIRECT_MODE_CREATE_XLSX_OUTPUT", "True" +) +DIRECT_MODE_CREATE_TOPICS_CSV = get_or_create_env_var( + "DIRECT_MODE_CREATE_TOPICS_CSV", "True" +) +DIRECT_MODE_SAMPLE_FRACTION = get_or_create_env_var("DIRECT_MODE_SAMPLE_FRACTION", "20") +DIRECT_MODE_S3_UPLOAD_ONLY_XLSX = convert_string_to_boolean( + get_or_create_env_var("DIRECT_MODE_S3_UPLOAD_ONLY_XLSX", "False") +) + +MAX_QUEUE_SIZE = int(get_or_create_env_var("MAX_QUEUE_SIZE", "5")) + +MAX_FILE_SIZE = get_or_create_env_var("MAX_FILE_SIZE", "250mb") + +GRADIO_SERVER_PORT = int(get_or_create_env_var("GRADIO_SERVER_PORT", "7860")) + +ROOT_PATH = get_or_create_env_var("ROOT_PATH", "") + +DEFAULT_CONCURRENCY_LIMIT = get_or_create_env_var("DEFAULT_CONCURRENCY_LIMIT", "3") + +FILE_INPUT_HEIGHT = int(get_or_create_env_var("FILE_INPUT_HEIGHT", "125")) + +SHOW_EXAMPLES = get_or_create_env_var("SHOW_EXAMPLES", "True") + +ALL_IN_ONE_USE_LLM_DEDUP = convert_string_to_boolean( + get_or_create_env_var("ALL_IN_ONE_USE_LLM_DEDUP", "False") +) + +### +# COST CODE OPTIONS +### + +SHOW_COSTS = get_or_create_env_var("SHOW_COSTS", "False") + +GET_COST_CODES = get_or_create_env_var("GET_COST_CODES", "False") + +DEFAULT_COST_CODE = get_or_create_env_var("DEFAULT_COST_CODE", "") + +DIRECT_MODE_DEFAULT_COST_CODE = get_or_create_env_var( + "DIRECT_MODE_DEFAULT_COST_CODE", "" +) + +COST_CODES_PATH = get_or_create_env_var( + "COST_CODES_PATH", "" +) # 'config/COST_CENTRES.csv' # file should be a csv file with a single table in it that has two columns with a header. First column should contain cost codes, second column should contain a name or description for the cost code + +S3_COST_CODES_PATH = get_or_create_env_var( + "S3_COST_CODES_PATH", "" +) # COST_CENTRES.csv # This is a path within the DOCUMENT_SUMMARISATION_BUCKET + +# A default path in case s3 cost code location is provided but no local cost code location given +if COST_CODES_PATH: + OUTPUT_COST_CODES_PATH = COST_CODES_PATH +else: + OUTPUT_COST_CODES_PATH = "config/cost_codes.csv" + +ENFORCE_COST_CODES = get_or_create_env_var( + "ENFORCE_COST_CODES", "False" +) # If you have cost codes listed, is it compulsory to choose one before redacting? + +if ENFORCE_COST_CODES == "True": + GET_COST_CODES = "True" + +### +# VALIDATE FOLDERS AND CONFIG OPTIONS +### + + +# Convert string environment variables to string or list +if CSV_ACCESS_LOG_HEADERS: + CSV_ACCESS_LOG_HEADERS = _get_env_list(CSV_ACCESS_LOG_HEADERS) +if CSV_FEEDBACK_LOG_HEADERS: + CSV_FEEDBACK_LOG_HEADERS = _get_env_list(CSV_FEEDBACK_LOG_HEADERS) +if CSV_USAGE_LOG_HEADERS: + CSV_USAGE_LOG_HEADERS = _get_env_list(CSV_USAGE_LOG_HEADERS) + +if DYNAMODB_ACCESS_LOG_HEADERS: + DYNAMODB_ACCESS_LOG_HEADERS = _get_env_list(DYNAMODB_ACCESS_LOG_HEADERS) +if DYNAMODB_FEEDBACK_LOG_HEADERS: + DYNAMODB_FEEDBACK_LOG_HEADERS = _get_env_list(DYNAMODB_FEEDBACK_LOG_HEADERS) +if DYNAMODB_USAGE_LOG_HEADERS: + DYNAMODB_USAGE_LOG_HEADERS = _get_env_list(DYNAMODB_USAGE_LOG_HEADERS) diff --git a/tools/custom_csvlogger.py b/tools/custom_csvlogger.py new file mode 100644 index 0000000000000000000000000000000000000000..79e2ab5a2e90f1b39cceebc47947c0db9fe7f815 --- /dev/null +++ b/tools/custom_csvlogger.py @@ -0,0 +1,333 @@ +from __future__ import annotations + +import csv +import os +import re +import time +import uuid +from collections.abc import Sequence +from datetime import datetime + +# from multiprocessing import Lock +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import boto3 +import botocore +from gradio import utils +from gradio_client import utils as client_utils + +from tools.config import AWS_ACCESS_KEY, AWS_REGION, AWS_SECRET_KEY, RUN_AWS_FUNCTIONS + +if TYPE_CHECKING: + from gradio.components import Component +from threading import Lock + +from gradio.flagging import FlaggingCallback + + +class CSVLogger_custom(FlaggingCallback): + """ + The default implementation of the FlaggingCallback abstract class in gradio>=5.0. Each flagged + sample (both the input and output data) is logged to a CSV file with headers on the machine running + the gradio app. Unlike ClassicCSVLogger, this implementation is concurrent-safe and it creates a new + dataset file every time the headers of the CSV (derived from the labels of the components) change. It also + only creates columns for "username" and "flag" if the flag_option and username are provided, respectively. + + Example: + import gradio as gr + def image_classifier(inp): + return {'cat': 0.3, 'dog': 0.7} + demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label", + flagging_callback=CSVLogger()) + Guides: using-flagging + """ + + def __init__( + self, + simplify_file_data: bool = True, + verbose: bool = True, + dataset_file_name: str | None = None, + ): + """ + Parameters: + simplify_file_data: If True, the file data will be simplified before being written to the CSV file. If CSVLogger is being used to cache examples, this is set to False to preserve the original FileData class + verbose: If True, prints messages to the console about the dataset file creation + dataset_file_name: The name of the dataset file to be created (should end in ".csv"). If None, the dataset file will be named "dataset1.csv" or the next available number. + """ + self.simplify_file_data = simplify_file_data + self.verbose = verbose + self.dataset_file_name = dataset_file_name + self.lock = Lock() + + def setup( + self, + components: Sequence[Component], + flagging_dir: str | Path, + ): + self.components = components + self.flagging_dir = Path(flagging_dir) + self.first_time = True + + def _create_dataset_file( + self, + additional_headers: list[str] | None = None, + replacement_headers: list[str] | None = None, + ): + os.makedirs(self.flagging_dir, exist_ok=True) + + if replacement_headers: + if additional_headers is None: + additional_headers = [] + + if len(replacement_headers) != len(self.components): + raise ValueError( + f"replacement_headers must have the same length as components " + f"({len(replacement_headers)} provided, {len(self.components)} expected)" + ) + headers = replacement_headers + additional_headers + ["timestamp"] + else: + if additional_headers is None: + additional_headers = [] + headers = ( + [ + getattr(component, "label", None) or f"component {idx}" + for idx, component in enumerate(self.components) + ] + + additional_headers + + ["timestamp"] + ) + + headers = utils.sanitize_list_for_csv(headers) + dataset_files = list(Path(self.flagging_dir).glob("dataset*.csv")) + + if self.dataset_file_name: + self.dataset_filepath = self.flagging_dir / self.dataset_file_name + elif dataset_files: + try: + latest_file = max( + dataset_files, key=lambda f: int(re.findall(r"\d+", f.stem)[0]) + ) + latest_num = int(re.findall(r"\d+", latest_file.stem)[0]) + + with open(latest_file, newline="", encoding="utf-8-sig") as csvfile: + reader = csv.reader(csvfile) + existing_headers = next(reader, None) + + if existing_headers != headers: + new_num = latest_num + 1 + self.dataset_filepath = self.flagging_dir / f"dataset{new_num}.csv" + else: + self.dataset_filepath = latest_file + except Exception: + self.dataset_filepath = self.flagging_dir / "dataset1.csv" + else: + self.dataset_filepath = self.flagging_dir / "dataset1.csv" + + if not Path(self.dataset_filepath).exists(): + with open( + self.dataset_filepath, "w", newline="", encoding="utf-8-sig" + ) as csvfile: + writer = csv.writer(csvfile) + writer.writerow(utils.sanitize_list_for_csv(headers)) + if self.verbose: + print("Created dataset file at:", self.dataset_filepath) + elif self.verbose: + print("Using existing dataset file at:", self.dataset_filepath) + + def flag( + self, + flag_data: list[Any], + flag_option: str | None = None, + username: str | None = None, + save_to_csv: bool = True, + save_to_dynamodb: bool = False, + dynamodb_table_name: str | None = None, + dynamodb_headers: list[str] | None = None, # New: specify headers for DynamoDB + replacement_headers: list[str] | None = None, + ) -> int: + if self.first_time: + # print("First time creating log file") + additional_headers = [] + if flag_option is not None: + additional_headers.append("flag") + if username is not None: + additional_headers.append("username") + additional_headers.append("id") + # additional_headers.append("timestamp") + self._create_dataset_file( + additional_headers=additional_headers, + replacement_headers=replacement_headers, + ) + self.first_time = False + + csv_data = [] + for idx, (component, sample) in enumerate( + zip(self.components, flag_data, strict=False) + ): + save_dir = ( + self.flagging_dir + / client_utils.strip_invalid_filename_characters( + getattr(component, "label", None) or f"component {idx}" + ) + ) + if utils.is_prop_update(sample): + csv_data.append(str(sample)) + else: + data = ( + component.flag(sample, flag_dir=save_dir) + if sample is not None + else "" + ) + if self.simplify_file_data: + data = utils.simplify_file_data_in_str(data) + csv_data.append(data) + + if flag_option is not None: + csv_data.append(flag_option) + if username is not None: + csv_data.append(username) + + generated_id = str(uuid.uuid4()) + csv_data.append(generated_id) + + timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[ + :-3 + ] # Correct format for Amazon Athena + csv_data.append(timestamp) + + # Build the headers + headers = [ + getattr(component, "label", None) or f"component {idx}" + for idx, component in enumerate(self.components) + ] + if flag_option is not None: + headers.append("flag") + if username is not None: + headers.append("username") + headers.append("id") + headers.append("timestamp") + + line_count = -1 + + if save_to_csv: + with self.lock: + with open( + self.dataset_filepath, "a", newline="", encoding="utf-8-sig" + ) as csvfile: + writer = csv.writer(csvfile) + writer.writerow(utils.sanitize_list_for_csv(csv_data)) + with open(self.dataset_filepath, encoding="utf-8-sig") as csvfile: + line_count = len(list(csv.reader(csvfile))) - 1 + + if save_to_dynamodb is True: + print("Saving to DynamoDB") + + if RUN_AWS_FUNCTIONS == "1": + try: + print("Connecting to DynamoDB via existing SSO connection") + dynamodb = boto3.resource("dynamodb", region_name=AWS_REGION) + # client = boto3.client('dynamodb') + + dynamodb.meta.client.list_tables() + + except Exception as e: + print("No SSO credentials found:", e) + if AWS_ACCESS_KEY and AWS_SECRET_KEY: + print("Trying DynamoDB credentials from environment variables") + dynamodb = boto3.resource( + "dynamodb", + aws_access_key_id=AWS_ACCESS_KEY, + aws_secret_access_key=AWS_SECRET_KEY, + region_name=AWS_REGION, + ) + # client = boto3.client('dynamodb',aws_access_key_id=AWS_ACCESS_KEY, + # aws_secret_access_key=AWS_SECRET_KEY, region_name=AWS_REGION) + else: + raise Exception( + "AWS credentials for DynamoDB logging not found" + ) + else: + raise Exception("AWS credentials for DynamoDB logging not found") + + if dynamodb_table_name is None: + raise ValueError( + "You must provide a dynamodb_table_name if save_to_dynamodb is True" + ) + + if dynamodb_headers: + dynamodb_headers = dynamodb_headers + if not dynamodb_headers and replacement_headers: + dynamodb_headers = replacement_headers + elif headers: + dynamodb_headers = headers + elif not dynamodb_headers: + raise ValueError( + "Headers not found. You must provide dynamodb_headers or replacement_headers to create a new table." + ) + + if flag_option is not None: + if "flag" not in dynamodb_headers: + dynamodb_headers.append("flag") + if username is not None: + if "username" not in dynamodb_headers: + dynamodb_headers.append("username") + if "timestamp" not in dynamodb_headers: + dynamodb_headers.append("timestamp") + if "id" not in dynamodb_headers: + dynamodb_headers.append("id") + + # Table doesn't exist — create it + try: + table = dynamodb.Table(dynamodb_table_name) + table.load() + except botocore.exceptions.ClientError as e: + if e.response["Error"]["Code"] == "ResourceNotFoundException": + + attribute_definitions = [ + { + "AttributeName": "id", + "AttributeType": "S", + } # Only define key attributes here + ] + + table = dynamodb.create_table( + TableName=dynamodb_table_name, + KeySchema=[ + {"AttributeName": "id", "KeyType": "HASH"} # Partition key + ], + AttributeDefinitions=attribute_definitions, + BillingMode="PAY_PER_REQUEST", + ) + # Wait until the table exists + table.meta.client.get_waiter("table_exists").wait( + TableName=dynamodb_table_name + ) + time.sleep(5) + print(f"Table '{dynamodb_table_name}' created successfully.") + else: + raise + + # Prepare the DynamoDB item to upload + try: + item = { + "id": str(generated_id), # UUID primary key + #'created_by': username if username else "unknown", + "timestamp": timestamp, + } + + # Map the headers to values + item.update( + { + header: str(value) + for header, value in zip(dynamodb_headers, csv_data) + } + ) + + table.put_item(Item=item) + + print("Successfully uploaded log to DynamoDB") + except Exception as e: + print("Could not upload log to DynamobDB due to", e) + + return line_count diff --git a/tools/dedup_summaries.py b/tools/dedup_summaries.py new file mode 100644 index 0000000000000000000000000000000000000000..dbf26a2637d085a3c02659830613bbadcc5e4c55 --- /dev/null +++ b/tools/dedup_summaries.py @@ -0,0 +1,4677 @@ +import os +import re +import time +from io import StringIO +from typing import List, Tuple, Union + +import boto3 +import gradio as gr +import markdown +import pandas as pd +import spaces +from rapidfuzz import fuzz, process +from tqdm import tqdm + +from tools.aws_functions import connect_to_bedrock_runtime +from tools.config import ( + BATCH_SIZE_DEFAULT, + CHOSEN_LOCAL_MODEL_TYPE, + DEDUPLICATION_THRESHOLD, + DEFAULT_SAMPLED_SUMMARIES, + INFERENCE_SERVER_SAMPLER_TOKEN_BUDGET_FRACTION, + LLM_CONTEXT_HEADROOM_FRACTION, + LLM_CONTEXT_LENGTH, + LLM_MAX_NEW_TOKENS, + LLM_SEED, + MAX_COMMENT_CHARS, + MAX_GROUPS, + MAX_SPACES_GPU_RUN_TIME, + MAX_TIME_FOR_LOOP, + MAXIMUM_ALLOWED_TOPICS, + NUMBER_OF_RETRY_ATTEMPTS, + OUTPUT_DEBUG_FILES, + OUTPUT_FOLDER, + REASONING_SUFFIX, + RUN_LOCAL_MODEL, + TIMEOUT_WAIT, + model_name_map, +) + + +def _effective_context_limit() -> int: + """Return max context length with safety headroom applied.""" + try: + frac = float(LLM_CONTEXT_HEADROOM_FRACTION) + except Exception: + frac = 0.05 + frac = min(max(frac, 0.0), 0.5) + return max(1, int(LLM_CONTEXT_LENGTH * (1.0 - frac))) + + +from tools.helper_functions import ( + clean_column_name, + convert_reference_table_to_pivot_table, + create_batch_file_path_details, + create_topic_summary_df_from_reference_table, + ensure_model_in_map, + generate_zero_shot_topics_df, + get_basic_response_data, + get_file_name_no_ext, + initial_clean, + load_in_data_file, + normalize_topic_name_for_llm, + read_file, + wrap_text, +) +from tools.llm_funcs import ( + calculate_tokens_from_metadata, + call_llm_with_markdown_table_checks, + construct_azure_client, + construct_gemini_generative_model, + get_assistant_model, + get_model, + get_tokenizer, + process_requests, +) +from tools.prompts import ( + comprehensive_summary_format_prompt, + comprehensive_summary_format_prompt_by_group, + llm_deduplication_prompt, + llm_deduplication_prompt_with_candidates, + llm_deduplication_system_prompt, + summarise_everything_prompt, + summarise_everything_system_prompt, + summarise_topic_descriptions_prompt, + summarise_topic_descriptions_system_prompt, + summary_assistant_prefill, + system_prompt, +) + +max_tokens = LLM_MAX_NEW_TOKENS +timeout_wait = TIMEOUT_WAIT +number_of_api_retry_attempts = NUMBER_OF_RETRY_ATTEMPTS +max_time_for_loop = MAX_TIME_FOR_LOOP +batch_size_default = BATCH_SIZE_DEFAULT +deduplication_threshold = DEDUPLICATION_THRESHOLD +max_comment_character_length = MAX_COMMENT_CHARS +reasoning_suffix = REASONING_SUFFIX +output_debug_files = OUTPUT_DEBUG_FILES +default_number_of_sampled_summaries = DEFAULT_SAMPLED_SUMMARIES +max_text_length = 500 +max_number_of_topics = MAXIMUM_ALLOWED_TOPICS + + +# DEDUPLICATION/SUMMARISATION FUNCTIONS +def deduplicate_categories( + category_series: pd.Series, + join_series: pd.Series, + reference_df: pd.DataFrame, + general_topic_series: pd.Series = None, + merge_general_topics="No", + merge_sentiment: str = "No", + threshold: float = 90, +) -> pd.DataFrame: + """ + Deduplicates similar category names in a pandas Series based on a fuzzy matching threshold, + merging smaller topics into larger topics. + + Parameters: + category_series (pd.Series): Series containing category names to deduplicate. + join_series (pd.Series): Additional series used for joining back to original results. + reference_df (pd.DataFrame): DataFrame containing the reference data to count occurrences. + threshold (float): Similarity threshold for considering two strings as duplicates. + + Returns: + pd.DataFrame: DataFrame with columns ['old_category', 'deduplicated_category']. + """ + # Count occurrences of each category in the reference_df + category_counts = reference_df["Subtopic"].value_counts().to_dict() + + # Initialize dictionaries for both category mapping and scores + deduplication_map = {} + match_scores = {} # New dictionary to store match scores + + # First pass: Handle exact matches + for category in category_series.unique(): + if category in deduplication_map: + continue + + # Find all exact matches + exact_matches = category_series[ + category_series.str.lower() == category.lower() + ].index.tolist() + if len(exact_matches) > 1: + # Find the variant with the highest count + match_counts = { + match: category_counts.get(category_series[match], 0) + for match in exact_matches + } + most_common = max(match_counts.items(), key=lambda x: x[1])[0] + most_common_category = category_series[most_common] + + # Map all exact matches to the most common variant and store score + for match in exact_matches: + deduplication_map[category_series[match]] = most_common_category + match_scores[category_series[match]] = ( + 100 # Exact matches get score of 100 + ) + + # Second pass: Handle fuzzy matches for remaining categories + # Create a DataFrame to maintain the relationship between categories and general topics + categories_df = pd.DataFrame( + {"category": category_series, "general_topic": general_topic_series} + ).drop_duplicates() + + for _, row in categories_df.iterrows(): + category = row["category"] + if category in deduplication_map: + continue + + current_general_topic = row["general_topic"] + + # Filter potential matches to only those within the same General topic if relevant + if merge_general_topics == "No": + potential_matches = categories_df[ + (categories_df["category"] != category) + & (categories_df["general_topic"] == current_general_topic) + ]["category"].tolist() + else: + potential_matches = categories_df[(categories_df["category"] != category)][ + "category" + ].tolist() + + matches = process.extract( + category, potential_matches, scorer=fuzz.WRatio, score_cutoff=threshold + ) + + if matches: + best_match = max(matches, key=lambda x: x[1]) + match, score, _ = best_match + + if category_counts.get(category, 0) < category_counts.get(match, 0): + deduplication_map[category] = match + match_scores[category] = score + else: + deduplication_map[match] = category + match_scores[match] = score + else: + deduplication_map[category] = category + match_scores[category] = 100 + + # Create the result DataFrame with scores + result_df = pd.DataFrame( + { + "old_category": category_series + " | " + join_series, + "deduplicated_category": category_series.map( + lambda x: deduplication_map.get(x, x) + ), + "match_score": category_series.map( + lambda x: match_scores.get(x, 100) + ), # Add scores column + } + ) + + # print(result_df) + + return result_df + + +def deduplicate_topics( + reference_df: pd.DataFrame, + topic_summary_df: pd.DataFrame, + reference_table_file_name: str, + unique_topics_table_file_name: str, + in_excel_sheets: str = "", + merge_sentiment: str = "No", + merge_general_topics: str = "No", + score_threshold: int = 90, + in_data_files: Union[List[str], pd.DataFrame] = list(), + chosen_cols: List[str] = "", + output_folder: str = OUTPUT_FOLDER, + deduplicate_topics: str = "Yes", + output_files: str = "True", + total_number_of_batches: int = None, + data_file_names_textbox: str = None, + sentiment_checkbox: str = "Negative, Neutral, or Positive", +): + """ + Deduplicate topics based on a reference and unique topics table, merging similar topics. + + Args: + reference_df (pd.DataFrame): DataFrame containing reference data with topics. + topic_summary_df (pd.DataFrame): DataFrame summarizing unique topics. + reference_table_file_name (str): Base file name for the output reference table. + unique_topics_table_file_name (str): Base file name for the output unique topics table. + in_excel_sheets (str, optional): Comma-separated list of Excel sheet names to load. Defaults to "". + merge_sentiment (str, optional): Whether to merge topics regardless of sentiment ("Yes" or "No"). Defaults to "No". + merge_general_topics (str, optional): Whether to merge topics across different general topics ("Yes" or "No"). Defaults to "No". + score_threshold (int, optional): Fuzzy matching score threshold for deduplication. Defaults to 90. + in_data_files (Union[List[str], pd.DataFrame], optional): List of input data file paths or a pandas DataFrame. If a DataFrame is provided, it will be used directly without loading from file. Defaults to []. + chosen_cols (List[str], optional): List of chosen columns from the input data files. Defaults to "". + output_folder (str, optional): Folder path to save output files. Defaults to OUTPUT_FOLDER. + deduplicate_topics (str, optional): Whether to perform topic deduplication ("Yes" or "No"). Defaults to "Yes". + output_files (str, optional): Whether to output files ("True" or "False"). Defaults to "True". + total_number_of_batches (int, optional): Total number of batches when in_data_files is a DataFrame. If None and in_data_files is a DataFrame, defaults to 1. Defaults to None. + data_file_names_textbox (str, optional): File name when in_data_files is a DataFrame. If None and in_data_files is a DataFrame, defaults to "dataframe". Defaults to None. + """ + # Save the parameter value before it gets overwritten + should_output_files = output_files + output_files = list() # This is now the list of output file paths + log_output_files = list() + file_data = pd.DataFrame() + deduplicated_unique_table_markdown = "" + # Use provided values if available, otherwise use defaults + if data_file_names_textbox is None: + data_file_names_textbox = "dataframe" + if total_number_of_batches is None: + total_number_of_batches = 1 + + # Validate that required columns exist in reference_df + if "Response ID" not in reference_df.columns: + raise ValueError( + "reference_df must contain 'Response ID' column. " + f"Available columns: {list(reference_df.columns)}" + ) + + # Many downstream dedup steps assume a Sentiment column exists. + # When sentiment assessment is disabled upstream, we normalise to a single placeholder value. + if sentiment_checkbox == "Do not assess sentiment": + if "Sentiment" not in reference_df.columns: + reference_df["Sentiment"] = "Not assessed" + if "Sentiment" not in topic_summary_df.columns: + topic_summary_df["Sentiment"] = "Not assessed" + + # Add 'Topic number' column if it doesn't exist (row number starting from 1) + if "Topic number" not in topic_summary_df.columns: + if not topic_summary_df.empty: + topic_summary_df["Topic number"] = range(1, len(topic_summary_df) + 1) + else: + # If empty, create empty column - it will be populated when DataFrame is recreated + topic_summary_df["Topic number"] = pd.Series(dtype="int64") + + if "Response ID" in reference_df.columns: + # Convert float or str to int + reference_df["Response ID"] = ( + reference_df["Response ID"].astype(float).astype(int) + ) + if "Start row of group" in reference_df.columns: + # Convert float or str to int + reference_df["Start row of group"] = ( + reference_df["Start row of group"].astype(float).astype(int) + ) + + if (len(reference_df["Response ID"].unique()) == 1) | ( + len(topic_summary_df["Topic number"].unique()) == 1 + ): + print( + "Data file outputs are too short for deduplicating. Returning original data." + ) + + # Get file name without extension and create proper output paths + reference_table_file_name_no_ext = get_file_name_no_ext( + reference_table_file_name + ) + unique_topics_table_file_name_no_ext = get_file_name_no_ext( + unique_topics_table_file_name + ) + + # Create output paths with _dedup suffix to match normal path + reference_file_out_path = ( + output_folder + reference_table_file_name_no_ext + "_dedup.csv" + ) + unique_topics_file_out_path = ( + output_folder + unique_topics_table_file_name_no_ext + "_dedup.csv" + ) + + if should_output_files == "True": + # Save the DataFrames to CSV files + reference_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + reference_file_out_path, index=None, encoding="utf-8-sig" + ) + topic_summary_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + unique_topics_file_out_path, index=None, encoding="utf-8-sig" + ) + output_files.append(reference_file_out_path) + output_files.append(unique_topics_file_out_path) + + # Create markdown output for display + topic_summary_df_revised_display = topic_summary_df.apply( + lambda col: col.map(lambda x: wrap_text(x, max_text_length=max_text_length)) + ) + deduplicated_unique_table_markdown = ( + topic_summary_df_revised_display.to_markdown(index=False) + ) + + return ( + reference_df, + topic_summary_df, + output_files, + log_output_files, + deduplicated_unique_table_markdown, + ) + + # For checking that data is not lost during the process + initial_unique_references = len(reference_df["Response ID"].unique()) + + if topic_summary_df.empty: + topic_summary_df = create_topic_summary_df_from_reference_table( + reference_df, sentiment_checkbox=sentiment_checkbox + ) + + # Then merge the topic numbers back to the original dataframe + reference_df = reference_df.merge( + topic_summary_df[ + ["General topic", "Subtopic", "Sentiment", "Topic number"] + ], + on=["General topic", "Subtopic", "Sentiment"], + how="left", + ) + + # Check if in_data_files is not empty (handles both DataFrame and list) + has_data_files = ( + (isinstance(in_data_files, pd.DataFrame) and not in_data_files.empty) + or (isinstance(in_data_files, list) and len(in_data_files) > 0) + or (not isinstance(in_data_files, (pd.DataFrame, list)) and in_data_files) + ) + if has_data_files and chosen_cols: + # Check if in_data_files is already a DataFrame + if isinstance(in_data_files, pd.DataFrame): + # Use the DataFrame directly + file_data = in_data_files.copy() + # Filter to chosen columns if specified + if chosen_cols and isinstance(chosen_cols, list): + # Ensure all chosen columns exist in the DataFrame + available_cols = [ + col for col in chosen_cols if col in file_data.columns + ] + if available_cols: + file_data = file_data[available_cols] + else: + print( + f"Warning: None of the chosen columns {chosen_cols} found in DataFrame. Using all columns." + ) + + else: + # Load from file path(s) as before + file_data, data_file_names_textbox, total_number_of_batches = ( + load_in_data_file( + in_data_files, chosen_cols, total_number_of_batches, in_excel_sheets + ) + ) + else: + out_message = "No file data found, pivot table output will not be created." + print(out_message) + # raise Exception(out_message) + + # Run through this x times to try to get all duplicate topics + if deduplicate_topics == "Yes": + if "Group" not in reference_df.columns: + reference_df["Group"] = "All" + for i in range(0, 8): + if merge_sentiment == "No": + if merge_general_topics == "No": + reference_df["old_category"] = ( + reference_df["Subtopic"] + " | " + reference_df["Sentiment"] + ) + reference_df_unique = reference_df.drop_duplicates("old_category") + + # Create an empty list to store results from each group + results = list() + # Iterate over each group instead of using .apply() + for name, group in reference_df_unique.groupby( + ["General topic", "Sentiment", "Group"] + ): + # Run your function on the 'group' DataFrame + result = deduplicate_categories( + group["Subtopic"], + group["Sentiment"], + reference_df, + general_topic_series=group["General topic"], + merge_general_topics="No", + threshold=score_threshold, + ) + results.append(result) + + # Concatenate all the results into a single DataFrame + deduplicated_topic_map_df = pd.concat(results).reset_index( + drop=True + ) + + else: + # This case should allow cross-topic matching but is still grouping by Sentiment + reference_df["old_category"] = ( + reference_df["Subtopic"] + " | " + reference_df["Sentiment"] + ) + reference_df_unique = reference_df.drop_duplicates("old_category") + + results = list() + for name, group in reference_df_unique.groupby("Sentiment"): + result = deduplicate_categories( + group["Subtopic"], + group["Sentiment"], + reference_df, + general_topic_series=None, + merge_general_topics="Yes", + threshold=score_threshold, + ) + results.append(result) + deduplicated_topic_map_df = pd.concat(results).reset_index( + drop=True + ) + + else: + if merge_general_topics == "No": + reference_df["old_category"] = ( + reference_df["Subtopic"] + " | " + reference_df["Sentiment"] + ) + reference_df_unique = reference_df.drop_duplicates("old_category") + + results = list() + for name, group in reference_df_unique.groupby("General topic"): + result = deduplicate_categories( + group["Subtopic"], + group["Sentiment"], + reference_df, + general_topic_series=group["General topic"], + merge_general_topics="No", + merge_sentiment=merge_sentiment, + threshold=score_threshold, + ) + results.append(result) + deduplicated_topic_map_df = pd.concat(results).reset_index( + drop=True + ) + + else: + reference_df["old_category"] = ( + reference_df["Subtopic"] + " | " + reference_df["Sentiment"] + ) + reference_df_unique = reference_df.drop_duplicates("old_category") + + deduplicated_topic_map_df = deduplicate_categories( + reference_df_unique["Subtopic"], + reference_df_unique["Sentiment"], + reference_df, + general_topic_series=None, + merge_general_topics="Yes", + merge_sentiment=merge_sentiment, + threshold=score_threshold, + ).reset_index(drop=True) + + if deduplicated_topic_map_df["deduplicated_category"].isnull().all(): + print("No deduplicated categories found, skipping the following code.") + + else: + # Remove rows where 'deduplicated_category' is blank or NaN + deduplicated_topic_map_df = deduplicated_topic_map_df.loc[ + ( + deduplicated_topic_map_df["deduplicated_category"].str.strip() + != "" + ) + & ~(deduplicated_topic_map_df["deduplicated_category"].isnull()), + ["old_category", "deduplicated_category", "match_score"], + ] + + reference_df = reference_df.merge( + deduplicated_topic_map_df, on="old_category", how="left" + ) + + reference_df.rename( + columns={"Subtopic": "Subtopic_old", "Sentiment": "Sentiment_old"}, + inplace=True, + ) + # Extract subtopic and sentiment from deduplicated_category + reference_df["Subtopic"] = reference_df[ + "deduplicated_category" + ].str.extract(r"^(.*?) \|")[ + 0 + ] # Extract subtopic + reference_df["Sentiment"] = reference_df[ + "deduplicated_category" + ].str.extract(r"\| (.*)$")[ + 0 + ] # Extract sentiment + + # Combine with old values to ensure no data is lost + reference_df["Subtopic"] = reference_df[ + "deduplicated_category" + ].combine_first(reference_df["Subtopic_old"]) + reference_df["Sentiment"] = reference_df["Sentiment"].combine_first( + reference_df["Sentiment_old"] + ) + + reference_df = reference_df.rename( + columns={"General Topic": "General topic"}, errors="ignore" + ) + reference_df = reference_df[ + [ + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + "Summary", + "Start row of group", + "Group", + ] + ] + + if merge_general_topics == "Yes": + # Replace General topic names for each Subtopic with that for the Subtopic with the most responses + # Step 1: Count the number of occurrences for each General topic and Subtopic combination + count_df = ( + reference_df.groupby(["Subtopic", "General topic"]) + .size() + .reset_index(name="Count") + ) + + # Step 2: Find the General topic with the maximum count for each Subtopic + max_general_topic = count_df.loc[ + count_df.groupby("Subtopic")["Count"].idxmax() + ] + + # Step 3: Map the General topic back to the original DataFrame + reference_df = reference_df.merge( + max_general_topic[["Subtopic", "General topic"]], + on="Subtopic", + suffixes=("", "_max"), + how="left", + ) + + reference_df["General topic"] = reference_df[ + "General topic_max" + ].combine_first(reference_df["General topic"]) + + if merge_sentiment == "Yes": + # Step 1: Count the number of occurrences for each General topic and Subtopic combination + count_df = ( + reference_df.groupby(["Subtopic", "Sentiment"]) + .size() + .reset_index(name="Count") + ) + + # Step 2: Determine the number of unique Sentiment values for each Subtopic + unique_sentiments = ( + count_df.groupby("Subtopic")["Sentiment"] + .nunique() + .reset_index(name="UniqueCount") + ) + + # Step 3: Update Sentiment to 'Mixed' where there is more than one unique sentiment + reference_df = reference_df.merge( + unique_sentiments, on="Subtopic", how="left" + ) + reference_df["Sentiment"] = reference_df.apply( + lambda row: "Mixed" if row["UniqueCount"] > 1 else row["Sentiment"], + axis=1, + ) + + # Clean up the DataFrame by dropping the UniqueCount column + reference_df.drop(columns=["UniqueCount"], inplace=True) + + # print("reference_df:", reference_df) + reference_df = reference_df[ + [ + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + "Summary", + "Start row of group", + "Group", + ] + ] + + # Update reference summary column with all summaries + reference_df["Summary"] = reference_df.groupby( + ["Response ID", "General topic", "Subtopic", "Sentiment"] + )["Summary"].transform("
".join) + + # Check that we have not inadvertantly removed some data during the above process + end_unique_references = len(reference_df["Response ID"].unique()) + + if initial_unique_references != end_unique_references: + raise Exception( + f"Number of unique references changed during processing: Initial={initial_unique_references}, Final={end_unique_references}" + ) + + # Drop duplicates in the reference table - each comment should only have the same topic referred to once + reference_df.drop_duplicates( + ["Response ID", "General topic", "Subtopic", "Sentiment"], + inplace=True, + ) + + # Before recreating topic_summary_df, check if input had Group information + # If input topic_summary_df doesn't have Group or all have same Group, normalize reference_df Group + input_has_group = "Group" in topic_summary_df.columns + input_unique_groups = ( + topic_summary_df["Group"].nunique() if input_has_group else 0 + ) + + # Check reference_df Group values before normalization + ref_df_has_group = "Group" in reference_df.columns + ref_df_unique_groups_before = ( + reference_df["Group"].nunique() if ref_df_has_group else 0 + ) + + # If input didn't have meaningful Group distinction, normalise to a single Group + # ONLY when the reference table also doesn't have meaningful Group distinction. + # If reference_df contains multiple Groups, collapsing them here would silently drop + # group-level outputs downstream (e.g. only one Group appears in the XLSX). + if (not input_has_group or input_unique_groups <= 1) and ( + ref_df_unique_groups_before <= 1 + ): + # Get the most common Group value from reference_df, or use "All" if not present + if ref_df_has_group and not reference_df["Group"].empty: + most_common_group = ( + reference_df["Group"].mode()[0] + if len(reference_df["Group"].mode()) > 0 + else "All" + ) + else: + most_common_group = "All" + # Normalize all Groups to the most common one to prevent topic count increase + reference_df["Group"] = most_common_group + + # Verify normalization worked + ref_df_unique_groups_after = reference_df["Group"].nunique() + if ref_df_unique_groups_before > 1 and ref_df_unique_groups_after > 1: + print( + f"Warning: Group normalization may have failed. " + f"reference_df had {ref_df_unique_groups_before} unique Groups before normalization, " + f"and {ref_df_unique_groups_after} after. Forcing all to '{most_common_group}'." + ) + reference_df["Group"] = most_common_group + elif (not input_has_group or input_unique_groups <= 1) and ( + ref_df_unique_groups_before > 1 + ): + print( + "Warning: topic_summary_df had no/one Group but reference_df has multiple Groups. " + "Preserving reference_df Group values to avoid dropping grouped outputs." + ) + + # Remake topic_summary_df based on new reference_df + # Rebuild topic_summary_df from updated reference_df to ensure consistency + # This ensures the topic count reflects the actual merged state + if not reference_df.empty: + # Check if Response ID are present and not all empty + if "Response ID" in reference_df.columns: + non_empty_refs = reference_df["Response ID"].notna() & ( + reference_df["Response ID"].astype(str).str.strip() != "" + ) + if not non_empty_refs.any(): + print( + "Warning: All Response ID are empty in reference_df. " + "Number of responses will be 0 for all topics." + ) + topic_summary_df = create_topic_summary_df_from_reference_table( + reference_df, sentiment_checkbox=sentiment_checkbox + ) + else: + print( + "Warning: reference_df is empty after deduplication. " + "Cannot recreate topic_summary_df. Number of responses will be 0 for all topics." + ) + # Create an empty topic_summary_df with the expected structure + topic_summary_df = pd.DataFrame() + + # Then merge the topic numbers back to the original dataframe + # Only merge if both dataframes are not empty + if not reference_df.empty and not topic_summary_df.empty: + on_cols = ["General topic", "Subtopic"] + if ( + sentiment_checkbox != "Do not assess sentiment" + and "Sentiment" in reference_df.columns + and "Sentiment" in topic_summary_df.columns + ): + on_cols.append("Sentiment") + if "Group" in reference_df.columns and "Group" in topic_summary_df.columns: + on_cols.append("Group") + on_cols = [ + c + for c in on_cols + if c in reference_df.columns and c in topic_summary_df.columns + ] + + right_cols = list(on_cols) + if "Topic number" in topic_summary_df.columns: + right_cols.append("Topic number") + right_cols = [c for c in right_cols if c in topic_summary_df.columns] + + reference_df = reference_df.merge( + topic_summary_df[right_cols], + on=on_cols, + how="left", + ) + + else: + print("Topics have not beeen deduplicated") + + reference_table_file_name_no_ext = get_file_name_no_ext(reference_table_file_name) + unique_topics_table_file_name_no_ext = get_file_name_no_ext( + unique_topics_table_file_name + ) + + if not file_data.empty: + basic_response_data = get_basic_response_data(file_data, chosen_cols) + reference_df_pivot = convert_reference_table_to_pivot_table( + reference_df, basic_response_data + ) + + reference_pivot_file_path = ( + output_folder + reference_table_file_name_no_ext + "_pivot_dedup.csv" + ) + if should_output_files == "True": + reference_df_pivot.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + reference_pivot_file_path, index=None, encoding="utf-8-sig" + ) + log_output_files.append(reference_pivot_file_path) + + reference_file_out_path = ( + output_folder + reference_table_file_name_no_ext + "_dedup.csv" + ) + unique_topics_file_out_path = ( + output_folder + unique_topics_table_file_name_no_ext + "_dedup.csv" + ) + + if should_output_files == "True": + reference_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + reference_file_out_path, index=None, encoding="utf-8-sig" + ) + topic_summary_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + unique_topics_file_out_path, index=None, encoding="utf-8-sig" + ) + + output_files.append(reference_file_out_path) + output_files.append(unique_topics_file_out_path) + + # Outputs for markdown table output + topic_summary_df_revised_display = topic_summary_df.apply( + lambda col: col.map(lambda x: wrap_text(x, max_text_length=max_text_length)) + ) + deduplicated_unique_table_markdown = topic_summary_df_revised_display.to_markdown( + index=False + ) + + print("Deduplication task successfully completed") + + return ( + reference_df, + topic_summary_df, + output_files, + log_output_files, + deduplicated_unique_table_markdown, + ) + + +def deduplicate_topics_llm( + reference_df: pd.DataFrame, + topic_summary_df: pd.DataFrame, + reference_table_file_name: str, + unique_topics_table_file_name: str, + model_choice: str, + in_api_key: str, + temperature: float, + model_source: str, + local_model=None, + tokenizer=None, + assistant_model=None, + in_excel_sheets: str = "", + merge_sentiment: str = "No", + merge_general_topics: str = "No", + in_data_files: Union[List[str], pd.DataFrame] = list(), + chosen_cols: List[str] = "", + output_folder: str = OUTPUT_FOLDER, + candidate_topics=None, + azure_endpoint: str = "", + output_debug_files: str = "False", + api_url: str = None, + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + azure_api_key_textbox: str = "", + model_name_map: dict = model_name_map, + output_files: str = "True", + sentiment_checkbox: str = "Negative or Positive", +): + """ + Deduplicate topics using LLM semantic understanding to identify and merge similar topics. + + Args: + reference_df (pd.DataFrame): DataFrame containing reference data with topics. + topic_summary_df (pd.DataFrame): DataFrame summarizing unique topics. + reference_table_file_name (str): Base file name for the output reference table. + unique_topics_table_file_name (str): Base file name for the output unique topics table. + model_choice (str): The LLM model to use for deduplication. + in_api_key (str): Google API key for the LLM service (for Gemini models). + temperature (float): Temperature setting for the LLM. + model_source (str): Source of the model (AWS, Gemini, Local, etc.). + local_model: Local model instance (if using local model). + tokenizer: Tokenizer for local model. + assistant_model: Assistant model for speculative decoding. + in_excel_sheets (str, optional): Comma-separated list of Excel sheet names to load. Defaults to "". + merge_sentiment (str, optional): Whether to merge topics regardless of sentiment ("Yes" or "No"). Defaults to "No". + merge_general_topics (str, optional): Whether to merge topics across different general topics ("Yes" or "No"). Defaults to "No". + in_data_files (Union[List[str], pd.DataFrame], optional): List of input data file paths or a pandas DataFrame. If a DataFrame is provided, it will be used directly without loading from file. Defaults to []. + chosen_cols (List[str], optional): List of chosen columns from the input data files. Defaults to "". + output_folder (str, optional): Folder path to save output files. Defaults to OUTPUT_FOLDER. + candidate_topics (optional): Candidate topics file for zero-shot guidance. Defaults to None. + azure_endpoint (str, optional): Azure endpoint for the LLM. Defaults to "". + output_debug_files (str, optional): Whether to output debug files. Defaults to "False". + api_url (str, optional): API URL for inference-server models. Defaults to None. + aws_access_key_textbox (str, optional): AWS access key for Bedrock. Defaults to "". + aws_secret_key_textbox (str, optional): AWS secret key for Bedrock. Defaults to "". + aws_region_textbox (str, optional): AWS region for Bedrock. Defaults to "". + azure_api_key_textbox (str, optional): Azure API key for Azure/OpenAI models. Defaults to "". + model_name_map (dict, optional): Mapping of model names to their configurations. Defaults to model_name_map from config. + output_files (str, optional): Whether to output files ("True" or "False"). Defaults to "True". + sentiment_checkbox (str, optional): Sentiment analysis option ("Negative or Positive", "Negative, Neutral, or Positive", or "Do not assess sentiment"). Defaults to "Negative or Positive". + output_files (str, optional): Whether to output files ("True" or "False"). Defaults to "True". + """ + + tic = time.perf_counter() + + # Save the parameter value before it gets overwritten + should_output_files = output_files + output_files = list() # This is now the list of output file paths + log_output_files = list() + file_data = pd.DataFrame() + deduplicated_unique_table_markdown = "" + + if "Response ID" in reference_df.columns: + # Convert float or str to int + reference_df["Response ID"] = ( + reference_df["Response ID"].astype(float).astype(int) + ) + if "Start row of group" in reference_df.columns: + # Convert float or str to int + reference_df["Start row of group"] = ( + reference_df["Start row of group"].astype(float).astype(int) + ) + + # Check if data is too short for deduplication + if (len(reference_df["Response ID"].unique()) == 1) | ( + len(topic_summary_df["Topic number"].unique()) == 1 + ): + print( + "Data file outputs are too short for deduplicating. Returning original data." + ) + + # Get file name without extension and create proper output paths + reference_table_file_name_no_ext = get_file_name_no_ext( + reference_table_file_name + ) + unique_topics_table_file_name_no_ext = get_file_name_no_ext( + unique_topics_table_file_name + ) + + # Create output paths with _dedup suffix to match normal path + reference_file_out_path = ( + output_folder + reference_table_file_name_no_ext + "_dedup.csv" + ) + unique_topics_file_out_path = ( + output_folder + unique_topics_table_file_name_no_ext + "_dedup.csv" + ) + + if should_output_files == "True": + # Save the DataFrames to CSV files + reference_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + reference_file_out_path, index=None, encoding="utf-8-sig" + ) + topic_summary_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + unique_topics_file_out_path, index=None, encoding="utf-8-sig" + ) + + output_files.append(reference_file_out_path) + output_files.append(unique_topics_file_out_path) + + # Create markdown output for display + topic_summary_df_revised_display = topic_summary_df.apply( + lambda col: col.map(lambda x: wrap_text(x, max_text_length=max_text_length)) + ) + deduplicated_unique_table_markdown = ( + topic_summary_df_revised_display.to_markdown(index=False) + ) + + # Compare input vs output for early return case + if not topic_summary_df.empty: + input_unique_topics = topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + output_unique_topics = topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + print( + f"Topic count comparison (early return): Input had {input_unique_topics} unique 'General topic' | 'Subtopic' combinations, " + f"Output has {output_unique_topics} unique 'General topic' | 'Subtopic' combinations. " + f"No deduplication performed (data too short)." + ) + + # Return with token counts set to 0 for early return + return ( + reference_df, + topic_summary_df, + output_files, + log_output_files, + deduplicated_unique_table_markdown, + 0, # input_tokens + 0, # output_tokens + 0, # number_of_calls + 0.0, # estimated_time_taken + ) + + # For checking that data is not lost during the process + initial_unique_references = len(reference_df["Response ID"].unique()) + + # Capture input topic count for comparison + if not topic_summary_df.empty: + input_unique_topics = topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + else: + input_unique_topics = 0 + + # Create topic summary if it doesn't exist + if topic_summary_df.empty: + topic_summary_df = create_topic_summary_df_from_reference_table(reference_df) + # Update input count if we just created the topic summary + input_unique_topics = topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + + topic_summary_df.drop( + ["1", "2", "3", "Response ID"], + axis=1, + errors="ignore", + inplace=True, + ) + + if "Topic number" not in reference_df.columns: + # Merge topic numbers back to the original dataframe + reference_df = reference_df.merge( + topic_summary_df[ + ["General topic", "Subtopic", "Sentiment", "Topic number"] + ], + on=["General topic", "Subtopic", "Sentiment"], + how="left", + ) + + orig_unique_topics_file_out_path = ( + output_folder + + get_file_name_no_ext(unique_topics_table_file_name) + + "_orig_pre_dedup.csv" + ) + if should_output_files == "True": + topic_summary_df.drop( + ["1", "2", "3", "Response ID"], axis=1, errors="ignore" + ).to_csv(orig_unique_topics_file_out_path, index=None, encoding="utf-8-sig") + + # Load data files if provided + # Check if in_data_files is not empty (handles both DataFrame and list) + has_data_files = ( + (isinstance(in_data_files, pd.DataFrame) and not in_data_files.empty) + or (isinstance(in_data_files, list) and len(in_data_files) > 0) + or (not isinstance(in_data_files, (pd.DataFrame, list)) and in_data_files) + ) + if has_data_files and chosen_cols: + # Check if in_data_files is already a DataFrame + if isinstance(in_data_files, pd.DataFrame): + # Use the DataFrame directly + file_data = in_data_files.copy() + # Filter to chosen columns if specified + if chosen_cols and isinstance(chosen_cols, list): + # Ensure all chosen columns exist in the DataFrame + available_cols = [ + col for col in chosen_cols if col in file_data.columns + ] + if available_cols: + file_data = file_data[available_cols] + else: + print( + f"Warning: None of the chosen columns {chosen_cols} found in DataFrame. Using all columns." + ) + data_file_names_textbox = "dataframe" + total_number_of_batches = 1 + else: + # Load from file path(s) as before + file_data, data_file_names_textbox, total_number_of_batches = ( + load_in_data_file(in_data_files, chosen_cols, 1, in_excel_sheets) + ) + else: + out_message = "No file data found, pivot table output will not be created." + print(out_message) + + # Process candidate topics if provided + candidate_topics_table = "" + if candidate_topics is not None: + try: + + # Read and process candidate topics + # Handle both string paths (CLI) and gr.FileData objects (Gradio) + candidate_topics_path = ( + candidate_topics + if isinstance(candidate_topics, str) + else getattr(candidate_topics, "name", None) + ) + if candidate_topics_path is None: + raise ValueError( + "candidate_topics must be a file path string or a FileData object with a 'name' attribute" + ) + candidate_topics_df = read_file(candidate_topics_path) + candidate_topics_df = candidate_topics_df.fillna("") + candidate_topics_df = candidate_topics_df.astype(str) + + # Generate zero-shot topics DataFrame + zero_shot_topics_df = generate_zero_shot_topics_df( + candidate_topics_df, "No", False + ) + + if not zero_shot_topics_df.empty: + candidate_topics_table = zero_shot_topics_df[ + ["General topic", "Subtopic"] + ].to_markdown(index=False) + print( + f"Found {len(zero_shot_topics_df)} candidate topics to consider during deduplication" + ) + except Exception as e: + print(f"Error processing candidate topics: {e}") + candidate_topics_table = "" + + # Determine if sentiment should be included based on sentiment_checkbox + include_sentiment = sentiment_checkbox != "Do not assess sentiment" + + # Normalize topic names for LLM comparison to avoid capitalization-only "merges" + # Create a copy of topic_summary_df with normalized topic names for the LLM + # This prevents the LLM from seeing capitalization differences as different topics + # Use the shared normalize_topic_name_for_llm function from helper_functions + # Create normalized version for LLM (but keep original for mapping back) + topics_for_llm = topic_summary_df.copy() + topics_for_llm["General topic"] = topics_for_llm["General topic"].apply( + normalize_topic_name_for_llm + ) + topics_for_llm["Subtopic"] = topics_for_llm["Subtopic"].apply( + normalize_topic_name_for_llm + ) + + # Remove duplicates after normalization (these are the same topic with different capitalization) + # This prevents the LLM from seeing them as separate topics + if include_sentiment: + topics_for_llm = topics_for_llm.drop_duplicates( + subset=["General topic", "Subtopic", "Sentiment"] + ) + else: + topics_for_llm = topics_for_llm.drop_duplicates( + subset=["General topic", "Subtopic"] + ) + + # Prepare topics table for LLM analysis (conditionally include sentiment) + if include_sentiment: + topics_table = topics_for_llm[ + ["General topic", "Subtopic", "Sentiment", "Number of responses"] + ].to_markdown(index=False) + sentiment_text = ", and Sentiment classifications" + sentiment_columns = "\n3. 'Original Sentiment' - The current sentiment" + merged_sentiment_columns = "\n4. 'Merged Sentiment' - The consolidated sentiment (use 'Mixed' if sentiments differ)" + else: + topics_table = topics_for_llm[ + ["General topic", "Subtopic", "Number of responses"] + ].to_markdown(index=False) + sentiment_text = "" + sentiment_columns = "" + merged_sentiment_columns = "" + + # Format the prompt with candidate topics if available + if candidate_topics_table: + formatted_prompt = llm_deduplication_prompt_with_candidates.format( + topics_table=topics_table, + candidate_topics_table=candidate_topics_table, + max_number_of_topics=max_number_of_topics, + sentiment_text=sentiment_text, + sentiment_columns=sentiment_columns, + merged_sentiment_columns=merged_sentiment_columns, + ) + else: + formatted_prompt = llm_deduplication_prompt.format( + topics_table=topics_table, + max_number_of_topics=max_number_of_topics, + sentiment_text=sentiment_text, + sentiment_columns=sentiment_columns, + merged_sentiment_columns=merged_sentiment_columns, + ) + + # Initialise conversation history + conversation_history = list() + whole_conversation = list() + whole_conversation_metadata = list() + + # Set up model clients based on model source + if "Gemini" in model_source: + print("Using Gemini model:", model_choice) + client, config = construct_gemini_generative_model( + in_api_key, + temperature, + model_choice, + llm_deduplication_system_prompt, + max_tokens, + LLM_SEED, + ) + bedrock_runtime = None + elif "Azure/OpenAI" in model_source: + print("Using Azure/OpenAI AI Inference model:", model_choice) + if azure_api_key_textbox: + os.environ["AZURE_INFERENCE_CREDENTIAL"] = azure_api_key_textbox + client, config = construct_azure_client( + in_api_key=azure_api_key_textbox, endpoint=azure_endpoint + ) + bedrock_runtime = None + elif "AWS" in model_source: + print("Using AWS Bedrock model:", model_choice) + bedrock_runtime = connect_to_bedrock_runtime( + model_name_map, + model_choice, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + ) + client = None + config = None + elif "Local" in model_source: + print("Using local model:", model_choice) + client = None + config = None + bedrock_runtime = None + elif "inference-server" in model_source: + print("Using inference-server model:", model_choice) + client = None + config = None + bedrock_runtime = None + # api_url is already passed to call_llm_with_markdown_table_checks + if api_url is None: + raise ValueError( + "api_url is required when model_source is 'inference-server'" + ) + else: + raise ValueError(f"Unsupported model source: {model_source}") + + # Call LLM to get deduplication suggestions + print("Calling LLM for topic deduplication analysis...") + + # Use the existing call_llm_with_markdown_table_checks function + ( + responses, + conversation_history, + whole_conversation, + whole_conversation_metadata, + response_text, + ) = call_llm_with_markdown_table_checks( + batch_prompts=[formatted_prompt], + system_prompt=llm_deduplication_system_prompt, + conversation_history=conversation_history, + whole_conversation=whole_conversation, + whole_conversation_metadata=whole_conversation_metadata, + client=client, + client_config=config, + model_choice=model_choice, + temperature=temperature, + reported_batch_no=1, + local_model=local_model, + tokenizer=tokenizer, + bedrock_runtime=bedrock_runtime, + model_source=model_source, + MAX_OUTPUT_VALIDATION_ATTEMPTS=3, + assistant_prefill="", + master=False, + CHOSEN_LOCAL_MODEL_TYPE=CHOSEN_LOCAL_MODEL_TYPE, + random_seed=LLM_SEED, + api_url=api_url, + ) + + # Generate debug files if enabled + if should_output_files == "True": + try: + # Create batch file path details for debug files + batch_file_path_details = ( + get_file_name_no_ext(reference_table_file_name) + "_llm_dedup" + ) + model_choice_clean_short = ( + model_choice.replace("/", "_").replace(":", "_").replace(".", "_") + ) + + # Create full prompt for debug output + full_prompt = llm_deduplication_system_prompt + "\n" + formatted_prompt + + # Write debug files + ( + current_prompt_content_logged, + current_summary_content_logged, + current_conversation_content_logged, + current_metadata_content_logged, + ) = process_debug_output_iteration( + OUTPUT_DEBUG_FILES, + output_folder, + batch_file_path_details, + model_choice_clean_short, + full_prompt, + response_text, + whole_conversation, + whole_conversation_metadata, + log_output_files, + task_type="llm_deduplication", + ) + + print("Debug files written for LLM deduplication analysis") + + except Exception as e: + print(f"Error writing debug files for LLM deduplication: {e}") + + # Parse the LLM response to extract merge suggestions + merge_suggestions_df = ( + pd.DataFrame() + ) # Initialize empty DataFrame for analysis results + num_merges_applied = 0 + # Track unique topic combinations being removed and added + # This helps explain the discrepancy between merge operations and actual topic reductions + unique_original_combinations = set() + unique_merged_combinations = set() + + try: + # Extract the markdown table from the response + table_match = re.search( + r"\|.*\|.*\n\|.*\|.*\n(\|.*\|.*\n)*", response_text, re.MULTILINE + ) + if table_match: + table_text = table_match.group(0) + + merge_suggestions_df = pd.read_csv( + StringIO(table_text), sep="|", skipinitialspace=True + ) + + # Clean up the DataFrame + # Remove empty columns (created from leading/trailing pipes in markdown) + merge_suggestions_df = merge_suggestions_df.dropna( + axis=1, how="all" + ) # Remove empty columns + merge_suggestions_df.columns = merge_suggestions_df.columns.str.strip() + + # Remove columns that are empty strings or just whitespace (from markdown table parsing) + # This fixes issues where leading/trailing pipes create empty column names + # Also remove columns with invalid names like 'end', '>' that can appear from parsing errors + expected_column_keywords = [ + "Original General topic", + "Original Subtopic", + "Merged General topic", + "Merged Subtopic", + "Merge Reason", + "Original Sentiment", + "Merged Sentiment", + ] + + valid_columns = [] + found_expected_column = False + for col in merge_suggestions_df.columns: + col_stripped = str(col).strip() + # Check if this column matches an expected column name + is_expected = any( + keyword.lower() in col_stripped.lower() + for keyword in expected_column_keywords + ) + + if is_expected: + found_expected_column = True + valid_columns.append(col) + elif found_expected_column: + # Once we've found expected columns, keep subsequent columns too + # (in case there are additional valid columns) + if col_stripped and col_stripped not in ["", "nan", "None"]: + valid_columns.append(col) + elif not found_expected_column: + # Before finding expected columns, filter out obvious artifacts + # Exclude single character columns and common parsing artifacts + if ( + col_stripped + and col_stripped not in ["", "nan", "None", "end", ">"] + and len(col_stripped) > 1 + ): + valid_columns.append(col) + + # If we found expected columns, use only valid columns; otherwise keep all non-empty + if found_expected_column and valid_columns: + merge_suggestions_df = merge_suggestions_df[valid_columns] + elif not found_expected_column: + # If no expected columns found, filter out obvious artifacts + valid_columns = [ + col + for col in merge_suggestions_df.columns + if str(col).strip() not in ["", "nan", "None", "end", ">"] + and len(str(col).strip()) > 1 + ] + if valid_columns: + merge_suggestions_df = merge_suggestions_df[valid_columns] + + # Also remove any columns where all values are empty strings after stripping + cols_to_drop = [] + for col in merge_suggestions_df.columns: + if merge_suggestions_df[col].astype(str).str.strip().eq("").all(): + cols_to_drop.append(col) + if cols_to_drop: + merge_suggestions_df = merge_suggestions_df.drop(columns=cols_to_drop) + + # Remove rows where all values are NaN + merge_suggestions_df = merge_suggestions_df.dropna(how="all") + + # Convert all columns to string to avoid float/NaN issues when calling .strip() + for col in merge_suggestions_df.columns: + merge_suggestions_df[col] = ( + merge_suggestions_df[col] + .astype(str) + .replace("nan", "") + .replace("NaN", "") + .replace("None", "") + ) + + # Filter out markdown table divider rows (rows that are primarily dashes/hyphens) + # These are false positives from parsing the markdown table structure + def is_divider_row(row): + """Check if a row is a markdown table divider (contains mostly dashes/hyphens).""" + row_str = " ".join(str(val) for val in row.values if pd.notna(val)) + # Remove whitespace and check if it's mostly dashes/hyphens + row_clean = row_str.replace(" ", "").replace("|", "").strip() + # If the row is empty or consists mostly of dashes/hyphens, it's a divider + if not row_clean or len(row_clean) == 0: + return True + # Check if more than 80% of characters are dashes/hyphens + dash_count = sum(1 for c in row_clean if c in ["-", "_", "="]) + return dash_count > 0.8 * len(row_clean) if len(row_clean) > 0 else True + + # Remove divider rows + merge_suggestions_df = merge_suggestions_df[ + ~merge_suggestions_df.apply(is_divider_row, axis=1) + ] + + if not merge_suggestions_df.empty: + print( + f"LLM identified {len(merge_suggestions_df)} potential topic merges" + ) + + # Deduplicate merge suggestions to avoid processing the same merge multiple times + # Normalize the original topic names for deduplication + def normalize_for_dedup(name): + """Normalize topic name for deduplication comparison.""" + if pd.isna(name) or name is None: + return "" + normalized = str(name) + normalized = initial_clean(normalized) + normalized = normalized.strip() + normalized = normalized.replace("\n", " ") + normalized = normalized.replace("\r", " ") + normalized = normalized.replace("/", " or ") + normalized = normalized.replace("&", " and ") + normalized = normalized.replace(" s ", "s ") + normalized = normalized.lower() + return normalized + + # Create normalized columns for deduplication + if "Original General topic" in merge_suggestions_df.columns: + merge_suggestions_df["_orig_gen_norm"] = merge_suggestions_df[ + "Original General topic" + ].apply(normalize_for_dedup) + else: + merge_suggestions_df["_orig_gen_norm"] = "" + + if "Original Subtopic" in merge_suggestions_df.columns: + merge_suggestions_df["_orig_sub_norm"] = merge_suggestions_df[ + "Original Subtopic" + ].apply(normalize_for_dedup) + else: + merge_suggestions_df["_orig_sub_norm"] = "" + + if include_sentiment: + # Deduplicate based on normalized original topics and sentiment + if "Original Sentiment" in merge_suggestions_df.columns: + merge_suggestions_df["_orig_sent"] = merge_suggestions_df[ + "Original Sentiment" + ].astype(str) + else: + merge_suggestions_df["_orig_sent"] = "" + initial_count = len(merge_suggestions_df) + merge_suggestions_df = merge_suggestions_df.drop_duplicates( + subset=["_orig_gen_norm", "_orig_sub_norm", "_orig_sent"], + keep="first", + ) + else: + # Deduplicate based on normalized original topics only + initial_count = len(merge_suggestions_df) + merge_suggestions_df = merge_suggestions_df.drop_duplicates( + subset=["_orig_gen_norm", "_orig_sub_norm"], keep="first" + ) + + # Remove the temporary normalization columns + merge_suggestions_df = merge_suggestions_df.drop( + columns=["_orig_gen_norm", "_orig_sub_norm", "_orig_sent"], + errors="ignore", + ) + + genuine_merges_count = len(merge_suggestions_df) + duplicates_removed = initial_count - genuine_merges_count + + if duplicates_removed > 0: + print( + f"Removed {duplicates_removed} duplicate merge suggestion(s). " + f"Processing {genuine_merges_count} genuine merge(s)." + ) + else: + print(f"Processing {genuine_merges_count} genuine merge(s).") + + # Apply the merges to the reference_df + # Helper function to normalize topic names for comparison + # Uses the same transformations applied to topic names in reference_df + def normalize_topic_name(name): + """Normalize topic name using the same transformations as reference_df topics for comparison.""" + if pd.isna(name) or name is None: + return "" + # Apply the same cleaning and transformations used in llm_api_call.py + normalized = str(name) + normalized = initial_clean(normalized) + normalized = normalized.strip() + normalized = normalized.replace("\n", " ") + normalized = normalized.replace("\r", " ") + normalized = normalized.replace("/", " or ") + normalized = normalized.replace("&", " and ") + normalized = normalized.replace(" s ", "s ") + normalized = normalized.lower() + return normalized + + # OPTIMIZATION: Pre-compute normalized topic mappings for efficient lookup + # Instead of normalizing every row for every merge suggestion, normalize unique combinations once + # Create a mapping from (actual_general, actual_subtopic, sentiment) -> normalized_general, normalized_subtopic + # This allows us to quickly find matching rows without repeated normalization + + # Get unique topic combinations from reference_df + if include_sentiment: + unique_topics = reference_df[ + ["General topic", "Subtopic", "Sentiment"] + ].drop_duplicates() + # Create normalized lookup dictionaries + # Map: (actual_general, actual_subtopic, sentiment) -> (normalized_general, normalized_subtopic) + topic_normalization_map = {} + for _, row in unique_topics.iterrows(): + key = ( + str(row["General topic"]), + str(row["Subtopic"]), + str(row["Sentiment"]), + ) + topic_normalization_map[key] = ( + normalize_topic_name(row["General topic"]), + normalize_topic_name(row["Subtopic"]), + ) + + # Create reverse lookup: (normalized_general, normalized_subtopic, sentiment) -> set of (actual_general, actual_subtopic, sentiment) + normalized_to_actual_map = {} + for key, (norm_gen, norm_sub) in topic_normalization_map.items(): + norm_key = ( + norm_gen, + norm_sub, + key[2], + ) # (normalized_gen, normalized_sub, sentiment) + if norm_key not in normalized_to_actual_map: + normalized_to_actual_map[norm_key] = set() + normalized_to_actual_map[norm_key].add(key) + + # Create a mask lookup: for each row index, store its normalized topic combination + # This allows us to quickly find all rows matching a normalized topic + row_normalized_map = {} + for idx, row in reference_df.iterrows(): + key = ( + str(row["General topic"]), + str(row["Subtopic"]), + str(row["Sentiment"]), + ) + if key in topic_normalization_map: + norm_gen, norm_sub = topic_normalization_map[key] + norm_key = (norm_gen, norm_sub, key[2]) + if norm_key not in row_normalized_map: + row_normalized_map[norm_key] = [] + row_normalized_map[norm_key].append(idx) + else: + unique_topics = reference_df[ + ["General topic", "Subtopic"] + ].drop_duplicates() + # Create normalized lookup dictionaries + topic_normalization_map = {} + for _, row in unique_topics.iterrows(): + key = (str(row["General topic"]), str(row["Subtopic"])) + topic_normalization_map[key] = ( + normalize_topic_name(row["General topic"]), + normalize_topic_name(row["Subtopic"]), + ) + + # Create reverse lookup: (normalized_general, normalized_subtopic) -> set of (actual_general, actual_subtopic) + normalized_to_actual_map = {} + for key, (norm_gen, norm_sub) in topic_normalization_map.items(): + norm_key = (norm_gen, norm_sub) + if norm_key not in normalized_to_actual_map: + normalized_to_actual_map[norm_key] = set() + normalized_to_actual_map[norm_key].add(key) + + # Create a mask lookup: for each normalized topic, store list of row indices + row_normalized_map = {} + for idx, row in reference_df.iterrows(): + key = (str(row["General topic"]), str(row["Subtopic"])) + if key in topic_normalization_map: + norm_gen, norm_sub = topic_normalization_map[key] + norm_key = (norm_gen, norm_sub) + if norm_key not in row_normalized_map: + row_normalized_map[norm_key] = [] + row_normalized_map[norm_key].append(idx) + + for _, row in merge_suggestions_df.iterrows(): + # Safely extract and convert values to strings, handling NaN/float values + def safe_get_strip(row, key, default=""): + value = row.get(key, default) + if ( + pd.isna(value) + or value is None + or str(value).lower() in ["nan", "none", ""] + ): + return default + return str(value).strip() + + original_general = safe_get_strip(row, "Original General topic", "") + original_subtopic = safe_get_strip(row, "Original Subtopic", "") + merged_general = safe_get_strip(row, "Merged General topic", "") + merged_subtopic = safe_get_strip(row, "Merged Subtopic", "") + + # Conditionally handle sentiment based on include_sentiment flag + if include_sentiment: + original_sentiment = safe_get_strip( + row, "Original Sentiment", "" + ) + merged_sentiment = safe_get_strip(row, "Merged Sentiment", "") + + # Check all required fields including sentiment + if all( + [ + original_general, + original_subtopic, + original_sentiment, + merged_general, + merged_subtopic, + merged_sentiment, + ] + ): + # Find matching rows using pre-computed normalized lookup + normalized_orig_gen = normalize_topic_name(original_general) + normalized_orig_sub = normalize_topic_name( + original_subtopic + ) + lookup_key = ( + normalized_orig_gen, + normalized_orig_sub, + original_sentiment, + ) + + # Get row indices that match this normalized topic combination + matching_indices = row_normalized_map.get(lookup_key, []) + # Create boolean mask using index.isin for efficiency + mask = reference_df.index.isin(matching_indices) + + if mask.any(): + # Normalize merged topics (original topics already normalized above) + normalized_merged_gen = normalize_topic_name( + merged_general + ) + normalized_merged_sub = normalize_topic_name( + merged_subtopic + ) + + # Skip if normalized versions are identical (only capitalization changed) + if ( + normalized_orig_gen == normalized_merged_gen + and normalized_orig_sub == normalized_merged_sub + ): + # This is just a capitalization change or self-merge, skip it + print( + f"Skipped self-merge: '{original_general}' | '{original_subtopic}' | '{original_sentiment}' -> " + f"'{merged_general}' | '{merged_subtopic}' | '{merged_sentiment}' " + f"(normalized versions are identical)" + ) + continue + + # Check if the merged combination already exists in reference_df (before applying the merge) + # This indicates a real merge (consolidation) vs just a rename + # Use pre-computed lookup for efficiency + merged_lookup_key = ( + normalized_merged_gen, + normalized_merged_sub, + merged_sentiment, + ) + merged_exists_before = ( + merged_lookup_key in row_normalized_map + and len(row_normalized_map[merged_lookup_key]) > 0 + ) + + # Check if the merged combo is different from the original (not just renaming to itself) + is_different_combo = ( + normalized_orig_gen != normalized_merged_gen + or normalized_orig_sub != normalized_merged_sub + or original_sentiment != merged_sentiment + ) + + # It's a real consolidation if the merged combo already exists and is different from original + is_real_merge = ( + merged_exists_before and is_different_combo + ) + + # Update the matching rows (including sentiment) + # Use merged values which already have whitespace stripped and preserve original case + reference_df.loc[mask, "General topic"] = merged_general + reference_df.loc[mask, "Subtopic"] = merged_subtopic + reference_df.loc[mask, "Sentiment"] = merged_sentiment + num_merges_applied += 1 + + # Track unique combinations for reporting + orig_combo = ( + normalized_orig_gen, + normalized_orig_sub, + original_sentiment, + ) + merged_combo = ( + normalized_merged_gen, + normalized_merged_sub, + merged_sentiment, + ) + unique_original_combinations.add(orig_combo) + unique_merged_combinations.add(merged_combo) + + merge_type = ( + "consolidated" if is_real_merge else "renamed" + ) + print( + f"Merged ({merge_type}): {original_general} | {original_subtopic} | {original_sentiment} -> {merged_general} | {merged_subtopic} | {merged_sentiment}" + ) + else: + # Debug: show why merge failed + # Use pre-computed normalized topics from lookup + # Get available topics from pre-computed map for debugging + available_gen_topics = { + norm_gen + for ( + norm_gen, + norm_sub, + _, + ) in row_normalized_map.keys() + } + available_sub_topics = { + norm_sub + for ( + norm_gen, + norm_sub, + _, + ) in row_normalized_map.keys() + } + # Check if the normalized general topic exists + gen_exists = normalized_orig_gen in available_gen_topics + sub_exists = normalized_orig_sub in available_sub_topics + print( + f"Failed to merge: '{original_general}' | '{original_subtopic}' | '{original_sentiment}' " + f"(normalized: '{normalized_orig_gen}' | '{normalized_orig_sub}'). " + f"General topic exists: {gen_exists}, Subtopic exists: {sub_exists}. " + f"No matching rows found in reference_df." + ) + + else: + # Check required fields without sentiment + if all( + [ + original_general, + original_subtopic, + merged_general, + merged_subtopic, + ] + ): + # Find matching rows using pre-computed normalized lookup + normalized_orig_gen = normalize_topic_name(original_general) + normalized_orig_sub = normalize_topic_name( + original_subtopic + ) + lookup_key = (normalized_orig_gen, normalized_orig_sub) + + # Get row indices that match this normalized topic combination + matching_indices = row_normalized_map.get(lookup_key, []) + # Create boolean mask using index.isin for efficiency + mask = reference_df.index.isin(matching_indices) + + if mask.any(): + # Normalize merged topics (original topics already normalized above) + normalized_merged_gen = normalize_topic_name( + merged_general + ) + normalized_merged_sub = normalize_topic_name( + merged_subtopic + ) + + # Skip if normalized versions are identical (only capitalization changed) + if ( + normalized_orig_gen == normalized_merged_gen + and normalized_orig_sub == normalized_merged_sub + ): + # This is just a capitalization change or self-merge, skip it + print( + f"Skipped self-merge: '{original_general}' | '{original_subtopic}' -> " + f"'{merged_general}' | '{merged_subtopic}' " + f"(normalized versions are identical)" + ) + continue + + # Check if the merged combination already exists in reference_df (before applying the merge) + # This indicates a real merge (consolidation) vs just a rename + # Use pre-computed lookup for efficiency + merged_lookup_key = ( + normalized_merged_gen, + normalized_merged_sub, + ) + merged_exists_before = ( + merged_lookup_key in row_normalized_map + and len(row_normalized_map[merged_lookup_key]) > 0 + ) + + # Check if the merged combo is different from the original (not just renaming to itself) + is_different_combo = ( + normalized_orig_gen != normalized_merged_gen + or normalized_orig_sub != normalized_merged_sub + ) + + # It's a real consolidation if the merged combo already exists and is different from original + is_real_merge = ( + merged_exists_before and is_different_combo + ) + + # Update the matching rows (without sentiment) + # Use merged values which already have whitespace stripped and preserve original case + reference_df.loc[mask, "General topic"] = merged_general + reference_df.loc[mask, "Subtopic"] = merged_subtopic + num_merges_applied += 1 + + # Track unique combinations for reporting + orig_combo = (normalized_orig_gen, normalized_orig_sub) + merged_combo = ( + normalized_merged_gen, + normalized_merged_sub, + ) + unique_original_combinations.add(orig_combo) + unique_merged_combinations.add(merged_combo) + + merge_type = ( + "consolidated" if is_real_merge else "renamed" + ) + print( + f"Merged ({merge_type}): {original_general} | {original_subtopic} -> {merged_general} | {merged_subtopic}" + ) + else: + # Debug: show why merge failed + # Use pre-computed normalized topics from lookup + # Get available topics from pre-computed map for debugging + available_gen_topics = { + norm_gen + for ( + norm_gen, + norm_sub, + ) in row_normalized_map.keys() + } + available_sub_topics = { + norm_sub + for ( + norm_gen, + norm_sub, + ) in row_normalized_map.keys() + } + # Check if the normalized general topic exists + gen_exists = normalized_orig_gen in available_gen_topics + sub_exists = normalized_orig_sub in available_sub_topics + print( + f"Failed to merge: '{original_general}' | '{original_subtopic}' " + f"(normalized: '{normalized_orig_gen}' | '{normalized_orig_sub}'). " + f"General topic exists: {gen_exists}, Subtopic exists: {sub_exists}. " + f"No matching rows found in reference_df." + ) + else: + print("No merge suggestions found in LLM response") + else: + print("No markdown table found in LLM response") + + if num_merges_applied == 0: + print("No duplicate topics found to merge") + else: + # Report on unique topic combinations affected + unique_originals_count = len(unique_original_combinations) + unique_merged_count = len(unique_merged_combinations) + unique_reduction = unique_originals_count - unique_merged_count + + print( + f"Merge summary: {num_merges_applied} merge operation(s) processed, " + f"affecting {unique_originals_count} unique original topic combination(s), " + f"resulting in {unique_merged_count} unique merged topic combination(s). " + f"Net reduction: {unique_reduction} unique topic combination(s)." + ) + + except Exception as e: + print(f"Error parsing LLM response: {e}") + print("Continuing with original data...") + + # Update reference summary column with all summaries + if include_sentiment: + reference_df["Summary"] = reference_df.groupby( + ["Response ID", "General topic", "Subtopic", "Sentiment"] + )["Summary"].transform("
".join) + else: + reference_df["Summary"] = reference_df.groupby( + ["Response ID", "General topic", "Subtopic"] + )["Summary"].transform("
".join) + + # Check that we have not inadvertently removed some data during the process + end_unique_references = len(reference_df["Response ID"].unique()) + + if initial_unique_references != end_unique_references: + raise Exception( + f"Number of unique references changed during processing: Initial={initial_unique_references}, Final={end_unique_references}" + ) + + # Drop duplicates in the reference table + if include_sentiment: + reference_df.drop_duplicates( + ["Response ID", "General topic", "Subtopic", "Sentiment"], + inplace=True, + ) + else: + reference_df.drop_duplicates( + ["Response ID", "General topic", "Subtopic"], inplace=True + ) + + # Before recreating topic_summary_df, check if input had Group information + # If input topic_summary_df doesn't have Group or all have same Group, normalize reference_df Group + input_has_group = "Group" in topic_summary_df.columns + input_unique_groups = topic_summary_df["Group"].nunique() if input_has_group else 0 + + # If input didn't have meaningful Group distinction, normalize to single Group + if not input_has_group or input_unique_groups <= 1: + # Get the most common Group value from reference_df, or use "All" if not present + if "Group" in reference_df.columns and not reference_df["Group"].empty: + most_common_group = ( + reference_df["Group"].mode()[0] + if len(reference_df["Group"].mode()) > 0 + else "All" + ) + else: + most_common_group = "All" + # Normalize all Groups to the most common one to prevent topic count increase + reference_df["Group"] = most_common_group + + # Remake topic_summary_df based on new reference_df + topic_summary_df = create_topic_summary_df_from_reference_table( + reference_df, sentiment_checkbox=sentiment_checkbox + ) + + # Normalize topic names in both dataframes to ensure consistent formatting + # Use the same normalization function as defined earlier in this function + # Apply normalization to topic_summary_df + if "General topic" in topic_summary_df.columns: + topic_summary_df["General topic"] = topic_summary_df["General topic"].apply( + normalize_topic_name_for_llm + ) + if "Subtopic" in topic_summary_df.columns: + topic_summary_df["Subtopic"] = topic_summary_df["Subtopic"].apply( + normalize_topic_name_for_llm + ) + + # Apply normalization to reference_df + if "General topic" in reference_df.columns: + reference_df["General topic"] = reference_df["General topic"].apply( + normalize_topic_name_for_llm + ) + if "Subtopic" in reference_df.columns: + reference_df["Subtopic"] = reference_df["Subtopic"].apply( + normalize_topic_name_for_llm + ) + + if "Topic number" not in reference_df.columns: + + # Merge the topic numbers back to the original dataframe + reference_df = reference_df.merge( + topic_summary_df[ + ["General topic", "Subtopic", "Sentiment", "Group", "Topic number"] + ], + on=["General topic", "Subtopic", "Sentiment", "Group"], + how="left", + ) + + # Create pivot table if file data is available + if not file_data.empty: + basic_response_data = get_basic_response_data(file_data, chosen_cols) + reference_df_pivot = convert_reference_table_to_pivot_table( + reference_df, basic_response_data + ) + + reference_pivot_file_path = ( + output_folder + + get_file_name_no_ext(reference_table_file_name) + + "_pivot_dedup.csv" + ) + if should_output_files == "True": + reference_df_pivot.to_csv( + reference_pivot_file_path, index=None, encoding="utf-8-sig" + ) + log_output_files.append(reference_pivot_file_path) + + # Save analysis results CSV if merge suggestions were found + if not merge_suggestions_df.empty: + analysis_results_file_path = ( + output_folder + + get_file_name_no_ext(reference_table_file_name) + + "_dedup_llm_analysis_results.csv" + ) + if should_output_files == "True": + merge_suggestions_df.to_csv( + analysis_results_file_path, index=None, encoding="utf-8-sig" + ) + log_output_files.append(analysis_results_file_path) + print(f"Analysis results saved to: {analysis_results_file_path}") + + # Save output files + reference_file_out_path = ( + output_folder + get_file_name_no_ext(reference_table_file_name) + "_dedup.csv" + ) + unique_topics_file_out_path = ( + output_folder + + get_file_name_no_ext(unique_topics_table_file_name) + + "_dedup.csv" + ) + if should_output_files == "True": + reference_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + reference_file_out_path, index=None, encoding="utf-8-sig" + ) + topic_summary_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + unique_topics_file_out_path, index=None, encoding="utf-8-sig" + ) + + output_files.append(reference_file_out_path) + output_files.append(unique_topics_file_out_path) + + # Outputs for markdown table output + topic_summary_df_revised_display = topic_summary_df.apply( + lambda col: col.map(lambda x: wrap_text(x, max_text_length=max_text_length)) + ) + deduplicated_unique_table_markdown = topic_summary_df_revised_display.to_markdown( + index=False + ) + + # Calculate token usage and timing information for logging + total_input_tokens = 0 + total_output_tokens = 0 + number_of_calls = 1 # Single LLM call for deduplication + + # Extract token usage from conversation metadata + if whole_conversation_metadata: + for metadata in whole_conversation_metadata: + if "input_tokens:" in metadata and "output_tokens:" in metadata: + try: + input_tokens = int( + metadata.split("input_tokens: ")[1].split(" ")[0] + ) + output_tokens = int( + metadata.split("output_tokens: ")[1].split(" ")[0] + ) + total_input_tokens += input_tokens + total_output_tokens += output_tokens + except (ValueError, IndexError): + pass + + # Compare input vs output unique topic combinations + if not topic_summary_df.empty: + output_unique_topics = topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + print( + f"Topic count comparison: Input had {input_unique_topics} unique 'General topic' | 'Subtopic' combinations, " + f"Output has {output_unique_topics} unique 'General topic' | 'Subtopic' combinations. " + f"Reduction: {input_unique_topics - output_unique_topics} topics merged." + ) + else: + print( + f"Topic count comparison: Input had {input_unique_topics} unique 'General topic' | 'Subtopic' combinations, " + f"Output has 0 unique 'General topic' | 'Subtopic' combinations." + ) + + # Calculate estimated time taken (rough estimate based on token usage) + toc = time.perf_counter() + estimated_time_taken = toc - tic + + print("LLM deduplication task successfully completed") + + return ( + reference_df, + topic_summary_df, + output_files, + log_output_files, + deduplicated_unique_table_markdown, + total_input_tokens, + total_output_tokens, + number_of_calls, + estimated_time_taken, + ) # , num_merges_applied + + +def join_unique_summaries(x): + unique_summaries = [] + seen = set() + + for s in x: + if pd.isna(s): + continue + + # 1. Normalize whitespace and split lines + s_str = str(s).strip() + # Split on both newlines and
separators to avoid duplicating segments + # that were previously concatenated in other stages. + lines = re.split(r"\n|
|
", s_str) + + for line in lines: + # 2. Aggressive Cleaning + # Remove "Rows X to Y:" prefix + line = re.sub( + r"^Rows\s+\d+\s+to\s+\d+:\s*", "", line, flags=re.IGNORECASE + ).strip() + + # Remove generic "Prefix:" if it exists (e.g., "Summary: ...") + if ": " in line: + parts = line.split(": ", 1) + if len(parts[0]) < 50 and " " not in parts[0]: + line = parts[1].strip() + + # 3. Handle Invisible Characters (Crucial) + # Replace non-breaking spaces (\xa0) and multiple spaces with a single standard space + normalized_line = re.sub(r"\s+", " ", line).strip() + + # 4. Check against Seen + if normalized_line and normalized_line not in seen: + unique_summaries.append(normalized_line) + seen.add(normalized_line) + + return "\n".join(unique_summaries) + + +def _rows_for_topic_summarisation(reference_df: pd.DataFrame) -> pd.DataFrame: + """Reference rows to feed into per-topic summarisation (excludes 'Not Mentioned').""" + if reference_df.empty: + return reference_df.copy() + filtered = reference_df.loc[reference_df["Sentiment"] != "Not Mentioned"].copy() + return filtered if not filtered.empty else reference_df.copy() + + +def sample_reference_table_summaries( + reference_df: pd.DataFrame, + random_seed: int, + no_of_sampled_summaries: int = default_number_of_sampled_summaries, + sample_reference_table_checkbox: bool = False, + available_tokens_for_summaries: int = None, + tokenizer=None, + model_source: str = "Local", +): + """ + Sample x number of summaries from which to produce summaries, so that the input token length is not too long. + """ + + if sample_reference_table_checkbox: + # If caller didn't provide a budget, use a conservative default. + # (We avoid truncating mid-passage; the budget controls how many whole samples we keep.) + if available_tokens_for_summaries is None: + available_tokens_for_summaries = int(_effective_context_limit() * 0.6) + + all_summaries = pd.DataFrame( + columns=[ + "General topic", + "Subtopic", + "Sentiment", + "Group", + "Response ID", + "Summary", + ] + ) + + if "Group" not in reference_df.columns: + reference_df["Group"] = "All" + + reference_df_grouped = reference_df.groupby( + ["General topic", "Subtopic", "Sentiment", "Group"] + ) + + # Sampling diagnostics + sampling_stats = { + "topic_combos_seen": 0, + "topics_budget_reduced": 0, + "total_passages_dropped": 0, + "topics_dropped_entirely": 0, + } + + if "Revised summary" in reference_df.columns: + out_message = "Summary has already been created for this file" + print(out_message) + raise Exception(out_message) + + for group_keys, reference_df_group in reference_df_grouped: + if len(reference_df_group["General topic"]) > 1: + sampling_stats["topic_combos_seen"] += 1 + + filtered_reference_df = reference_df_group.reset_index() + + filtered_reference_df_unique = filtered_reference_df.drop_duplicates( + [ + "General topic", + "Subtopic", + "Sentiment", + "Group", + "Start row of group", + ] + ) + + # Sample n of the unique topic summaries PER GROUP. To limit the length of the text going into the summarisation tool + # This ensures each group gets up to no_of_sampled_summaries summaries, not the total across all groups + number_of_summaries_to_sample = min( + no_of_sampled_summaries, len(filtered_reference_df_unique) + ) + + # Sample once at the maximum size, then shrink by taking a prefix if needed. + sampled_max_df = filtered_reference_df_unique.sample( + number_of_summaries_to_sample, random_state=random_seed + ).reset_index(drop=True) + + # Never allow the same underlying response to be sampled twice for a topic combo + if "Response ID" in sampled_max_df.columns: + sampled_max_df = sampled_max_df.drop_duplicates( + subset=["Response ID"] + ).reset_index(drop=True) + + def _split_passages_for_budget(text: str) -> List[str]: + if not isinstance(text, str) or not text.strip(): + return [] + # Prefer newline (join_unique_summaries uses "\n"), then
+ parts = [p.strip() for p in str(text).split("\n") if p.strip()] + if len(parts) <= 1 and "
" in str(text): + parts = [ + p.strip() for p in str(text).split("
") if p.strip() + ] + return parts if parts else [str(text).strip()] + + def _trim_joined_text_to_budget(text: str) -> str: + """Trim joined summaries by dropping whole passages to fit token budget.""" + passages = _split_passages_for_budget(text) + if not passages: + return "" + + # Fast path: already fits + joined = "\n".join(passages) + if ( + count_tokens_in_text(joined, tokenizer, model_source) + <= available_tokens_for_summaries + ): + return joined + + # Binary search max number of passages that fit + lo, hi = 1, len(passages) + best = None + while lo <= hi: + mid = (lo + hi) // 2 + candidate = "\n".join(passages[:mid]) + cand_tokens = count_tokens_in_text( + candidate, tokenizer, model_source + ) + if cand_tokens <= available_tokens_for_summaries: + best = candidate + lo = mid + 1 + else: + hi = mid - 1 + + if best is None: + # Even one passage doesn't fit: return empty rather than crash the pipeline. + one_tokens = count_tokens_in_text( + passages[0], tokenizer, model_source + ) + print( + "Warning: A single passage is too large to fit token budget during sampling. " + f"topic-sentiment-group={group_keys}, passage_tokens={one_tokens}, " + f"budget={available_tokens_for_summaries}. Dropping this summary from sampling." + ) + sampling_stats["topics_dropped_entirely"] += 1 + return "" + + # Track how much we dropped + try: + kept = len([p for p in best.splitlines() if p.strip()]) + except Exception: + kept = 0 + dropped = max(0, len(passages) - kept) + if dropped > 0: + sampling_stats["topics_budget_reduced"] += 1 + sampling_stats["total_passages_dropped"] += dropped + + return best + + def _joined_summary_token_count(df: pd.DataFrame) -> int: + joined = join_unique_summaries(df["Summary"]) + joined = _trim_joined_text_to_budget(joined) + return count_tokens_in_text(joined, tokenizer, model_source) + + # Ensure the joined summary fits within the token budget by dropping whole samples + # (never truncating a passage mid-way). + if len(sampled_max_df) > 0: + current_tokens = _joined_summary_token_count(sampled_max_df) + if current_tokens > available_tokens_for_summaries: + lo, hi = 1, len(sampled_max_df) + best_k = 0 + while lo <= hi: + mid = (lo + hi) // 2 + test_df = sampled_max_df.iloc[:mid] + test_tokens = _joined_summary_token_count(test_df) + if test_tokens <= available_tokens_for_summaries: + best_k = mid + lo = mid + 1 + else: + hi = mid - 1 + + if best_k == 0: + # With passage-level trimming, we should always be able to fit *something* unless + # even one passage is too large. In that case we just keep 1 row and allow the + # downstream summariser to skip/handle empty summaries. + best_k = 1 + + print( + f"Warning: Joined summaries for topic-sentiment-group combination {group_keys} exceed token budget. " + f"Reducing samples from {len(sampled_max_df)} to {best_k} to fit." + ) + sampled_max_df = sampled_max_df.iloc[:best_k] + + all_summaries = pd.concat([all_summaries, sampled_max_df]) + + # If no responses/topics qualify, just go ahead with the original reference dataframe + if all_summaries.empty: + sampled_reference_table_df = _rows_for_topic_summarisation(reference_df) + else: + # Deduplicate summaries within each group before joining to prevent repeated summaries + + def join_unique_summaries_with_budget(x): + joined = join_unique_summaries(x) + return _trim_joined_text_to_budget(joined) + + sampled_reference_table_df = ( + all_summaries.groupby( + ["General topic", "Subtopic", "Sentiment", "Group"] + ) + .agg( + { + "Response ID": "size", # Count the number of references + "Summary": join_unique_summaries_with_budget, # Join unique summaries only (budgeted) + } + ) + .reset_index() + ) + sampled_reference_table_df = sampled_reference_table_df.loc[ + sampled_reference_table_df["Sentiment"] != "Not Mentioned" + ] + + # Groups whose topic combos only had one response each never enter all_summaries. + # Keep them by appending their full reference rows so they are not dropped. + if not sampled_reference_table_df.empty: + groups_in_sampled = set(sampled_reference_table_df["Group"].unique()) + else: + groups_in_sampled = set() + missing_group_rows = _rows_for_topic_summarisation( + reference_df.loc[~reference_df["Group"].isin(groups_in_sampled)] + ) + if not missing_group_rows.empty: + sampled_reference_table_df = pd.concat( + [sampled_reference_table_df, missing_group_rows], + ignore_index=True, + ) + + print( + "Sampling stats: " + f"topic_combos_seen={sampling_stats['topic_combos_seen']}, " + f"topics_budget_reduced={sampling_stats['topics_budget_reduced']}, " + f"total_passages_dropped={sampling_stats['total_passages_dropped']}, " + f"topics_dropped_entirely={sampling_stats['topics_dropped_entirely']}" + ) + else: + sampled_reference_table_df = reference_df + + summarised_references_markdown = sampled_reference_table_df.to_markdown(index=False) + + return sampled_reference_table_df, summarised_references_markdown + + +def count_tokens_in_text(text: str, tokenizer=None, model_source: str = "Local") -> int: + """ + Count the number of tokens in the given text. + + Args: + text (str): The text to count tokens for + tokenizer (object, optional): Tokenizer object for local models. Defaults to None. + model_source (str): Source of the model to determine tokenization method. Defaults to "Local". + + Returns: + int: Number of tokens in the text + """ + if not text: + return 0 + + try: + if model_source == "Local" and tokenizer and len(tokenizer) > 0: + # Use local tokenizer if available + tokens = tokenizer[0].encode(text, add_special_tokens=False) + return len(tokens) + else: + # Fallback: rough estimation using word count (approximately 1.3 tokens per word) + word_count = len(text.split()) + return int(word_count * 1.3) + except Exception as e: + print(f"Error counting tokens: {e}. Using word count estimation.") + # Fallback: rough estimation using word count + word_count = len(text.split()) + return int(word_count * 1.3) + + +def _inference_server_prompt_token_count( + api_url: str, messages: List[dict], chat_template_kwargs: dict | None = None +) -> int | None: + """ + Best-effort server-side token counting for llama.cpp / OpenAI-compatible inference servers. + + We try to: + 1) Apply the server's chat template (so counts match what the server will tokenize) + 2) Tokenize that rendered prompt + + If the server doesn't expose these endpoints, returns None. + """ + if not api_url: + return None + + try: + import requests + + # 1) Try apply-template (llama.cpp exposes this on some builds) + prompt_text = None + for path in ("/apply-template", "/v1/apply-template"): + try: + r = requests.post( + f"{api_url}{path}", + json={ + "messages": messages, + "chat_template_kwargs": chat_template_kwargs or {}, + }, + timeout=10, + ) + if r.status_code == 404: + continue + r.raise_for_status() + data = ( + r.json() + if "application/json" in r.headers.get("content-type", "") + else r.text + ) + if isinstance(data, dict): + prompt_text = ( + data.get("prompt") or data.get("content") or data.get("text") + ) + elif isinstance(data, str): + prompt_text = data + if prompt_text: + break + except Exception: + continue + + # If no apply-template endpoint, we can't reliably match server tokenization. + if not prompt_text: + return None + + # 2) Tokenize the rendered prompt + for path in ("/tokenize", "/v1/tokenize"): + try: + r = requests.post( + f"{api_url}{path}", + json={"content": prompt_text}, + timeout=10, + ) + if r.status_code == 404: + continue + r.raise_for_status() + data = r.json() + if isinstance(data, dict): + if "tokens" in data and isinstance(data["tokens"], list): + return len(data["tokens"]) + if "n_tokens" in data: + return int(data["n_tokens"]) + break + except Exception: + continue + except Exception: + return None + + return None + + +def clean_markdown_table_whitespace(markdown_text: str) -> str: + if not markdown_text: + return markdown_text + + lines = markdown_text.splitlines() + cleaned_lines = [] + + for line in lines: + # 1. Clean all types of whitespace (including non-breaking spaces \u00A0) + # This turns every cell into a single-spaced string + cells = [re.sub(r"[\s\u00A0]+", " ", cell.strip()) for cell in line.split("|")] + + # 2. Check if the row is effectively empty (only pipes or whitespace) + # We join the content; if nothing is left, it's a "ghost" row. + if not "".join(cells).strip(): + continue + + # 3. Handle the separator row specifically (e.g., |:---|---:|) + # We reset these to a small fixed width so they don't stretch the table. + if re.match(r"^[|\s\-:]+$", line): + new_separator = [] + for cell in cells: + if not cell: # Outer pipes + new_separator.append("") + elif ":" in cell: # Alignment markers + left = ":" if cell.startswith(":") else "-" + right = ":" if cell.endswith(":") else "-" + new_separator.append(f"{left}---{right}") + else: + new_separator.append("---") + cleaned_lines.append("|".join(new_separator)) + continue + + # 4. Standard data row: Rejoin with single padding + # We filter out empty outer parts caused by leading/trailing pipes + formatted_row = ( + "| " + + " | ".join( + c for c in cells if c or cells.index(c) not in [0, len(cells) - 1] + ) + + " |" + ) + + # Simple fallback if the logic above is too aggressive for your specific table style: + # formatted_row = "|".join(f" {c} " if c else "" for c in cells) + + cleaned_lines.append(formatted_row) + + return "\n".join(cleaned_lines) + + +def summarise_output_topics_query( + model_choice: str, + in_api_key: str, + temperature: float, + formatted_summary_prompt: str, + summarise_topic_descriptions_system_prompt: str, + model_source: str, + bedrock_runtime: boto3.Session.client, + local_model=list(), + tokenizer=list(), + assistant_model=list(), + azure_endpoint: str = "", + api_url: str = None, +): + """ + Query an LLM to generate a summary of topics based on the provided prompts. + + Args: + model_choice (str): The name/type of model to use for generation + in_api_key (str): API key for accessing the model service + temperature (float): Temperature parameter for controlling randomness in generation + formatted_summary_prompt (str): The formatted prompt containing topics to summarize + summarise_topic_descriptions_system_prompt (str): System prompt providing context and instructions + model_source (str): Source of the model (e.g. "AWS", "Gemini", "Local") + bedrock_runtime (boto3.Session.client): AWS Bedrock runtime client for AWS models + local_model (object, optional): Local model object if using local inference. Defaults to empty list. + tokenizer (object, optional): Tokenizer object if using local inference. Defaults to empty list. + Returns: + tuple: Contains: + - response_text (str): The generated summary text + - conversation_history (list): History of the conversation with the model + - whole_conversation_metadata (list): Metadata about the conversation + """ + conversation_history = list() + whole_conversation_metadata = list() + client = list() + client_config = {} + + # Combine system prompt and user prompt for token counting + full_input_text = ( + summarise_topic_descriptions_system_prompt + "\n" + formatted_summary_prompt[0] + if isinstance(formatted_summary_prompt, list) + else summarise_topic_descriptions_system_prompt + + "\n" + + formatted_summary_prompt + ) + + # Count tokens in the input text + if "inference-server" in str(model_source).lower() and api_url: + messages = [ + {"role": "system", "content": summarise_topic_descriptions_system_prompt}, + { + "role": "user", + "content": ( + formatted_summary_prompt[0] + if isinstance(formatted_summary_prompt, list) + else formatted_summary_prompt + ), + }, + ] + input_token_count = _inference_server_prompt_token_count(api_url, messages) + if input_token_count is None: + input_token_count = count_tokens_in_text( + full_input_text, tokenizer, model_source + ) + print( + "Warning: Using approximate token count for inference-server; " + "server-side tokenization endpoint not available." + ) + else: + input_token_count = count_tokens_in_text( + full_input_text, tokenizer, model_source + ) + + effective_limit = _effective_context_limit() + + # Check if input exceeds context length (with headroom) + if input_token_count > effective_limit: + error_message = ( + f"Input text exceeds LLM context length (with headroom). " + f"Input tokens: {input_token_count}, Max (effective): {effective_limit} " + f"(raw max: {LLM_CONTEXT_LENGTH}). Please reduce the input text size." + ) + print(error_message) + raise ValueError(error_message) + + print(f"Input token count: {input_token_count} (Max effective: {effective_limit})") + + # Prepare Gemini models before query + if "Gemini" in model_source: + # print("Using Gemini model:", model_choice) + client, config = construct_gemini_generative_model( + in_api_key=in_api_key, + temperature=temperature, + model_choice=model_choice, + system_prompt=system_prompt, + max_tokens=max_tokens, + ) + elif "Azure/OpenAI" in model_source: + client, config = construct_azure_client( + in_api_key=os.environ.get("AZURE_INFERENCE_CREDENTIAL", ""), + endpoint=azure_endpoint, + ) + elif "Local" in model_source: + pass + # print("Using local model: ", model_choice) + elif "AWS" in model_source: + pass + # print("Using AWS Bedrock model:", model_choice) + + whole_conversation = [summarise_topic_descriptions_system_prompt] + + # Process requests to large language model + ( + responses, + conversation_history, + whole_conversation, + whole_conversation_metadata, + response_text, + ) = process_requests( + formatted_summary_prompt, + summarise_topic_descriptions_system_prompt, + conversation_history, + whole_conversation, + whole_conversation_metadata, + client, + client_config, + model_choice, + temperature, + bedrock_runtime=bedrock_runtime, + model_source=model_source, + local_model=local_model, + tokenizer=tokenizer, + assistant_model=assistant_model, + assistant_prefill=summary_assistant_prefill, + api_url=api_url, + ) + + summarised_output = re.sub( + r"\n{2,}", "\n", response_text + ) # Replace multiple line breaks with a single line break + summarised_output = re.sub( + r"^\n{1,}", "", summarised_output + ) # Remove one or more line breaks at the start + summarised_output = re.sub( + r"\n", "
", summarised_output + ) # Replace \n with more html friendly
tags + summarised_output = summarised_output.strip() + + print("Finished summary query") + + # Ensure the system prompt is included in the conversation history + try: + if isinstance(conversation_history, list): + has_system_prompt = False + + if conversation_history: + first_entry = conversation_history[0] + if isinstance(first_entry, dict): + role_is_system = first_entry.get("role") == "system" + parts = first_entry.get("parts") + content_matches = ( + parts == summarise_topic_descriptions_system_prompt + or ( + isinstance(parts, list) + and summarise_topic_descriptions_system_prompt in parts + ) + ) + has_system_prompt = role_is_system and content_matches + elif isinstance(first_entry, str): + has_system_prompt = ( + first_entry.strip().lower().startswith("system:") + ) + + if not has_system_prompt: + conversation_history.insert( + 0, + { + "role": "system", + "parts": [summarise_topic_descriptions_system_prompt], + }, + ) + except Exception as _e: + # Non-fatal: if anything goes wrong, return the original conversation history + pass + + return ( + summarised_output, + conversation_history, + whole_conversation_metadata, + response_text, + ) + + +def process_debug_output_iteration( + output_debug_files: str, + output_folder: str, + batch_file_path_details: str, + model_choice_clean_short: str, + final_system_prompt: str, + summarised_output: str, + conversation_history: list, + metadata: list, + log_output_files: list, + task_type: str, +) -> tuple[str, str, str, str]: + """ + Writes debug files for summary generation if output_debug_files is "True", + and returns the content of the prompt, summary, conversation, and metadata for the current iteration. + + Args: + output_debug_files (str): Flag to indicate if debug files should be written. + output_folder (str): The folder where output files are saved. + batch_file_path_details (str): Details for the batch file path. + model_choice_clean_short (str): Shortened cleaned model choice. + final_system_prompt (str): The system prompt content. + summarised_output (str): The summarised output content. + conversation_history (list): The full conversation history. + metadata (list): The metadata for the conversation. + log_output_files (list): A list to append paths of written log files. This list is modified in-place. + task_type (str): The type of task being performed. + Returns: + tuple[str, str, str, str]: A tuple containing the content of the prompt, + summarised output, conversation history (as string), + and metadata (as string) for the current iteration. + """ + current_prompt_content = final_system_prompt + current_summary_content = summarised_output + + if isinstance(conversation_history, list): + + # Handle both list of strings and list of dicts + if conversation_history and isinstance(conversation_history[0], dict): + # Convert list of dicts to list of strings + conversation_strings = list() + for entry in conversation_history: + if "role" in entry and "parts" in entry: + role = entry["role"].capitalize() + message = ( + " ".join(entry["parts"]) + if isinstance(entry["parts"], list) + else str(entry["parts"]) + ) + conversation_strings.append(f"{role}: {message}") + else: + # Fallback for unexpected dict format + conversation_strings.append(str(entry)) + current_conversation_content = "\n".join(conversation_strings) + else: + # Handle list of strings + current_conversation_content = "\n".join(conversation_history) + else: + current_conversation_content = str(conversation_history) + current_metadata_content = str(metadata) + current_task_type = task_type + + if output_debug_files == "True": + try: + formatted_prompt_output_path = ( + output_folder + + batch_file_path_details + + "_full_prompt_" + + model_choice_clean_short + + "_" + + current_task_type + + ".txt" + ) + final_table_output_path = ( + output_folder + + batch_file_path_details + + "_full_response_" + + model_choice_clean_short + + "_" + + current_task_type + + ".txt" + ) + whole_conversation_path = ( + output_folder + + batch_file_path_details + + "_full_conversation_" + + model_choice_clean_short + + "_" + + current_task_type + + ".txt" + ) + whole_conversation_path_meta = ( + output_folder + + batch_file_path_details + + "_metadata_" + + model_choice_clean_short + + "_" + + current_task_type + + ".txt" + ) + + with open( + formatted_prompt_output_path, + "w", + encoding="utf-8-sig", + errors="replace", + ) as f: + f.write(current_prompt_content) + with open( + final_table_output_path, "w", encoding="utf-8-sig", errors="replace" + ) as f: + f.write(current_summary_content) + with open( + whole_conversation_path, "w", encoding="utf-8-sig", errors="replace" + ) as f: + f.write(current_conversation_content) + with open( + whole_conversation_path_meta, + "w", + encoding="utf-8-sig", + errors="replace", + ) as f: + f.write(current_metadata_content) + + log_output_files.append(formatted_prompt_output_path) + log_output_files.append(final_table_output_path) + log_output_files.append(whole_conversation_path) + log_output_files.append(whole_conversation_path_meta) + except Exception as e: + print(f"Error in writing debug files for summary: {e}") + + # Return the content of the objects for the current iteration. + # The caller can then append these to separate lists if accumulation is desired. + return ( + current_prompt_content, + current_summary_content, + current_conversation_content, + current_metadata_content, + ) + + +@spaces.GPU(duration=MAX_SPACES_GPU_RUN_TIME) +def summarise_output_topics( + sampled_reference_table_df: pd.DataFrame, + topic_summary_df: pd.DataFrame, + reference_table_df: pd.DataFrame, + model_choice: str, + in_api_key: str, + temperature: float, + reference_data_file_name: str, + summarised_outputs: list = list(), + latest_summary_completed: int = 0, + out_metadata_str: str = "", + in_data_files: List[str] = list(), + in_excel_sheets: str = "", + chosen_cols: List[str] = list(), + log_output_files: list[str] = list(), + summarise_format_radio: str = "Return a summary up to two paragraphs long that includes as much detail as possible from the original text", + output_folder: str = OUTPUT_FOLDER, + context_textbox: str = "", + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + model_name_map: dict = model_name_map, + hf_api_key_textbox: str = "", + azure_endpoint_textbox: str = "", + existing_logged_content: list = list(), + additional_summary_instructions_provided: str = "", + output_debug_files: str = "False", + group_value: str = "All", + reasoning_suffix: str = reasoning_suffix, + local_model: object = None, + tokenizer: object = None, + assistant_model: object = None, + summarise_topic_descriptions_prompt: str = summarise_topic_descriptions_prompt, + summarise_topic_descriptions_system_prompt: str = summarise_topic_descriptions_system_prompt, + do_summaries: str = "Yes", + api_url: str = None, + progress=gr.Progress(track_tqdm=True), +): + """ + Create improved summaries of topics by consolidating raw batch-level summaries from the initial model run. Works on a single group of summaries at a time (called from wrapper function summarise_output_topics_by_group). + + Args: + sampled_reference_table_df (pd.DataFrame): DataFrame containing sampled reference data with summaries + topic_summary_df (pd.DataFrame): DataFrame containing topic summary information + reference_table_df (pd.DataFrame): DataFrame mapping response references to topics + model_choice (str): Name of the LLM model to use + in_api_key (str): API key for model access + temperature (float): Temperature parameter for model generation + reference_data_file_name (str): Name of the reference data file + summarised_outputs (list, optional): List to store generated summaries. Defaults to empty list. + latest_summary_completed (int, optional): Index of last completed summary. Defaults to 0. + out_metadata_str (str, optional): String for metadata output. Defaults to empty string. + in_data_files (List[str], optional): List of input data file paths. Defaults to empty list. + in_excel_sheets (str, optional): Excel sheet names if using Excel files. Defaults to empty string. + chosen_cols (List[str], optional): List of columns selected for analysis. Defaults to empty list. + log_output_files (list[str], optional): List of log file paths. Defaults to empty list. + summarise_format_radio (str, optional): Format instructions for summary generation. Defaults to two paragraph format. + output_folder (str, optional): Folder path for outputs. Defaults to OUTPUT_FOLDER. + context_textbox (str, optional): Additional context for summarization. Defaults to empty string. + aws_access_key_textbox (str, optional): AWS access key. Defaults to empty string. + aws_secret_key_textbox (str, optional): AWS secret key. Defaults to empty string. + model_name_map (dict, optional): Dictionary mapping model choices to their properties. Defaults to model_name_map. + hf_api_key_textbox (str, optional): Hugging Face API key. Defaults to empty string. + azure_endpoint_textbox (str, optional): Azure endpoint. Defaults to empty string. + additional_summary_instructions_provided (str, optional): Additional summary instructions provided by the user. Defaults to empty string. + existing_logged_content (list, optional): List of existing logged content. Defaults to empty list. + output_debug_files (str, optional): Flag to indicate if debug files should be written. Defaults to "False". + group_value (str, optional): Value of the group to summarise. Defaults to "All". + reasoning_suffix (str, optional): Suffix for reasoning. Defaults to reasoning_suffix. + local_model (object, optional): Local model object if using local inference. Defaults to None. + tokenizer (object, optional): Tokenizer object if using local inference. Defaults to None. + assistant_model (object, optional): Assistant model object if using local inference. Defaults to None. + summarise_topic_descriptions_prompt (str, optional): Prompt template for topic summarization. + summarise_topic_descriptions_system_prompt (str, optional): System prompt for topic summarization. + do_summaries (str, optional): Flag to control summary generation. Defaults to "Yes". + progress (gr.Progress, optional): Gradio progress tracker. Defaults to track_tqdm=True. + + Returns: + Multiple outputs including summarized content, metadata, and file paths + """ + out_metadata = list() + summarised_output_markdown = "" + output_files = list() + acc_input_tokens = 0 + acc_output_tokens = 0 + acc_number_of_calls = 0 + time_taken = 0 + out_metadata_str = ( + "" # Output metadata is currently replaced on starting a summarisation task + ) + out_message = list() + task_type = "Topic summarisation" + topic_summary_df_revised = pd.DataFrame() + + all_prompts_content = list() + all_summaries_content = list() + all_metadata_content = list() + all_groups_content = list() + all_batches_content = list() + all_model_choice_content = list() + all_validated_content = list() + all_task_type_content = list() + all_logged_content = list() + all_file_names_content = list() + + tic = time.perf_counter() + + # Ensure custom model_choice is registered in model_name_map + ensure_model_in_map(model_choice, model_name_map) + + model_choice_clean = clean_column_name( + model_name_map[model_choice]["short_name"], + max_length=20, + front_characters=False, + ) + + if context_textbox and "The context of this analysis is" not in context_textbox: + context_textbox = "The context of this analysis is '" + context_textbox + "'." + + if log_output_files is None: + log_output_files = list() + + # Check for data for summarisations + if not topic_summary_df.empty and not reference_table_df.empty: + print("Unique table and reference table data found.") + else: + out_message = "Please upload a unique topic table and reference table file to continue with summarisation." + print(out_message) + raise Exception(out_message) + + if "Revised summary" in reference_table_df.columns: + out_message = "Summary has already been created for this file" + print(out_message) + raise Exception(out_message) + + # Load in data file and chosen columns if exists to create pivot table later + file_data = pd.DataFrame() + if in_data_files and chosen_cols: + file_data, data_file_names_textbox, total_number_of_batches = load_in_data_file( + in_data_files, chosen_cols, 1, in_excel_sheets=in_excel_sheets + ) + else: + out_message = "No file data found, pivot table output will not be created." + print(out_message) + # Use sys.stdout.write to avoid issues with progress bars + # sys.stdout.write(out_message + "\n") + # sys.stdout.flush() + # Note: file_data will remain empty, pivot tables will not be created + + reference_table_df = reference_table_df.rename( + columns={"General Topic": "General topic"}, errors="ignore" + ) + topic_summary_df = topic_summary_df.rename( + columns={"General Topic": "General topic"}, errors="ignore" + ) + + # Sentiment is optional upstream. Normalise to a placeholder value so the + # rest of this function (joins/grouping/output schemas) stays consistent. + for _df in (reference_table_df, topic_summary_df, sampled_reference_table_df): + if "Sentiment" not in _df.columns: + _df["Sentiment"] = "Not assessed" + + if "Number of responses" not in topic_summary_df.columns: + topic_summary_df["Number of responses"] = 0 + + if "Group" not in reference_table_df.columns: + reference_table_df["Group"] = "All" + if "Group" not in topic_summary_df.columns: + topic_summary_df["Group"] = "All" + if "Group" not in sampled_reference_table_df.columns: + sampled_reference_table_df["Group"] = "All" + + # Use the Summary column if it exists, otherwise use the Revised summary column + if "Summary" in sampled_reference_table_df.columns: + all_summaries = sampled_reference_table_df["Summary"].tolist() + else: + all_summaries = sampled_reference_table_df["Revised summary"].tolist() + + # Normalise separators in sampled summaries to avoid wasting context window + # with repeated blank lines /
tags. + def _normalise_summary_separators(text: str) -> str: + if not isinstance(text, str): + return "" + t = text + # Collapse repeated spaces/tabs to a single space (keep newlines intact) + t = re.sub(r"[ \t]{2,}", " ", t) + t = re.sub(r"(
\s*){2,}", "
", t, flags=re.IGNORECASE) + t = re.sub(r"\n{2,}", "\n", t) + t = re.sub(r"(\s*
\s*){2,}", "
", t, flags=re.IGNORECASE) + # Trim stray separators at ends + t = t.strip().strip("
").strip() + return t + + all_summaries = [_normalise_summary_separators(s) for s in all_summaries] + + all_groups = sampled_reference_table_df["Group"].tolist() + + if not group_value: + group_value = str(all_groups[0]) + else: + group_value = str(group_value) + + length_all_summaries = len(all_summaries) + + model_source = model_name_map[model_choice]["source"] + + if (model_source == "Local") & (RUN_LOCAL_MODEL == "1") & (not local_model): + progress(0.1, f"Using global model: {CHOSEN_LOCAL_MODEL_TYPE}") + local_model = get_model() + tokenizer = get_tokenizer() + assistant_model = get_assistant_model() + + ( + "Revising topic-level summaries. " + + str(latest_summary_completed) + + " summaries completed so far." + ) + summary_loop = progress.tqdm( + range(latest_summary_completed, length_all_summaries), + desc="Revising topic-level summaries", + unit="summaries", + ) + + if do_summaries == "Yes": + + bedrock_runtime = connect_to_bedrock_runtime( + model_name_map, + model_choice, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + ) + + create_batch_file_path_details(reference_data_file_name) + model_choice_clean_short = clean_column_name( + model_choice_clean, max_length=20, front_characters=False + ) + file_name_clean = f"{clean_column_name(reference_data_file_name, max_length=15)}_{clean_column_name(str(group_value), max_length=15).replace(' ','_')}" + # file_name_clean = clean_column_name(reference_data_file_name, max_length=20, front_characters=True) + in_column_cleaned = clean_column_name(chosen_cols, max_length=20) + + combined_summary_instructions = ( + summarise_format_radio + ". " + additional_summary_instructions_provided + ) + + for summary_no in summary_loop: + print("Current summary number is:", summary_no) + + batch_file_path_details = f"{file_name_clean}_batch_{latest_summary_completed + 1}_size_1_col_{in_column_cleaned}" + + formatted_summarise_topic_descriptions_system_prompt = ( + summarise_topic_descriptions_system_prompt.format( + column_name=chosen_cols, consultation_context=context_textbox + ) + ) + + # Apply reasoning suffix for GPT-OSS models (Local, inference-server, or AWS) + is_gpt_oss_model = ( + "gpt-oss" in model_choice.lower() or "gpt_oss" in model_choice.lower() + ) + + if is_gpt_oss_model: + # Use default reasoning suffix if not set + effective_reasoning_suffix = ( + reasoning_suffix if reasoning_suffix else "Reasoning: low" + ) + if effective_reasoning_suffix: + formatted_summarise_topic_descriptions_system_prompt = ( + formatted_summarise_topic_descriptions_system_prompt + + "\n" + + effective_reasoning_suffix + ) + elif "Local" in model_source and reasoning_suffix: + # For other local models, use reasoning_suffix if provided + formatted_summarise_topic_descriptions_system_prompt = ( + formatted_summarise_topic_descriptions_system_prompt + + "\n" + + reasoning_suffix + ) + + # ---- Context-length safety (never cut passages mid-way) ---- + # We prefer dropping whole sampled passages over truncating. + # Passages are expected to be joined with newlines by join_unique_summaries(). + summary_text = all_summaries[summary_no] + + # Compute base token overhead (system prompt + prompt template without summaries) + base_prompt = summarise_topic_descriptions_prompt.format( + summaries="", + summary_format=combined_summary_instructions, + ) + base_token_count = count_tokens_in_text( + formatted_summarise_topic_descriptions_system_prompt + + "\n" + + base_prompt, + tokenizer, + model_source, + ) + # Use effective context limit (with headroom) to reduce risk of server-side mismatch + effective_limit = _effective_context_limit() + available_tokens_for_summaries = max(0, effective_limit - base_token_count) + + def _full_prompt_token_count_for_summary_text(text: str) -> int: + """Token count for the full system+user prompt for this summary.""" + user_prompt = summarise_topic_descriptions_prompt.format( + summaries=text, summary_format=combined_summary_instructions + ) + full_text = ( + formatted_summarise_topic_descriptions_system_prompt + + "\n" + + user_prompt + ) + + if "inference-server" in str(model_source).lower() and api_url: + messages = [ + { + "role": "system", + "content": formatted_summarise_topic_descriptions_system_prompt, + }, + {"role": "user", "content": user_prompt}, + ] + server_tokens = _inference_server_prompt_token_count( + api_url, messages + ) + if server_tokens is not None: + return int(server_tokens) + + return int(count_tokens_in_text(full_text, tokenizer, model_source)) + + def _split_passages(text: str) -> List[str]: + if not isinstance(text, str) or not text.strip(): + return [] + # Prefer newline-split first (join_unique_summaries uses "\n") + parts = [p.strip() for p in str(text).split("\n") if p.strip()] + if len(parts) <= 1 and "
" in str(text): + parts = [p.strip() for p in str(text).split("
") if p.strip()] + # If we still can't split, treat the entire text as one passage + return parts if parts else [str(text).strip()] + + passages = _split_passages(summary_text) + if passages: + joined_text = "\n".join(passages) + passage_tokens = count_tokens_in_text( + joined_text, tokenizer, model_source + ) + if passage_tokens > available_tokens_for_summaries: + # Binary-search the maximum number of whole passages that fit. + lo, hi = 1, len(passages) + best = None + while lo <= hi: + mid = (lo + hi) // 2 + candidate = "\n".join(passages[:mid]) + cand_tokens = count_tokens_in_text( + candidate, tokenizer, model_source + ) + if cand_tokens <= available_tokens_for_summaries: + best = candidate + lo = mid + 1 + else: + hi = mid - 1 + + if best is None: + # Even a single passage is too large. Don't truncate mid-passage: skip this topic. + single_tokens = count_tokens_in_text( + passages[0], tokenizer, model_source + ) + print( + "Warning: A single sampled passage is too large to summarise within context. " + f"Skipping this topic summary. Passage tokens: {single_tokens}, " + f"Available for passages: {available_tokens_for_summaries}, " + f"Max effective context length: {effective_limit}." + ) + summarised_outputs.append(None) + # Keep metadata aligned + out_metadata_str = ". ".join(out_metadata) + latest_summary_completed += 1 + continue + + print( + "Warning: Sampled summaries exceed context limit for this topic. " + f"Reducing from {len(passages)} passages to {len(best.splitlines())} passages to fit." + ) + summary_text = best + + # Final guard: ensure the FULL prompt is within context (with headroom). + # If not, drop whole passages until it fits. This prevents constructing a request + # that the server will reject. + if passages: + prompt_tokens = _full_prompt_token_count_for_summary_text(summary_text) + if prompt_tokens > effective_limit: + print( + "Warning: Full prompt still exceeds context after passage trimming. " + f"Prompt tokens: {prompt_tokens}, Max effective: {effective_limit}. " + "Dropping more passages." + ) + lo, hi = 1, len(passages) + best_text = None + while lo <= hi: + mid = (lo + hi) // 2 + candidate_text = "\n".join(passages[:mid]) + cand_tokens = _full_prompt_token_count_for_summary_text( + candidate_text + ) + if cand_tokens <= effective_limit: + best_text = candidate_text + lo = mid + 1 + else: + hi = mid - 1 + + if best_text is None: + print( + "Warning: Even one passage cannot fit within context for this summary prompt. " + "Skipping this topic summary." + ) + summarised_outputs.append(None) + out_metadata_str = ( + ". ".join(out_metadata) if out_metadata else "" + ) + latest_summary_completed += 1 + continue + + summary_text = best_text + + formatted_summary_prompt = [ + summarise_topic_descriptions_prompt.format( + summaries=summary_text, summary_format=combined_summary_instructions + ) + ] + + # Ensure locals exist even if the call fails (e.g. context-length error) + conversation_history = list() + metadata = [] + try: + response, conversation_history, metadata, response_text = ( + summarise_output_topics_query( + model_choice, + in_api_key, + temperature, + formatted_summary_prompt, + formatted_summarise_topic_descriptions_system_prompt, + model_source, + bedrock_runtime, + local_model, + tokenizer=tokenizer, + assistant_model=assistant_model, + azure_endpoint=azure_endpoint_textbox, + api_url=api_url, + ) + ) + summarised_output = response_text + except Exception as e: + print("Creating summary failed:", e) + summarised_output = "" + conversation_history = list() + metadata = [] + + summarised_outputs.append(summarised_output) + if metadata: + out_metadata.extend(metadata) + out_metadata_str = ". ".join(out_metadata) if out_metadata else "" + + # Call the new function to process and log debug outputs for the current iteration. + # The returned values are the contents of the prompt, summary, conversation, and metadata + + full_prompt = ( + formatted_summarise_topic_descriptions_system_prompt + + "\n" + + formatted_summary_prompt[0] + ) + + # Coerce toggle to string expected by debug writer (accepts True/False or "True"/"False") + output_debug_files_str = ( + "True" + if ( + (isinstance(output_debug_files, bool) and output_debug_files) + or (str(output_debug_files) == "True") + ) + else "False" + ) + + ( + current_prompt_content_logged, + current_summary_content_logged, + current_conversation_content_logged, + current_metadata_content_logged, + ) = process_debug_output_iteration( + output_debug_files_str, + output_folder, + batch_file_path_details, + model_choice_clean_short, + full_prompt, + summarised_output, + conversation_history, + metadata, + log_output_files, + task_type=task_type, + ) + + all_prompts_content.append(current_prompt_content_logged) + all_summaries_content.append(current_summary_content_logged) + # all_conversation_content.append(current_conversation_content_logged) + all_metadata_content.append(current_metadata_content_logged) + all_groups_content.append(all_groups[summary_no]) + all_batches_content.append(f"{summary_no}:") + all_model_choice_content.append(model_choice_clean_short) + all_validated_content.append("No") + all_task_type_content.append(task_type) + all_file_names_content.append(reference_data_file_name) + latest_summary_completed += 1 + + toc = time.perf_counter() + time_taken = toc - tic + + if time_taken > max_time_for_loop: + print( + "Time taken for loop is greater than maximum time allowed. Exiting and restarting loop" + ) + summary_loop.close() + tqdm._instances.clear() + break + + # If all summaries completed, make final outputs + if latest_summary_completed >= length_all_summaries: + print("All summaries completed. Creating outputs.") + + sampled_reference_table_df["Revised summary"] = summarised_outputs + + join_cols = ["General topic", "Subtopic", "Sentiment"] + join_plus_summary_cols = [ + "General topic", + "Subtopic", + "Sentiment", + "Revised summary", + ] + + summarised_references_j = sampled_reference_table_df[ + join_plus_summary_cols + ].drop_duplicates(join_plus_summary_cols) + + topic_summary_df_revised = topic_summary_df.merge( + summarised_references_j, on=join_cols, how="left" + ) + + # If no new summary is available, keep the original + # But prefer the version without "Rows X to Y" prefix to avoid duplication + def clean_summary_text(text): + if pd.isna(text): + return text + # Remove "Rows X to Y:" prefix if present (both at start and after
tags) + # First remove from the beginning + cleaned = re.sub(r"^Rows\s+\d+\s+to\s+\d+:\s*", "", str(text)) + # Then remove from after
tags + cleaned = re.sub(r"
\s*Rows\s+\d+\s+to\s+\d+:\s*", "
", cleaned) + return cleaned + + def deduplicate_summary_text(text): + """Remove duplicate summary segments separated by
or newlines.""" + if pd.isna(text): + return text + text_str = str(text) + # Split by
tags and newlines + segments = re.split(r"
|\n", text_str) + # Remove empty segments and strip whitespace + segments = [s.strip() for s in segments if s.strip()] + # Remove duplicates while preserving order + seen = set() + unique_segments = [] + for segment in segments: + if segment not in seen: + seen.add(segment) + unique_segments.append(segment) + # Join back with
(preserving original format) + return "
".join(unique_segments) if unique_segments else text_str + + # Deduplicate the original Summary before using combine_first + topic_summary_df_revised["Summary"] = topic_summary_df_revised["Summary"].apply( + deduplicate_summary_text + ) + + topic_summary_df_revised["Revised summary"] = topic_summary_df_revised[ + "Revised summary" + ].combine_first(topic_summary_df_revised["Summary"]) + # Clean the revised summary to remove "Rows X to Y" prefixes + topic_summary_df_revised["Revised summary"] = topic_summary_df_revised[ + "Revised summary" + ].apply(clean_summary_text) + topic_summary_df_revised = topic_summary_df_revised[ + [ + "General topic", + "Subtopic", + "Sentiment", + "Group", + "Number of responses", + "Revised summary", + ] + ] + + # Note: "Rows X to Y:" prefixes are now cleaned by the clean_summary_text function above + topic_summary_df_revised["Topic number"] = range( + 1, len(topic_summary_df_revised) + 1 + ) + + # If no new summary is available, keep the original. Also join on topic number to ensure consistent topic number assignment + reference_table_df_revised = reference_table_df.copy() + reference_table_df_revised = reference_table_df_revised.drop( + "Topic number", axis=1, errors="ignore" + ) + + # Ensure reference table has Topic number column + if ( + "Topic number" not in reference_table_df_revised.columns + or "Revised summary" not in reference_table_df_revised.columns + ): + if ( + "Topic number" in topic_summary_df_revised.columns + and "Revised summary" in topic_summary_df_revised.columns + ): + reference_table_df_revised = reference_table_df_revised.merge( + topic_summary_df_revised[ + [ + "General topic", + "Subtopic", + "Sentiment", + "Group", + "Topic number", + "Revised summary", + ] + ], + on=["General topic", "Subtopic", "Sentiment", "Group"], + how="left", + ) + + # Deduplicate the original Summary before using combine_first to prevent repeated summaries + reference_table_df_revised["Summary"] = reference_table_df_revised[ + "Summary" + ].apply(deduplicate_summary_text) + + reference_table_df_revised["Revised summary"] = reference_table_df_revised[ + "Revised summary" + ].combine_first(reference_table_df_revised["Summary"]) + # Clean the revised summary to remove "Rows X to Y" prefixes + reference_table_df_revised["Revised summary"] = reference_table_df_revised[ + "Revised summary" + ].apply(clean_summary_text) + reference_table_df_revised = reference_table_df_revised.drop( + "Summary", axis=1, errors="ignore" + ) + + # Remove topics that are tagged as 'Not Mentioned' + topic_summary_df_revised = topic_summary_df_revised.loc[ + topic_summary_df_revised["Sentiment"] != "Not Mentioned", : + ] + reference_table_df_revised = reference_table_df_revised.loc[ + reference_table_df_revised["Sentiment"] != "Not Mentioned", : + ] + + # Combine the logged content into a list of dictionaries + all_logged_content = [ + { + "prompt": prompt, + "response": summary, + "metadata": metadata, + "batch": batch, + "model_choice": model_choice, + "validated": validated, + "group": group, + "task_type": task_type, + "file_name": file_name, + } + for prompt, summary, metadata, batch, model_choice, validated, group, task_type, file_name in zip( + all_prompts_content, + all_summaries_content, + all_metadata_content, + all_batches_content, + all_model_choice_content, + all_validated_content, + all_groups_content, + all_task_type_content, + all_file_names_content, + ) + ] + + if isinstance(existing_logged_content, pd.DataFrame): + existing_logged_content = existing_logged_content.to_dict(orient="records") + + out_logged_content = existing_logged_content + all_logged_content + + ### Save output files + + if output_debug_files == "True": + + if not file_data.empty: + basic_response_data = get_basic_response_data(file_data, chosen_cols) + reference_table_df_revised_pivot = ( + convert_reference_table_to_pivot_table( + reference_table_df_revised, basic_response_data + ) + ) + + ### Save pivot file to log area + reference_table_df_revised_pivot_path = ( + output_folder + + file_name_clean + + "_summ_reference_table_pivot_" + + model_choice_clean + + ".csv" + ) + reference_table_df_revised_pivot.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv( + reference_table_df_revised_pivot_path, + index=None, + encoding="utf-8-sig", + ) + log_output_files.append(reference_table_df_revised_pivot_path) + + # Save to file + topic_summary_df_revised_path = ( + output_folder + + file_name_clean + + "_summ_unique_topics_table_" + + model_choice_clean + + ".csv" + ) + topic_summary_df_revised.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv(topic_summary_df_revised_path, index=None, encoding="utf-8-sig") + + reference_table_df_revised_path = ( + output_folder + + file_name_clean + + "_summ_reference_table_" + + model_choice_clean + + ".csv" + ) + reference_table_df_revised.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv(reference_table_df_revised_path, index=None, encoding="utf-8-sig") + + log_output_files.extend( + [reference_table_df_revised_path, topic_summary_df_revised_path] + ) + + ### + topic_summary_df_revised_display = topic_summary_df_revised.apply( + lambda col: col.map(lambda x: wrap_text(x, max_text_length=max_text_length)) + ) + summarised_output_markdown = topic_summary_df_revised_display.to_markdown( + index=False + ) + + # Ensure same file name not returned twice + output_files = list(set(output_files)) + log_output_files = list(set(log_output_files)) + + acc_input_tokens, acc_output_tokens, acc_number_of_calls = ( + calculate_tokens_from_metadata( + out_metadata_str, model_choice, model_name_map + ) + ) + + toc = time.perf_counter() + time_taken = toc - tic + + if isinstance(out_message, list): + out_message = "\n".join(out_message) + else: + out_message = out_message + + out_message = ( + out_message + + f"\nTopic summarisation finished processing. Total time: {round(float(time_taken), 1)}s" + ) + print(out_message) + + return ( + sampled_reference_table_df, + topic_summary_df_revised, + reference_table_df_revised, + output_files, + summarised_outputs, + latest_summary_completed, + out_metadata_str, + summarised_output_markdown, + log_output_files, + output_files, + acc_input_tokens, + acc_output_tokens, + acc_number_of_calls, + time_taken, + out_message, + out_logged_content, + ) + + +@spaces.GPU(duration=MAX_SPACES_GPU_RUN_TIME) +def wrapper_summarise_output_topics_per_group( + grouping_col: str, + sampled_reference_table_df: pd.DataFrame, + topic_summary_df: pd.DataFrame, + reference_table_df: pd.DataFrame, + model_choice: str, + in_api_key: str, + temperature: float, + reference_data_file_name: str, + summarised_outputs: list = list(), + latest_summary_completed: int = 0, + out_metadata_str: str = "", + in_data_files: List[str] = list(), + in_excel_sheets: str = "", + chosen_cols: List[str] = list(), + log_output_files: list[str] = list(), + summarise_format_radio: str = "Return a summary up to two paragraphs long that includes as much detail as possible from the original text", + output_folder: str = OUTPUT_FOLDER, + context_textbox: str = "", + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + model_name_map: dict = model_name_map, + hf_api_key_textbox: str = "", + azure_endpoint_textbox: str = "", + existing_logged_content: list = list(), + sample_reference_table: bool = False, + no_of_sampled_summaries: int = default_number_of_sampled_summaries, + random_seed: int = 42, + api_url: str = None, + additional_summary_instructions_provided: str = "", + output_debug_files: str = OUTPUT_DEBUG_FILES, + reasoning_suffix: str = reasoning_suffix, + local_model: object = None, + tokenizer: object = None, + assistant_model: object = None, + summarise_topic_descriptions_prompt: str = summarise_topic_descriptions_prompt, + summarise_topic_descriptions_system_prompt: str = summarise_topic_descriptions_system_prompt, + do_summaries: str = "Yes", + progress=gr.Progress(track_tqdm=True), +) -> Tuple[ + pd.DataFrame, + pd.DataFrame, + pd.DataFrame, + List[str], + List[str], + int, + str, + str, + List[str], + List[str], + int, + int, + int, + float, + str, + List[dict], +]: + """ + A wrapper function that iterates through unique values in a specified grouping column + and calls the `summarise_output_topics` function for each group of summaries. + It accumulates results from each call and returns a consolidated output. + + :param grouping_col: The name of the column to group the data by. + :param sampled_reference_table_df: DataFrame containing sampled reference data with summaries + :param topic_summary_df: DataFrame containing topic summary information + :param reference_table_df: DataFrame mapping response references to topics + :param model_choice: Name of the LLM model to use + :param in_api_key: API key for model access + :param temperature: Temperature parameter for model generation + :param reference_data_file_name: Name of the reference data file + :param summarised_outputs: List to store generated summaries + :param latest_summary_completed: Index of last completed summary + :param out_metadata_str: String for metadata output + :param in_data_files: List of input data file paths + :param in_excel_sheets: Excel sheet names if using Excel files + :param chosen_cols: List of columns selected for analysis + :param log_output_files: List of log file paths + :param summarise_format_radio: Format instructions for summary generation + :param output_folder: Folder path for outputs + :param context_textbox: Additional context for summarization + :param aws_access_key_textbox: AWS access key + :param aws_secret_key_textbox: AWS secret key + :param model_name_map: Dictionary mapping model choices to their properties + :param hf_api_key_textbox: Hugging Face API key + :param azure_endpoint_textbox: Azure endpoint + :param existing_logged_content: List of existing logged content + :param additional_summary_instructions_provided: Additional summary instructions + :param output_debug_files: Flag to indicate if debug files should be written + :param reasoning_suffix: Suffix for reasoning + :param local_model: Local model object if using local inference + :param tokenizer: Tokenizer object if using local inference + :param assistant_model: Assistant model object if using local inference + :param summarise_topic_descriptions_prompt: Prompt template for topic summarization + :param summarise_topic_descriptions_system_prompt: System prompt for topic summarization + :param do_summaries: Flag to control summary generation + :param sample_reference_table: If True, sample the reference table at the top of the function + :param no_of_sampled_summaries: Number of summaries to sample per group (default 100) + :param random_seed: Random seed for reproducible sampling (default 42) + :param progress: Gradio progress tracker + :return: A tuple containing consolidated results, mimicking the return structure of `summarise_output_topics` + """ + + acc_input_tokens = 0 + acc_output_tokens = 0 + acc_number_of_calls = 0 + out_message = list() + + # Logged content + all_groups_logged_content = existing_logged_content + + # Check if we have data to process + # Allow empty sampled_reference_table_df if sample_reference_table is True (it will be created from reference_table_df) + if ( + (sampled_reference_table_df.empty and not sample_reference_table) + or topic_summary_df.empty + or reference_table_df.empty + ): + out_message = "Please upload reference table, topic summary, and sampled reference table files to continue with summarisation." + print(out_message) + raise Exception(out_message) + + # Ensure Group column exists + if "Group" not in sampled_reference_table_df.columns: + sampled_reference_table_df["Group"] = "All" + if "Group" not in topic_summary_df.columns: + topic_summary_df["Group"] = "All" + if "Group" not in reference_table_df.columns: + reference_table_df["Group"] = "All" + + # Sample reference table if requested + if sample_reference_table: + print( + f"Sampling reference table with {no_of_sampled_summaries} summaries per group..." + ) + # Compute a token budget for the sampled summaries so downstream prompts always fit. + model_source = model_name_map[model_choice]["source"] + combined_summary_instructions = ( + summarise_format_radio + ". " + additional_summary_instructions_provided + ) + formatted_sys = summarise_topic_descriptions_system_prompt.format( + column_name=chosen_cols, consultation_context=context_textbox + ) + # Apply reasoning suffix for GPT-OSS models (Local, inference-server, or AWS) + is_gpt_oss_model = ( + "gpt-oss" in model_choice.lower() or "gpt_oss" in model_choice.lower() + ) + if is_gpt_oss_model: + effective_reasoning_suffix = ( + reasoning_suffix if reasoning_suffix else "Reasoning: low" + ) + if effective_reasoning_suffix: + formatted_sys = formatted_sys + "\n" + effective_reasoning_suffix + elif "Local" in model_source and reasoning_suffix: + formatted_sys = formatted_sys + "\n" + reasoning_suffix + + base_prompt = summarise_topic_descriptions_prompt.format( + summaries="", summary_format=combined_summary_instructions + ) + base_tokens = count_tokens_in_text( + formatted_sys + "\n" + base_prompt, tokenizer, model_source + ) + available_tokens_for_summaries = max( + 0, _effective_context_limit() - base_tokens + ) + + # Extra safety for inference-server when we don't have server-side token counting. + # If the server applies a chat template, it can push us over the limit even if our + # local estimate says we're OK. + if "inference-server" in str(model_source).lower(): + try: + frac = float(INFERENCE_SERVER_SAMPLER_TOKEN_BUDGET_FRACTION) + except Exception: + frac = 0.9 + frac = min(max(frac, 0.1), 1.0) + + # If the server can tokenize/apply-template, we can rely on accurate limits elsewhere. + # If not, be conservative here. + can_server_count = False + if api_url: + try: + from tools.dedup_summaries import ( + _inference_server_prompt_token_count, + ) + + test_messages = [ + {"role": "system", "content": formatted_sys}, + {"role": "user", "content": base_prompt}, + ] + can_server_count = ( + _inference_server_prompt_token_count(api_url, test_messages) + is not None + ) + except Exception: + can_server_count = False + + if not can_server_count: + available_tokens_for_summaries = int( + available_tokens_for_summaries * frac + ) + + sampled_reference_table_df, _ = sample_reference_table_summaries( + reference_table_df, + random_seed=random_seed, + no_of_sampled_summaries=no_of_sampled_summaries, + sample_reference_table_checkbox=sample_reference_table, + available_tokens_for_summaries=available_tokens_for_summaries, + tokenizer=tokenizer, + model_source=model_source, + ) + print( + f"Sampling complete. {len(sampled_reference_table_df)} summaries selected." + ) + + # Summarise every group present in the reference table, not only those that survived sampling + unique_values = reference_table_df["Group"].unique() + + if len(unique_values) > MAX_GROUPS: + print( + f"Warning: More than {MAX_GROUPS} unique values found in '{grouping_col}'. Processing only the first {MAX_GROUPS}." + ) + unique_values = unique_values[:MAX_GROUPS] + + # Initialize accumulators for results across all groups + acc_sampled_reference_table_df = pd.DataFrame() + acc_topic_summary_df_revised = pd.DataFrame() + acc_reference_table_df_revised = pd.DataFrame() + acc_output_files = list() + acc_log_output_files = list() + acc_summarised_outputs = list() + acc_latest_summary_completed = latest_summary_completed + acc_out_metadata_str = out_metadata_str + acc_summarised_output_markdown = "" + acc_total_time_taken = 0.0 + acc_logged_content = list() + + if len(unique_values) == 1: + # If only one unique value, no need for progress bar, iterate directly + loop_object = unique_values + else: + # If multiple unique values, use tqdm progress bar + loop_object = progress.tqdm( + unique_values, desc="Summarising group", unit="groups" + ) + + for i, group_value in enumerate(loop_object): + print( + f"\nProcessing summary group: {grouping_col} = {group_value} ({i+1}/{len(unique_values)})" + ) + + # Filter data for current group + filtered_sampled_reference_table_df = sampled_reference_table_df[ + sampled_reference_table_df["Group"] == group_value + ].copy() + filtered_topic_summary_df = topic_summary_df[ + topic_summary_df["Group"] == group_value + ].copy() + filtered_reference_table_df = reference_table_df[ + reference_table_df["Group"] == group_value + ].copy() + + if filtered_sampled_reference_table_df.empty: + # Sampling can exclude whole groups when other groups have multi-response topics. + # Fall back to the full per-group reference table so small groups still get summarised. + filtered_sampled_reference_table_df = _rows_for_topic_summarisation( + filtered_reference_table_df + ) + if filtered_sampled_reference_table_df.empty: + print(f"No data for {grouping_col} = {group_value}. Skipping.") + continue + print( + f"No sampled summaries for {grouping_col} = {group_value}; " + "using full reference table for this group." + ) + + # Create unique file name for this group's outputs + group_file_name = f"{reference_data_file_name}_{clean_column_name(str(group_value), max_length=15).replace(' ','_')}" + + # Call summarise_output_topics for the current group + try: + ( + seg_sampled_reference_table_df, + seg_topic_summary_df_revised, + seg_reference_table_df_revised, + seg_output_files, + seg_summarised_outputs, + seg_latest_summary_completed, + seg_out_metadata_str, + seg_summarised_output_markdown, + seg_log_output_files, + seg_output_files_2, + seg_acc_input_tokens, + seg_acc_output_tokens, + seg_acc_number_of_calls, + seg_time_taken, + seg_out_message, + seg_logged_content, + ) = summarise_output_topics( + sampled_reference_table_df=filtered_sampled_reference_table_df, + topic_summary_df=filtered_topic_summary_df, + reference_table_df=filtered_reference_table_df, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + reference_data_file_name=group_file_name, + summarised_outputs=list(), # Fresh for each call + latest_summary_completed=0, # Reset for each group + out_metadata_str="", # Fresh for each call + in_data_files=in_data_files, + in_excel_sheets=in_excel_sheets, + chosen_cols=chosen_cols, + log_output_files=list(), # Fresh for each call + summarise_format_radio=summarise_format_radio, + output_folder=output_folder, + context_textbox=context_textbox, + aws_access_key_textbox=aws_access_key_textbox, + aws_secret_key_textbox=aws_secret_key_textbox, + aws_region_textbox=aws_region_textbox, + model_name_map=model_name_map, + hf_api_key_textbox=hf_api_key_textbox, + azure_endpoint_textbox=azure_endpoint_textbox, + existing_logged_content=all_groups_logged_content, + additional_summary_instructions_provided=additional_summary_instructions_provided, + output_debug_files=output_debug_files, + group_value=group_value, + reasoning_suffix=reasoning_suffix, + local_model=local_model, + tokenizer=tokenizer, + assistant_model=assistant_model, + summarise_topic_descriptions_prompt=summarise_topic_descriptions_prompt, + summarise_topic_descriptions_system_prompt=summarise_topic_descriptions_system_prompt, + do_summaries=do_summaries, + api_url=api_url, + ) + + # Aggregate results + acc_sampled_reference_table_df = pd.concat( + [acc_sampled_reference_table_df, seg_sampled_reference_table_df] + ) + acc_topic_summary_df_revised = pd.concat( + [acc_topic_summary_df_revised, seg_topic_summary_df_revised] + ) + acc_reference_table_df_revised = pd.concat( + [acc_reference_table_df_revised, seg_reference_table_df_revised] + ) + + # For lists, extend + acc_output_files.extend( + f for f in seg_output_files if f not in acc_output_files + ) + acc_log_output_files.extend( + f for f in seg_log_output_files if f not in acc_log_output_files + ) + acc_summarised_outputs.extend(seg_summarised_outputs) + + acc_latest_summary_completed = seg_latest_summary_completed + acc_out_metadata_str += ( + ("\n---\n" if acc_out_metadata_str else "") + + f"Group {grouping_col}={group_value}:\n" + + seg_out_metadata_str + ) + acc_summarised_output_markdown = ( + seg_summarised_output_markdown # Keep the latest markdown + ) + acc_total_time_taken += float(seg_time_taken) + acc_logged_content.extend(seg_logged_content) + + # Accumulate token counts + acc_input_tokens += seg_acc_input_tokens + acc_output_tokens += seg_acc_output_tokens + acc_number_of_calls += seg_acc_number_of_calls + + print( + f"Group {grouping_col} = {group_value} summarised. Time: {seg_time_taken:.2f}s" + ) + + except Exception as e: + print(f"Error processing summary group {grouping_col} = {group_value}: {e}") + # Optionally, decide if you want to continue with other groups or stop + # For now, it will continue + continue + + # Ensure custom model_choice is registered in model_name_map + ensure_model_in_map(model_choice, model_name_map) + + # Create consolidated output files + overall_file_name = clean_column_name(reference_data_file_name, max_length=20) + model_choice_clean = model_name_map[model_choice]["short_name"] + model_choice_clean_short = clean_column_name( + model_choice_clean, max_length=20, front_characters=False + ) + + # Save consolidated outputs + if ( + not acc_topic_summary_df_revised.empty + and not acc_reference_table_df_revised.empty + ): + # Sort the dataframes + if "General topic" in acc_topic_summary_df_revised.columns: + acc_topic_summary_df_revised["Number of responses"] = ( + acc_topic_summary_df_revised["Number of responses"].astype(int) + ) + acc_topic_summary_df_revised.sort_values( + [ + "Group", + "Number of responses", + "General topic", + "Subtopic", + "Sentiment", + ], + ascending=[True, False, True, True, True], + inplace=True, + ) + elif "Main heading" in acc_topic_summary_df_revised.columns: + acc_topic_summary_df_revised["Number of responses"] = ( + acc_topic_summary_df_revised["Number of responses"].astype(int) + ) + acc_topic_summary_df_revised.sort_values( + [ + "Group", + "Number of responses", + "Main heading", + "Subheading", + "Topic number", + ], + ascending=[True, False, True, True, True], + inplace=True, + ) + + # Save consolidated files + consolidated_topic_summary_path = ( + output_folder + + overall_file_name + + "_all_final_summ_unique_topics_" + + model_choice_clean_short + + ".csv" + ) + consolidated_reference_table_path = ( + output_folder + + overall_file_name + + "_all_final_summ_reference_table_" + + model_choice_clean_short + + ".csv" + ) + + acc_topic_summary_df_revised.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv(consolidated_topic_summary_path, index=None, encoding="utf-8-sig") + acc_reference_table_df_revised.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv(consolidated_reference_table_path, index=None, encoding="utf-8-sig") + + acc_output_files.extend( + [consolidated_topic_summary_path, consolidated_reference_table_path] + ) + + # Create markdown output for display + topic_summary_df_revised_display = acc_topic_summary_df_revised.apply( + lambda col: col.map(lambda x: wrap_text(x, max_text_length=max_text_length)) + ) + acc_summarised_output_markdown = topic_summary_df_revised_display.to_markdown( + index=False + ) + + out_message = "\n".join(out_message) + out_message = ( + out_message + + " " + + f"Topic summarisation finished processing all groups. Total time: {acc_total_time_taken:.2f}s" + ) + print(out_message) + + # The return signature should match summarise_output_topics + return ( + acc_sampled_reference_table_df, + acc_topic_summary_df_revised, + acc_reference_table_df_revised, + acc_output_files, + acc_summarised_outputs, + acc_latest_summary_completed, + acc_out_metadata_str, + acc_summarised_output_markdown, + acc_log_output_files, + acc_output_files, # Duplicate for compatibility + acc_input_tokens, + acc_output_tokens, + acc_number_of_calls, + acc_total_time_taken, + out_message, + acc_logged_content, + ) + + +def convert_markdown_headers_to_excel_format(text: str) -> str: + """ + Convert markdown headers to Excel-friendly format that preserves hierarchy. + + Converts: + - # Header (H1) -> === HEADER === (most prominent) + - ## Header (H2) -> --- Header --- (medium) + - ### Header (H3) -> ── Header ── (less prominent) + - #### Header (H4) -> • Header (with bullet) + - ##### Header (H5) -> • Header (indented) + - ###### Header (H6) -> • Header (more indented) + + Args: + text (str): Text containing markdown headers + + Returns: + str: Text with markdown headers converted to Excel-friendly format + """ + if not text: + return text + + lines = text.split("\n") + converted_lines = [] + + for line in lines: + # Match markdown headers (# through ######) + header_match = re.match(r"^(#{1,6})\s+(.+)$", line) + if header_match: + header_level = len(header_match.group(1)) # Number of # characters + header_text = header_match.group(2).strip() + + if header_level == 1: + # H1: Most prominent - uppercase with double equals + converted_line = f"=== {header_text.upper()} ===" + elif header_level == 2: + # H2: Medium prominence - title case with dashes + converted_line = f"--- {header_text.title()} ---" + elif header_level == 3: + # H3: Less prominent - title case with single dashes + converted_line = f"── {header_text.title()} ──" + elif header_level == 4: + # H4: Bullet with no indentation + converted_line = f"• {header_text}" + elif header_level == 5: + # H5: Bullet with indentation + converted_line = f" • {header_text}" + else: # header_level == 6 + # H6: Bullet with more indentation + converted_line = f" • {header_text}" + + converted_lines.append(converted_line) + else: + converted_lines.append(line) + + return "\n".join(converted_lines) + + +@spaces.GPU(duration=MAX_SPACES_GPU_RUN_TIME) +def overall_summary( + topic_summary_df: pd.DataFrame, + model_choice: str, + in_api_key: str, + temperature: float, + reference_data_file_name: str, + output_folder: str = OUTPUT_FOLDER, + chosen_cols: List[str] = list(), + context_textbox: str = "", + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + model_name_map: dict = model_name_map, + hf_api_key_textbox: str = "", + azure_endpoint_textbox: str = "", + existing_logged_content: list = list(), + api_url: str = None, + output_debug_files: str = output_debug_files, + log_output_files: list = list(), + reasoning_suffix: str = reasoning_suffix, + local_model: object = None, + tokenizer: object = None, + assistant_model: object = None, + summarise_everything_prompt: str = summarise_everything_prompt, + comprehensive_summary_format_prompt: str = comprehensive_summary_format_prompt, + comprehensive_summary_format_prompt_by_group: str = comprehensive_summary_format_prompt_by_group, + summarise_everything_system_prompt: str = summarise_everything_system_prompt, + do_summaries: str = "Yes", + progress=gr.Progress(track_tqdm=True), +) -> Tuple[ + List[str], + List[str], + int, + str, + List[str], + List[str], + int, + int, + int, + float, + List[dict], +]: + """ + Create an overall summary of all responses based on a topic summary table. + + Args: + topic_summary_df (pd.DataFrame): DataFrame containing topic summaries + model_choice (str): Name of the LLM model to use + in_api_key (str): API key for model access + temperature (float): Temperature parameter for model generation + reference_data_file_name (str): Name of reference data file + output_folder (str, optional): Folder to save outputs. Defaults to OUTPUT_FOLDER. + chosen_cols (List[str], optional): Columns to analyze. Defaults to empty list. + context_textbox (str, optional): Additional context. Defaults to empty string. + aws_access_key_textbox (str, optional): AWS access key. Defaults to empty string. + aws_secret_key_textbox (str, optional): AWS secret key. Defaults to empty string. + aws_region_textbox (str, optional): AWS region. Defaults to empty string. + model_name_map (dict, optional): Mapping of model names. Defaults to model_name_map. + hf_api_key_textbox (str, optional): Hugging Face API key. Defaults to empty string. + existing_logged_content (list, optional): List of existing logged content. Defaults to empty list. + output_debug_files (str, optional): Flag to indicate if debug files should be written. Defaults to "False". + log_output_files (list, optional): List of existing logged content. Defaults to empty list. + api_url (str, optional): API URL for inference-server models. Defaults to None. + reasoning_suffix (str, optional): Suffix for reasoning. Defaults to reasoning_suffix. + local_model (object, optional): Local model object. Defaults to None. + tokenizer (object, optional): Tokenizer object. Defaults to None. + assistant_model (object, optional): Assistant model object. Defaults to None. + summarise_everything_prompt (str, optional): Prompt for overall summary + comprehensive_summary_format_prompt (str, optional): Prompt for comprehensive summary format + comprehensive_summary_format_prompt_by_group (str, optional): Prompt for group summary format + summarise_everything_system_prompt (str, optional): System prompt for overall summary + do_summaries (str, optional): Whether to generate summaries. Defaults to "Yes". + progress (gr.Progress, optional): Progress tracker. Defaults to gr.Progress(track_tqdm=True). + + Returns: + Tuple containing: + List[str]: Output files + List[str]: Text summarized outputs + int: Latest summary completed + str: Output metadata + List[str]: Summarized outputs + List[str]: Summarized outputs for DataFrame + int: Number of input tokens + int: Number of output tokens + int: Number of API calls + float: Time taken + List[dict]: List of logged content + """ + + out_metadata = list() + latest_summary_completed = 0 + output_files = list() + txt_summarised_outputs = list() + summarised_outputs = list() + summarised_outputs_for_df = list() + input_tokens_num = 0 + output_tokens_num = 0 + number_of_calls_num = 0 + time_taken = 0 + out_message = list() + all_logged_content = list() + all_prompts_content = list() + all_summaries_content = list() + all_metadata_content = list() + all_groups_content = list() + all_batches_content = list() + all_model_choice_content = list() + all_validated_content = list() + task_type = "Overall summary" + all_task_type_content = list() + log_output_files = list() + all_logged_content = list() + all_file_names_content = list() + out_metadata_str = "" + tic = time.perf_counter() + + if "Group" not in topic_summary_df.columns: + topic_summary_df["Group"] = "All" + + if "Number of responses" not in topic_summary_df.columns: + topic_summary_df["Number of responses"] = 0 + + topic_summary_df = topic_summary_df.sort_values( + by=["Group", "Number of responses"], ascending=[True, False] + ) + + unique_groups = sorted(topic_summary_df["Group"].unique()) + + length_groups = len(unique_groups) + + if context_textbox and "The context of this analysis is" not in context_textbox: + context_textbox = "The context of this analysis is '" + context_textbox + "'." + + if length_groups > 1: + comprehensive_summary_format_prompt = ( + comprehensive_summary_format_prompt_by_group + ) + else: + comprehensive_summary_format_prompt = comprehensive_summary_format_prompt + + # Ensure custom model_choice is registered in model_name_map + ensure_model_in_map(model_choice, model_name_map) + + batch_file_path_details = create_batch_file_path_details(reference_data_file_name) + model_choice_clean = model_name_map[model_choice]["short_name"] + model_choice_clean_short = clean_column_name( + model_choice_clean, max_length=20, front_characters=False + ) + + tic = time.perf_counter() + + if ( + (model_choice == CHOSEN_LOCAL_MODEL_TYPE) + & (RUN_LOCAL_MODEL == "1") + & (not local_model) + ): + progress(0.1, f"Using model: {CHOSEN_LOCAL_MODEL_TYPE}") + local_model = get_model() + tokenizer = get_tokenizer() + assistant_model = get_assistant_model() + + summary_loop = tqdm( + unique_groups, desc="Creating overall summary for groups", unit="groups" + ) + + if do_summaries == "Yes": + model_source = model_name_map[model_choice]["source"] + bedrock_runtime = connect_to_bedrock_runtime( + model_name_map, + model_choice, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + ) + + for summary_group in summary_loop: + + print("Creating overall summary for group:", summary_group) + + # Get the group-specific DataFrame + group_df = topic_summary_df.loc[ + topic_summary_df["Group"] == summary_group + ].copy() + + # Prepare the system prompt first (needed for token counting) + formatted_summarise_everything_system_prompt = ( + summarise_everything_system_prompt.format( + column_name=chosen_cols, consultation_context=context_textbox + ) + ) + + # Apply reasoning suffix for GPT-OSS models (Local, inference-server, or AWS) + is_gpt_oss_model = ( + "gpt-oss" in model_choice.lower() or "gpt_oss" in model_choice.lower() + ) + + if is_gpt_oss_model: + # Use default reasoning suffix if not set + effective_reasoning_suffix = ( + reasoning_suffix if reasoning_suffix else "Reasoning: low" + ) + if effective_reasoning_suffix: + formatted_summarise_everything_system_prompt = ( + formatted_summarise_everything_system_prompt + + "\n" + + effective_reasoning_suffix + ) + elif "Local" in model_source and reasoning_suffix: + # For other local models, use reasoning_suffix if provided + formatted_summarise_everything_system_prompt = ( + formatted_summarise_everything_system_prompt + + "\n" + + reasoning_suffix + ) + + # Create a test prompt with empty table to get base token count + test_summary_text = "" + test_formatted_summary_prompt = [ + summarise_everything_prompt.format( + topic_summary_table=test_summary_text, + summary_format=comprehensive_summary_format_prompt, + ) + ] + + # Calculate base token count (system prompt + prompt template without table) + full_test_text = ( + formatted_summarise_everything_system_prompt + + "\n" + + test_formatted_summary_prompt[0] + ) + count_tokens_in_text(full_test_text, tokenizer, model_source) + + # Fit-to-context by dropping whole rows (never truncating text). + # We enforce the limit against the FULL prompt (system + formatted prompt), + # not just the table text, to avoid undercount edge cases. + def _prompt_tokens_for_df(df: pd.DataFrame) -> int: + table_text = clean_markdown_table_whitespace( + df.to_markdown(index=False) + ) + prompt_text = summarise_everything_prompt.format( + topic_summary_table=table_text, + summary_format=comprehensive_summary_format_prompt, + ) + full_text = ( + formatted_summarise_everything_system_prompt + "\n" + prompt_text + ) + return count_tokens_in_text(full_text, tokenizer, model_source) + + if len(group_df) > 0: + effective_limit = _effective_context_limit() + full_tokens = _prompt_tokens_for_df(group_df) + if full_tokens > effective_limit: + print( + f"Warning: Overall summary prompt for group '{summary_group}' exceeds context limit. " + f"Dropping rows. Prompt tokens: {full_tokens}, Max (effective): {effective_limit} (raw max: {LLM_CONTEXT_LENGTH})" + ) + + num_rows = len(group_df) + lo, hi = 1, num_rows + best_df = group_df.iloc[:0] + + while lo <= hi: + mid = (lo + hi) // 2 + test_df = group_df.iloc[:mid] + test_tokens = _prompt_tokens_for_df(test_df) + if test_tokens <= effective_limit: + best_df = test_df + lo = mid + 1 + else: + hi = mid - 1 + + if best_df.empty: + # Even one row doesn't fit (usually means system prompt itself is too large). + raise ValueError( + "Overall summary prompt is too large even with a single table row. " + f"Max effective context length: {effective_limit} (raw max: {LLM_CONTEXT_LENGTH}). " + "Reduce prompt size/format or use a model with a larger context window." + ) + + group_df = best_df + print( + f"Truncated to {len(group_df)} rows (from {num_rows} original rows) to fit context." + ) + + # Create summary_text from (possibly truncated) DataFrame + summary_text = group_df.to_markdown(index=False) + # Clean extraneous whitespace from markdown table cells + summary_text = clean_markdown_table_whitespace(summary_text) + + formatted_summary_prompt = [ + summarise_everything_prompt.format( + topic_summary_table=summary_text, + summary_format=comprehensive_summary_format_prompt, + ) + ] + + # Ensure locals are always defined even if the LLM call fails early + conversation_history = list() + metadata = list() + try: + response, conversation_history, metadata, response_text = ( + summarise_output_topics_query( + model_choice, + in_api_key, + temperature, + formatted_summary_prompt, + formatted_summarise_everything_system_prompt, + model_source, + bedrock_runtime, + local_model, + tokenizer=tokenizer, + assistant_model=assistant_model, + azure_endpoint=azure_endpoint_textbox, + api_url=api_url, + ) + ) + summarised_output_for_df = response_text + summarised_output = response + except Exception as e: + print( + "Cannot create overall summary for group:", + summary_group, + "due to:", + e, + ) + summarised_output = "" + summarised_output_for_df = "" + conversation_history = list() + metadata = list() + + # Remove multiple consecutive line breaks (2 or more) and replace with single line break + if summarised_output_for_df: + summarised_output_for_df = re.sub( + r"\n{2,}", "\n", summarised_output_for_df + ) + # Convert markdown headers to Excel-friendly format + summarised_output_for_df = convert_markdown_headers_to_excel_format( + summarised_output_for_df + ) + if summarised_output: + summarised_output = re.sub(r"\n{2,}", "\n", summarised_output) + + summarised_outputs_for_df.append(summarised_output_for_df) + summarised_outputs.append(summarised_output) + txt_summarised_outputs.append( + f"""Group name: {summary_group}\n""" + summarised_output + ) + + if metadata: + out_metadata.extend(metadata) + out_metadata_str = ". ".join(out_metadata) if out_metadata else "" + + full_prompt = ( + formatted_summarise_everything_system_prompt + + "\n" + + formatted_summary_prompt[0] + ) + + ( + current_prompt_content_logged, + current_summary_content_logged, + current_conversation_content_logged, + current_metadata_content_logged, + ) = process_debug_output_iteration( + output_debug_files, + output_folder, + batch_file_path_details, + model_choice_clean_short, + full_prompt, + summarised_output, + conversation_history, + metadata, + log_output_files, + task_type=task_type, + ) + + all_prompts_content.append(current_prompt_content_logged) + all_summaries_content.append(current_summary_content_logged) + # all_conversation_content.append(current_conversation_content_logged) + all_metadata_content.append(current_metadata_content_logged) + all_groups_content.append(summary_group) + all_batches_content.append("1") + all_model_choice_content.append(model_choice_clean_short) + all_validated_content.append("No") + all_task_type_content.append(task_type) + all_file_names_content.append(reference_data_file_name) + latest_summary_completed += 1 + clean_column_name(summary_group) + + # Write overall outputs to csv + overall_summary_output_csv_path = ( + output_folder + + batch_file_path_details + + "_overall_summary_" + + model_choice_clean_short + + ".csv" + ) + summarised_outputs_df = pd.DataFrame( + data={"Group": unique_groups, "Summary": summarised_outputs_for_df} + ) + summarised_outputs_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + overall_summary_output_csv_path, index=None, encoding="utf-8-sig" + ) + output_files.append(overall_summary_output_csv_path) + + summarised_outputs_df_for_display = pd.DataFrame( + data={"Group": unique_groups, "Summary": summarised_outputs} + ) + summarised_outputs_df_for_display["Summary"] = ( + summarised_outputs_df_for_display["Summary"].apply( + lambda x: markdown.markdown(x) if isinstance(x, str) else "" + ) + ) + summarised_outputs_df_for_display["Summary"] = ( + summarised_outputs_df_for_display["Summary"] + .astype(str) + .str.replace(r"\n", "
", regex=False) + .str.replace(r"(
\s*){2,}", "
", regex=True) + ) + html_output_table = summarised_outputs_df_for_display.to_html( + index=False, escape=False + ) + + output_files = list(set(output_files)) + + input_tokens_num, output_tokens_num, number_of_calls_num = ( + calculate_tokens_from_metadata( + out_metadata_str, model_choice, model_name_map + ) + ) + + # Check if beyond max time allowed for processing and break if necessary + toc = time.perf_counter() + time_taken = toc - tic + + out_message = "\n".join(out_message) + out_message = ( + out_message + + " " + + f"Overall summary finished processing. Total time: {time_taken:.2f}s" + ) + print(out_message) + + # Combine the logged content into a list of dictionaries + all_logged_content = [ + { + "prompt": prompt, + "response": summary, + "metadata": metadata, + "batch": batch, + "model_choice": model_choice, + "validated": validated, + "group": group, + "task_type": task_type, + "file_name": file_name, + } + for prompt, summary, metadata, batch, model_choice, validated, group, task_type, file_name in zip( + all_prompts_content, + all_summaries_content, + all_metadata_content, + all_batches_content, + all_model_choice_content, + all_validated_content, + all_groups_content, + all_task_type_content, + all_file_names_content, + ) + ] + + if isinstance(existing_logged_content, pd.DataFrame): + existing_logged_content = existing_logged_content.to_dict(orient="records") + + out_logged_content = existing_logged_content + all_logged_content + + return ( + output_files, + html_output_table, + summarised_outputs_df, + out_metadata_str, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + time_taken, + out_message, + out_logged_content, + ) diff --git a/tools/example_table_outputs.py b/tools/example_table_outputs.py new file mode 100644 index 0000000000000000000000000000000000000000..dc0392723083704425f8b9176c1dd1ecd679817a --- /dev/null +++ b/tools/example_table_outputs.py @@ -0,0 +1,94 @@ +dummy_consultation_table = """| General topic | Subtopic | Sentiment | Group | Number of responses | Revised summary | +|:---------------------|:-----------------------|:------------|:--------|----------------------:|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| Development proposal | Affordable housing | Positive | All | 6 | The proposed development is overwhelmingly viewed favorably by the local community, primarily due to
its potential to address significant needs. residents express strong support for the provision of
amenities, particularly much-needed housing for young people and families. crucially, the
development is also considered vital for increasing green space, with some suggesting this could
even contribute to further housing opportunities.
furthermore, a key theme emerging from the
responses is the des... | +| Development proposal | Environmental damage | Negative | All | 5 | A primary concern expressed across the dataset relates to the potential negative consequences of a
proposed development on the local environment and wildlife. multiple respondents highlighted worries
about environmental damage, suggesting a significant apprehension regarding the project’s impact.
specifically, there’s a shared concern about the detrimental effects on local wildlife, indicating a
potential disruption to the natural ecosystem.
furthermore, the proposed development is perceived
to ... | +| Community impact | Loss of local business | Negative | All | 4 | A significant concern expressed within the dataset relates to the potential negative consequences of
a development project on local businesses and the wider community. multiple respondents voiced
worries about increased traffic congestion, directly impacting the viability of local businesses and
leading to apprehension about their future. specifically, there is a palpable sadness surrounding
the possibility of a beloved local cafe closing, emphasizing the detrimental effect on community
connect... | +| Development proposal | Architectural style | Negative | All | 4 | The primary concern regarding the proposed development is its incompatibility with the established
character of the area. residents express a strong feeling that the design is fundamentally at odds
with the existing aesthetic and atmosphere, suggesting a lack of sensitivity to the local context.
this sentiment highlights a significant apprehension about disrupting the area’s unique
identity.
furthermore, significant anxieties center on the potential negative impact of the
development on the surr... | +| Economic impact | Investment and jobs | Positive | All | 4 | The proposed development is widely anticipated to generate significant positive economic impacts
within the local community. residents believe it will lead to substantial investment and the
creation of numerous job opportunities, directly boosting the local economy and revitalizing the
town centre. there’s a strong consensus that this development represents a key step towards economic
growth and prosperity for the area.
specifically, the anticipated benefits include the creation
of jobs for loca... | +| Development proposal | Height of building | Negative | All | 3 | Residents expressed significant concerns regarding the proposed development’s height, specifically
highlighting the five-storey structure as a major issue. this height was perceived as excessively
tall and likely to cause overshadowing of existing buildings in the area, directly impacting the
views enjoyed by current residents. the potential for diminished sunlight and altered visual
landscapes was a central point of contention.
furthermore, the impact on the existing character
of the neighborho... | +| Development proposal | Infrastructure impact | Negative | All | 3 | Analysis of the provided text reveals significant concerns regarding the potential consequences of a
proposed project on existing infrastructure. specifically, there is a notable worry about the
detrimental effects on local infrastructure, with the possibility of widespread disruption as a key
consequence. this suggests a need for careful assessment and mitigation strategies to avoid
negatively impacting essential services and community operations.
furthermore, the repeated
emphasis on infrastru... | +| Development proposal | Noise pollution | Negative | All | 3 | The primary concern highlighted within the dataset relates to anticipated noise pollution stemming
from the proposed development. multiple responses explicitly express this worry, emphasizing it as a
“significant concern” and a key area of apprehension. there’s a clear understanding that the
development will likely exacerbate existing noise levels within the surrounding area, suggesting a
potential negative impact on residents and the local environment.
several respondents reiterate
this concer... | +| Housing needs | Supply of housing | Positive | All | 3 | The proposed development is viewed as a crucial solution to address the town’s significant housing
shortage, reflecting a clear desire for increased housing supply within the community. residents
express a strong need for more homes, and this development is seen as a key step towards meeting
that demand.
furthermore, the project is anticipated to alleviate existing parking issues, which
is considered a valuable contribution to the overall housing supply. the provision of additional
parking spac... | +| Development proposal | Height of building | Neutral | All | 2 | The analysis of the provided text reveals a nuanced perspective regarding a development project,
with no explicit sentiment expressed concerning the building's height itself. however, significant
concerns are raised about the potential impact of the development on local schools. these concerns
appear to be linked, at least in part, to the building's height, suggesting a worry that the
increased scale could strain existing resources and infrastructure within the school system.

further investigat... | +| Community impact | Community facilities | Negative | All | 1 | Concerns exist regarding the negative impact on local amenities. | +| Community impact | Community facilities | Positive | All | 1 | The development will provide much-needed community facilities, enhancing the local area. | +| Development proposal | Architectural style | Neutral | All | 1 | The development is expected to provide facilities for young people, but no specific architectural
concerns. | +| Development proposal | Noise pollution | Neutral | All | 1 | Potential for increased noise pollution due to the development is a concern. | +| Economic impact | Economic decline | Negative | All | 1 | Worries about a negative impact on the local economy are expressed, suggesting potential harm. |""" + +dummy_consultation_table_zero_shot = """| General topic | Subtopic | Sentiment | Group | Number of responses | Revised summary | +|:---------------------------|:------------------------------------|:------------|:--------|----------------------:|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| Planning & development | Impact on the character of the area | Negative | All | 10 | Residents overwhelmingly express strong objections to the proposed development, primarily focusing
on its incompatibility with the established character of the area. A central concern is the
development's height and design, which they believe clashes significantly with the existing
aesthetic and creates a sense of being overshadowed by taller structures, leading to a feeling of
crampedness. Many respondents specifically highlighted the potential for the development to
negatively impact Main Stre... | +| Environmental impact | Impact on the local environment | Negative | All | 8 | Several concerns have been raised regarding the potential negative impacts of a development on the
local environment. Multiple respondents expressed worry about the development’s possible detrimental
effects on the surrounding environment and quality of life, highlighting a significant area of
concern. These anxieties include potential damage to the environment and a general feeling of unease
about the development’s consequences.

Despite a single positive note regarding the provision
of green s... | +| Infrastructure & transport | Traffic congestion | Negative | All | 7 | Concerns regarding increased traffic congestion are prevalent in the dataset, largely stemming from
the anticipated impact of the proposed development. Specifically, Main Street is predicted to
experience heightened congestion due to the increased volume of traffic it will attract. Multiple
responses repeatedly highlight this anticipation as a key issue associated with the
project.

Despite the consistent apprehension about traffic congestion, no direct responses
offer specific solutions or miti... | +| Planning & development | Need for family housing | Positive | All | 7 | The proposed development is overwhelmingly viewed as a crucial solution to the need for family
housing within the community. Multiple sources highlight its significance in providing much-needed
homes, particularly for families, and specifically addressing the demand for affordable family
housing options. Several respondents emphasized the beneficial impact on local residents, with the
development also anticipated to create jobs and offer facilities geared towards young people
alongside housing. ... | +| Quality of life | Impact on quality of life | Negative | All | 7 | Analysis of the provided text reveals significant concerns regarding a proposed development's
potential negative impact on the quality of life within the area. Residents are particularly worried
that the development will overshadow existing buildings, creating a sense of crampedness and
diminishing their living experience. Furthermore, anxieties extend beyond immediate residential
impacts, encompassing broader concerns about the development’s effects on local businesses, schools,
and crucial inf... | +| Economic impact | Investment and job creation | Positive | All | 6 | The proposed development is overwhelmingly viewed positively, with significant anticipation for its
economic impact on the area. Residents and observers alike believe it will stimulate considerable
investment and generate numerous job opportunities, particularly for local residents. Furthermore,
the project is expected to revitalize the town center and provide crucial affordable housing,
potentially benefiting young people seeking to establish themselves in the
community.

Specifically, the deve... | +| Infrastructure & transport | Parking | Negative | All | 6 | Analysis of the '{column_name}' column reveals significant concerns regarding the potential impact
of a new development on Main Street. The primary issue identified is increased traffic congestion,
directly linked to the development’s activity. Furthermore, there is widespread apprehension that
the project will worsen existing parking problems, with multiple respondents explicitly stating a
lack of adequate parking provisions as a key worry.

Specifically, numerous individuals
expressed concern... | +| Community & local life | Amenities for the local community | Positive | All | 5 | The proposed development is anticipated to significantly benefit the local community, offering a
range of amenities and a positive contribution to the area. Specifically, the project will deliver
crucial green space alongside facilities designed to cater to the needs of young people and the
broader community.

Furthermore, the development is expected to address critical social needs
by providing much-needed community facilities and social housing, indicating a commitment to
supporting local resi... | +| Environmental impact | Impact on local wildlife | Neutral | All | 4 | No specific responses were provided, and the dataset contained no information relevant to the
specified consultation context. Consequently, a summary cannot be generated based on the provided
data.

Due to the absence of any textual data within the dataset, there is no content to
consolidate and summarize. | +| Improvement of main street | Improvement of main street | Positive | All | 4 | This development is being hailed as a positive step for the revitalization of Main Street, primarily
due to its anticipated improvement in the street’s appearance. Stakeholders view this initiative as
a crucial element in breathing new life into the area, suggesting a significant upgrade to the
existing landscape.

Specifically, the project aims to enhance the visual appeal of Main
Street, representing a tangible advancement in its overall attractiveness and desirability. The
development is wide... | +| Planning & development | Impact on views | Negative | All | 4 | A primary concern expressed regarding the proposed development is its potential negative impact on
existing views. Multiple respondents voiced worries about how the development might obstruct or
diminish the current vistas, alongside specific concerns about its effect on views from neighboring
properties. This suggests a significant sensitivity to the visual landscape and its value within the
community.

Furthermore, the potential aesthetic consequences of the development are
highlighted, with s... | +| Community & local life | Amenities for the local community | Negative | All | 2 | Residents are voicing significant concerns regarding a proposed development, primarily focusing on
its anticipated detrimental effects on local amenities. A key point of contention is the planned
removal of the existing cafe, which is being viewed as a substantial loss to the community’s social
fabric and a vital local resource.

The overall sentiment suggests a strong apprehension that
the development will diminish the quality of life for those living nearby, highlighting a desire to
preserve c... | +| Impact on local businesses | Impact on local businesses | Negative | All | 2 | A primary concern expressed relates to the potential detrimental effects of the development on local
businesses. There’s a clear worry that the project will negatively impact these businesses,
suggesting a potential loss of revenue, customer base, or even business closure. The repeated
emphasis on a “negative impact” highlights a significant apprehension regarding the economic
repercussions for the existing business community.

The sentiment underscores a desire to
mitigate potential harm and li... | +| Impact on local heritage | Impact on local heritage | Negative | All | 2 | There are growing concerns regarding the potential negative impact of the development on the local
heritage. While specific details and references haven’t been explicitly stated, the underlying
sentiment suggests a worry about the development’s effects on historically significant elements
within the area. This implies a recognition that the proposed project could, perhaps inadvertently,
threaten or diminish the cultural value and character of the local environment.

The presence
of these concern... | +| Environmental impact | Impact on local wildlife | Negative | All | 1 | Concerns regarding the negative impact of the development on local wildlife. | +| Impact on local heritage | Impact on local heritage | Neutral | All | 1 | No specific responses mention this topic. | +| Impact on local schools | Impact on local schools | Negative | All | 1 | Concerns about the negative impact on the local schools. | +| Impact on local schools | Impact on local schools | Neutral | All | 1 | No specific responses mention this topic. | +| Infrastructure & transport | Parking | Positive | All | 1 | The development is expected to provide much-needed parking spaces. |""" + +case_notes_table = """| General topic | Subtopic | Sentiment | Group | Number of responses | Revised summary | +|:------------------|:----------------------------|:------------|:--------|----------------------:|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| Family dynamics | Parental conflict | Negative | All | 6 | Several parents expressed significant concerns regarding the well-being of their children, primarily
focusing on escalating aggression and withdrawal. alex’s mother specifically highlighted a pattern
of arguments at home and attributed the aggressive behavior to external provocation, suggesting a
destabilizing family environment. furthermore, parents voiced a lack of confidence in existing
interventions for their children, particularly jamie, indicating a perceived need for supplemental
support ... | +| Mental health | Feelings of isolation | Negative | All | 6 | Several individuals expressed significant emotional distress, primarily centered around feelings of
isolation and hopelessness. jamie’s withdrawal and reported emptiness suggest a deep-seated sense of
disconnection, while alex powerfully articulated his feeling of being misunderstood with the
statement “no one gets me.” these experiences appear to be impacting daily functioning, as evidenced
by jamie’s struggles and disruption of his sleep patterns.
the observations from parents further
indicat... | +| Overall mood | Agitation | Negative | All | 6 | Several individuals expressed concerns regarding heightened emotional distress within the subject
group, primarily focusing on alex. alex repeatedly demonstrated significant frustration and
aggression, necessitating anger management support, and appeared increasingly agitated in recent
meetings, suggesting a worsening emotional state. parents specifically highlighted jamie’s ongoing
struggles and agitation, emphasizing the need for a comprehensive assessment and subsequent
intervention.
while so... | +| Family dynamics | Peer influence | Negative | All | 5 | Concerns centered around alex’s social circle and behavior highlighted potential negative peer
influence, specifically due to his new friends and frequent late-night activities. further
investigation revealed troubling admissions from alex himself, including alcohol use and
participation in a physical altercation with a fellow student, indicating a concerning pattern of
risk-taking behavior.
simultaneously, jamie’s situation presented a separate area of concern,
characterized by isolation and l... | +| Substance use | Potential substance abuse | Negative | All | 5 | Alex has disclosed alcohol use, raising significant concerns about potential ongoing substance
abuse. the situation is further complicated by reports from alex’s mother, who has observed
potential signs of substance abuse and expressed her worries regarding this matter. these
observations highlight a need for further assessment and support to address the individual’s
substance use patterns and ensure their well-being.
the situation requires careful monitoring and
intervention. the mother’s repor... | +| Mental health | Depression | Positive | All | 3 | Jamie is diagnosed with major depressive disorder and initiated on antidepressant medication, with
initial positive feedback on mood and energy. | +| Mental health | Self-harm | Negative | All | 3 | The assessment revealed a complex picture regarding the individual’s mental state. while initial
observations did not indicate any active self-harm, a thorough evaluation is strongly recommended to
identify potential underlying issues contributing to the risk. this proactive approach is crucial
for a complete understanding of the individual’s needs.
more concerningly, alex presented with
visible self-harm indicators on his arms and explicitly communicated thoughts of self-harm,
signifying a sign... | +| School engagement | Absenteeism | Negative | All | 3 | Recent reports highlight a concerning trend of declining student engagement within the school
environment. specifically, there has been an increase in absences alongside a decrease in academic
performance, suggesting a fundamental lack of connection with schoolwork and learning. several
students, including jamie, are exhibiting problematic behaviors that further underscore this issue,
such as consistent tardiness and reduced participation in classroom activities.
furthermore,
observations indica... | +| Mental health | Depression | Negative | All | 2 | Jamie exhibits symptoms of moderate depression, requiring further evaluation and intervention. | +| Mental health | Self-harm | Neutral | All | 2 | The psychiatrist’s assessment centered on the potential advantages and drawbacks of antidepressant
medication, with a notable emphasis on evaluating the possibility of self-harm risk. this indicates
a proactive approach to patient safety and a recognition of the complex interplay between medication
and mental health. the discussion highlights a careful consideration of the potential for increased
suicidal ideation, suggesting a thorough risk assessment was undertaken.
furthermore, the
analysis o... | +| School engagement | Academic performance | Negative | All | 2 | Analysis of the provided text reveals concerns regarding student engagement and academic
performance. specifically, jamie’s reduced involvement in class is flagged as a potential indicator
of negative consequences, with declining grades reported as a direct result. this suggests a
concerning downward trend in alex’s academic progress, highlighting a need for further investigation
into the underlying causes of this shift.
the combined observations point to a possible
correlation between decreased... | +| Substance use | Substance use (unspecified) | Negative | All | 2 | Concerns regarding ongoing substance use prompted discussion about the possibility of a short-term
residential treatment program. alex’s involvement highlighted a potential issue, as they reported
occasional substance use, though the specific substances involved were not detailed during the
consultation. this lack of specificity regarding the substances used raises a need for further
investigation into the nature and frequency of alex’s substance use.
the consultation focused on
assessing the ri... | +| Family dynamics | Stepfather relationship | Negative | All | 1 | Alex displayed sudden outbursts of anger when discussing his new stepfather, indicating significant
distress related to this family change. | +| School engagement | Academic performance | Positive | All | 1 | Jamie's academic performance has slightly improved, indicating a potential positive change. |""" + +case_notes_table_grouped = """| General topic | Subtopic | Sentiment | Group | Number of responses | Revised summary | +|:--------------------|:---------------------------|:------------|:---------|----------------------:|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| Trends over time | Trends over time | Negative | Alex D. | 7 | Alex’s case note reveals a troubling deterioration in his well-being marked by a gradual escalation
of issues. Initially, the record details an incident involving a physical altercation, which quickly
spiraled into increasingly concerning behaviours at home, specifically escalating aggression. Over
subsequent meetings, observations consistently pointed towards heightened agitation and expressions
of hopelessness, indicating a worsening emotional state and a significant decline in his overall
con... | +| Physical health | Substance misuse | Negative | Alex D. | 6 | Alex’s substance use remains a significant concern, necessitating continued vigilance and support
despite recent positive developments in group therapy. While Alex has acknowledged instances of
substance use, the details surrounding these occurrences have not been shared, raising questions
about the extent and nature of the problem. Concerns were specifically noted regarding potential
substance abuse, highlighting a need for further investigation and assessment.

Ongoing
monitoring is crucial to... | +| Behaviour at school | Behaviour at school | Negative | Alex D. | 3 | A recent case note details a troubling incident involving a physical altercation at school,
alongside concerning admissions from Alex regarding alcohol use. This event has sparked worries
about potential behavioural issues within the school setting, suggesting a need for further
investigation and support. Alex’s demeanor was notably problematic, characterized by sullen behavior
and a deliberate avoidance of eye contact, indicating a possible struggle with emotional
regulation.

Furthermore, Alex... | +| Mental health | Anger | Negative | Alex D. | 3 | Alex exhibits a pronounced anger issue, characterized by frustration and a tendency to blame others
for triggering his aggressive behavior. He demonstrated this significantly when discussing his
personal life, particularly relating to his new stepfather, suggesting a volatile emotional response
to this change. The observed outbursts highlight a need for immediate intervention to manage his
escalating anger.

Further investigation reveals that Alex’s anger is closely linked to his
home environmen... | +| Mental health | Self-harm | Negative | Alex D. | 3 | The analysis reveals significant concerns regarding Alex’s mental health, centering around potential
self-harm behaviors. Indications suggest a possible diagnosis of Oppositional Defiant Disorder
alongside a co-occurring substance use disorder, warranting a comprehensive treatment plan. Alex
demonstrated visible signs of self-harm and openly confessed to experiencing thoughts of self-harm,
highlighting a critical need for immediate intervention.

Following this disclosure, an
immediate referral ... | +| Mental health | Social issues | Negative | Alex D. | 3 | Alex exhibits a pattern of blaming others for his problematic behavior, indicating underlying
challenges in social interaction and conflict resolution. This behavior appears to be contributing
to further instability in his life. Specifically, his mother voiced concerns regarding his new
social circle and increasingly frequent late-night activities, suggesting she perceives these
relationships and outings as potentially risky.

The mother’s observations highlight a
potential area of concern for A... | +| Mental health | Depression | Negative | Jamie L. | 6 | Jamie is currently experiencing concerning symptoms indicative of depression, as noted by both
Jamie’s behavior and parental observations. Specifically, he demonstrates limited social
interaction, struggles with his mood, and has difficulty engaging with his schoolwork. These
difficulties appear persistent, with parents reporting ongoing struggles despite occasional positive
moments.

Further assessment suggests a more pronounced picture, with indications of moderate
depression characterized by... | +| Mental health | Social isolation | Negative | Jamie L. | 4 | Jamie is experiencing significant social isolation, which is negatively affecting both his academic
performance and his general well-being. He has expressed feelings of loneliness and difficulty
sleeping, strongly suggesting a core social issue is contributing to his distress. Current efforts
are focused on promoting increased social interaction to address these challenges.

The report
highlights the urgency of this situation, emphasizing the need for intervention to mitigate Jamie’s
isolation a... | +| Mental health | Medication | Neutral | Jamie L. | 3 | Consideration is being given to medication as a potential intervention alongside therapy to manage
depressive symptoms. Initial feedback on the antidepressant is positive. | +| Mental health | Withdrawal & sadness | Negative | Jamie L. | 3 | Jamie is experiencing a significant downturn in his emotional state, characterized by withdrawal,
sadness, and a pervasive sense of emptiness and hopelessness. These negative feelings appear to be
triggered by recent reports of tardiness and decreased participation, suggesting a possible link
between his behavior and external pressures or expectations. The combination of these symptoms
points to a low mood and a feeling of struggle, indicating a potentially serious situation requiring
attention.... | +| Mental health | Low self-worth | Negative | Jamie L. | 2 | Parents are increasingly concerned about Jamie’s well-being due to observed difficulties and a
potential lack of self-worth. These concerns are primarily fueled by Jamie’s own statements, where
he articulated feelings of low self-esteem and a significant struggle to find
motivation.

Further investigation revealed a direct link between Jamie’s emotional state and
recent family financial hardships. The pressures of these struggles appear to have deeply impacted
his self-perception and ability to ... | +| Trends over time | Increasing withdrawal | Negative | Jamie L. | 2 | A significant and worrying trend is emerging regarding withdrawal, necessitating continuous
observation and targeted intervention strategies. Specifically, Jamie is exhibiting a noticeable
decline in engagement with family activities, representing a key indicator of this broader issue.
This withdrawal suggests a potential underlying problem requiring careful assessment and proactive
support.

The observed pattern of withdrawal highlights the importance of sustained monitoring
to understand its p... | +| Behaviour at school | Attendance issues | Negative | Jamie L. | 1 | Jamie’s consistent tardiness was a concern leading to a meeting. | +| Behaviour at school | Reduced participation | Negative | Jamie L. | 1 | Jamie’s decreased participation in class was noted. | +| Behaviour at school | Social engagement | Negative | Jamie L. | 1 | Jamie's withdrawal from family activities and hobbies was highlighted. | +| Behaviour at school | Social engagement | Positive | Jamie L. | 1 | Encouraging Jamie to join school clubs and groups is a strategy to foster social connection and
improve his social engagement. | +| Family & social | Family communication | Negative | Jamie L. | 1 | Parents expressed concerns about Jamie’s withdrawal and lack of communication within the family. | +| Family & social | Family communication | Neutral | Jamie L. | 1 | Parents are actively involved in Jamie's care and are communicating their observations to the care
team. | +| Family & social | Family financial struggles | Negative | Jamie L. | 1 | Jamie's low motivation is attributed to recent family financial difficulties. |""" + +case_notes_table_structured_summary = """| Main heading | Subheading | Summary | Group | +|:--------------------|:--------------------|:--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:---------| +| Behaviour at school | Behaviour at school | Several cases involved disruptions at school, including increased absences, declining grades, and a
physical altercation. Alex displayed sullenness, avoidance, and agitation, sometimes reacting with
frustration. A key theme was isolation and a lack of connection with peers and school staff. | Alex D. | +| Mental health | Anger | Anger was a prominent feature across multiple cases, particularly when discussing home life and
family dynamics. Outbursts of anger were observed, especially related to a new stepfather, and Alex
displayed defensiveness when questioned about his actions. | Alex D. | +| Mental health | Social issues | Alex experienced feelings of isolation and difficulty connecting with others. He had a new group of
friends and engaged in late-night outings, which raised concerns about potential risky behaviours
and social influences. | Alex D. | +| Physical health | General | Signs of self-harm were present on Alex’s arms, indicating a heightened level of distress and
potentially a need for immediate support. He displayed visible agitation and defensive behaviour
during questioning. | Alex D. | +| Physical health | Substance misuse | Substance use was a recurring concern, with Alex admitting to occasional substance use and his
mother reporting potential signs of abuse. Alcohol use was noted in several instances, leading to
recommendations for assessment and potential intervention. | Alex D. | +| Trends over time | Trends over time | There was a gradual escalation of concerning behaviours over time. Early interventions focused on
initial meetings and observation, progressing to more intensive interventions like referrals to
mental health professionals, residential treatment programs, and family counseling. | Alex D. | +| Behaviour at school | Behaviour at school | Jamie exhibited concerning behaviours at school, including consistent tardiness and decreased
participation in class. This was accompanied by withdrawn behaviour and signs of sadness, suggesting
a need for immediate intervention to address potential underlying issues impacting his academic
performance. | Jamie L. | +| Mental health | Anger | There is no direct indication of anger in Jamie's case notes. | Jamie L. | +| Mental health | Mental health | Jamie displayed concerning signs of mental health difficulties, including feelings of emptiness,
hopelessness, low self-worth, and isolation. He reported difficulty sleeping and a lack of
motivation. The need for a comprehensive mental health assessment was highlighted to fully
understand the nature and severity of his condition. | Jamie L. | +| Mental health | Social issues | Jamie experienced significant social difficulties, including limited social interactions, feelings
of isolation, and a lack of engagement with family activities and hobbies. He spends a lot of time
alone in his room. Recommendations focused on fostering connection through school clubs and family
therapy were made. | Jamie L. | +| Physical health | General | While no direct physical health concerns were explicitly stated, Jamie's emotional state and
associated symptoms (difficulty sleeping) warrant consideration of his overall well-being and
potential physical manifestations of his mental health challenges. | Jamie L. | +| Physical health | Substance misuse | There is no indication of substance misuse in the provided case notes. | Jamie L. | +| Trends over time | Trends over time | Jamie’s case demonstrates fluctuating progress. Initial feedback indicated slight improvements in
mood on some days, but overall he continues to struggle. A shift occurred with the commencement of
antidepressant medication, showing initial positive feedback in terms of mood and energy levels,
requiring continued monitoring and adjustment. | Jamie L. |""" diff --git a/tools/helper_functions.py b/tools/helper_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..7a71514180e621f4558f6d3b136427dcbf7838c7 --- /dev/null +++ b/tools/helper_functions.py @@ -0,0 +1,1686 @@ +import codecs +import math +import os +import re +from typing import List + +import boto3 +import gradio as gr +import numpy as np +import pandas as pd +from botocore.exceptions import ClientError + +from tools.config import ( + AWS_USER_POOL_ID, + CUSTOM_HEADER, + CUSTOM_HEADER_VALUE, + INPUT_FOLDER, + MAXIMUM_ALLOWED_TOPICS, + OUTPUT_FOLDER, + SESSION_OUTPUT_FOLDER, + model_full_names, + model_name_map, +) + + +class ColumnNotFoundError(ValueError): + """Raised when a requested column name is not present in the loaded dataset.""" + + +def empty_output_vars_extract_topics(): + # Empty output objects before processing a new file + + master_topic_df_state = pd.DataFrame() + master_topic_summary_df_state = pd.DataFrame() + master_reference_df_state = pd.DataFrame() + text_output_file = list() + text_output_file_list_state = list() + latest_batch_completed = 0 + log_files_output = list() + log_files_output_list_state = list() + conversation_metadata_textbox = "" + estimated_time_taken_number = 0 + file_data_state = pd.DataFrame() + reference_data_file_name_textbox = "" + display_topic_table_markdown = "" + summary_output_file_list = list() + summary_input_file_list = list() + overall_summarisation_input_files = list() + overall_summary_output_files = list() + + return ( + master_topic_df_state, + master_topic_summary_df_state, + master_reference_df_state, + text_output_file, + text_output_file_list_state, + latest_batch_completed, + log_files_output, + log_files_output_list_state, + conversation_metadata_textbox, + estimated_time_taken_number, + file_data_state, + reference_data_file_name_textbox, + display_topic_table_markdown, + summary_output_file_list, + summary_input_file_list, + overall_summarisation_input_files, + overall_summary_output_files, + ) + + +def empty_output_vars_summarise(): + # Empty output objects before summarising files + + summary_reference_table_sample_state = pd.DataFrame() + master_topic_summary_df_revised_summaries_state = pd.DataFrame() + master_reference_df_revised_summaries_state = pd.DataFrame() + summary_output_files = list() + summarised_outputs_list = list() + latest_summary_completed_num = 0 + overall_summarisation_input_files = list() + + return ( + summary_reference_table_sample_state, + master_topic_summary_df_revised_summaries_state, + master_reference_df_revised_summaries_state, + summary_output_files, + summarised_outputs_list, + latest_summary_completed_num, + overall_summarisation_input_files, + ) + + +def get_or_create_env_var(var_name: str, default_value: str): + # Get the environment variable if it exists + value = os.environ.get(var_name) + + # If it doesn't exist, set it to the default value + if value is None: + os.environ[var_name] = default_value + value = default_value + + return value + + +def get_file_path_with_extension(file_path: str): + # First, get the basename of the file (e.g., "example.txt" from "/path/to/example.txt") + basename = os.path.basename(file_path) + + # Return the basename with its extension + return basename + + +def get_file_name_no_ext(file_path: str): + # First, get the basename of the file (e.g., "example.txt" from "/path/to/example.txt") + basename = os.path.basename(file_path) + + # Then, split the basename and its extension and return only the basename without the extension + filename_without_extension, _ = os.path.splitext(basename) + + # print(filename_without_extension) + + return filename_without_extension + + +def detect_file_type(filename: str): + """Detect the file type based on its extension.""" + + # Strip quotes and whitespace that might have been accidentally included + filename = filename.strip().strip("'\"") + + if ( + (filename.endswith(".csv")) + | (filename.endswith(".csv.gz")) + | (filename.endswith(".zip")) + ): + return "csv" + elif filename.endswith(".xlsx"): + return "xlsx" + elif filename.endswith(".parquet"): + return "parquet" + elif filename.endswith(".pdf"): + return "pdf" + elif filename.endswith(".jpg"): + return "jpg" + elif filename.endswith(".jpeg"): + return "jpeg" + elif filename.endswith(".png"): + return "png" + else: + raise ValueError("Unsupported file type.") + + +def read_file(filename: str, sheet: str = ""): + """Read the file based on its detected type.""" + # Strip quotes and whitespace that might have been accidentally included + filename = filename.strip().strip("'\"") + file_type = detect_file_type(filename) + + if file_type == "csv": + # CSVs in UK local gov contexts are often Windows-1252 (smart quotes like 0x92), + # UTF-8, or UTF-8 with BOM. Try the common encodings first. + encodings_to_try = ["utf-8", "utf-8-sig", "cp1252", "latin1"] + last_err: Exception | None = None + for enc in encodings_to_try: + try: + return pd.read_csv(filename, low_memory=False, encoding=enc) + except UnicodeDecodeError as e: + last_err = e + continue + except Exception: + # Non-encoding errors should surface immediately + raise + + # Final fallback: force a load with replacement chars so the user can still proceed. + # (Better to show � than crash the whole app.) + try: + return pd.read_csv( + filename, + low_memory=False, + encoding="utf-8", + encoding_errors="replace", + ) + except Exception: + # Re-raise the original decode error with context + if last_err is not None: + raise last_err + raise UnicodeDecodeError("utf-8", b"", 0, 1, "Could not decode CSV") + elif file_type == "xlsx": + if sheet: + return pd.read_excel(filename, sheet_name=sheet) + else: + return pd.read_excel(filename) + elif file_type == "parquet": + return pd.read_parquet(filename) + + +def load_in_file(file_path: str, colnames: List[str] = "", excel_sheet: str = ""): + """ + Loads in a tabular data file and returns data and file name. + + Parameters: + - file_path (str): The path to the file to be processed. + - colnames (List[str], optional): list of colnames to load in + """ + + file_name = get_file_name_no_ext(file_path) + file_data = read_file(file_path, excel_sheet) + + if colnames and isinstance(colnames, list): + col_list = colnames + else: + col_list = list(file_data.columns) + + if not isinstance(col_list, List): + col_list = [col_list] + + col_list = [item for item in col_list if item not in ["", "NA"]] + + missing_cols = [col for col in col_list if col not in file_data.columns] + if missing_cols: + missing_label = ", ".join(f"'{col}'" for col in missing_cols) + available_cols = ", ".join(f"'{col}'" for col in file_data.columns) + raise ColumnNotFoundError( + "Column not found in data - aborting task: " + f"{missing_label} not found in file '{file_name}'. " + f"Available columns: {available_cols}" + ) + + for col in col_list: + file_data[col] = file_data[col].fillna("") + file_data[col] = ( + file_data[col].astype(str).str.replace("\bnan\b", "", regex=True) + ) + + # print(file_data[colnames]) + + return file_data, file_name + + +def load_in_data_file( + file_paths: List[str], + in_colnames: List[str], + batch_size: int = 5, + in_excel_sheets: str = "", +): + """Load in data table, work out how many batches needed.""" + + if not isinstance(in_colnames, list): + in_colnames = [in_colnames] + + # print("in_colnames:", in_colnames) + + try: + file_data, file_name = load_in_file( + file_paths[0], colnames=in_colnames, excel_sheet=in_excel_sheets + ) + num_batches = math.ceil(len(file_data) / batch_size) + print( + f"File {file_name} loaded successfully. Number of rows: {len(file_data)}. Total number of batches: {num_batches}" + ) + + except ColumnNotFoundError: + raise + except Exception as e: + print("Could not load data file due to:", e) + file_data = pd.DataFrame() + file_name = "" + num_batches = 1 + + return file_data, file_name, num_batches + + +def clean_column_name( + column_name: str, max_length: int = 20, front_characters: bool = True +): + # Convert to string + column_name = str(column_name) + # Replace non-alphanumeric characters (except underscores) with underscores + column_name = re.sub(r"\W+", "_", column_name) + # Remove leading/trailing underscores + column_name = column_name.strip("_") + # Ensure the result is not empty; fall back to "column" if necessary + column_name = column_name if column_name else "column" + # Truncate to max_length + if front_characters is True: + output_text = column_name[:max_length] + else: + output_text = column_name[-max_length:] + return output_text + + +def load_in_previous_reference_file(file: str): + """Load in data table from a partially completed consultation summary to continue it.""" + + reference_file_data = pd.DataFrame() + reference_file_name = "" + out_message = "" + + # for file in file_paths: + + print("file:", file) + + # If reference table + if "reference_table" in file: + try: + reference_file_data, reference_file_name = load_in_file(file) + # print("reference_file_data:", reference_file_data.head(2)) + out_message = out_message + " Response ID file load successful." + except Exception as e: + out_message = "Could not load reference file data:" + str(e) + raise Exception("Could not load reference file data:", e) + + if reference_file_data.empty: + out_message = out_message + " No reference data table provided." + raise Exception(out_message) + + print(out_message) + + return reference_file_data, reference_file_name + + +def load_in_previous_data_files( + file_paths_partial_output: List[str], for_modified_table: bool = False +): + """Load in data table from a partially completed consultation summary to continue it.""" + + reference_file_data = pd.DataFrame() + reference_file_name = "" + unique_file_data = pd.DataFrame() + unique_file_name = "" + out_message = "" + latest_batch = 0 + + if not file_paths_partial_output: + out_message = out_message + " No reference or unique data table provided." + return ( + reference_file_data, + unique_file_data, + latest_batch, + out_message, + reference_file_name, + unique_file_name, + ) + + if not isinstance(file_paths_partial_output, list): + file_paths_partial_output = [file_paths_partial_output] + + for file in file_paths_partial_output: + + if isinstance(file, gr.FileData): + name = file.name + else: + name = file + + # If reference table + if "reference_table" in name: + try: + reference_file_data, reference_file_name = load_in_file(file) + # print("reference_file_data:", reference_file_data.head(2)) + out_message = out_message + " Response ID file load successful." + + except Exception as e: + out_message = "Could not load reference file data:" + str(e) + raise Exception("Could not load reference file data:", e) + # If unique table + if "unique_topic" in name: + try: + unique_file_data, unique_file_name = load_in_file(file) + # print("unique_topics_file:", unique_file_data.head(2)) + out_message = out_message + " Unique table file load successful." + except Exception as e: + out_message = "Could not load unique table file data:" + str(e) + raise Exception("Could not load unique table file data:", e) + if "batch_" in name: + latest_batch = re.search(r"batch_(\d+)", file.name).group(1) + print("latest batch:", latest_batch) + latest_batch = int(latest_batch) + + if latest_batch == 0: + out_message = out_message + " Latest batch number not found." + if reference_file_data.empty: + out_message = out_message + " No reference data table provided." + # raise Exception(out_message) + if unique_file_data.empty: + out_message = out_message + " No unique data table provided." + + print(out_message) + + # Return all data if using for deduplication task. Return just modified unique table if using just for table modification + if for_modified_table is False: + return ( + reference_file_data, + unique_file_data, + latest_batch, + out_message, + reference_file_name, + unique_file_name, + ) + else: + reference_file_data.drop("Topic number", axis=1, inplace=True, errors="ignore") + + unique_file_data = create_topic_summary_df_from_reference_table( + reference_file_data + ) + + unique_file_data.drop("Summary", axis=1, inplace=True) + + # Then merge the topic numbers back to the original dataframe + reference_file_data = reference_file_data.merge( + unique_file_data[ + ["General topic", "Subtopic", "Sentiment", "Topic number"] + ], + on=["General topic", "Subtopic", "Sentiment"], + how="left", + ) + + out_file_names = [reference_file_name + ".csv"] + out_file_names.append(unique_file_name + ".csv") + + return ( + unique_file_data, + reference_file_data, + unique_file_data, + reference_file_name, + unique_file_name, + out_file_names, + ) # gr.Dataframe(value=unique_file_data, headers=None, column_count=(unique_file_data.shape[1], "fixed"), row_count = (unique_file_data.shape[0], "fixed"), visible=True, type="pandas") + + +def join_cols_onto_reference_df( + reference_df: pd.DataFrame, + original_data_df: pd.DataFrame, + join_columns: List[str], + original_file_name: str, + output_folder: str = OUTPUT_FOLDER, +): + + # print("original_data_df columns:", original_data_df.columns) + # print("original_data_df:", original_data_df) + + original_data_df.reset_index(names="Response ID", inplace=True) + original_data_df["Response ID"] += 1 + + # print("reference_df columns:", reference_df.columns) + # print("reference_df:", reference_df) + + join_columns.append("Response ID") + + reference_df["Response ID"] = reference_df["Response ID"].fillna("-1").astype(int) + + save_file_name = output_folder + original_file_name + "_j.csv" + + out_reference_df = reference_df.merge( + original_data_df[join_columns], on="Response ID", how="left" + ) + out_reference_df.to_csv(save_file_name, index=None) + + file_data_outputs = [save_file_name] + + return out_reference_df, file_data_outputs + + +def get_basic_response_data( + file_data: pd.DataFrame, chosen_cols: List[str], verify_titles: bool = False +) -> pd.DataFrame: + + if not isinstance(chosen_cols, list): + chosen_cols = [chosen_cols] + + if chosen_cols[0] not in file_data.columns: + error_msg = ( + f"Column '{chosen_cols[0]}' not found in file_data columns. " + f"Available columns: {list(file_data.columns)}" + ) + print(error_msg) + raise KeyError(error_msg) + + # If verify_titles is True, we need to check and include the second column + if verify_titles is True: + if len(chosen_cols) < 2: + error_msg = ( + "verify_titles is True but only one column provided. " + "Need at least 2 columns: one for response text and one for title." + ) + print(error_msg) + raise ValueError(error_msg) + if chosen_cols[1] not in file_data.columns: + error_msg = ( + f"Column '{chosen_cols[1]}' not found in file_data columns for title. " + f"Available columns: {list(file_data.columns)}" + ) + print(error_msg) + raise KeyError(error_msg) + # Include both columns when verify_titles is True + basic_response_data = file_data[[chosen_cols[0], chosen_cols[1]]] + basic_response_data = basic_response_data.rename( + columns={ + basic_response_data.columns[0]: "Response", + basic_response_data.columns[1]: "Title", + } + ) + else: + basic_response_data = file_data[[chosen_cols[0]]] + basic_response_data = basic_response_data.rename( + columns={basic_response_data.columns[0]: "Response"} + ) + basic_response_data = basic_response_data.reset_index( + names="Original Response ID" + ) # .reset_index(drop=True) # + # Try to convert to int, if it fails, return a range of 1 to last row + 1 + try: + basic_response_data["Original Response ID"] = ( + basic_response_data["Original Response ID"].astype(int) + 1 + ) + except (ValueError, TypeError): + basic_response_data["Original Response ID"] = range( + 1, len(basic_response_data) + 1 + ) + + basic_response_data["Response ID"] = basic_response_data.index.astype(int) + 1 + + if verify_titles is True: + basic_response_data["Title"] = basic_response_data["Title"].str.strip() + basic_response_data["Title"] = basic_response_data["Title"].apply(initial_clean) + else: + basic_response_data = basic_response_data[ + ["Response ID", "Response", "Original Response ID"] + ] + + basic_response_data["Response"] = basic_response_data["Response"].str.strip() + basic_response_data["Response"] = basic_response_data["Response"].apply( + initial_clean + ) + + return basic_response_data + + +def convert_reference_table_to_pivot_table( + df: pd.DataFrame, basic_response_data: pd.DataFrame = pd.DataFrame() +): + df = df.copy() + if "Sentiment" not in df.columns: + df["Sentiment"] = "Not assessed" + + df_in = df[["Response ID", "General topic", "Subtopic", "Sentiment"]].copy() + + # Convert to numeric first (handles float strings like '1.0'), then to int + df_in["Response ID"] = pd.to_numeric(df_in["Response ID"], errors="coerce").astype( + "Int64" + ) + + # Create a combined category column + df_in["Category"] = ( + df_in["General topic"] + " - " + df_in["Subtopic"] + " - " + df_in["Sentiment"] + ) + + # Create pivot table counting occurrences of each unique combination + pivot_table = pd.crosstab( + index=df_in["Response ID"], + columns=[df_in["General topic"], df_in["Subtopic"], df_in["Sentiment"]], + margins=True, + ) + + # Flatten column names to make them more readable + pivot_table.columns = [" - ".join(col) for col in pivot_table.columns] + + pivot_table.reset_index(inplace=True) + + if not basic_response_data.empty: + pivot_table = basic_response_data.merge( + pivot_table, right_on="Response ID", left_on="Response ID", how="left" + ) + + pivot_table.drop("Response ID", axis=1, inplace=True) + + pivot_table.columns = pivot_table.columns.str.replace( + "Not assessed - ", "" + ).str.replace("- Not assessed", "") + + leading_cols = [ + col + for col in ["Original Response ID", "Response"] + if col in pivot_table.columns + ] + if leading_cols: + other_cols = [col for col in pivot_table.columns if col not in leading_cols] + pivot_table = pivot_table[leading_cols + other_cols] + + return pivot_table + + +def create_topic_summary_df_from_reference_table( + reference_df: pd.DataFrame, + sentiment_checkbox: str = "Negative, Neutral, or Positive", +): + + if "Group" not in reference_df.columns: + reference_df["Group"] = "All" + + # Ensure 'Start row of group' column is numeric to avoid comparison errors + if "Start row of group" in reference_df.columns: + reference_df["Start row of group"] = pd.to_numeric( + reference_df["Start row of group"], errors="coerce" + ) + + # Determine if sentiment should be included in grouping + include_sentiment = sentiment_checkbox != "Do not assess sentiment" + + # Define grouping columns based on whether sentiment is being assessed + # When sentiment is assessed: group by General topic, Subtopic, Sentiment, and Group + # When sentiment is NOT assessed: group by General topic, Subtopic, and Group + # In both cases, duplicate rows with the same combination will be merged: + # - Response ID will be concatenated into a comma-separated list + # - Summaries will be concatenated with "
" separators + if include_sentiment: + groupby_cols = ["General topic", "Subtopic", "Sentiment", "Group"] + else: + groupby_cols = ["General topic", "Subtopic", "Group"] + + # Helper function to collect and join Response ID + # This aggregates all unique Response ID from rows with the same topic combination + def aggregate_response_references(x): + # Collect all unique Response ID, convert to string, sort numerically + refs = set() + for ref in x: + if pd.notna(ref) and str(ref).strip(): + # Handle both single numbers and comma-separated lists + ref_str = str(ref).strip() + # Split by comma if it's a comma-separated list + for r in ref_str.split(","): + r = r.strip() + if r and r.isdigit(): + refs.add(int(r)) + # Sort and join with ", " + sorted_refs = sorted(refs) + return ", ".join(map(str, sorted_refs)) if sorted_refs else "" + + # Helper function to concatenate summaries + # This aggregates all unique summaries from rows with the same topic combination + def aggregate_summaries(x): + # First, collect all summary segments (handling both single summaries and already-aggregated summaries) + all_segments = [] + + for summary in x: + if pd.notna(summary): + summary_str = str(summary).strip() + if summary_str: + # If summary already contains
separators, split it into segments + if "
" in summary_str or "
" in summary_str: + segments = re.split(r"
|
", summary_str) + segments = [seg.strip() for seg in segments if seg.strip()] + all_segments.extend(segments) + else: + # Single summary + all_segments.append(summary_str) + + # Remove duplicate segments while preserving order + unique_segments = [] + seen_segments = set() + for segment in all_segments: + if segment and segment not in seen_segments: + seen_segments.add(segment) + unique_segments.append(segment) + + # Sort by minimum Start row of group if available (try to match segments to original summaries) + if "Start row of group" in reference_df.columns and unique_segments: + try: + # Create a mapping of segments to their minimum start row + segment_to_start_row = {} + for segment in unique_segments: + # Find rows where Summary contains this segment + matching_rows = reference_df[ + reference_df["Summary"].str.contains( + segment, na=False, regex=False + ) + ] + if ( + not matching_rows.empty + and "Start row of group" in matching_rows.columns + ): + min_start_row = matching_rows["Start row of group"].min() + segment_to_start_row[segment] = min_start_row + else: + segment_to_start_row[segment] = float("inf") + + # Sort by minimum start row + unique_segments = sorted( + unique_segments, + key=lambda seg: segment_to_start_row.get(seg, float("inf")), + ) + except Exception: + # If sorting fails, just use the order we have + pass + + return "
".join(unique_segments) if unique_segments else "" + + out_topic_summary_df = ( + reference_df.groupby(groupby_cols) + .agg( + { + "Response ID": aggregate_response_references, + "Summary": aggregate_summaries, + } + ) + .reset_index() + ) + + # Calculate number of responses from the Response ID string + # (count comma-separated values) + def count_responses(ref_str): + if not ref_str or str(ref_str).strip() == "": + return 0 + return len([r for r in str(ref_str).split(",") if r.strip()]) + + # Store Response ID before potentially renaming + response_refs_col = out_topic_summary_df["Response ID"].copy() + + # Calculate Number of responses from Response ID + out_topic_summary_df["Number of responses"] = response_refs_col.apply( + count_responses + ) + + # For backward compatibility with code that expects "Number of responses" instead of "Response ID", + # we keep both columns. The "Response ID" column contains the concatenated reference numbers, + # and "Number of responses" contains the count. + + # Sort the dataframe first + sort_cols = ["Group", "Number of responses", "General topic", "Subtopic"] + if include_sentiment: + sort_cols.append("Sentiment") + out_topic_summary_df = out_topic_summary_df.sort_values( + sort_cols, + ascending=( + [True, False, True, True, True] + if include_sentiment + else [True, False, True, True] + ), + ) + + # Then assign Topic number based on the final sorted order + out_topic_summary_df = out_topic_summary_df.assign( + Topic_number=lambda df: np.arange(1, len(df) + 1) + ) + + out_topic_summary_df.rename(columns={"Topic_number": "Topic number"}, inplace=True) + + out_topic_summary_df.drop( + ["1", "2", "3", "Response ID"], axis=1, errors="ignore", inplace=True + ) + + return out_topic_summary_df + + +# Wrap text in each column to the specified max width, including whole words +def wrap_text(text: str, max_width=50, max_text_length=None): + if not isinstance(text, str): + return text + + # If max_text_length is set, truncate the text and add ellipsis + if max_text_length and len(text) > max_text_length: + text = text[:max_text_length] + "..." + + text = text.replace("\r\n", "
").replace("\n", "
") + + words = text.split() + if not words: + return text + + # First pass: initial word wrapping + wrapped_lines = list() + current_line = list() + current_length = 0 + + def add_line(): + if current_line: + wrapped_lines.append(" ".join(current_line)) + current_line.clear() + + for i, word in enumerate(words): + word_length = len(word) + + # Handle words longer than max_width + if word_length > max_width: + add_line() + wrapped_lines.append(word) + current_length = 0 + continue + + # Calculate space needed for this word + space_needed = word_length if not current_line else word_length + 1 + + # Check if adding this word would exceed max_width + if current_length + space_needed > max_width: + add_line() + current_line.append(word) + current_length = word_length + else: + current_line.append(word) + current_length += space_needed + + add_line() # Add any remaining text + + # Second pass: redistribute words from lines following single-word lines + def can_fit_in_previous_line(prev_line, word): + return len(prev_line) + 1 + len(word) <= max_width + + i = 0 + while i < len(wrapped_lines) - 1: + words_in_line = wrapped_lines[i].split() + next_line_words = wrapped_lines[i + 1].split() + + # If current line has only one word and isn't too long + if len(words_in_line) == 1 and len(words_in_line[0]) < max_width * 0.8: + # Try to bring words back from the next line + words_to_bring_back = list() + remaining_words = list() + current_length = len(words_in_line[0]) + + for word in next_line_words: + if current_length + len(word) + 1 <= max_width: + words_to_bring_back.append(word) + current_length += len(word) + 1 + else: + remaining_words.append(word) + + if words_to_bring_back: + # Update current line with additional words + wrapped_lines[i] = " ".join(words_in_line + words_to_bring_back) + + # Update next line with remaining words + if remaining_words: + wrapped_lines[i + 1] = " ".join(remaining_words) + else: + wrapped_lines.pop(i + 1) + continue # Don't increment i if we removed a line + i += 1 + + return "
".join(wrapped_lines) + + +def normalize_topic_name_for_llm(name): + """ + Normalize topic name for LLM presentation and consistent formatting. + + This function applies consistent cleaning and normalization to topic names + across the codebase to ensure uniform formatting. It handles: + - Cleaning with initial_clean + - Removing newlines and carriage returns + - Replacing special characters (/ with " or ", & with " and ") + - Normalizing apostrophes to standard apostrophe (') + - Normalizing whitespace + - Converting to title case for better readability + + Args: + name: Topic name to normalize (can be string, NaN, or None) + + Returns: + Normalized topic name string, or empty string if input is NaN/None + """ + if pd.isna(name) or name is None: + return "" + normalized = str(name) + normalized = initial_clean(normalized) + normalized = normalized.strip() + normalized = normalized.replace("\n", " ") + normalized = normalized.replace("\r", " ") + normalized = normalized.replace("/", " or ") + normalized = normalized.replace("&", " and ") + normalized = normalized.replace(" s ", "s ") + # Normalize different types of apostrophes to standard apostrophe (') + # Replace various Unicode apostrophe/quotation mark characters with standard apostrophe + normalized = normalized.replace("'", "'") # U+2019 right single quotation mark + normalized = normalized.replace("'", "'") # U+2018 left single quotation mark + normalized = normalized.replace( + "`", "'" + ) # U+2018 left single quotation mark (backtick style) + normalized = normalized.replace("´", "'") # U+00B4 acute accent + normalized = normalized.replace("'", "'") # U+02BC modifier letter apostrophe + normalized = normalized.lower() + # Capitalize first letter of each word for better readability + normalized = normalized.title() + return normalized + + +def initial_clean(text: str): + #### Some of my cleaning functions + html_pattern_regex = r"<.*?>|&([a-z0-9]+|#[0-9]{1,6}|#x[0-9a-f]{1,6});|\xa0| " + html_start_pattern_end_dots_regex = r"<(.*?)\.\." + non_ascii_pattern = r"[^\x00-\x7F]+" + multiple_spaces_regex = r"\s{2,}" + + # Markdown formatting patterns - remove formatting but preserve text content + # Bold: **text** or __text__ + markdown_bold_double_asterisk = r"\*\*([^*]+?)\*\*" + markdown_bold_double_underscore = r"__([^_]+?)__" + # Italic: *text* or _text_ (but be careful not to match asterisks/underscores in words) + markdown_italic_asterisk = r"(? 0: + return ( + gr.Dropdown(choices=concat_choices, value=concat_choices[0]), + gr.Dropdown( + choices=all_sheet_names, + value=all_sheet_names[0], + visible=True, + interactive=True, + ), + file_end, + gr.Dropdown(choices=concat_choices), + gr.Dropdown(choices=concat_choices), + ) + else: + return ( + gr.Dropdown(choices=concat_choices, value=concat_choices[0]), + gr.Dropdown(visible=False), + file_end, + gr.Dropdown(choices=concat_choices), + gr.Dropdown(choices=concat_choices), + ) + + +# Following function is only relevant for locally-created executable files based on this app (when using pyinstaller it creates a _internal folder that contains tesseract and poppler. These need to be added to the system path to enable the app to run) +def add_folder_to_path(folder_path: str): + """ + Check if a folder exists on your system. If so, get the absolute path and then add it to the system Path variable if it doesn't already exist. + """ + + if os.path.exists(folder_path) and os.path.isdir(folder_path): + print(folder_path, "folder exists.") + + # Resolve relative path to absolute path + absolute_path = os.path.abspath(folder_path) + + current_path = os.environ["PATH"] + if absolute_path not in current_path.split(os.pathsep): + full_path_extension = absolute_path + os.pathsep + current_path + os.environ["PATH"] = full_path_extension + # print(f"Updated PATH with: ", full_path_extension) + else: + print(f"Directory {folder_path} already exists in PATH.") + else: + print(f"Folder not found at {folder_path} - not added to PATH") + + +# Upon running a process, the feedback buttons are revealed +def reveal_feedback_buttons(): + return ( + gr.Radio(visible=True), + gr.Textbox(visible=True), + gr.Button(visible=True), + gr.Markdown(visible=True), + ) + + +def wipe_logs(feedback_logs_loc: str, usage_logs_loc: str): + try: + os.remove(feedback_logs_loc) + except Exception as e: + print("Could not remove feedback logs file", e) + try: + os.remove(usage_logs_loc) + except Exception as e: + print("Could not remove usage logs file", e) + + +async def get_connection_params( + request: gr.Request, + output_folder_textbox: str = OUTPUT_FOLDER, + input_folder_textbox: str = INPUT_FOLDER, + session_output_folder: str = SESSION_OUTPUT_FOLDER, +): + + # print("Session hash:", request.session_hash) + + if CUSTOM_HEADER and CUSTOM_HEADER_VALUE: + if CUSTOM_HEADER in request.headers: + supplied_custom_header_value = request.headers[CUSTOM_HEADER] + if supplied_custom_header_value == CUSTOM_HEADER_VALUE: + print("Custom header supplied and matches CUSTOM_HEADER_VALUE") + else: + print("Custom header value does not match expected value.") + raise ValueError("Custom header value does not match expected value.") + else: + print("Custom header value not found.") + raise ValueError("Custom header value not found.") + + # Get output save folder from 1 - username passed in from direct Cognito login, 2 - Cognito ID header passed through a Lambda authenticator, 3 - the session hash. + + if request.username: + out_session_hash = request.username + # print("Request username found:", out_session_hash) + + elif "x-cognito-id" in request.headers: + out_session_hash = request.headers["x-cognito-id"] + # print("Cognito ID found:", out_session_hash) + + elif "x-amzn-oidc-identity" in request.headers: + out_session_hash = request.headers["x-amzn-oidc-identity"] + + # Fetch email address using Cognito client + cognito_client = boto3.client("cognito-idp") + try: + response = cognito_client.admin_get_user( + UserPoolId=AWS_USER_POOL_ID, # Replace with your User Pool ID + Username=out_session_hash, + ) + email = next( + attr["Value"] + for attr in response["UserAttributes"] + if attr["Name"] == "email" + ) + # print("Email address found:", email) + + out_session_hash = email + except ClientError as e: + print("Error fetching user details:", e) + email = None + + print("Cognito ID found:", out_session_hash) + + else: + out_session_hash = request.session_hash + + if session_output_folder == "True" or session_output_folder is True: + output_folder = output_folder_textbox + out_session_hash + "/" + input_folder = input_folder_textbox + out_session_hash + "/" + else: + output_folder = output_folder_textbox + input_folder = input_folder_textbox + + if not os.path.exists(output_folder): + os.mkdir(output_folder) + if not os.path.exists(input_folder): + os.mkdir(input_folder) + + return out_session_hash, output_folder, out_session_hash, input_folder + + +def load_in_default_cost_codes(cost_codes_path: str, default_cost_code: str = ""): + """ + Load in the cost codes list from file. + """ + cost_codes_df = pd.read_csv(cost_codes_path) + dropdown_choices = cost_codes_df.iloc[:, 0].astype(str).tolist() + + # Avoid inserting duplicate or empty cost code values + if default_cost_code and default_cost_code not in dropdown_choices: + dropdown_choices.insert(0, default_cost_code) + + # Always have a blank option at the top + if "" not in dropdown_choices: + dropdown_choices.insert(0, "") + + out_dropdown = gr.Dropdown( + value=default_cost_code if default_cost_code in dropdown_choices else "", + label="Choose cost code for analysis", + choices=dropdown_choices, + allow_custom_value=False, + ) + + return cost_codes_df, cost_codes_df, out_dropdown + + +def df_select_callback_cost(df: pd.DataFrame, evt: gr.SelectData): + row_value_code = evt.row_value[0] # This is the value for cost code + + return row_value_code + + +def update_cost_code_dataframe_from_dropdown_select( + cost_dropdown_selection: str, cost_code_df: pd.DataFrame +): + cost_code_df = cost_code_df.loc[ + cost_code_df.iloc[:, 0] == cost_dropdown_selection, : + ] + return cost_code_df + + +def reset_base_dataframe(df: pd.DataFrame): + return df + + +def enforce_cost_codes( + enforce_cost_code_textbox: str, + cost_code_choice: str, + cost_code_df: pd.DataFrame, + verify_cost_codes: bool = True, +): + """ + Check if the enforce cost codes variable is set to true, and then check that a cost cost has been chosen. If not, raise an error. Then, check against the values in the cost code dataframe to ensure that the cost code exists. + """ + + if enforce_cost_code_textbox == "True": + if not cost_code_choice: + raise Exception("Please choose a cost code before continuing") + + if verify_cost_codes is True: + if cost_code_df.empty: + # Warn but don't block - cost code is still required above + print( + "Warning: Cost code dataframe is empty. Verification skipped. Please ensure cost codes are loaded for full validation." + ) + else: + valid_cost_codes_list = list(cost_code_df.iloc[:, 0].unique()) + + if cost_code_choice not in valid_cost_codes_list: + raise Exception( + "Selected cost code not found in list. Please contact Finance if you cannot find the correct cost code from the given list of suggestions." + ) + return + + +def _get_env_list(env_var_name: str, strip_strings: bool = True) -> List[str]: + """Parses a comma-separated environment variable into a list of strings.""" + value = env_var_name[1:-1].strip().replace('"', "").replace("'", "") + if not value: + return [] + # Split by comma and filter out any empty strings that might result from extra commas + if strip_strings: + return [s.strip() for s in value.split(",") if s.strip()] + else: + return [codecs.decode(s, "unicode_escape") for s in value.split(",") if s] + + +def create_batch_file_path_details( + reference_data_file_name: str, + latest_batch_completed: int = None, + batch_size_number: int = None, + in_column: str = None, +) -> str: + """ + Creates a standardised batch file path detail string from a reference data filename. + + Args: + reference_data_file_name (str): Name of the reference data file + latest_batch_completed (int, optional): Latest batch completed. Defaults to None. + batch_size_number (int, optional): Batch size number. Defaults to None. + in_column (str, optional): In column. Defaults to None. + Returns: + str: Formatted batch file path detail string + """ + + # Extract components from filename using regex + file_name = ( + re.search( + r"(.*?)(?:_all_|_final_|_batch_|_col_)", reference_data_file_name + ).group(1) + if re.search(r"(.*?)(?:_all_|_final_|_batch_|_col_)", reference_data_file_name) + else reference_data_file_name + ) + latest_batch_completed = ( + int(re.search(r"batch_(\d+)_", reference_data_file_name).group(1)) + if "batch_" in reference_data_file_name + else latest_batch_completed + ) + batch_size_number = ( + int(re.search(r"size_(\d+)_", reference_data_file_name).group(1)) + if "size_" in reference_data_file_name + else batch_size_number + ) + in_column = ( + re.search(r"col_(.*?)_reference", reference_data_file_name).group(1) + if "col_" in reference_data_file_name + else in_column + ) + + # Clean the extracted names + file_name_cleaned = clean_column_name(file_name, max_length=20) + in_column_cleaned = clean_column_name(in_column, max_length=20) + + # Create batch file path details string + if latest_batch_completed: + return f"{file_name_cleaned}_batch_{latest_batch_completed}_size_{batch_size_number}_col_{in_column_cleaned}" + return f"{file_name_cleaned}_col_{in_column_cleaned}" + + +def move_overall_summary_output_files_to_front_page( + overall_summary_output_files_xlsx: List[str], +): + return overall_summary_output_files_xlsx + + +def generate_zero_shot_topics_df( + zero_shot_topics: pd.DataFrame, + force_zero_shot_radio: str = "No", + create_revised_general_topics: bool = False, + max_topic_no: int = MAXIMUM_ALLOWED_TOPICS, +): + """ + Preprocesses a DataFrame of zero-shot topics, cleaning and formatting them + for use with a large language model. It handles different column configurations + (e.g., only subtopics, general topics and subtopics, or subtopics with descriptions) + and enforces a maximum number of topics. + + Args: + zero_shot_topics (pd.DataFrame): A DataFrame containing the initial zero-shot topics. + Expected columns can vary, but typically include + "General topic", "Subtopic", and/or "Description". + force_zero_shot_radio (str, optional): A string indicating whether to force + the use of zero-shot topics. Defaults to "No". + (Currently not used in the function logic, but kept for signature consistency). + create_revised_general_topics (bool, optional): A boolean indicating whether to + create revised general topics. Defaults to False. + (Currently not used in the function logic, but kept for signature consistency). + max_topic_no (int, optional): The maximum number of topics allowed to fit within + LLM context limits. If `zero_shot_topics` exceeds this, + it will be truncated. Defaults to 120. + + Returns: + tuple: A tuple containing: + - zero_shot_topics_gen_topics_list (list): A list of cleaned general topics. + - zero_shot_topics_subtopics_list (list): A list of cleaned subtopics. + - zero_shot_topics_description_list (list): A list of cleaned topic descriptions. + """ + + def _has_column_starting_with(columns, prefix): + """Check if any column starts with the given prefix.""" + return any(str(col).lower().startswith(prefix.lower()) for col in columns) + + def _get_column_starting_with(columns, prefix): + """Get the first column that starts with the given prefix, or None if not found.""" + for col in columns: + if str(col).lower().startswith(prefix.lower()): + return col + return None + + def _has_all_columns_starting_with(columns, prefixes): + """Check if there are columns starting with all given prefixes.""" + return all(_has_column_starting_with(columns, prefix) for prefix in prefixes) + + zero_shot_topics_gen_topics_list = list() + zero_shot_topics_subtopics_list = list() + zero_shot_topics_description_list = list() + + # Check if input exceeds maximum number of topics allowed + if zero_shot_topics.shape[0] > max_topic_no: + out_message = ( + "Maximum " + + str(max_topic_no) + + " zero-shot topics allowed according to application configuration." + ) + print(out_message) + raise Exception(out_message) + + # Forward slashes in the topic names seems to confuse the model + if zero_shot_topics.shape[1] >= 1: # Check if there is at least one column + for x in zero_shot_topics.columns: + if not zero_shot_topics[x].isnull().all(): + zero_shot_topics[x] = zero_shot_topics[x].apply( + normalize_topic_name_for_llm + ) + + zero_shot_topics.columns = zero_shot_topics.columns.str.lower() + + # If number of columns is 1, keep only subtopics + if zero_shot_topics.shape[1] == 1 and not _has_column_starting_with( + zero_shot_topics.columns, "general topic" + ): + print("Found only Subtopic in zero shot topics") + zero_shot_topics_gen_topics_list = [""] * zero_shot_topics.shape[0] + zero_shot_topics_subtopics_list = list(zero_shot_topics.iloc[:, 0]) + # Allow for possibility that the user only wants to set general topics and not subtopics + elif zero_shot_topics.shape[1] == 1 and _has_column_starting_with( + zero_shot_topics.columns, "general topic" + ): + print("Found only General topic in zero shot topics") + gen_topic_col = _get_column_starting_with( + zero_shot_topics.columns, "general topic" + ) + zero_shot_topics_gen_topics_list = list(zero_shot_topics[gen_topic_col]) + zero_shot_topics_subtopics_list = [""] * zero_shot_topics.shape[0] + # If general topic, subtopic and description are specified + elif _has_all_columns_starting_with( + zero_shot_topics.columns, ["general topic", "subtopic", "description"] + ): + print("Found General topic, Subtopic and Description in zero shot topics") + gen_topic_col = _get_column_starting_with( + zero_shot_topics.columns, "general topic" + ) + subtopic_col = _get_column_starting_with( + zero_shot_topics.columns, "subtopic" + ) + desc_col = _get_column_starting_with( + zero_shot_topics.columns, "description" + ) + zero_shot_topics_gen_topics_list = list(zero_shot_topics[gen_topic_col]) + zero_shot_topics_subtopics_list = list(zero_shot_topics[subtopic_col]) + zero_shot_topics_description_list = list(zero_shot_topics[desc_col]) + # If general topic and subtopic are specified + elif _has_all_columns_starting_with( + zero_shot_topics.columns, ["general topic", "subtopic"] + ) and not _has_column_starting_with(zero_shot_topics.columns, "description"): + print("Found General topic and Subtopic in zero shot topics") + gen_topic_col = _get_column_starting_with( + zero_shot_topics.columns, "general topic" + ) + subtopic_col = _get_column_starting_with( + zero_shot_topics.columns, "subtopic" + ) + zero_shot_topics_gen_topics_list = list(zero_shot_topics[gen_topic_col]) + zero_shot_topics_subtopics_list = list(zero_shot_topics[subtopic_col]) + # If subtopic and description are specified + elif _has_all_columns_starting_with( + zero_shot_topics.columns, ["subtopic", "description"] + ) and not _has_column_starting_with(zero_shot_topics.columns, "general topic"): + print("Found Subtopic and Description in zero shot topics") + zero_shot_topics_gen_topics_list = [""] * zero_shot_topics.shape[0] + subtopic_col = _get_column_starting_with( + zero_shot_topics.columns, "subtopic" + ) + desc_col = _get_column_starting_with( + zero_shot_topics.columns, "description" + ) + zero_shot_topics_subtopics_list = list(zero_shot_topics[subtopic_col]) + zero_shot_topics_description_list = list(zero_shot_topics[desc_col]) + elif _has_all_columns_starting_with( + zero_shot_topics.columns, ["general topic", "description"] + ) and not _has_column_starting_with(zero_shot_topics.columns, "subtopic"): + print("Found General topic and Description in zero shot topics") + gen_topic_col = _get_column_starting_with( + zero_shot_topics.columns, "general topic" + ) + desc_col = _get_column_starting_with( + zero_shot_topics.columns, "description" + ) + zero_shot_topics_gen_topics_list = list(zero_shot_topics[gen_topic_col]) + zero_shot_topics_subtopics_list = [""] * zero_shot_topics.shape[0] + zero_shot_topics_description_list = list(zero_shot_topics[desc_col]) + + # If number of columns is at least 2, keep general topics and subtopics + # (only if named columns don't exist) + elif ( + zero_shot_topics.shape[1] == 2 + and not _has_column_starting_with(zero_shot_topics.columns, "description") + and not _has_column_starting_with(zero_shot_topics.columns, "general topic") + and not _has_column_starting_with(zero_shot_topics.columns, "subtopic") + ): + zero_shot_topics_gen_topics_list = list(zero_shot_topics.iloc[:, 0]) + zero_shot_topics_subtopics_list = list(zero_shot_topics.iloc[:, 1]) + # If number of columns is at least 3, assume general topics, subtopics and descriptions were intended for the first three columns + elif ( + zero_shot_topics.shape[1] >= 3 + and not _has_column_starting_with(zero_shot_topics.columns, "general topic") + and not _has_column_starting_with(zero_shot_topics.columns, "subtopic") + and not _has_column_starting_with(zero_shot_topics.columns, "description") + ): + zero_shot_topics_gen_topics_list = list(zero_shot_topics.iloc[:, 0]) + zero_shot_topics_subtopics_list = list(zero_shot_topics.iloc[:, 1]) + zero_shot_topics_description_list = list(zero_shot_topics.iloc[:, 2]) + else: + # If there are more columns, just assume that the first column was meant to be a subtopic + zero_shot_topics_gen_topics_list = [""] * zero_shot_topics.shape[0] + zero_shot_topics_subtopics_list = list(zero_shot_topics.iloc[:, 0]) + zero_shot_topics_description_list = [""] * zero_shot_topics.shape[0] + + # Add a description if column is present and not already added above + if not zero_shot_topics_description_list: + zero_shot_topics_description_list = [""] * zero_shot_topics.shape[0] + + # If the responses are being forced into zero shot topics, allow an option for nothing relevant + if force_zero_shot_radio == "Yes": + zero_shot_topics_gen_topics_list.append("") + zero_shot_topics_subtopics_list.append("No relevant topic") + zero_shot_topics_description_list.append("") + + # Add description or not + zero_shot_topics_df = pd.DataFrame( + data={ + "General topic": zero_shot_topics_gen_topics_list, + "Subtopic": zero_shot_topics_subtopics_list, + "Description": zero_shot_topics_description_list, + } + ) + + # Filter out duplicate General topic and subtopic names + zero_shot_topics_df = zero_shot_topics_df.drop_duplicates( + ["General topic", "Subtopic"], keep="first" + ) + + # Sort the dataframe by General topic and subtopic + zero_shot_topics_df = zero_shot_topics_df.sort_values( + ["General topic", "Subtopic"], ascending=[True, True] + ) + + return zero_shot_topics_df + + +def create_candidate_topics_df_from_topic_summary( + topic_summary_df: pd.DataFrame, +) -> pd.DataFrame: + """ + Extract unique General topic / Subtopic pairs from a topic summary dataframe. + + Deduplicates across Group and Sentiment so the output can be reused as a + zero-shot topics input file (same column format as generate_zero_shot_topics_df). + """ + if topic_summary_df is None or topic_summary_df.empty: + return pd.DataFrame(columns=["General topic", "Subtopic"]) + + df = topic_summary_df.copy() + + if "General topic" not in df.columns and "Main heading" in df.columns: + df = df.rename(columns={"Main heading": "General topic"}) + if "Subtopic" not in df.columns and "Subheading" in df.columns: + df = df.rename(columns={"Subheading": "Subtopic"}) + + if "General topic" not in df.columns and "Subtopic" not in df.columns: + print( + "Warning: topic summary has no General topic / Subtopic columns; " + "cannot create suggested topics CSV." + ) + return pd.DataFrame(columns=["General topic", "Subtopic"]) + + if "General topic" not in df.columns: + df["General topic"] = "" + if "Subtopic" not in df.columns: + df["Subtopic"] = "" + + out_df = df[["General topic", "Subtopic"]].copy() + out_df = out_df.drop_duplicates(subset=["General topic", "Subtopic"], keep="first") + out_df = out_df.sort_values(["General topic", "Subtopic"], ascending=[True, True]) + return out_df.reset_index(drop=True) + + +def write_candidate_topics_csv( + topic_summary_df: pd.DataFrame, + output_path: str, +) -> str: + """Write a suggested-topics CSV from a topic summary dataframe.""" + topics_df = create_candidate_topics_df_from_topic_summary(topic_summary_df) + if topics_df.empty: + return "" + + topics_df.to_csv(output_path, index=False, encoding="utf-8-sig") + print(f"Suggested topics CSV saved as '{output_path}'") + return output_path + + +def subsample_responses_for_topic_discovery( + file_data: pd.DataFrame, + sample_fraction: float, + random_seed: int, + group_col: str | None = None, + min_rows_per_group: int = 1, +) -> tuple[pd.DataFrame, dict]: + """ + Randomly subsample rows for topic discovery. + + When group_col is set and present in file_data, sampling is stratified per group. + Otherwise a single random sample is taken from the full dataset. + """ + if file_data is None or file_data.empty: + raise ValueError("Cannot subsample empty data.") + + if sample_fraction <= 0 or sample_fraction > 1: + raise ValueError( + f"sample_fraction must be between 0 and 1 (exclusive of 0), got {sample_fraction}" + ) + + original_rows = len(file_data) + df = file_data.copy() + df["_discovery_original_row_index"] = df.index + per_group_counts: dict = {} + + use_group_col = ( + group_col + and str(group_col).strip() + and group_col in df.columns + and group_col != "group_col" + ) + + if use_group_col: + sampled_parts = [] + for group_value, group_df in df.groupby(group_col, dropna=False): + n_sample = max( + min_rows_per_group, math.ceil(len(group_df) * sample_fraction) + ) + n_sample = min(n_sample, len(group_df)) + part = group_df.sample(n=n_sample, random_state=random_seed) + per_group_counts[str(group_value)] = { + "original": len(group_df), + "sampled": n_sample, + } + sampled_parts.append(part) + sampled_df = pd.concat(sampled_parts).sort_index() + else: + n_sample = max(1, math.ceil(original_rows * sample_fraction)) + n_sample = min(n_sample, original_rows) + sampled_df = df.sample(n=n_sample, random_state=random_seed).sort_index() + per_group_counts = {"All": {"original": original_rows, "sampled": n_sample}} + + if sampled_df.empty: + raise ValueError("Subsample produced no rows.") + + metadata = { + "original_rows": original_rows, + "sampled_rows": len(sampled_df), + "sample_fraction": sample_fraction, + "random_seed": random_seed, + "per_group_counts": per_group_counts, + } + return sampled_df, metadata + + +def write_topic_discovery_manifest_csv( + sampled_df: pd.DataFrame, + output_path: str, + group_col: str | None = None, +) -> str: + """Write a manifest CSV listing which rows were included in topic discovery.""" + if sampled_df is None or sampled_df.empty: + return "" + + manifest_cols = ["_discovery_original_row_index"] + if ( + group_col + and str(group_col).strip() + and group_col in sampled_df.columns + and group_col != "group_col" + ): + manifest_cols.append(group_col) + + manifest_df = sampled_df[manifest_cols].copy() + manifest_df.rename( + columns={"_discovery_original_row_index": "Original row index"}, + inplace=True, + ) + manifest_df.to_csv(output_path, index=False, encoding="utf-8-sig") + print(f"Topic discovery manifest saved as '{output_path}'") + return output_path + + +def update_model_choice(model_source): + # Filter models by source and return the first matching model name + matching_models = [ + model_name + for model_name, model_info in model_name_map.items() + if model_info["source"] == model_source + ] + + output_model = matching_models[0] if matching_models else model_full_names[0] + + return gr.Dropdown( + value=output_model, + choices=matching_models, + label="Large language model for topic extraction and summarisation", + multiselect=False, + ) + + +def ensure_model_in_map(model_choice: str, model_name_map_dict: dict = None) -> dict: + """ + Ensures that a model_choice is registered in model_name_map. + If the model_choice is not found, it assumes it's an inference-server model + and adds it to the map with source "inference-server". + + Args: + model_choice (str): The model name to check/register + model_name_map_dict (dict, optional): The model_name_map dictionary to update. + If None, uses the global model_name_map from config. + + Returns: + dict: The model_name_map dictionary (updated if needed) + """ + # Use provided dict or global one + if model_name_map_dict is None: + from tools.config import model_name_map + + model_name_map_dict = model_name_map + + # If model_choice is not in the map, assume it's an inference-server model + if model_choice not in model_name_map_dict: + model_name_map_dict[model_choice] = { + "short_name": model_choice, + "source": "inference-server", + } + print(f"Registered custom model '{model_choice}' as inference-server model") + + return model_name_map_dict diff --git a/tools/llm_api_call.py b/tools/llm_api_call.py new file mode 100644 index 0000000000000000000000000000000000000000..94aafb197dd1246b7cfc3c77b942ca41ebdf37a1 --- /dev/null +++ b/tools/llm_api_call.py @@ -0,0 +1,6648 @@ +import json +import os +import re +import string +import time +from io import StringIO +from typing import Any, List, Tuple + +import gradio as gr +import markdown +import pandas as pd +import spaces +from gradio import Progress +from tqdm import tqdm + +from tools.aws_functions import connect_to_bedrock_runtime +from tools.config import ( + ALL_IN_ONE_USE_LLM_DEDUP, + BATCH_DEDUPLICATION_EVERY_N_BATCHES, + BATCH_SIZE_DEFAULT, + CHOSEN_LOCAL_MODEL_TYPE, + DEDUPLICATION_THRESHOLD, + ENABLE_BATCH_DEDUPLICATION, + ENABLE_VALIDATION, + LLM_CONTEXT_HEADROOM_FRACTION, + LLM_CONTEXT_LENGTH, + LLM_MAX_NEW_TOKENS, + LLM_SEED, + MAX_COMMENT_CHARS, + MAX_GROUPS, + MAX_OUTPUT_VALIDATION_ATTEMPTS, + MAX_ROWS, + MAX_SPACES_GPU_RUN_TIME, + MAX_TIME_FOR_LOOP, + MAXIMUM_ALLOWED_TOPICS, + NUMBER_OF_RETRY_ATTEMPTS, + OUTPUT_DEBUG_FILES, + OUTPUT_FOLDER, + REASONING_SUFFIX, + RUN_LOCAL_MODEL, + TIMEOUT_WAIT, + model_name_map, +) +from tools.dedup_summaries import ( + deduplicate_topics, + deduplicate_topics_llm, + overall_summary, + process_debug_output_iteration, + wrapper_summarise_output_topics_per_group, +) +from tools.helper_functions import ( + clean_column_name, + convert_reference_table_to_pivot_table, + create_topic_summary_df_from_reference_table, + ensure_model_in_map, + generate_zero_shot_topics_df, + get_basic_response_data, + load_in_data_file, + load_in_previous_data_files, + normalize_topic_name_for_llm, + put_columns_in_df, + read_file, + subsample_responses_for_topic_discovery, + wrap_text, + write_candidate_topics_csv, + write_topic_discovery_manifest_csv, +) +from tools.llm_funcs import ( + calculate_tokens_from_metadata, + call_llm_with_markdown_table_checks, + construct_azure_client, + construct_gemini_generative_model, + create_missing_references_df, + get_assistant_model, + get_model, + get_tokenizer, +) +from tools.prompts import ( + add_existing_topics_assistant_prefill, + add_existing_topics_prompt, + add_existing_topics_system_prompt, + allow_new_topics_prompt, + default_response_reference_format, + default_sentiment_prompt, + force_existing_topics_prompt, + force_single_topic_prompt, + initial_table_assistant_prefill, + initial_table_prompt, + initial_table_system_prompt, + negative_neutral_positive_sentiment_prompt, + negative_or_positive_sentiment_prompt, + previous_table_introduction_default, + structured_summary_prompt, + validation_prompt_prefix_default, + validation_prompt_suffix_default, + validation_prompt_suffix_struct_summary_default, + validation_system_prompt, +) + +max_tokens = LLM_MAX_NEW_TOKENS +timeout_wait = TIMEOUT_WAIT +number_of_api_retry_attempts = NUMBER_OF_RETRY_ATTEMPTS +max_time_for_loop = MAX_TIME_FOR_LOOP +batch_size_default = BATCH_SIZE_DEFAULT +deduplication_threshold = DEDUPLICATION_THRESHOLD +max_comment_character_length = MAX_COMMENT_CHARS +random_seed = LLM_SEED +reasoning_suffix = REASONING_SUFFIX +max_rows = MAX_ROWS +maximum_allowed_topics = MAXIMUM_ALLOWED_TOPICS +output_debug_files = OUTPUT_DEBUG_FILES + +max_text_length = 500 + +### HELPER FUNCTIONS + + +def remove_markdown_emphasis(text: str) -> str: + """ + Remove markdown emphasis characters from text that are not used in general English. + Removes: ** (bold), * (italic), __ (bold), _ (italic), ~~ (strikethrough), ` (code), # (headers) + """ + if not isinstance(text, str): + return text + + # Remove markdown emphasis patterns in order of specificity + + # Remove strikethrough: ~~text~~ + text = re.sub(r"~~([^~]+)~~", r"\1", text) + + # Remove bold with double asterisks: **text** (handles cases with or without spaces) + text = re.sub(r"\*\*([^*]+?)\*\*", r"\1", text) + # Remove any remaining double asterisks + text = text.replace("**", "") + + # Remove bold with double underscores: __text__ + text = re.sub(r"__([^_]+?)__", r"\1", text) + # Remove any remaining double underscores + text = text.replace("__", "") + + # Remove italic with single asterisks: *text* (but preserve asterisks that are part of words) + # Only match when asterisk is at word boundary or surrounded by non-word characters + text = re.sub(r"(? tuple[str, pd.DataFrame]: + """ + Reconstructs a markdown table from reference_df data when all_responses_content is missing. + Filters to only include rows from the current batch if start_row and end_row are provided. + + Parameters: + - reference_df (pd.DataFrame): The reference dataframe containing topic analysis data + - start_row (int, optional): The starting row number for the current batch + - end_row (int, optional): The ending row number for the current batch + + Returns: + - tuple[str, pd.DataFrame]: A tuple containing: + - str: A markdown table string in the required format + - pd.DataFrame: A pandas DataFrame with the same data as the markdown table + """ + if reference_df.empty: + return "", pd.DataFrame() + + # Filter reference_df to current batch if start_row and end_row are provided + filtered_df = reference_df.copy() + if start_row is not None and end_row is not None: + # Convert Response ID to numeric for filtering + filtered_df["Response ID"] = pd.to_numeric( + filtered_df["Response ID"], errors="coerce" + ) + # Filter to only include rows where Response ID fall within the current batch range + filtered_df = filtered_df[ + (filtered_df["Response ID"] >= start_row + 1) + & (filtered_df["Response ID"] <= end_row + 1) + ] + + if filtered_df.empty: + return "", pd.DataFrame() + + if ( + "Revised summary" in filtered_df.columns + and "Summary" not in filtered_df.columns + ): + filtered_df = filtered_df.rename(columns={"Revised summary": "Summary"}) + + # Group by General topic, Subtopic, and Sentiment to aggregate response references + grouped_df = ( + filtered_df.groupby(["General topic", "Subtopic", "Sentiment"]) + .agg( + { + "Response ID": lambda x: ", ".join(map(str, sorted(x.unique()))), + "Summary": "first", # Take the first summary for each group + } + ) + .reset_index() + ) + + # Adjust response references to be relative to the batch (subtract start_row if provided) + if start_row is not None: + # Convert response references to relative numbers by subtracting start_row + def adjust_references(refs_str): + if not refs_str or refs_str == "": + return refs_str + try: + # Split by comma, convert to int, subtract start_row, convert back to string + refs = [ + str(int(ref.strip()) - start_row) + for ref in refs_str.split(",") + if ref.strip().isdigit() + ] + return ", ".join(refs) + except (ValueError, TypeError): + return refs_str + + grouped_df["Response ID"] = grouped_df["Response ID"].apply(adjust_references) + + # Clean up the data to handle any NaN values and remove "Rows x to y: " prefix from summary + cleaned_df = grouped_df.copy() + for col in [ + "General topic", + "Subtopic", + "Sentiment", + "Response ID", + "Summary", + ]: + cleaned_df[col] = cleaned_df[col].fillna("").astype(str) + + # Remove "Rows x to y: " prefix from summary if present + cleaned_df["Summary"] = cleaned_df["Summary"].apply( + lambda x: ( + re.sub(r"^Rows\s+\d+\s+to\s+\d+:\s*", "", x) if isinstance(x, str) else x + ) + ) + + cleaned_df.drop_duplicates( + ["General topic", "Subtopic", "Sentiment", "Response ID"], inplace=True + ) + + # Create the markdown table + markdown_table = ( + "| General topic | Subtopic | Sentiment | Response ID | Summary |\n" + ) + markdown_table += "|---|---|---|---|---|\n" + + for _, row in cleaned_df.iterrows(): + general_topic = row["General topic"] + subtopic = row["Subtopic"] + sentiment = row["Sentiment"] + response_refs = row["Response ID"] + summary = row["Summary"] + + # Add row to markdown table + markdown_table += f"| {general_topic} | {subtopic} | {sentiment} | {response_refs} | {summary} |\n" + + return markdown_table, cleaned_df + + +def _assess_sentiment(sentiment_checkbox: str) -> bool: + return sentiment_checkbox != "Do not assess sentiment" + + +def _topic_summary_dedup_columns( + sentiment_checkbox: str, topic_summary_df: pd.DataFrame +) -> List[str]: + cols = ["General topic", "Subtopic"] + if ( + _assess_sentiment(sentiment_checkbox) + and "Sentiment" in topic_summary_df.columns + ): + cols.append("Sentiment") + if "Group" in topic_summary_df.columns: + cols.append("Group") + return cols + + +def _merge_topic_number_keys( + sentiment_checkbox: str, + reference_df: pd.DataFrame, + topic_summary_df: pd.DataFrame, +) -> Tuple[List[str], List[str]]: + """Return (on_cols, right_table_cols) to join Topic number onto reference_df.""" + on_cols = ["General topic", "Subtopic"] + if ( + _assess_sentiment(sentiment_checkbox) + and "Sentiment" in reference_df.columns + and "Sentiment" in topic_summary_df.columns + ): + on_cols.append("Sentiment") + if "Group" in reference_df.columns and "Group" in topic_summary_df.columns: + on_cols.append("Group") + on_cols = [ + c + for c in on_cols + if c in reference_df.columns and c in topic_summary_df.columns + ] + summary_cols = list(on_cols) + if "Topic number" in topic_summary_df.columns: + summary_cols.append("Topic number") + summary_cols = [c for c in summary_cols if c in topic_summary_df.columns] + return on_cols, summary_cols + + +def validate_topics( + file_data: pd.DataFrame, + reference_df: pd.DataFrame, + topic_summary_df: pd.DataFrame, + file_name: str, + chosen_cols: List[str], + batch_size: int, + model_choice: str, + in_api_key: str, + temperature: float, + max_tokens: int, + azure_api_key_textbox: str = "", + azure_endpoint_textbox: str = "", + reasoning_suffix: str = "", + group_name: str = "All", + produce_structured_summary_radio: str = "No", + force_zero_shot_radio: str = "No", + force_single_topic_radio: str = "No", + context_textbox: str = "", + additional_instructions_summary_format: str = "", + output_folder: str = OUTPUT_FOLDER, + output_debug_files: str = "False", + original_full_file_name: str = "", + additional_validation_issues_provided: str = "", + max_time_for_loop: int = MAX_TIME_FOR_LOOP, + sentiment_checkbox: str = "Negative or Positive", + logged_content: list = None, + show_previous_table: str = "Yes", + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + api_url: str = None, + max_topics_number: int = MAXIMUM_ALLOWED_TOPICS, + progress=gr.Progress(track_tqdm=True), +) -> Tuple[pd.DataFrame, pd.DataFrame, list, str, int, int, int]: + """ + Validates topics by re-running the topic extraction process on all batches + using the consolidated topics from the original run. + + Parameters: + - file_data (pd.DataFrame): The input data to validate + - reference_df (pd.DataFrame): The reference dataframe from the original run + - topic_summary_df (pd.DataFrame): The topic summary dataframe from the original run + - file_name (str): Name of the file being processed + - chosen_cols (List[str]): Columns to process + - batch_size (int): Size of each batch + - model_choice (str): The model to use for validation + - in_api_key (str): API key for the model + - temperature (float): Temperature for the model + - max_tokens (int): Maximum tokens for the model + - azure_api_key_textbox (str): Azure API key if using Azure + - azure_endpoint_textbox (str): Azure endpoint if using Azure + - reasoning_suffix (str): Suffix for reasoning + - group_name (str): Name of the group + - produce_structured_summary_radio (str): Whether to produce structured summaries + - force_zero_shot_radio (str): Whether to force zero-shot + - force_single_topic_radio (str): Whether to force single topic + - context_textbox (str): Context for the validation + - additional_instructions_summary_format (str): Additional instructions + - output_folder (str): Output folder for files + - output_debug_files (str): Whether to output debug files + - original_full_file_name (str): Original file name + - additional_validation_issues_provided (str): Additional validation issues provided + - max_time_for_loop (int): Maximum time for the loop + - logged_content (list, optional): The logged content from the original run. If None, tables will be reconstructed from reference_df + - show_previous_table (str): Whether to show the previous table ("Yes" or "No"). + - aws_access_key_textbox (str): AWS access key. + - aws_secret_key_textbox (str): AWS secret key. + - progress: Progress bar object + + Returns: + - Tuple[pd.DataFrame, pd.DataFrame, list, str, int, int, int]: Updated reference_df, topic_summary_df, logged_content, conversation_metadata_str, total_input_tokens, total_output_tokens, total_llm_calls + """ + print("Starting validation process...") + + # Ensure custom model_choice is registered in model_name_map + ensure_model_in_map(model_choice) + + # Calculate number of batches + num_batches = (len(file_data) + batch_size - 1) // batch_size + + # Initialize model components + model_choice_clean = model_name_map[model_choice]["short_name"] + model_source = model_name_map[model_choice]["source"] + + if context_textbox and "The context of this analysis is" not in context_textbox: + context_textbox = "The context of this analysis is '" + context_textbox + "'." + + # Initialize model objects + local_model = None + tokenizer = None + bedrock_runtime = None + + # Load local model if needed + if (model_name_map[model_choice]["source"] == "Local") & (RUN_LOCAL_MODEL == "1"): + local_model = get_model() + tokenizer = get_tokenizer() + + # Set up bedrock runtime if needed + if model_source == "AWS": + bedrock_runtime = connect_to_bedrock_runtime( + model_name_map, + model_choice, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + ) + + # Clean file name for output + file_name_clean = clean_column_name( + file_name, max_length=20, front_characters=False + ) + in_column_cleaned = clean_column_name(chosen_cols, max_length=20) + model_choice_clean_short = clean_column_name( + model_choice_clean, max_length=20, front_characters=False + ) + + # Create validation-specific logged content lists + validation_all_prompts_content = list() + validation_all_summaries_content = list() + validation_all_conversation_content = list() + validation_all_metadata_content = list() + validation_all_groups_content = list() + validation_all_batches_content = list() + validation_all_model_choice_content = list() + validation_all_validated_content = list() + validation_all_task_type_content = list() + validation_all_file_names_content = list() + + # Extract previous summaries from logged content for validation + if logged_content is None: + logged_content = list() + all_responses_content = [ + item.get("response", "") for item in logged_content if "response" in item + ] + + # Initialize validation dataframes + validation_reference_df = reference_df.copy() + validation_topic_summary_df = topic_summary_df.copy() + + # Clean topic names in initial dataframes to ensure consistent formatting throughout processing + for col_name in ["General topic", "Subtopic"]: + if ( + col_name in validation_reference_df.columns + and not validation_reference_df[col_name].isnull().all() + ): + validation_reference_df[col_name] = validation_reference_df[col_name].apply( + normalize_topic_name_for_llm + ) + if ( + col_name in validation_topic_summary_df.columns + and not validation_topic_summary_df[col_name].isnull().all() + ): + validation_topic_summary_df[col_name] = validation_topic_summary_df[ + col_name + ].apply(normalize_topic_name_for_llm) + + sentiment_prefix = "In the next column named 'Sentiment', " + sentiment_suffix = "." + if sentiment_checkbox == "Negative, Neutral, or Positive": + sentiment_prompt = ( + sentiment_prefix + + negative_neutral_positive_sentiment_prompt + + sentiment_suffix + ) + elif sentiment_checkbox == "Negative or Positive": + sentiment_prompt = ( + sentiment_prefix + negative_or_positive_sentiment_prompt + sentiment_suffix + ) + elif sentiment_checkbox == "Do not assess sentiment": + sentiment_prompt = "" # Just remove line completely. Previous: sentiment_prefix + do_not_assess_sentiment_prompt + sentiment_suffix + else: + sentiment_prompt = ( + sentiment_prefix + default_sentiment_prompt + sentiment_suffix + ) + + # Validation loop through all batches + validation_latest_batch_completed = 0 + validation_loop = progress.tqdm( + range(num_batches), + total=num_batches, + desc="Validating topic extraction batches", + unit="validation batches", + ) + + tic = time.perf_counter() + + for validation_i in validation_loop: + validation_reported_batch_no = validation_latest_batch_completed + 1 + print("Running validation batch:", validation_reported_batch_no) + + # Call the function to prepare the input table for validation + ( + validation_simplified_csv_table_path, + validation_normalised_simple_markdown_table, + validation_start_row, + validation_end_row, + validation_batch_basic_response_df, + ) = data_file_to_markdown_table( + file_data, + file_name, + chosen_cols, + validation_latest_batch_completed, + batch_size, + ) + + if validation_batch_basic_response_df.shape[0] == 1: + validation_response_reference_format = "" + else: + validation_response_reference_format = ( + "\n" + default_response_reference_format + ) + + if validation_normalised_simple_markdown_table: + validation_response_table_prompt = ( + "Response table:\n" + validation_normalised_simple_markdown_table + ) + else: + validation_response_table_prompt = "" + + # If the validation batch of responses contains at least one instance of text. The function will first try to get the previous table from logged outputs, and will reconstruct the table from reference_df data if not available. + if not reference_df.empty: + validation_latest_batch_completed = int(validation_latest_batch_completed) + validation_start_row = int(validation_start_row) + validation_end_row = int(validation_end_row) + + # Get the previous table from all_responses_content for this batch + if validation_latest_batch_completed < len(all_responses_content): + previous_table_content = all_responses_content[ + validation_latest_batch_completed + ] + _, previous_topic_df = reconstruct_markdown_table_from_reference_df( + reference_df, validation_start_row, validation_end_row + ) + else: + # Try to reconstruct markdown table from reference_df data + previous_table_content, previous_topic_df = ( + reconstruct_markdown_table_from_reference_df( + reference_df, validation_start_row, validation_end_row + ) + ) + + # Always use the consolidated topics from the first run for validation + validation_formatted_system_prompt = validation_system_prompt.format( + consultation_context=context_textbox, column_name=chosen_cols + ) + + # Use the accumulated topic summary from previous validation batches (or initial if first batch) + validation_existing_topic_summary_df = validation_topic_summary_df.copy() + validation_existing_topic_summary_df["Number of responses"] = "" + # validation_existing_topic_summary_df.fillna("", inplace=True) + validation_existing_topic_summary_df["General topic"] = ( + validation_existing_topic_summary_df["General topic"].str.replace( + "(?i)^Nan$", "", regex=True + ) + ) + validation_existing_topic_summary_df["Subtopic"] = ( + validation_existing_topic_summary_df["Subtopic"].str.replace( + "(?i)^Nan$", "", regex=True + ) + ) + validation_existing_topic_summary_df = ( + validation_existing_topic_summary_df.drop_duplicates() + ) + + # Create topics table to be presented to LLM for validation + validation_keep_cols = [ + col + for col in ["General topic", "Subtopic", "Description"] + if col in validation_existing_topic_summary_df.columns + and not validation_existing_topic_summary_df[col] + .replace(r"^\s*$", pd.NA, regex=True) + .isna() + .all() + ] + + validation_topics_df_for_markdown = validation_existing_topic_summary_df[ + validation_keep_cols + ].drop_duplicates(validation_keep_cols) + if ( + "General topic" in validation_topics_df_for_markdown.columns + and "Subtopic" in validation_topics_df_for_markdown.columns + ): + validation_topics_df_for_markdown = ( + validation_topics_df_for_markdown.sort_values( + ["General topic", "Subtopic"] + ) + ) + + if "Description" in validation_existing_topic_summary_df: + if validation_existing_topic_summary_df["Description"].isnull().all(): + validation_existing_topic_summary_df.drop( + "Description", axis=1, inplace=True + ) + + if produce_structured_summary_radio == "Yes": + if "General topic" in validation_topics_df_for_markdown.columns: + validation_topics_df_for_markdown = ( + validation_topics_df_for_markdown.rename( + columns={"General topic": "Main heading"} + ) + ) + if "Subtopic" in validation_topics_df_for_markdown.columns: + validation_topics_df_for_markdown = ( + validation_topics_df_for_markdown.rename( + columns={"Subtopic": "Subheading"} + ) + ) + + # Clean topic names before converting to markdown to ensure consistent formatting + if "General topic" in validation_topics_df_for_markdown.columns: + validation_topics_df_for_markdown["General topic"] = ( + validation_topics_df_for_markdown["General topic"].apply( + normalize_topic_name_for_llm + ) + ) + if "Subtopic" in validation_topics_df_for_markdown.columns: + validation_topics_df_for_markdown["Subtopic"] = ( + validation_topics_df_for_markdown["Subtopic"].apply( + normalize_topic_name_for_llm + ) + ) + if "Main heading" in validation_topics_df_for_markdown.columns: + validation_topics_df_for_markdown["Main heading"] = ( + validation_topics_df_for_markdown["Main heading"].apply( + normalize_topic_name_for_llm + ) + ) + if "Subheading" in validation_topics_df_for_markdown.columns: + validation_topics_df_for_markdown["Subheading"] = ( + validation_topics_df_for_markdown["Subheading"].apply( + normalize_topic_name_for_llm + ) + ) + + validation_unique_topics_markdown = ( + validation_topics_df_for_markdown.to_markdown(index=False) + ) + validation_unique_topics_markdown = normalise_string( + validation_unique_topics_markdown + ) + + if force_zero_shot_radio == "Yes": + validation_topic_assignment_prompt = force_existing_topics_prompt + else: + validation_topic_assignment_prompt = allow_new_topics_prompt + + # Should the outputs force only one single topic assignment per response? + if force_single_topic_radio != "Yes": + validation_force_single_topic_prompt = "" + else: + validation_topic_assignment_prompt = ( + validation_topic_assignment_prompt.replace( + "Assign topics", "Assign a topic" + ) + .replace("assign Subtopics", "assign a Subtopic") + .replace("Subtopics", "Subtopic") + .replace("Topics", "Topic") + .replace("topics", "a topic") + ) + + # Provide new validation issues on a new line if provided + # if additional_validation_issues_provided: + # additional_validation_issues_provided = "\n" + additional_validation_issues_provided + + # Format the validation prompt with the response table and topics + if produce_structured_summary_radio != "Yes": + validation_formatted_summary_prompt = add_existing_topics_prompt.format( + validate_prompt_prefix=validation_prompt_prefix_default, + response_table=validation_response_table_prompt, + topics=validation_unique_topics_markdown, + topic_assignment=validation_topic_assignment_prompt, + force_single_topic=validation_force_single_topic_prompt, + sentiment_choices=sentiment_prompt, + response_reference_format=validation_response_reference_format, + add_existing_topics_summary_format=additional_instructions_summary_format, + previous_table_introduction=previous_table_introduction_default, + previous_table=( + previous_table_content if show_previous_table == "Yes" else "" + ), + validate_prompt_suffix=validation_prompt_suffix_default.format( + additional_validation_issues=additional_validation_issues_provided + ), + ) + else: + # Ensure the validation wrapper is applied even for structured summaries + structured_summary_instructions = structured_summary_prompt.format( + response_table=validation_response_table_prompt, + topics=validation_unique_topics_markdown, + summary_format=additional_instructions_summary_format, + ) + + validation_formatted_summary_prompt = ( + f"{validation_prompt_prefix_default}" + f"{structured_summary_instructions}" + f"{previous_table_introduction_default}" + f"{previous_table_content if show_previous_table == 'Yes' else ''}" + f"{validation_prompt_suffix_struct_summary_default.format(additional_validation_issues=additional_validation_issues_provided)}" + ) + + validation_batch_file_path_details = f"{file_name_clean}_val_batch_{validation_latest_batch_completed + 1}_size_{batch_size}_col_{in_column_cleaned}" + + # Use the helper function to process the validation batch + ( + validation_new_topic_df, + validation_new_reference_df, + validation_new_topic_summary_df, + validation_is_error, + validation_current_prompt_content_logged, + validation_current_summary_content_logged, + validation_current_conversation_content_logged, + validation_current_metadata_content_logged, + validation_topic_table_out_path, + validation_reference_table_out_path, + validation_topic_summary_df_out_path, + ) = process_batch_with_llm( + is_first_batch=False, + formatted_system_prompt=validation_formatted_system_prompt, + formatted_prompt=validation_formatted_summary_prompt, + batch_file_path_details=validation_batch_file_path_details, + model_source=model_source, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + max_tokens=max_tokens, + azure_api_key_textbox=azure_api_key_textbox, + azure_endpoint_textbox=azure_endpoint_textbox, + reasoning_suffix=reasoning_suffix, + local_model=local_model, + tokenizer=tokenizer, + bedrock_runtime=bedrock_runtime, + reported_batch_no=validation_reported_batch_no, + response_text="", + whole_conversation=list(), + all_metadata_content=list(), + start_row=validation_start_row, + end_row=validation_end_row, + model_choice_clean=model_choice_clean, + log_files_output_paths=list(), + existing_reference_df=validation_reference_df, + existing_topic_summary_df=validation_existing_topic_summary_df, + batch_size=batch_size, + batch_basic_response_df=validation_batch_basic_response_df, + group_name=group_name, + produce_structured_summary_radio=produce_structured_summary_radio, + output_folder=output_folder, + output_debug_files=output_debug_files, + task_type="Validation", + assistant_prefill=add_existing_topics_assistant_prefill, + api_url=api_url, + ) + + if validation_new_topic_df.empty: + validation_new_topic_df = previous_topic_df + # print("Validation new topic df is empty, using previous topic df:", validation_new_topic_df) + # print("Validation new topic df columns:", validation_new_topic_df.columns) + + # Collect conversation metadata from validation batch + if validation_current_metadata_content_logged: + validation_all_metadata_content.append( + validation_current_metadata_content_logged + ) + + validation_all_prompts_content.append( + validation_current_prompt_content_logged + ) + validation_all_summaries_content.append( + validation_current_summary_content_logged + ) + validation_all_conversation_content.append( + validation_current_conversation_content_logged + ) + validation_all_groups_content.append(group_name) + validation_all_batches_content.append( + f"Validation {validation_reported_batch_no}:" + ) + validation_all_model_choice_content.append(model_choice_clean_short) + validation_all_validated_content.append("Yes") + validation_all_task_type_content.append("Validation") + validation_all_file_names_content.append(original_full_file_name) + + # print("Appended to logs") + + # Update validation dataframes with validation results + # For validation, we need to accumulate results from each batch, not overwrite them + # The validation_new_* dataframes contain the results for the current batch + # We need to concatenate them with the existing validation dataframes + + # For reference_df, we need to be careful about duplicates + if not validation_new_reference_df.empty: + # Check if the new reference_df is the same as the existing one (indicating "no change" response) + # This happens when the LLM responds with "no change" and returns the existing data + if validation_new_reference_df.equals(validation_reference_df): + print( + "Validation new reference df is identical to existing df (no change response), skipping concatenation" + ) + else: + # print("Validation new reference df is not empty, appending new table to validation reference df") + # Remove any existing entries for this batch range to avoid duplicates + start_row_reported = int(validation_start_row) + 1 + end_row_reported = int(validation_end_row) + 1 + validation_reference_df["Start row of group"] = ( + validation_reference_df["Start row of group"].astype(int) + ) + + # Remove existing entries for this batch range from validation_reference_df + if "Start row of group" in validation_reference_df.columns: + validation_reference_df = validation_reference_df[ + ~( + ( + validation_reference_df["Start row of group"] + >= start_row_reported + ) + & ( + validation_reference_df["Start row of group"] + <= end_row_reported + ) + ) + ] + + # Concatenate the new results + validation_reference_df = pd.concat( + [validation_reference_df, validation_new_reference_df] + ).dropna(how="all") + + # Rebuild topic_summary_df from accumulated reference_df to ensure consistency + # This properly aggregates response counts and handles all dimensions (General topic, Subtopic, Sentiment, Group) + if not validation_reference_df.empty: + validation_topic_summary_df = ( + create_topic_summary_df_from_reference_table( + validation_reference_df, sentiment_checkbox=sentiment_checkbox + ) + ) + + else: + print( + "Current validation batch of responses contains no text, moving onto next. Batch number:", + str(validation_latest_batch_completed + 1), + ". Start row:", + validation_start_row, + ". End row:", + validation_end_row, + ) + + # Increase validation batch counter + validation_latest_batch_completed += 1 + + # Check if we've exceeded max time for validation loop + validation_toc = time.perf_counter() + validation_final_time = validation_toc - tic + + if validation_final_time > max_time_for_loop: + print("Max time reached during validation, breaking validation loop.") + if progress: + validation_loop.close() + tqdm._instances.clear() + break + + # Deduplicate topics after each batch if enabled and conditions are met + if ( + ENABLE_BATCH_DEDUPLICATION + and force_zero_shot_radio == "No" + and produce_structured_summary_radio == "No" + and not validation_reference_df.empty + and not validation_topic_summary_df.empty + and validation_latest_batch_completed % BATCH_DEDUPLICATION_EVERY_N_BATCHES + == 0 + ): + print( + f"Deduplicating topics after validation batch {validation_latest_batch_completed}..." + ) + # Create temporary file names for deduplication + reference_table_file_name = f"{file_name_clean}_val_batch_{validation_latest_batch_completed}_reference" + unique_topics_table_file_name = f"{file_name_clean}_val_batch_{validation_latest_batch_completed}_unique_topics" + + # Prepare file_data for deduplication if available + in_data_files_for_dedup = None + validation_num_batches = None + validation_data_file_names_textbox = None + if not file_data.empty and chosen_cols: + # Pass file_data as DataFrame and calculate num_batches + in_data_files_for_dedup = file_data + validation_num_batches = (len(file_data) + batch_size - 1) // batch_size + validation_data_file_names_textbox = file_name + + # Clean General topic and Subtopic columns using the same process as zero-shot topics + for col_name in ["General topic", "Subtopic"]: + for df in [validation_reference_df, validation_topic_summary_df]: + if col_name in df.columns and not df[col_name].isnull().all(): + df[col_name] = df[col_name].apply(normalize_topic_name_for_llm) + + try: + topics_before = validation_topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + # Check Group values before deduplication for better diagnostics + input_groups_before = ( + validation_topic_summary_df["Group"].nunique() + if "Group" in validation_topic_summary_df.columns + else 0 + ) + input_ref_groups_before = ( + validation_reference_df["Group"].nunique() + if "Group" in validation_reference_df.columns + else 0 + ) + # Call deduplicate_topics function + ( + validation_reference_df, + validation_topic_summary_df, + _, + _, + _, + ) = deduplicate_topics( + reference_df=validation_reference_df, + topic_summary_df=validation_topic_summary_df, + reference_table_file_name=reference_table_file_name, + unique_topics_table_file_name=unique_topics_table_file_name, + in_excel_sheets="", # in_excel_sheets not available in validate_topics + merge_sentiment="No", + merge_general_topics="No", + score_threshold=95, + in_data_files=in_data_files_for_dedup, + chosen_cols=( + chosen_cols + if isinstance(chosen_cols, list) + else [chosen_cols] if chosen_cols else [] + ), + output_folder=output_folder, + deduplicate_topics="Yes", + output_files="False", + total_number_of_batches=validation_num_batches, + data_file_names_textbox=validation_data_file_names_textbox, + ) + topics_after = validation_topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + # Check Group values after deduplication + output_groups_after = ( + validation_topic_summary_df["Group"].nunique() + if "Group" in validation_topic_summary_df.columns + else 0 + ) + output_ref_groups_after = ( + validation_reference_df["Group"].nunique() + if "Group" in validation_reference_df.columns + else 0 + ) + print( + f"Topics deduplicated successfully after validation batch {validation_latest_batch_completed}. " + f"Topics before: {topics_before}, Topics after: {topics_after}" + ) + if topics_after > topics_before: + print( + f"WARNING: Number of topics increased from {topics_before} to {topics_after}. " + f"This may be due to Group column differences or new topic combinations created during deduplication. " + f"Input topic_summary_df had {input_groups_before} unique Group values, " + f"Input reference_df had {input_ref_groups_before} unique Group values. " + f"Output topic_summary_df has {output_groups_after} unique Group values, " + f"Output reference_df has {output_ref_groups_after} unique Group values." + ) + except Exception as e: + print( + f"Warning: Deduplication failed after validation batch {validation_latest_batch_completed}: {str(e)}. Continuing with original topics." + ) + + # Check if number of topics exceeds max_topics_number and use LLM deduplication if needed + topics_after_dedup = validation_topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + if topics_after_dedup > max_topics_number: + print( + f"Number of topics ({topics_after_dedup}) exceeds max_topics_number ({max_topics_number}). " + f"Using LLM-based deduplication after validation batch {validation_latest_batch_completed}..." + ) + try: + # Prepare file names for LLM deduplication + reference_table_file_name_llm = f"{file_name_clean}_val_batch_{validation_latest_batch_completed}_reference_llm" + unique_topics_table_file_name_llm = f"{file_name_clean}_val_batch_{validation_latest_batch_completed}_unique_topics_llm" + + # Prepare in_data_files - use file_data if available + in_data_files_for_llm_dedup = ( + file_data if not file_data.empty else None + ) + + # Store topics count before LLM deduplication + topics_before_llm_dedup = ( + validation_topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + ) + + # Call deduplicate_topics_llm function + ( + validation_reference_df, + validation_topic_summary_df, + _, + _, + _, + _, + _, + _, + _, + ) = deduplicate_topics_llm( + reference_df=validation_reference_df, + topic_summary_df=validation_topic_summary_df, + reference_table_file_name=reference_table_file_name_llm, + unique_topics_table_file_name=unique_topics_table_file_name_llm, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + model_source=model_source, + local_model=( + local_model if "local_model" in locals() else None + ), + tokenizer=tokenizer if "tokenizer" in locals() else None, + assistant_model=None, # assistant_model not typically available in this scope + in_excel_sheets="", # in_excel_sheets not available in validate_topics + merge_sentiment="No", + merge_general_topics="Yes", + in_data_files=in_data_files_for_llm_dedup, + chosen_cols=( + chosen_cols + if isinstance(chosen_cols, list) + else [chosen_cols] if chosen_cols else [] + ), + output_folder=output_folder, + candidate_topics=None, + azure_endpoint=azure_endpoint_textbox, + output_debug_files=output_debug_files, + api_url=api_url, + aws_access_key_textbox=aws_access_key_textbox, + aws_secret_key_textbox=aws_secret_key_textbox, + aws_region_textbox=aws_region_textbox, + azure_api_key_textbox=azure_api_key_textbox, + model_name_map=model_name_map, + output_files="False", + sentiment_checkbox=sentiment_checkbox, + ) + # Rebuild topic_summary_df from updated reference_df to ensure consistency + # This ensures the topic count reflects the actual merged state + if not validation_reference_df.empty: + validation_topic_summary_df = ( + create_topic_summary_df_from_reference_table( + validation_reference_df, + sentiment_checkbox=sentiment_checkbox, + ) + ) + + topics_after_llm_dedup = ( + validation_topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + ) + print( + f"LLM-based deduplication completed successfully after validation batch {validation_latest_batch_completed}. " + f"Number of topics before LLM deduplication: {topics_before_llm_dedup}, Number of topics after LLM deduplication: {topics_after_llm_dedup}" + ) + except Exception as e: + print( + f"Warning: LLM-based deduplication failed after validation batch {validation_latest_batch_completed}: {str(e)}. " + f"Continuing with existing topics." + ) + + # Combine validation logged content + validation_all_logged_content = [ + { + "prompt": prompt, + "response": summary, + "metadata": metadata, + "batch": batch, + "model_choice": model_choice, + "validated": validated, + "group": group, + "task_type": task_type, + "file_name": file_name, + } + for prompt, summary, metadata, batch, model_choice, validated, group, task_type, file_name in zip( + validation_all_prompts_content, + validation_all_summaries_content, + validation_all_metadata_content, + validation_all_batches_content, + validation_all_model_choice_content, + validation_all_validated_content, + validation_all_groups_content, + validation_all_task_type_content, + validation_all_file_names_content, + ) + ] + + # Append validation content to original logged content + updated_logged_content = list(logged_content) + list(validation_all_logged_content) + + # Combine validation conversation metadata + validation_conversation_metadata_str = " ".join(validation_all_metadata_content) + + # Ensure consistent Topic number assignment by recreating topic_summary_df from reference_df + if not validation_reference_df.empty: + validation_topic_summary_df = create_topic_summary_df_from_reference_table( + validation_reference_df, sentiment_checkbox=sentiment_checkbox + ) + + # Sort output dataframes + validation_reference_df["Response ID"] = ( + validation_reference_df["Response ID"].astype(float).astype(int) + ) + validation_reference_df["Start row of group"] = validation_reference_df[ + "Start row of group" + ].astype(int) + validation_reference_df.drop_duplicates( + ["Response ID", "General topic", "Subtopic", "Sentiment"], inplace=True + ) + validation_reference_df.sort_values( + [ + "Group", + "Start row of group", + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + ], + inplace=True, + ) + + validation_topic_summary_df["Number of responses"] = validation_topic_summary_df[ + "Number of responses" + ].astype(int) + validation_topic_summary_df.drop_duplicates( + _topic_summary_dedup_columns(sentiment_checkbox, validation_topic_summary_df), + inplace=True, + ) + topic_sort_cols = [ + "Group", + "Number of responses", + "General topic", + "Subtopic", + ] + topic_sort_asc = [True, False, True, True] + if ( + _assess_sentiment(sentiment_checkbox) + and "Sentiment" in validation_topic_summary_df.columns + ): + topic_sort_cols.append("Sentiment") + topic_sort_asc.append(True) + validation_topic_summary_df.sort_values( + topic_sort_cols, + ascending=topic_sort_asc, + inplace=True, + ) + + print("Validation process completed.") + + return ( + validation_reference_df, + validation_topic_summary_df, + updated_logged_content, + validation_conversation_metadata_str, + ) + + +# Define validation wrapper function +def validate_topics_wrapper( + file_data: pd.DataFrame, + reference_df: pd.DataFrame, + topic_summary_df: pd.DataFrame, + file_name: str, + chosen_cols: List[str], + batch_size: int, + model_choice: str, + in_api_key: str, + temperature: float, + max_tokens: int, + azure_api_key_textbox: str, + azure_endpoint_textbox: str, + reasoning_suffix: str, + group_name: str, + produce_structured_summary_radio: str, + force_zero_shot_radio: str, + force_single_topic_radio: str, + context_textbox: str, + additional_instructions_summary_format: str, + output_folder: str, + output_debug_files: str, + original_full_file_name: str, + additional_validation_issues_provided: str, + max_time_for_loop: int, + in_data_files: Any = None, + sentiment_checkbox: str = "Negative or Positive", + logged_content: List[dict] = None, + show_previous_table: str = "Yes", + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + api_url: str = None, + max_topics_number: int = MAXIMUM_ALLOWED_TOPICS, + progress=gr.Progress(track_tqdm=True), +) -> Tuple[pd.DataFrame, pd.DataFrame, List[dict], str, int, int, int, List[str]]: + """ + Wrapper function for validate_topics that processes data grouped by Group values and accumulates results, + similar to wrapper_extract_topics_per_column_value. + + Args: + file_data (pd.DataFrame): The input data to validate. + reference_df (pd.DataFrame): The reference dataframe from the original run. + topic_summary_df (pd.DataFrame): The topic summary dataframe from the original run. + file_name (str): Name of the file being processed. + chosen_cols (List[str]): Columns to process. + batch_size (int): Size of each batch. + model_choice (str): The model to use for validation. + in_api_key (str): API key for the model. + temperature (float): Temperature for the model. + max_tokens (int): Maximum tokens for the model. + azure_api_key_textbox (str): Azure API key if using Azure. + azure_endpoint_textbox (str): Azure endpoint if using Azure. + reasoning_suffix (str): Suffix for reasoning. + group_name (str): Name of the group. + produce_structured_summary_radio (str): Whether to produce structured summaries ("Yes" or "No"). + force_zero_shot_radio (str): Whether to force zero-shot ("Yes" or "No"). + force_single_topic_radio (str): Whether to force single topic ("Yes" or "No"). + context_textbox (str): Context for the validation. + additional_instructions_summary_format (str): Additional instructions for summary format. + output_folder (str): Output folder for files. + output_debug_files (str): Whether to output debug files ("True" or "False"). + original_full_file_name (str): Original file name. + additional_validation_issues_provided (str): Additional validation issues provided. + max_time_for_loop (int): Maximum time for the loop. + in_data_files (Any, optional): The input data files (e.g., Gradio FileData). If None, file_data must be provided. + sentiment_checkbox (str): Sentiment analysis option. + logged_content (List[dict], optional): The logged content from the original run. If None, tables will be reconstructed from reference_df. + show_previous_table (str): Whether to show the previous table ("Yes" or "No"). + aws_access_key_textbox (str): AWS access key. + aws_secret_key_textbox (str): AWS secret key. + progress (gr.Progress): Progress bar object. + + Returns: + Tuple[pd.DataFrame, pd.DataFrame, List[dict], str, int, int, int, List[str]]: + Accumulated reference_df, topic_summary_df, logged_content, conversation_metadata_str, + total_input_tokens, total_output_tokens, total_llm_calls, and a list of output file paths. + """ + + # Ensure custom model_choice is registered in model_name_map + ensure_model_in_map(model_choice) + + # Handle None logged_content + if logged_content is None: + logged_content = list() + + # If you have a file input but no file data it hasn't yet been loaded. Load it here. + if file_data.empty: + print("No data table found, loading from file") + try: + ( + in_colnames_drop, + in_excel_sheets, + file_name, + join_colnames, + join_colnames_drop, + ) = put_columns_in_df(in_data_files) + file_data, file_name, num_batches = load_in_data_file( + in_data_files, chosen_cols, batch_size_default, in_excel_sheets + ) + except Exception as e: + out_message = "Could not load in data file due to: " + str(e) + print(out_message) + raise Exception(out_message) + + if file_data.shape[0] > max_rows: + out_message = ( + "Your data has more than " + + str(max_rows) + + " rows, which has been set as the maximum in the application configuration." + ) + print(out_message) + raise Exception(out_message) + + if group_name is None: + print("No grouping column found") + file_data["group_col"] = "All" + group_name = "group_col" + + if group_name not in file_data.columns: + raise ValueError(f"Selected column '{group_name}' not found in file_data.") + + # Get unique Group values from the input dataframes + unique_groups = list() + if "Group" in reference_df.columns and not reference_df["Group"].isnull().all(): + unique_groups = reference_df["Group"].unique() + elif ( + "Group" in topic_summary_df.columns + and not topic_summary_df["Group"].isnull().all() + ): + unique_groups = topic_summary_df["Group"].unique() + else: + # If no Group column exists, use the provided group_name + unique_groups = [group_name] + + # Limit to MAX_GROUPS if there are too many + if len(unique_groups) > MAX_GROUPS: + print( + f"Warning: More than {MAX_GROUPS} unique groups found. Processing only the first {MAX_GROUPS}." + ) + unique_groups = unique_groups[:MAX_GROUPS] + + print(f"Processing validation for {len(unique_groups)} groups: {unique_groups}") + + # Initialise accumulators for results across all groups + acc_reference_df = pd.DataFrame() + acc_topic_summary_df = pd.DataFrame() + acc_logged_content = list() + acc_conversation_metadata = "" + acc_input_tokens = 0 + acc_output_tokens = 0 + acc_llm_calls = 0 + acc_output_files = list() + + if len(unique_groups) == 1: + # If only one unique value, no need for progress bar, iterate directly + loop_object = unique_groups + else: + # If multiple unique values, use tqdm progress bar + loop_object = progress.tqdm( + unique_groups, desc="Validating groups", unit="groups" + ) + + # Process each group separately + for i, current_group in enumerate(loop_object): + print( + f"\nProcessing validation for group: {current_group} ({i+1}/{len(unique_groups)})" + ) + + # Filter data for current group + if "Group" in reference_df.columns: + group_reference_df = reference_df[ + reference_df["Group"] == current_group + ].copy() + else: + group_reference_df = reference_df.copy() + + if "Group" in topic_summary_df.columns: + group_topic_summary_df = topic_summary_df[ + topic_summary_df["Group"] == current_group + ].copy() + else: + group_topic_summary_df = topic_summary_df.copy() + + # Filter file_data if it has a Group column + if "Group" in file_data.columns: + group_file_data = file_data[file_data["Group"] == current_group].copy() + else: + group_file_data = file_data.copy() + + # Skip if no data for this group + if group_reference_df.empty and group_topic_summary_df.empty: + print(f"No data for group {current_group}. Skipping.") + continue + + try: + # Call validate_topics for this group + ( + validation_reference_df, + validation_topic_summary_df, + updated_logged_content, + validation_conversation_metadata_str, + ) = validate_topics( + file_data=group_file_data, + reference_df=group_reference_df, + topic_summary_df=group_topic_summary_df, + file_name=file_name, + chosen_cols=chosen_cols, + batch_size=batch_size, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + max_tokens=max_tokens, + azure_api_key_textbox=azure_api_key_textbox, + azure_endpoint_textbox=azure_endpoint_textbox, + reasoning_suffix=reasoning_suffix, + group_name=current_group, + produce_structured_summary_radio=produce_structured_summary_radio, + force_zero_shot_radio=force_zero_shot_radio, + force_single_topic_radio=force_single_topic_radio, + context_textbox=context_textbox, + additional_instructions_summary_format=additional_instructions_summary_format, + output_folder=output_folder, + output_debug_files=output_debug_files, + original_full_file_name=original_full_file_name, + additional_validation_issues_provided=additional_validation_issues_provided, + max_time_for_loop=max_time_for_loop, + sentiment_checkbox=sentiment_checkbox, + logged_content=logged_content, + show_previous_table=show_previous_table, + aws_access_key_textbox=aws_access_key_textbox, + aws_secret_key_textbox=aws_secret_key_textbox, + aws_region_textbox=aws_region_textbox, + api_url=api_url, + max_topics_number=max_topics_number, + ) + + # Accumulate results + if not validation_reference_df.empty: + acc_reference_df = pd.concat( + [acc_reference_df, validation_reference_df], ignore_index=True + ) + acc_reference_df.drop_duplicates( + ["Response ID", "General topic", "Subtopic", "Sentiment"], + inplace=True, + ) + if not validation_topic_summary_df.empty: + acc_topic_summary_df = pd.concat( + [acc_topic_summary_df, validation_topic_summary_df], + ignore_index=True, + ) + acc_topic_summary_df.drop_duplicates( + _topic_summary_dedup_columns( + sentiment_checkbox, acc_topic_summary_df + ), + inplace=True, + ) + + acc_logged_content.extend(updated_logged_content) + acc_conversation_metadata += ( + ("\n---\n" if acc_conversation_metadata else "") + + f"Validation for group {current_group}:\n" + + validation_conversation_metadata_str + ) + + # Calculate token counts for this group + group_input_tokens, group_output_tokens, group_llm_calls = ( + calculate_tokens_from_metadata( + validation_conversation_metadata_str, model_choice, model_name_map + ) + ) + + acc_input_tokens += int(group_input_tokens) + acc_output_tokens += int(group_output_tokens) + acc_llm_calls += int(group_llm_calls) + + print(f"Group {current_group} validation completed.") + + except Exception as e: + print(f"Error processing validation for group {current_group}: {e}") + continue + + # Create consolidated output files + file_name_clean = clean_column_name(file_name, max_length=20) + clean_column_name(chosen_cols, max_length=20) + model_choice_clean = model_name_map[model_choice]["short_name"] + model_choice_clean_short = clean_column_name( + model_choice_clean, max_length=20, front_characters=False + ) + + # Create consolidated output file paths + validation_reference_table_path = f"{output_folder}{file_name_clean}_all_final_reference_table_{model_choice_clean_short}_valid.csv" + validation_unique_topics_path = f"{output_folder}{file_name_clean}_all_final_unique_topics_{model_choice_clean_short}_valid.csv" + + # Need to join "Topic number" onto acc_reference_df + # If any blanks, there is an issue somewhere, drop and redo + if "Topic number" in acc_reference_df.columns: + try: + # Convert to numeric first to avoid dtype issues with object arrays + topic_num_series = pd.to_numeric( + acc_reference_df["Topic number"], errors="coerce" + ) + if topic_num_series.isnull().any(): + acc_reference_df = acc_reference_df.drop("Topic number", axis=1) + except (TypeError, ValueError): + # If conversion fails, drop the column to be safe + acc_reference_df = acc_reference_df.drop("Topic number", axis=1) + + if "Topic number" not in acc_reference_df.columns: + if "Topic number" in acc_topic_summary_df.columns: + if "General topic" in acc_topic_summary_df.columns: + on_cols, summary_cols = _merge_topic_number_keys( + sentiment_checkbox, acc_reference_df, acc_topic_summary_df + ) + if on_cols and summary_cols: + acc_reference_df = acc_reference_df.merge( + acc_topic_summary_df[summary_cols], + on=on_cols, + how="left", + ) + # Sort output dataframes + acc_reference_df["Response ID"] = ( + acc_reference_df["Response ID"].astype(float).astype(int) + ) + acc_reference_df["Start row of group"] = acc_reference_df[ + "Start row of group" + ].astype(int) + ref_sort_cols = [ + "Group", + "Start row of group", + "Response ID", + "General topic", + "Subtopic", + ] + if ( + _assess_sentiment(sentiment_checkbox) + and "Sentiment" in acc_reference_df.columns + ): + ref_sort_cols.append("Sentiment") + acc_reference_df.sort_values(ref_sort_cols, inplace=True) + elif "Main heading" in acc_topic_summary_df.columns: + acc_reference_df = acc_reference_df.merge( + acc_topic_summary_df[ + ["Main heading", "Subheading", "Topic number"] + ], + on=["Main heading", "Subheading"], + how="left", + ) + # Sort output dataframes + acc_reference_df["Response ID"] = ( + acc_reference_df["Response ID"].astype(float).astype(int) + ) + acc_reference_df["Start row of group"] = acc_reference_df[ + "Start row of group" + ].astype(int) + acc_reference_df.sort_values( + [ + "Group", + "Start row of group", + "Response ID", + "Main heading", + "Subheading", + "Topic number", + ], + inplace=True, + ) + + if "General topic" in acc_topic_summary_df.columns: + acc_topic_summary_df["Number of responses"] = acc_topic_summary_df[ + "Number of responses" + ].astype(int) + v_topic_sort_cols = [ + "Group", + "Number of responses", + "General topic", + "Subtopic", + ] + v_topic_sort_asc = [True, False, True, True] + if ( + _assess_sentiment(sentiment_checkbox) + and "Sentiment" in acc_topic_summary_df.columns + ): + v_topic_sort_cols.append("Sentiment") + v_topic_sort_asc.append(True) + acc_topic_summary_df.sort_values( + v_topic_sort_cols, + ascending=v_topic_sort_asc, + inplace=True, + ) + elif "Main heading" in acc_topic_summary_df.columns: + acc_topic_summary_df["Number of responses"] = acc_topic_summary_df[ + "Number of responses" + ].astype(int) + acc_topic_summary_df.sort_values( + [ + "Group", + "Number of responses", + "Main heading", + "Subheading", + "Topic number", + ], + ascending=[True, False, True, True, True], + inplace=True, + ) + + # Save consolidated validation dataframes to CSV + if not acc_reference_df.empty: + acc_reference_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + validation_reference_table_path, index=None, encoding="utf-8-sig" + ) + acc_output_files.append(validation_reference_table_path) + + if not acc_topic_summary_df.empty: + acc_topic_summary_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + validation_unique_topics_path, index=None, encoding="utf-8-sig" + ) + acc_output_files.append(validation_unique_topics_path) + + if "Group" in acc_reference_df.columns: + # Create missing references dataframe using consolidated data from all groups + # This ensures we correctly identify missing references across all groups + # Get all basic_response_data from all groups + all_basic_response_data = list() + for logged_item in acc_logged_content: + if "basic_response_data" in logged_item: + all_basic_response_data.extend(logged_item["basic_response_data"]) + + if all_basic_response_data: + all_basic_response_df = pd.DataFrame(all_basic_response_data) + acc_missing_df = create_missing_references_df( + all_basic_response_df, acc_reference_df + ) + else: + # Fallback: if no logged content, create empty missing_df + acc_missing_df = pd.DataFrame( + columns=["Missing Response ID", "Response Character Count"] + ) + else: + # Fallback: if no logged content, create empty missing_df + acc_missing_df = pd.DataFrame( + columns=["Missing Response ID", "Response Character Count"] + ) + + # Create display table markdown for validation results + if not acc_topic_summary_df.empty: + validation_display_table = acc_topic_summary_df.copy() + if "Summary" in validation_display_table.columns: + validation_display_table = validation_display_table.drop("Summary", axis=1) + + # Apply text wrapping for display + validation_display_table = validation_display_table.apply( + lambda col: col.map(lambda x: wrap_text(x, max_text_length=max_text_length)) + ) + + # Handle structured summary format + if produce_structured_summary_radio == "Yes": + if "General topic" in validation_display_table.columns: + validation_display_table = validation_display_table.rename( + columns={"General topic": "Main heading"} + ) + if "Subtopic" in validation_display_table.columns: + validation_display_table = validation_display_table.rename( + columns={"Subtopic": "Subheading"} + ) + validation_display_table_markdown = validation_display_table.to_markdown( + index=False + ) + else: + validation_display_table_markdown = validation_display_table.to_markdown( + index=False + ) + else: + validation_display_table_markdown = "No validation results available." + + print( + f"Validation completed for all groups. Total tokens: {acc_input_tokens} input, {acc_output_tokens} output" + ) + + # Return the same format as wrapper_extract_topics_per_column_value + return ( + validation_display_table_markdown, # display_topic_table_markdown + acc_topic_summary_df, # master_unique_topics_df_state + acc_topic_summary_df, # master_unique_topics_df_state (duplicate for compatibility) + acc_reference_df, # master_reference_df_state + acc_output_files, # topic_extraction_output_files + acc_output_files, # text_output_file_list_state + 0, # latest_batch_completed (reset for validation) + [], # log_files_output (empty for validation) + [], # log_files_output_list_state (empty for validation) + acc_conversation_metadata, # conversation_metadata_textbox + 0.0, # estimated_time_taken_number (reset for validation) + acc_output_files, # deduplication_input_files + acc_output_files, # summarisation_input_files + acc_topic_summary_df, # modifiable_unique_topics_df_state + acc_output_files, # modification_input_files + [], # in_join_files (empty for validation) + acc_missing_df, # missing_df_state (empty for validation) + acc_input_tokens, # input_tokens_num + acc_output_tokens, # output_tokens_num + acc_llm_calls, # number_of_calls_num + f"Validation completed for {len(unique_groups)} groups", # output_messages_textbox + acc_logged_content, # logged_content_df + ) + + +def data_file_to_markdown_table( + file_data: pd.DataFrame, + file_name: str, + chosen_cols: List[str], + batch_number: int, + batch_size: int, + verify_titles: bool = False, +) -> Tuple[str, str, str]: + """ + Processes a file by simplifying its content based on chosen columns and saves the result to a specified output folder. + + Parameters: + - file_data (pd.DataFrame): Tabular data file with responses. + - file_name (str): File name with extension. + - chosen_cols (List[str]): A list of column names to include in the simplified file. + - batch_number (int): The current batch number for processing. + - batch_size (int): The number of rows to process in each batch. + + Returns: + - Tuple[str, str, str]: A tuple containing the path to the simplified CSV file, the simplified markdown table as a string, and the file path end (used for naming the output file). + """ + + normalised_simple_markdown_table = "" + simplified_csv_table_path = "" + + # Simplify table to just responses column and the Response reference number + basic_response_data = get_basic_response_data( + file_data, chosen_cols, verify_titles=verify_titles + ) + + file_len = int(len(basic_response_data["Response ID"])) + batch_size = int(batch_size) + batch_number = int(batch_number) + + # Subset the data for the current batch + start_row = int(batch_number * batch_size) + + if start_row > file_len + 1: + print("Start row greater than file row length") + return simplified_csv_table_path, normalised_simple_markdown_table, file_name + if start_row < 0: + raise Exception("Start row is below 0") + + if ((start_row + batch_size) - 1) <= file_len + 1: + end_row = int((start_row + batch_size) - 1) + else: + end_row = file_len + 1 + + batch_basic_response_data = basic_response_data.loc[ + start_row:end_row, ["Response ID", "Response", "Original Response ID"] + ] # Select the current batch + + # Now replace the reference numbers with numbers starting from 1 + batch_basic_response_data.loc[:, "Response ID"] = ( + batch_basic_response_data["Response ID"] - start_row + ) + + # Remove problematic characters including control characters, special characters, and excessive leading/trailing whitespace + batch_basic_response_data.loc[:, "Response"] = batch_basic_response_data[ + "Response" + ].str.replace( + r'[\x00-\x1F\x7F]|[""<>]|\\', "", regex=True + ) # Remove control and special characters + batch_basic_response_data.loc[:, "Response"] = batch_basic_response_data[ + "Response" + ].str.strip() # Remove leading and trailing whitespace + batch_basic_response_data.loc[:, "Response"] = batch_basic_response_data[ + "Response" + ].str.replace( + r"\s+", " ", regex=True + ) # Replace multiple spaces with a single space + batch_basic_response_data.loc[:, "Response"] = batch_basic_response_data[ + "Response" + ].str.replace( + r"\n{2,}", "\n", regex=True + ) # Replace multiple line breaks with a single line break + batch_basic_response_data.loc[:, "Response"] = batch_basic_response_data[ + "Response" + ].str.slice( + 0, max_comment_character_length + ) # Maximum 1,500 character responses + + # Remove blank and extremely short responses + batch_basic_response_data = batch_basic_response_data.loc[ + ~(batch_basic_response_data["Response"].isnull()) + & ~(batch_basic_response_data["Response"] == "None") + & ~(batch_basic_response_data["Response"] == " ") + & ~(batch_basic_response_data["Response"] == ""), + :, + ] # ~(batch_basic_response_data["Response"].str.len() < 5), :] + + simple_markdown_table = batch_basic_response_data[ + ["Response ID", "Response"] + ].to_markdown(index=None) + + normalised_simple_markdown_table = normalise_string(simple_markdown_table) + + # print("normalised_simple_markdown_table:", normalised_simple_markdown_table) + + return ( + simplified_csv_table_path, + normalised_simple_markdown_table, + start_row, + end_row, + batch_basic_response_data, + ) + + +def replace_punctuation_with_underscore(input_string: str): + # Create a translation table where each punctuation character maps to '_' + translation_table = str.maketrans(string.punctuation, "_" * len(string.punctuation)) + + # Translate the input string using the translation table + return input_string.translate(translation_table) + + +### INITIAL TOPIC MODEL DEVELOPMENT FUNCTIONS + + +def clean_markdown_table(text: str): + # Split text into lines + lines = text.splitlines() + + # Step 1: Identify table structure and process line continuations + table_rows = list() + current_row = None + + for line in lines: + stripped = line.strip() + + # Skip empty lines + if not stripped: + continue + + # Check if this is a table row or alignment row + is_table_row = "|" in stripped or stripped.startswith(":-") or ":-:" in stripped + + if is_table_row: + # If we have a current row being built, add it to our list + if current_row is not None: + table_rows.append(current_row) + + # Start a new row + current_row = stripped + elif current_row is not None: + # This must be a continuation of the previous row + current_row += " " + stripped + else: + # Not part of the table + current_row = stripped + + # Don't forget the last row + if current_row is not None: + table_rows.append(current_row) + + # Step 2: Properly format the table + # First, determine the maximum number of columns + max_columns = 0 + for row in table_rows: + cells = row.split("|") + # Account for rows that may not start/end with a pipe + if row.startswith("|"): + cells = cells[1:] + if row.endswith("|"): + cells = cells[:-1] + max_columns = max(max_columns, len(cells)) + + # Now format each row + formatted_rows = list() + for row in table_rows: + # Ensure the row starts and ends with pipes + if not row.startswith("|"): + row = "|" + row + if not row.endswith("|"): + row = row + "|" + + # Split into cells + cells = row.split("|")[1:-1] # Remove empty entries from split + + # Ensure we have the right number of cells + while len(cells) < max_columns: + cells.append("") + + # Rebuild the row + formatted_row = "|" + "|".join(cells) + "|" + formatted_rows.append(formatted_row) + + # Step 3: Remove duplicate rows while preserving header and separator rows + if len(formatted_rows) > 0: + # Keep the first row (header) + deduplicated_rows = [formatted_rows[0]] + seen_rows = {formatted_rows[0]} # Track seen rows + + # Check if second row is a separator (contains alignment markers) + separator_row = None + if len(formatted_rows) > 1: + second_row = formatted_rows[1] + # Check if it's a separator row (contains ":-", ":-:", or similar alignment markers) + if ":-" in second_row or ":-:" in second_row or "-:" in second_row: + separator_row = second_row + deduplicated_rows.append(separator_row) + seen_rows.add(separator_row) + + # Process data rows (skip header and separator) + start_idx = 2 if separator_row else 1 + for row in formatted_rows[start_idx:]: + # Only add if we haven't seen this exact row before + if row not in seen_rows: + deduplicated_rows.append(row) + seen_rows.add(row) + + formatted_rows = deduplicated_rows + + # Join everything back together + result = "\n".join(formatted_rows) + + return result + + +# Convert output table to markdown and then to a pandas dataframe to csv +def remove_before_last_term(input_string: str) -> str: + # Use regex to find the last occurrence of the term + match = re.search(r"(\| ?General topic)", input_string) + if match: + # Find the last occurrence by using rfind + last_index = input_string.rfind(match.group(0)) + return input_string[last_index:] # Return everything from the last match onward + return input_string # Return the original string if the term is not found + + +def convert_to_html_table(input_string: str, table_type: str = "Main table"): + # Remove HTML tags from input string + input_string = input_string.replace("

", "").replace("

", "") + + if " 1 else col + for i, col in enumerate(df.columns) + ] + + # Convert DataFrame to HTML + html_table = df.to_html(index=False, border=1) + + # Ensure that the HTML structure is correct + if table_type == "Main table": + if " + + General topic + Subtopic + Sentiment + Response ID + Summary + + {html_table} + + """ + elif table_type == "Revised topics table": + if " + + General topic + Subtopic + + {html_table} + + """ + elif table_type == "Verify titles table": + if " + + Response ID + Is this a suitable title + Explanation + Alternative title + + {html_table} + + """ + + return html_table + + +def _normalize_parsed_table_column_name(name: object) -> str: + return str(name).lower().strip().replace("_", " ") + + +def _find_parsed_table_column(df: pd.DataFrame, standard_name: str) -> str | None: + target = _normalize_parsed_table_column_name(standard_name) + for col in df.columns: + if _normalize_parsed_table_column_name(col) == target: + return col + return None + + +def _four_column_table_has_sentiment(df: pd.DataFrame) -> bool: + """Distinguish 4-column tables with Sentiment (batch_size==1) from those without.""" + if _find_parsed_table_column(df, "Sentiment") is not None: + return True + if _find_parsed_table_column(df, "Response ID") is not None: + return False + if df.shape[1] < 3: + return False + col2 = df.iloc[:, 2].astype(str).str.strip().str.lower() + sentiment_values = {"negative", "neutral", "positive", "not assessed"} + if col2.empty: + return False + return col2.isin(sentiment_values).mean() >= 0.5 + + +def convert_response_text_to_dataframe( + response_text: str, table_type: str = "Main table" +): + is_error = False + start_of_table_response = remove_before_last_term(response_text) + + cleaned_response = clean_markdown_table(start_of_table_response) + + # Add a space after commas between numbers (e.g., "1,2" -> "1, 2") + cleaned_response = re.sub(r"(\d),(\d)", r"\1, \2", cleaned_response) + + try: + string_html_table = markdown.markdown( + cleaned_response, extensions=["markdown.extensions.tables"] + ) + except Exception as e: + print("Unable to convert response to string_html_table due to", e) + string_html_table = "" + + html_table = convert_to_html_table(string_html_table) + + html_buffer = StringIO(html_table) + + try: + tables = pd.read_html(html_buffer) + if tables: + out_df = tables[0] # Use the first table if available + else: + raise ValueError("No tables found in the provided HTML.") + is_error = True + out_df = pd.DataFrame() + except Exception as e: + print("Error when trying to parse table:", e) + is_error = True + out_df = pd.DataFrame() + + # Fill down blank values in General Topic/General topic column if it exists + # Also remove markdown emphasis characters from topic columns + if not out_df.empty: + # Check for both possible column name variations + general_topic_col = None + for col in out_df.columns: + if col.lower().strip() in ["general topic", "general_topic"]: + general_topic_col = col + break + + if general_topic_col is not None: + # Fill down blank/NaN values in the General topic column + out_df[general_topic_col] = out_df[general_topic_col].ffill() + # Also fill empty strings if they appear after non-empty values + # Convert empty strings to NaN temporarily for ffill, then back + mask_empty = out_df[general_topic_col].astype(str).str.strip() == "" + out_df.loc[mask_empty, general_topic_col] = pd.NA + out_df[general_topic_col] = out_df[general_topic_col].ffill() + # Convert back any remaining NA to empty string if needed + out_df[general_topic_col] = out_df[general_topic_col].fillna("") + + # Remove markdown emphasis from General topic, Subtopic, and Sentiment columns + for col_name in ["General topic", "Subtopic", "Sentiment"]: + # Check for column name variations (case-insensitive, handle underscores) + matching_col = None + for col in out_df.columns: + if col.lower().strip().replace("_", " ") == col_name.lower(): + matching_col = col + break + + if matching_col is not None: + # Remove markdown emphasis characters + out_df[matching_col] = ( + out_df[matching_col].astype(str).apply(remove_markdown_emphasis) + ) + + return out_df, is_error + + +def write_llm_output_and_logs( + response_text: str, + whole_conversation: List[str], + all_metadata_content: List[str], + batch_file_path_details: str, + start_row: int, + end_row: int, + model_choice_clean: str, + log_files_output_paths: List[str], + existing_reference_df: pd.DataFrame, + existing_topics_df: pd.DataFrame, + batch_size_number: int, + batch_basic_response_df: pd.DataFrame, + group_name: str = "All", + produce_structured_summary_radio: str = "No", + return_logs: bool = False, + output_folder: str = OUTPUT_FOLDER, +) -> Tuple: + """ + Writes the output of the large language model requests and logs to files. + + Parameters: + - response_text (str): The text of the response from the model. + - whole_conversation (List[str]): A list of strings representing the complete conversation including prompts and responses. + - all_metadata_content (List[str]): A list of strings representing metadata about the whole conversation. + - batch_file_path_details (str): String containing details for constructing batch-specific file paths. + - start_row (int): Start row of the current batch. + - end_row (int): End row of the current batch. + - model_choice_clean (str): The cleaned model choice string. + - log_files_output_paths (List[str]): A list of paths to the log files. + - existing_reference_df (pd.DataFrame): The existing reference dataframe mapping response numbers to topics. + - existing_topics_df (pd.DataFrame): The existing unique topics dataframe. + - batch_size_number (int): The size of batches in terms of number of responses. + - batch_basic_response_df (pd.DataFrame): The dataframe that contains the response data. + - group_name (str, optional): The name of the current group. + - produce_structured_summary_radio (str, optional): Whether the option to produce structured summaries has been selected. + - return_logs (bool): A boolean indicating if logs should be returned. Defaults to False. + - output_folder (str): The name of the folder where output files are saved. + """ + topic_summary_df_out_path = list() + topic_table_out_path = "topic_table_error.csv" + reference_table_out_path = "reference_table_error.csv" + topic_summary_df_out_path = "unique_topic_table_error.csv" + topic_with_response_df = pd.DataFrame( + columns=[ + "General topic", + "Subtopic", + "Sentiment", + "Response ID", + "Summary", + ] + ) + out_reference_df = pd.DataFrame( + columns=[ + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + "Summary", + "Start row of group", + ] + ) + out_topic_summary_df = pd.DataFrame( + columns=["General topic", "Subtopic", "Sentiment"] + ) + is_error = False # If there was an error in parsing, return boolean saying error + + if produce_structured_summary_radio == "Yes": + existing_topics_df.rename( + columns={"Main heading": "General topic", "Subheading": "Subtopic"}, + inplace=True, + errors="ignore", + ) + existing_reference_df.rename( + columns={"Main heading": "General topic", "Subheading": "Subtopic"}, + inplace=True, + errors="ignore", + ) + topic_with_response_df.rename( + columns={"Main heading": "General topic", "Subheading": "Subtopic"}, + inplace=True, + errors="ignore", + ) + out_reference_df.rename( + columns={"Main heading": "General topic", "Subheading": "Subtopic"}, + inplace=True, + errors="ignore", + ) + out_topic_summary_df.rename( + columns={"Main heading": "General topic", "Subheading": "Subtopic"}, + inplace=True, + errors="ignore", + ) + + # Convert conversation to string and add to log outputs + whole_conversation_str = "\n".join(whole_conversation) + all_metadata_content_str = "\n".join(all_metadata_content) + start_row_reported = int(start_row) + 1 + + # Need to reduce output file names as full length files may be too long + model_choice_clean_short = clean_column_name( + model_choice_clean, max_length=20, front_characters=False + ) + + row_number_string_start = f"Rows {start_row_reported} to {end_row + 1}: " + + if output_debug_files == "True" and return_logs is True: + whole_conversation_path = ( + output_folder + + batch_file_path_details + + "_full_conversation_" + + model_choice_clean_short + + ".txt" + ) + whole_conversation_path_meta = ( + output_folder + + batch_file_path_details + + "_metadata_" + + model_choice_clean_short + + ".txt" + ) + with open( + whole_conversation_path, "w", encoding="utf-8-sig", errors="replace" + ) as f: + f.write(whole_conversation_str) + with open( + whole_conversation_path_meta, "w", encoding="utf-8-sig", errors="replace" + ) as f: + f.write(all_metadata_content_str) + log_files_output_paths.append(whole_conversation_path_meta) + + # Check if response is "No change" - if so, return input dataframes + stripped_response = response_text.strip() + + if stripped_response.lower().startswith("no change"): + print("LLM response indicates no changes needed, returning input dataframes") + + # For "No change" responses, we need to return the existing dataframes + # but we still need to process them through the same logic as normal processing + + # Create empty topic_with_response_df since no new topics were generated + if produce_structured_summary_radio == "Yes": + topic_with_response_df = pd.DataFrame( + columns=[ + "Main heading", + "Subheading", + "Sentiment", + "Response ID", + "Summary", + ] + ) + else: + topic_with_response_df = pd.DataFrame( + columns=[ + "General topic", + "Subtopic", + "Sentiment", + "Response ID", + "Summary", + ] + ) + + # For "No change", we return the existing dataframes as-is (they already contain all the data) + # This is equivalent to the normal processing where new_reference_df would be empty + out_reference_df = existing_reference_df.copy() + out_topic_summary_df = existing_topics_df.copy() + + # Set up output file paths + topic_table_out_path = ( + output_folder + + batch_file_path_details + + "_topic_table_" + + model_choice_clean_short + + ".csv" + ) + reference_table_out_path = ( + output_folder + + batch_file_path_details + + "_reference_table_" + + model_choice_clean_short + + ".csv" + ) + topic_summary_df_out_path = ( + output_folder + + batch_file_path_details + + "_unique_topics_" + + model_choice_clean_short + + ".csv" + ) + + # Return the existing dataframes (no changes needed) + return ( + topic_table_out_path, + reference_table_out_path, + topic_summary_df_out_path, + topic_with_response_df, + out_reference_df, + out_topic_summary_df, + batch_file_path_details, + is_error, + False, # has_incomplete_output + ) + + # Convert response text to a markdown table + try: + topic_with_response_df, is_error = convert_response_text_to_dataframe( + response_text + ) + except Exception as e: + print("Error in parsing markdown table from response text:", e) + + return ( + topic_table_out_path, + reference_table_out_path, + topic_summary_df_out_path, + topic_with_response_df, + out_reference_df, + out_topic_summary_df, + batch_file_path_details, + is_error, + False, # has_incomplete_output + ) + + # Check if the parsed dataframe has less than 3 columns (incomplete output) + # This will be used to trigger a retry + has_incomplete_output = False + if not is_error and not topic_with_response_df.empty: + # Check if dataframe has less than 3 columns (incomplete output format) + if topic_with_response_df.shape[1] < 3: + has_incomplete_output = True + + # If the table has 5 columns, rename them + # Rename columns to ensure consistent use of data frames later in code + if topic_with_response_df.shape[1] == 5: + new_column_names = { + topic_with_response_df.columns[0]: "General topic", + topic_with_response_df.columns[1]: "Subtopic", + topic_with_response_df.columns[2]: "Sentiment", + topic_with_response_df.columns[3]: "Response ID", + topic_with_response_df.columns[4]: "Summary", + } + + topic_with_response_df = topic_with_response_df.rename(columns=new_column_names) + elif topic_with_response_df.shape[1] == 4: + if _four_column_table_has_sentiment(topic_with_response_df): + # batch_size==1: General topic, Subtopic, Sentiment, Summary (no Response ID) + rename_map = {} + for standard_name in [ + "General topic", + "Subtopic", + "Sentiment", + "Summary", + ]: + found_col = _find_parsed_table_column( + topic_with_response_df, standard_name + ) + if found_col is not None: + rename_map[found_col] = standard_name + if len(rename_map) < 4: + rename_map.update( + { + topic_with_response_df.columns[0]: "General topic", + topic_with_response_df.columns[1]: "Subtopic", + topic_with_response_df.columns[2]: "Sentiment", + topic_with_response_df.columns[3]: "Summary", + } + ) + topic_with_response_df = topic_with_response_df.rename(columns=rename_map) + topic_with_response_df["Response ID"] = pd.Series( + ["1"] * len(topic_with_response_df), dtype=str + ) + else: + # 4-column case without Sentiment (Response ID present instead) + new_column_names = { + topic_with_response_df.columns[0]: "General topic", + topic_with_response_df.columns[1]: "Subtopic", + topic_with_response_df.columns[2]: "Response ID", + topic_with_response_df.columns[3]: "Summary", + } + topic_with_response_df = topic_with_response_df.rename( + columns=new_column_names + ) + topic_with_response_df["Sentiment"] = pd.Series( + ["Not assessed"] * len(topic_with_response_df), dtype=str + ) + topic_with_response_df = topic_with_response_df[ + ["General topic", "Subtopic", "Sentiment", "Response ID", "Summary"] + ] + else: + # Something went wrong with the table output, so add empty columns + print("Table output has wrong number of columns, adding with blank values") + + # Check for and handle duplicate column names which can cause DataFrame access issues + if topic_with_response_df.columns.duplicated().any(): + # Remove duplicate column names by adding suffix + topic_with_response_df.columns = [ + ( + f"{col}_{i}" + if topic_with_response_df.columns.tolist().count(col) > 1 + and topic_with_response_df.columns.tolist()[: i + 1].count(col) > 1 + else col + ) + for i, col in enumerate(topic_with_response_df.columns) + ] + + # First, rename first two columns that should always exist. + if len(topic_with_response_df.columns) >= 2: + new_column_names = { + topic_with_response_df.columns[0]: "General topic", + topic_with_response_df.columns[1]: "Subtopic", + } + topic_with_response_df.rename(columns=new_column_names, inplace=True) + else: + # If we don't have enough columns, create a minimal valid structure + topic_with_response_df = pd.DataFrame( + columns=[ + "General topic", + "Subtopic", + "Sentiment", + "Response ID", + "Summary", + ] + ) + + # Add empty columns if they are not present + # Ensure we're assigning to a Series, not accidentally creating a DataFrame + if "Sentiment" not in topic_with_response_df.columns: + topic_with_response_df["Sentiment"] = pd.Series( + ["Not assessed"] * len(topic_with_response_df), dtype=str + ) + if "Response ID" not in topic_with_response_df.columns: + if batch_size_number == 1: + topic_with_response_df["Response ID"] = pd.Series( + ["1"] * len(topic_with_response_df), dtype=str + ) + else: + topic_with_response_df["Response ID"] = pd.Series( + [""] * len(topic_with_response_df), dtype=str + ) + if "Summary" not in topic_with_response_df.columns: + topic_with_response_df["Summary"] = pd.Series( + [""] * len(topic_with_response_df), dtype=str + ) + + # Ensure we only have the expected columns and they're all Series + expected_cols = [ + "General topic", + "Subtopic", + "Sentiment", + "Response ID", + "Summary", + ] + topic_with_response_df = topic_with_response_df[expected_cols].copy() + + # Verify all columns are Series (not DataFrames) + for col in expected_cols: + if col in topic_with_response_df.columns: + if not isinstance(topic_with_response_df[col], pd.Series): + # If somehow it's a DataFrame, take the first column + topic_with_response_df[col] = ( + topic_with_response_df[col].iloc[:, 0] + if hasattr(topic_with_response_df[col], "iloc") + else topic_with_response_df[col] + ) + + # Fill in NA rows with values from above (topics seem to be included only on one row): + topic_with_response_df = topic_with_response_df.ffill() + + # Ensure we have a valid DataFrame with the expected columns before processing + if topic_with_response_df.empty: + # If DataFrame is empty, create a minimal valid structure + topic_with_response_df = pd.DataFrame( + columns=[ + "General topic", + "Subtopic", + "Sentiment", + "Response ID", + "Summary", + ] + ) + + # For instances where you end up with float values in Response ID + # Ensure we're working with a Series by using iloc if needed + if ( + "Response ID" in topic_with_response_df.columns + and not topic_with_response_df.empty + ): + # Handle case where column access might return DataFrame (e.g., duplicate column names) + try: + response_ref_series = topic_with_response_df["Response ID"] + # If it's a DataFrame (due to duplicate columns), take first column + if isinstance(response_ref_series, pd.DataFrame): + response_ref_series = response_ref_series.iloc[:, 0] + topic_with_response_df["Response ID"] = response_ref_series.astype( + str + ).str.replace(".0", "", regex=False) + except Exception as e: + print(f"Warning: Error processing Response ID column: {e}") + # Fallback: create a simple Series + topic_with_response_df["Response ID"] = topic_with_response_df[ + "Response ID" + ].astype(str) + + # Strip and lower case topic names to remove issues where model is randomly capitalising topics/sentiment + # Also remove markdown emphasis characters + # Ensure we're working with Series, not DataFrames + for col_name in ["General topic", "Subtopic", "Sentiment"]: + if ( + col_name in topic_with_response_df.columns + and not topic_with_response_df.empty + ): + try: + col_series = topic_with_response_df[col_name] + # If it's a DataFrame (due to duplicate columns), take first column + if isinstance(col_series, pd.DataFrame): + col_series = col_series.iloc[:, 0] + # Remove markdown emphasis, then strip, lower, and capitalize + topic_with_response_df[col_name] = ( + col_series.astype(str) + .apply(remove_markdown_emphasis) + .str.strip() + .str.lower() + .str.capitalize() + ) + except Exception as e: + print(f"Warning: Error processing {col_name} column: {e}") + # Fallback: just convert to string and remove markdown emphasis + topic_with_response_df[col_name] = ( + topic_with_response_df[col_name] + .astype(str) + .apply(remove_markdown_emphasis) + ) + + topic_table_out_path = ( + output_folder + + batch_file_path_details + + "_topic_table_" + + model_choice_clean_short + + ".csv" + ) + + # Table to map references to topics + reference_data = list() + existing_reference_numbers = False + + batch_basic_response_df["Response ID"] = batch_basic_response_df[ + "Response ID" + ].astype(str) + batch_size_number = int(batch_size_number) + + # Handle blank Response ID: remove rows if batch_size > 1, set to "1" if batch_size == 1 + # Skip this processing when producing structured summaries, as they don't have Response ID column + if ( + produce_structured_summary_radio != "Yes" + and "Response ID" in topic_with_response_df.columns + and not topic_with_response_df.empty + ): + # Convert Response ID to string and identify blank entries + topic_with_response_df["Response ID"] = topic_with_response_df[ + "Response ID" + ].astype(str) + + # Identify rows with invalid Response ID + # For batch_size > 1: blank, "nan", contains no digits, or contains "0" + # For batch_size == 1: only blank or "nan" (will be set to "1") + blank_or_invalid_mask = ( + topic_with_response_df["Response ID"].str.strip() == "" + ) | (topic_with_response_df["Response ID"].str.lower() == "nan") + + if batch_size_number > 1: + # Also check for non-numeric characters (no digits) or "0" when batch_size > 1 + no_digits_mask = ~topic_with_response_df["Response ID"].str.contains( + r"\d", regex=True, na=False + ) + contains_zero_mask = topic_with_response_df["Response ID"].str.contains( + r"\b0\b", regex=True, na=False + ) + + # Combine all invalid conditions for batch_size > 1 + invalid_mask = blank_or_invalid_mask | no_digits_mask | contains_zero_mask + + if invalid_mask.any(): + rows_removed = invalid_mask.sum() + print( + f"Removing {rows_removed} row(s) with invalid Response ID " + f"(blank, non-numeric, or '0' when batch_size > 1)" + ) + topic_with_response_df = topic_with_response_df[~invalid_mask].copy() + else: + # For batch_size == 1, only handle blank/nan (set to "1") + if blank_or_invalid_mask.any(): + rows_updated = blank_or_invalid_mask.sum() + print( + f"Setting {rows_updated} blank Response ID(s) to '1' (batch_size == 1)" + ) + topic_with_response_df.loc[blank_or_invalid_mask, "Response ID"] = "1" + + # Iterate through each row in the original DataFrame + for index, row in topic_with_response_df.iterrows(): + references_raw = str(row.iloc[3]) if pd.notna(row.iloc[3]) else "" + references = re.findall(r"\d+", references_raw) + + if batch_size_number == 1: + references = ["1"] + + # Filter out references that are outside the valid range + if references: + try: + # Convert all references to integers and keep only those within valid range + ref_numbers = [int(ref) for ref in references] + references = [ + ref + for ref in ref_numbers + if 1 <= int(ref) <= int(batch_size_number) + ] + except ValueError: + # If any reference can't be converted to int, skip this row + print("Response value could not be converted to number:", references) + continue + else: + references = [] + + topic = row.iloc[0] if pd.notna(row.iloc[0]) else "" + subtopic = row.iloc[1] if pd.notna(row.iloc[1]) else "" + sentiment = row.iloc[2] if pd.notna(row.iloc[2]) else "" + summary = row.iloc[4] if pd.notna(row.iloc[4]) else "" + + # If the reference response column is very long, and there's nothing in the summary column, assume that the summary was put in the reference column + if not summary and (len(str(row.iloc[3])) > 30): + summary = row.iloc[3] + + if produce_structured_summary_radio != "Yes": + summary = row_number_string_start + summary + + # Check if the 'references' list exists and is not empty + + if references: + existing_reference_numbers = True + + # We process one reference at a time to create one dictionary entry per reference. + for ref in references: + # This variable will hold the final reference number for the current 'ref' + response_ref_no = None + + # Now, we decide how to calculate 'response_ref_no' for the current 'ref' + if batch_basic_response_df.empty: + # --- Scenario 1: The DataFrame is empty, so we calculate the reference --- + try: + response_ref_no = int(ref) + int(start_row) + except ValueError: + print(f"Response ID '{ref}' is not a number and was skipped.") + continue # Skip to the next 'ref' in the loop + + else: + # --- Scenario 2: The DataFrame is NOT empty, so we look up the reference --- + matching_series = batch_basic_response_df.loc[ + batch_basic_response_df["Response ID"] == str(ref), + "Original Response ID", + ] + + if not matching_series.empty: + # If found, get the first match + response_ref_no = matching_series.iloc[0] + else: + # If not found, report it and skip this reference + print(f"Response ID '{ref}' not found in the DataFrame.") + continue # Skip to the next 'ref' in the loop + + # This code runs for every *valid* reference that wasn't skipped by 'continue'. + # It uses the 'response_ref_no' calculated in the if/else block above. + reference_data.append( + { + "Response ID": str(response_ref_no), + "General topic": topic, + "Subtopic": subtopic, + "Sentiment": sentiment, + "Summary": summary, + "Start row of group": start_row_reported, + } + ) + + # This 'else' corresponds to the 'if references:' at the top + else: + # This block runs only if the 'references' list was empty or None to begin with + existing_reference_numbers = False + response_ref_no = 0 # Default value when no references are provided + + reference_data.append( + { + "Response ID": str(response_ref_no), + "General topic": topic, + "Subtopic": subtopic, + "Sentiment": sentiment, + "Summary": summary, + "Start row of group": start_row_reported, + } + ) + + # Create a new DataFrame from the reference data + if reference_data: + new_reference_df = pd.DataFrame(reference_data) + else: + new_reference_df = pd.DataFrame( + columns=[ + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + "Summary", + "Start row of group", + ] + ) + + # Ensure new_reference_df has all required columns + required_cols = [ + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + "Summary", + "Start row of group", + ] + for col in required_cols: + if col not in new_reference_df.columns: + new_reference_df[col] = "" + + # Append on old reference data + if not new_reference_df.empty: + # Reset index to avoid reindexing errors with duplicate indices + new_reference_df = new_reference_df.reset_index(drop=True) + if not existing_reference_df.empty: + existing_reference_df = existing_reference_df.reset_index(drop=True) + # Ensure existing_reference_df also has all required columns + for col in required_cols: + if col not in existing_reference_df.columns: + existing_reference_df[col] = "" + out_reference_df = pd.concat( + [new_reference_df, existing_reference_df], ignore_index=True + ).dropna(how="all") + else: + out_reference_df = ( + existing_reference_df.copy() + if not existing_reference_df.empty + else new_reference_df + ).reset_index(drop=True) + # Ensure out_reference_df has all required columns + for col in required_cols: + if col not in out_reference_df.columns: + out_reference_df[col] = "" + + # Remove duplicate Response ID for the same topic + # Only if out_reference_df is not empty and has the required columns + if not out_reference_df.empty and "Response ID" in out_reference_df.columns: + out_reference_df.drop_duplicates( + ["Response ID", "General topic", "Subtopic", "Sentiment"], + inplace=True, + ) + + # Try converting response references column to int, keep as string if fails + if ( + existing_reference_numbers is True + and not out_reference_df.empty + and "Response ID" in out_reference_df.columns + ): + try: + out_reference_df["Response ID"] = ( + out_reference_df["Response ID"].astype(float).astype(int) + ) + except Exception as e: + print("Could not convert Response ID column to integer due to", e) + + # Only sort if out_reference_df is not empty and has the required columns + if not out_reference_df.empty and "Response ID" in out_reference_df.columns: + out_reference_df.sort_values( + [ + "Start row of group", + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + ], + inplace=True, + ) + + # Each topic should only be associated with each individual response once + out_reference_df.drop_duplicates( + ["Response ID", "General topic", "Subtopic", "Sentiment"], inplace=True + ) + out_reference_df["Group"] = group_name + + # Save the new DataFrame to CSV + reference_table_out_path = ( + output_folder + + batch_file_path_details + + "_reference_table_" + + model_choice_clean_short + + ".csv" + ) + + # Table of all unique topics with descriptions + new_topic_summary_df = topic_with_response_df[ + ["General topic", "Subtopic", "Sentiment"] + ] + + # new_topic_summary_df = new_topic_summary_df.rename( + # columns={ + # new_topic_summary_df.columns[0]: "General topic", + # new_topic_summary_df.columns[1]: "Subtopic", + # new_topic_summary_df.columns[2]: "Sentiment", + # } + # ) + + # Join existing and new unique topics + # Reset index to avoid reindexing errors with duplicate indices + new_topic_summary_df = new_topic_summary_df.reset_index(drop=True) + if not existing_topics_df.empty: + existing_topics_df = existing_topics_df.reset_index(drop=True) + out_topic_summary_df = pd.concat( + [new_topic_summary_df, existing_topics_df], ignore_index=True + ).dropna(how="all") + + out_topic_summary_df = out_topic_summary_df.rename( + columns={ + out_topic_summary_df.columns[0]: "General topic", + out_topic_summary_df.columns[1]: "Subtopic", + out_topic_summary_df.columns[2]: "Sentiment", + } + ) + + # print("out_topic_summary_df:", out_topic_summary_df) + + out_topic_summary_df = ( + out_topic_summary_df.drop_duplicates(["General topic", "Subtopic", "Sentiment"]) + .drop(["Number of responses", "Summary"], axis=1, errors="ignore") + .reset_index(drop=True) + ) + + # Get count of rows that refer to particular topics + # Only do this if out_reference_df is not empty and has the required columns + if not out_reference_df.empty and "Response ID" in out_reference_df.columns: + # Helper function to join unique summaries, preventing duplicates + def join_unique_summaries(x): + """Join unique summaries, handling both formats with and without '
' separators.""" + unique_summaries = [] + seen = set() + for s in x: + if pd.notna(s): # Skip NaN values + # Convert to string and strip whitespace + s_str = str(s).strip() + if s_str: # Skip empty strings + # If summary already contains
separators, split and deduplicate segments + if "
" in s_str or "
" in s_str: + # Split by
tags + segments = re.split(r"
|
", s_str) + segments = [seg.strip() for seg in segments if seg.strip()] + # Add unique segments + for segment in segments: + if segment and segment not in seen: + seen.add(segment) + unique_summaries.append(segment) + else: + # Single summary, add if not seen + if s_str not in seen: + seen.add(s_str) + unique_summaries.append(s_str) + return "
".join(unique_summaries) if unique_summaries else "" + + reference_counts = ( + out_reference_df.groupby(["General topic", "Subtopic", "Sentiment"]) + .agg( + { + "Response ID": "size", # Count the number of references + "Summary": join_unique_summaries, # Join unique summaries only + } + ) + .reset_index() + ) + # Rename the aggregated column to avoid merge conflicts + reference_counts = reference_counts.rename( + columns={"Response ID": "Number of responses"} + ) + + # Drop any existing "Response ID" or "Number of responses" columns from out_topic_summary_df + # to avoid duplicate columns after merge + out_topic_summary_df = out_topic_summary_df.drop( + ["Response ID", "Number of responses"], axis=1, errors="ignore" + ) + + # Join the counts to existing_topic_summary_df + out_topic_summary_df = out_topic_summary_df.merge( + reference_counts, how="left", on=["General topic", "Subtopic", "Sentiment"] + ) + + # Fill NaN values with 0 for Number of responses + if "Number of responses" in out_topic_summary_df.columns: + out_topic_summary_df["Number of responses"] = ( + out_topic_summary_df["Number of responses"].fillna(0).astype(int) + ) + out_topic_summary_df = out_topic_summary_df.sort_values( + "Number of responses", ascending=False + ) + else: + # If merge didn't add the column, add it with default values + out_topic_summary_df["Number of responses"] = 0 + else: + # If out_reference_df is empty or missing Response ID, add default column + out_topic_summary_df["Number of responses"] = 0 + if "Summary" not in out_topic_summary_df.columns: + out_topic_summary_df["Summary"] = "" + + out_topic_summary_df["Group"] = group_name + + topic_summary_df_out_path = ( + output_folder + + batch_file_path_details + + "_unique_topics_" + + model_choice_clean_short + + ".csv" + ) + + return ( + topic_table_out_path, + reference_table_out_path, + topic_summary_df_out_path, + topic_with_response_df, + out_reference_df, + out_topic_summary_df, + batch_file_path_details, + is_error, + has_incomplete_output, + ) + + +def process_batch_with_llm( + is_first_batch: bool, + formatted_system_prompt: str, + formatted_prompt: str, + batch_file_path_details: str, + model_source: str, + model_choice: str, + in_api_key: str, + temperature: float, + max_tokens: int, + azure_api_key_textbox: str, + azure_endpoint_textbox: str, + reasoning_suffix: str, + local_model: object, + tokenizer: object, + bedrock_runtime: object, + reported_batch_no: int, + response_text: str, + whole_conversation: list, + all_metadata_content: list, + start_row: int, + end_row: int, + model_choice_clean: str, + log_files_output_paths: list, + existing_reference_df: pd.DataFrame, + existing_topic_summary_df: pd.DataFrame, + batch_size: int, + batch_basic_response_df: pd.DataFrame, + group_name: str, + produce_structured_summary_radio: str, + output_folder: str, + output_debug_files: str, + task_type: str, + assistant_prefill: str = "", + api_url: str = None, +): + """Helper function to process a batch with LLM, handling the common logic between first and subsequent batches. + + This function orchestrates the interaction with various LLM providers (Gemini, Azure/OpenAI, AWS Bedrock, Local) + to process a given batch of data. It constructs the client, calls the LLM with the specified prompts, + and then processes the LLM's response to extract topics, references, and summaries, writing them to output files. + It also handles error conditions related to LLM output parsing. + + Args: + is_first_batch (bool): True if this is the first batch being processed, False otherwise. + formatted_system_prompt (str): The system prompt to be sent to the LLM. + formatted_prompt (str): The main user prompt for the LLM. + batch_file_path_details (str): String containing details for constructing batch-specific file paths. + model_source (str): The source of the LLM (e.g., "Gemini", "Azure/OpenAI", "AWS Bedrock", "Local"). + model_choice (str): The specific model chosen (e.g., "gemini-pro", "gpt-4", "anthropic.claude-v2"). + in_api_key (str): API key for the chosen model source (if applicable). + temperature (float): The sampling temperature for the LLM, controlling randomness. + max_tokens (int): The maximum number of tokens to generate in the LLM's response. + azure_api_key_textbox (str): API key for Azure OpenAI (if `model_source` is Azure/OpenAI). + azure_endpoint_textbox (str): Endpoint URL for Azure OpenAI (if `model_source` is Azure/OpenAI). + reasoning_suffix (str): Additional text to append to the system prompt for reasoning (primarily for local models). + local_model (object): The loaded local model object (if `model_source` is Local). + tokenizer (object): The tokenizer object for local models. + bedrock_runtime (object): AWS Bedrock runtime client object (if `model_source` is AWS Bedrock). + reported_batch_no (int): The current batch number being processed and reported. + response_text (str): The raw text response from the LLM (can be pre-filled or from a previous step). + whole_conversation (list): A list representing the entire conversation history. + all_metadata_content (list): Metadata associated with each turn in the conversation. + start_row (int): The starting row index of the current batch in the original dataset. + end_row (int): The ending row index of the current batch in the original dataset. + model_choice_clean (str): A cleaned, short name for the chosen model. + log_files_output_paths (list): A list of paths to log files generated during processing. + existing_reference_df (pd.DataFrame): DataFrame containing existing reference data. + existing_topic_summary_df (pd.DataFrame): DataFrame containing existing topic summary data. + batch_size (int): The number of items processed in each batch. + batch_basic_response_df (pd.DataFrame): DataFrame containing basic responses for the current batch. + group_name (str): The name of the group associated with the current batch. + produce_structured_summary_radio (str): Indicates whether to produce structured summaries ("Yes" or "No"). + output_folder (str): The directory where all output files will be saved. + output_debug_files (str): Flag indicating whether to output debug files ("Yes" or "No"). + task_type (str): The type of task being performed (e.g., "topic_extraction", "summarisation"). + assistant_prefill (str, optional): Optional prefill text for the assistant's response. Defaults to "". + + Returns: + tuple: A tuple containing various output paths and DataFrames after processing the batch: + - topic_table_out_path (str): Path to the output CSV for the topic table. + - reference_table_out_path (str): Path to the output CSV for the reference table. + - topic_summary_df_out_path (str): Path to the output CSV for the topic summary DataFrame. + - new_topic_df (pd.DataFrame): DataFrame of newly extracted topics. + - new_reference_df (pd.DataFrame): DataFrame of newly extracted references. + - new_topic_summary_df (pd.DataFrame): DataFrame of the updated topic summary. + - batch_file_path_details (str): The batch file path details used. + - is_error (bool): True if an error occurred during processing, False otherwise. + """ + client = list() + client_config = dict() + + # Prepare clients before query + if "Gemini" in model_source: + print("Using Gemini model:", model_choice) + client, client_config = construct_gemini_generative_model( + in_api_key=in_api_key, + temperature=temperature, + model_choice=model_choice, + system_prompt=formatted_system_prompt, + max_tokens=max_tokens, + ) + elif "Azure/OpenAI" in model_source: + print("Using Azure/OpenAI AI Inference model:", model_choice) + if azure_api_key_textbox: + os.environ["AZURE_INFERENCE_CREDENTIAL"] = azure_api_key_textbox + client, client_config = construct_azure_client( + in_api_key=azure_api_key_textbox, endpoint=azure_endpoint_textbox + ) + elif "AWS" in model_source: + print("Using AWS Bedrock model:", model_choice) + pass + elif "Local" in model_source: + print("Using local model:", model_choice) + pass + elif "inference-server" in model_source: + print("Using inference-server model:", model_choice) + pass + else: + raise ValueError(f"Unsupported model source: {model_source}") + + batch_prompts = [formatted_prompt] + + # Apply reasoning suffix for GPT-OSS models (Local, inference-server, or AWS) + is_gpt_oss_model = ( + "gpt-oss" in model_choice.lower() or "gpt_oss" in model_choice.lower() + ) + + if is_gpt_oss_model: + # Use default reasoning suffix if not set + effective_reasoning_suffix = ( + reasoning_suffix if reasoning_suffix else "Reasoning: low" + ) + if effective_reasoning_suffix: + formatted_system_prompt = ( + formatted_system_prompt + "\n" + effective_reasoning_suffix + ) + elif "Local" in model_source and reasoning_suffix: + # For other local models, use reasoning_suffix if provided + formatted_system_prompt = formatted_system_prompt + "\n" + reasoning_suffix + + # Combine system prompt and user prompt for token counting + full_input_text = formatted_system_prompt + "\n" + formatted_prompt + + # Count tokens in the input text + from tools.dedup_summaries import ( + _inference_server_prompt_token_count, + count_tokens_in_text, + ) + + if "inference-server" in str(model_source).lower() and api_url: + messages = [ + {"role": "system", "content": formatted_system_prompt}, + {"role": "user", "content": formatted_prompt}, + ] + input_token_count = _inference_server_prompt_token_count(api_url, messages) + if input_token_count is None: + input_token_count = count_tokens_in_text( + full_input_text, tokenizer, model_source + ) + print( + "Warning: Using approximate token count for inference-server; " + "server-side tokenization endpoint not available." + ) + else: + input_token_count = count_tokens_in_text( + full_input_text, tokenizer, model_source + ) + + # Apply context headroom to reduce risk of server-side tokenization mismatch + try: + _headroom = float(LLM_CONTEXT_HEADROOM_FRACTION) + except Exception: + _headroom = 0.05 + _headroom = min(max(_headroom, 0.0), 0.5) + effective_limit = max(1, int(LLM_CONTEXT_LENGTH * (1.0 - _headroom))) + + # Check if input exceeds context length (with headroom) + if input_token_count > effective_limit: + error_message = ( + f"Input text exceeds LLM context length (with headroom). " + f"Input tokens: {input_token_count}, Max (effective): {effective_limit} " + f"(raw max: {LLM_CONTEXT_LENGTH}). Please reduce the input text size." + ) + print(error_message) + raise ValueError(error_message) + + print(f"Input token count: {input_token_count} (Max effective: {effective_limit})") + + conversation_history = list() + whole_conversation = list() + + # Process requests to large language model with retry logic for incomplete outputs + max_retries = 3 + retry_count = 0 + retry_needed = True + original_temperature = temperature + current_temperature = temperature + + while retry_needed and retry_count < max_retries: + # Increase temperature by 0.1 on each retry (not on first attempt) + if retry_count > 0: + current_temperature = original_temperature + (retry_count * 0.1) + print( + f"Increasing temperature to {current_temperature:.1f} for retry attempt {retry_count + 1}" + ) + + # Recreate client with updated temperature if using Gemini + if "Gemini" in model_source: + client, client_config = construct_gemini_generative_model( + in_api_key=in_api_key, + temperature=current_temperature, + model_choice=model_choice, + system_prompt=formatted_system_prompt, + max_tokens=max_tokens, + ) + + ( + responses, + conversation_history, + whole_conversation, + all_metadata_content, + response_text, + ) = call_llm_with_markdown_table_checks( + batch_prompts, + formatted_system_prompt, + conversation_history, + whole_conversation, + all_metadata_content, + client, + client_config, + model_choice, + current_temperature, + reported_batch_no, + local_model, + tokenizer, + bedrock_runtime, + model_source, + MAX_OUTPUT_VALIDATION_ATTEMPTS, + assistant_prefill=assistant_prefill, + master=not is_first_batch, + api_url=api_url, + ) + + # print("Response text:", response_text) + + # Return output tables + ( + topic_table_out_path, + reference_table_out_path, + topic_summary_df_out_path, + new_topic_df, + new_reference_df, + new_topic_summary_df, + batch_file_path_details, + is_error, + has_incomplete_output, + ) = write_llm_output_and_logs( + response_text, + whole_conversation, + all_metadata_content, + batch_file_path_details, + start_row, + end_row, + model_choice_clean, + log_files_output_paths, + existing_reference_df, + existing_topic_summary_df, + batch_size, + batch_basic_response_df, + group_name, + produce_structured_summary_radio, + output_folder=output_folder, + ) + + # Check if output has less than 3 columns (incomplete output format) + # This indicates the LLM didn't follow the format properly + if has_incomplete_output: + retry_count += 1 + if retry_count < max_retries: + next_temperature = original_temperature + (retry_count * 0.1) + print( + f"LLM output table has less than 3 columns (incomplete format). " + f"Retrying LLM call with increased temperature {next_temperature:.1f} " + f"(attempt {retry_count + 1}/{max_retries})..." + ) + retry_needed = True + continue + else: + print( + f"LLM output still incomplete after {max_retries} attempts. Proceeding with incomplete data." + ) + retry_needed = False + else: + retry_needed = False + + # If error in table parsing, leave function + if is_error is True: + if is_first_batch: + raise Exception("Error in output table parsing") + else: + final_message_out = "Could not complete summary, error in LLM output." + raise Exception(final_message_out) + + # Write final output to text file and objects for logging purposes + full_prompt = formatted_system_prompt + "\n" + formatted_prompt + + ( + current_prompt_content_logged, + current_summary_content_logged, + current_conversation_content_logged, + current_metadata_content_logged, + ) = process_debug_output_iteration( + output_debug_files, + output_folder, + batch_file_path_details, + model_choice_clean, + full_prompt, + response_text, + whole_conversation, + all_metadata_content, + log_files_output_paths, + task_type=task_type, + ) + + # print("Finished processing batch with LLM") + + return ( + new_topic_df, + new_reference_df, + new_topic_summary_df, + is_error, + current_prompt_content_logged, + current_summary_content_logged, + current_conversation_content_logged, + current_metadata_content_logged, + topic_table_out_path, + reference_table_out_path, + topic_summary_df_out_path, + ) + + +def extract_topics( + in_data_file: gr.FileData, + file_data: pd.DataFrame, + existing_topics_table: pd.DataFrame, + existing_reference_df: pd.DataFrame, + existing_topic_summary_df: pd.DataFrame, + unique_table_df_display_table_markdown: str, + file_name: str, + num_batches: int, + in_api_key: str, + temperature: float, + chosen_cols: List[str], + model_choice: str, + candidate_topics: gr.FileData = None, + latest_batch_completed: int = 0, + out_message: List = list(), + out_file_paths: List = list(), + log_files_output_paths: List = list(), + first_loop_state: bool = False, + all_metadata_content_str: str = "", + initial_table_prompt: str = initial_table_prompt, + initial_table_system_prompt: str = initial_table_system_prompt, + add_existing_topics_system_prompt: str = add_existing_topics_system_prompt, + add_existing_topics_prompt: str = add_existing_topics_prompt, + number_of_prompts_used: int = 1, + batch_size: int = 5, + context_textbox: str = "", + time_taken: float = 0, + sentiment_checkbox: str = "Negative, Neutral, or Positive", + force_zero_shot_radio: str = "No", + in_excel_sheets: List[str] = list(), + force_single_topic_radio: str = "No", + output_folder: str = OUTPUT_FOLDER, + force_single_topic_prompt: str = force_single_topic_prompt, + group_name: str = "All", + produce_structured_summary_radio: str = "No", + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + hf_api_key_textbox: str = "", + azure_api_key_textbox: str = "", + azure_endpoint_textbox: str = "", + max_tokens: int = max_tokens, + model_name_map: dict = model_name_map, + existing_logged_content: list = list(), + max_time_for_loop: int = max_time_for_loop, + CHOSEN_LOCAL_MODEL_TYPE: str = CHOSEN_LOCAL_MODEL_TYPE, + reasoning_suffix: str = reasoning_suffix, + output_debug_files: str = output_debug_files, + model: object = list(), + tokenizer: object = list(), + assistant_model: object = list(), + max_rows: int = max_rows, + original_full_file_name: str = "", + additional_instructions_summary_format: str = "", + additional_validation_issues_provided: str = "", + api_url: str = None, + max_topics_number: int = MAXIMUM_ALLOWED_TOPICS, + progress=Progress(track_tqdm=True), +): + """ + Query an LLM (local, (Gemma/GPT-OSS if local, Gemini, AWS Bedrock or Azure/OpenAI AI Inference) with up to three prompts about a table of open text data. Up to 'batch_size' rows will be queried at a time. + + Parameters: + - in_data_file (gr.File): Gradio file object containing input data + - file_data (pd.DataFrame): Pandas dataframe containing the consultation response data. + - existing_topics_table (pd.DataFrame): Pandas dataframe containing the latest master topic table that has been iterated through batches. + - existing_reference_df (pd.DataFrame): Pandas dataframe containing the list of Response reference numbers alongside the derived topics and subtopics. + - existing_topic_summary_df (pd.DataFrame): Pandas dataframe containing the unique list of topics, subtopics, sentiment and summaries until this point. + - unique_table_df_display_table_markdown (str): Table for display in markdown format. + - file_name (str): File name of the data file. + - num_batches (int): Number of batches required to go through all the response rows. + - in_api_key (str): The API key for authentication (Google Gemini). + - temperature (float): The temperature parameter for the model. + - chosen_cols (List[str]): A list of chosen columns to process. + - candidate_topics (gr.FileData): File with a table of existing candidate topics files submitted by the user. + - model_choice (str): The choice of model to use. + - latest_batch_completed (int): The index of the latest file completed. + - out_message (list): A list to store output messages. + - out_file_paths (list): A list to store output file paths. + - log_files_output_paths (list): A list to store log file output paths. + - first_loop_state (bool): A flag indicating the first loop state. + - all_metadata_content_str (str): A string to store whole conversation metadata. + - initial_table_prompt (str): The first prompt for the model. + - initial_table_system_prompt (str): The system prompt for the model. + - add_existing_topics_system_prompt (str): The system prompt for the summary part of the model. + - add_existing_topics_prompt (str): The prompt for the model summary. + - number of requests (int): The number of prompts to send to the model. + - batch_size (int): The number of data rows to consider in each request. + - context_textbox (str, optional): A string giving some context to the consultation/task. + - time_taken (float, optional): The amount of time taken to process the responses up until this point. + - sentiment_checkbox (str, optional): What type of sentiment analysis should the topic modeller do? + - force_zero_shot_radio (str, optional): Should responses be forced into a zero shot topic or not. + - in_excel_sheets (List[str], optional): List of excel sheets to load from input file. + - force_single_topic_radio (str, optional): Should the model be forced to assign only one single topic to each response (effectively a classifier). + - produce_structured_summary_radio (str, optional): Should the model create a structured summary instead of extracting topics. + - output_folder (str, optional): Output folder where results will be stored. + - force_single_topic_prompt (str, optional): The prompt for forcing the model to assign only one single topic to each response. + - aws_access_key_textbox (str, optional): AWS access key for account with Bedrock permissions. + - aws_secret_key_textbox (str, optional): AWS secret key for account with Bedrock permissions. + - hf_api_key_textbox (str, optional): Hugging Face API key for account with Hugging Face permissions. + - max_tokens (int): The maximum number of tokens for the model. + - model_name_map (dict, optional): A dictionary mapping full model name to shortened. + - existing_logged_content (list, optional): A list of existing logged content. + - max_time_for_loop (int, optional): The number of seconds maximum that the function should run for before breaking (to run again, this is to avoid timeouts with some AWS services if deployed there). + - CHOSEN_LOCAL_MODEL_TYPE (str, optional): The name of the chosen local model. + - reasoning_suffix (str, optional): The suffix for the reasoning system prompt. + - output_debug_files (str, optional): Flag indicating whether to output debug files ("True" or "False"). + - model: Model object for local inference. + - tokenizer: Tokenizer object for local inference. + - assistant_model: Assistant model object for local inference. + - max_rows: The maximum number of rows to process. + - original_full_file_name: The original full file name. + - additional_instructions_summary_format: Initial instructions to guide the format for the initial summary of the topics. + - additional_validation_issues_provided: Additional validation issues provided by the user. + - progress (Progress): A progress tracker. + + """ + + # Ensure custom model_choice is registered in model_name_map + ensure_model_in_map(model_choice, model_name_map) + + tic = time.perf_counter() + + final_time = 0.0 + all_metadata_content = list() + create_revised_general_topics = False + local_model = None + tokenizer = None + zero_shot_topics_df = pd.DataFrame() + missing_df = pd.DataFrame() + new_reference_df = pd.DataFrame( + columns=[ + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + "Start row of group", + "Group", + "Topic number", + "Summary", + ] + ) + new_topic_summary_df = pd.DataFrame( + columns=[ + "General topic", + "Subtopic", + "Sentiment", + "Group", + "Number of responses", + "Summary", + ] + ) + if existing_topic_summary_df.empty: + existing_topic_summary_df = pd.DataFrame( + columns=[ + "General topic", + "Subtopic", + "Sentiment", + "Group", + "Number of responses", + "Summary", + ] + ) + if existing_reference_df.empty: + existing_reference_df = pd.DataFrame( + columns=[ + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + "Start row of group", + "Group", + "Topic number", + "Summary", + ] + ) + + # Clean topic names in initial dataframes to ensure consistent formatting throughout processing + if not existing_reference_df.empty: + for col_name in ["General topic", "Subtopic"]: + if ( + col_name in existing_reference_df.columns + and not existing_reference_df[col_name].isnull().all() + ): + existing_reference_df[col_name] = existing_reference_df[col_name].apply( + normalize_topic_name_for_llm + ) + if not existing_topic_summary_df.empty: + for col_name in ["General topic", "Subtopic"]: + if ( + col_name in existing_topic_summary_df.columns + and not existing_topic_summary_df[col_name].isnull().all() + ): + existing_topic_summary_df[col_name] = existing_topic_summary_df[ + col_name + ].apply(normalize_topic_name_for_llm) + if not existing_topics_table.empty: + for col_name in ["General topic", "Subtopic"]: + if ( + col_name in existing_topics_table.columns + and not existing_topics_table[col_name].isnull().all() + ): + existing_topics_table[col_name] = existing_topics_table[col_name].apply( + normalize_topic_name_for_llm + ) + + new_topic_df = pd.DataFrame( + columns=[ + "General topic", + "Subtopic", + "Sentiment", + "Group", + "Number of responses", + "Summary", + ] + ) + pd.DataFrame( + columns=[ + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + "Start row of group", + "Group", + "Topic number", + "Summary", + ] + ) + pd.DataFrame( + columns=[ + "General topic", + "Subtopic", + "Sentiment", + "Group", + "Number of responses", + "Summary", + ] + ) + task_type = "Topic extraction" + + # Logged content + all_prompts_content = list() + all_responses_content = list() + all_conversation_content = list() + all_metadata_content = list() + all_groups_content = list() + all_batches_content = list() + all_model_choice_content = list() + all_validated_content = list() + all_task_type_content = list() + all_file_names_content = list() + all_groups_logged_content = list() + # Need to reduce output file names as full length files may be too long + model_choice_clean = model_name_map[model_choice]["short_name"] + model_choice_clean_short = clean_column_name( + model_choice_clean, max_length=20, front_characters=False + ) + in_column_cleaned = clean_column_name(chosen_cols, max_length=20) + file_name_clean = clean_column_name( + file_name, max_length=20, front_characters=False + ) + + # If you have a file input but no file data it hasn't yet been loaded. Load it here. + if file_data.empty: + print("No data table found, loading from file") + try: + ( + in_colnames_drop, + in_excel_sheets, + file_name, + join_colnames, + join_colnames_drop, + ) = put_columns_in_df(in_data_file) + file_data, file_name, num_batches = load_in_data_file( + in_data_file, chosen_cols, batch_size_default, in_excel_sheets + ) + except Exception as e: + out_message = "Could not load in data file due to: " + str(e) + print(out_message) + raise Exception(out_message) + + if file_data.shape[0] > max_rows: + out_message = ( + "Your data has more than " + + str(max_rows) + + " rows, which has been set as the maximum in the application configuration." + ) + print(out_message) + raise Exception(out_message) + + model_choice_clean = model_name_map[model_choice]["short_name"] + model_source = model_name_map[model_choice]["source"] + + bedrock_runtime = connect_to_bedrock_runtime( + model_name_map, + model_choice, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + ) + + # If this is the first time around, set variables to 0/blank + if first_loop_state is True: + if (latest_batch_completed == 999) | (latest_batch_completed == 0): + latest_batch_completed = 0 + out_message = list() + out_file_paths = list() + final_time = 0 + + if (model_source == "Local") & (RUN_LOCAL_MODEL == "1") & (not model): + progress(0.1, f"Using local model: {model_choice_clean}") + local_model = get_model() + tokenizer = get_tokenizer() + get_assistant_model() + + if num_batches > 0: + progress_measure = round(latest_batch_completed / num_batches, 1) + progress(progress_measure, desc="Querying large language model") + else: + progress(0.1, desc="Querying large language model") + + latest_batch_completed = int(latest_batch_completed) + num_batches = int(num_batches) + + if latest_batch_completed < num_batches: + + # Load file + # If out message or out_file_paths are blank, change to a list so it can be appended to + if isinstance(out_message, str): + out_message = [out_message] + + if not out_file_paths: + out_file_paths = list() + + if "anthropic.claude-3-sonnet" in model_choice and file_data.shape[1] > 300: + out_message = "Your data has more than 300 rows, using the Sonnet model will be too expensive. Please choose the Haiku model instead." + print(out_message) + raise Exception(out_message) + + sentiment_prefix = "In the next column named 'Sentiment', " + sentiment_suffix = "." + if sentiment_checkbox == "Negative, Neutral, or Positive": + sentiment_prompt = ( + sentiment_prefix + + negative_neutral_positive_sentiment_prompt + + sentiment_suffix + ) + elif sentiment_checkbox == "Negative or Positive": + sentiment_prompt = ( + sentiment_prefix + + negative_or_positive_sentiment_prompt + + sentiment_suffix + ) + elif sentiment_checkbox == "Do not assess sentiment": + sentiment_prompt = "" # Just remove line completely. Previous: sentiment_prefix + do_not_assess_sentiment_prompt + sentiment_suffix + else: + sentiment_prompt = ( + sentiment_prefix + default_sentiment_prompt + sentiment_suffix + ) + + if context_textbox and "The context of this analysis is" not in context_textbox: + context_textbox = ( + "The context of this analysis is '" + context_textbox + "'." + ) + + topics_loop_description = ( + "Extracting topics from response batches (each batch of " + + str(batch_size) + + " responses)." + ) + total_batches_to_do = num_batches - latest_batch_completed + topics_loop = progress.tqdm( + range(total_batches_to_do), + desc=topics_loop_description, + unit="batches remaining", + ) + + for i in topics_loop: + reported_batch_no = latest_batch_completed + 1 + print("Running response batch:", reported_batch_no) + + # Call the function to prepare the input table + ( + simplified_csv_table_path, + normalised_simple_markdown_table, + start_row, + end_row, + batch_basic_response_df, + ) = data_file_to_markdown_table( + file_data, file_name, chosen_cols, latest_batch_completed, batch_size + ) + + if batch_basic_response_df.shape[0] == 1: + response_reference_format = "" # Blank, as the topics will always refer to the single response provided, '1' + else: + response_reference_format = "\n" + default_response_reference_format + + # If the response table is not empty, add it to the prompt with an intro line + if normalised_simple_markdown_table: + response_table_prompt = ( + "Response table:\n" + normalised_simple_markdown_table + ) + else: + response_table_prompt = "" + + existing_topic_summary_df.rename( + columns={"Main heading": "General topic", "Subheading": "Subtopic"}, + inplace=True, + errors="ignore", + ) + existing_reference_df.rename( + columns={"Main heading": "General topic", "Subheading": "Subtopic"}, + inplace=True, + errors="ignore", + ) + + # If the latest batch of responses contains at least one instance of text + if not batch_basic_response_df.empty: + + # If this is the second batch, the master table will refer back to the current master table when assigning topics to the new table. Also runs if there is an existing list of topics supplied by the user + if latest_batch_completed >= 1 or candidate_topics is not None: + + formatted_system_prompt = add_existing_topics_system_prompt.format( + consultation_context=context_textbox, column_name=chosen_cols + ) + + # Preparing candidate topics if no topics currently exist + if candidate_topics and existing_topic_summary_df.empty: + + # 'Zero shot topics' are those supplied by the user + # Handle both string paths (CLI) and gr.FileData objects (Gradio) + # Supports CSV, Excel (.xlsx), and Parquet files + candidate_topics_path = ( + candidate_topics + if isinstance(candidate_topics, str) + else getattr(candidate_topics, "name", None) + ) + if candidate_topics_path is None: + raise ValueError( + "candidate_topics must be a file path string or a FileData object with a 'name' attribute" + ) + + # Read the file (supports CSV, Excel .xlsx, and Parquet) + # For Excel files, reads the first sheet by default + try: + zero_shot_topics = read_file(candidate_topics_path) + except Exception as e: + raise ValueError( + f"Error reading candidate topics file '{candidate_topics_path}': {str(e)}. " + f"Supported formats: CSV (.csv), Excel (.xlsx), and Parquet (.parquet). " + f"For Excel files, the first sheet will be used." + ) from e + + zero_shot_topics = zero_shot_topics.fillna( + "" + ) # Replace NaN with empty string + zero_shot_topics = zero_shot_topics.astype(str) + + zero_shot_topics_df = generate_zero_shot_topics_df( + zero_shot_topics, + force_zero_shot_radio, + create_revised_general_topics, + max_topic_no=max_topics_number, + ) + + # This part concatenates all zero shot and new topics together, so that for the next prompt the LLM will have the full list available + if ( + not existing_topic_summary_df.empty + and force_zero_shot_radio != "Yes" + ): + existing_topic_summary_df = pd.concat( + [existing_topic_summary_df, zero_shot_topics_df] + ).drop_duplicates("Subtopic") + else: + existing_topic_summary_df = zero_shot_topics_df + + if candidate_topics and not zero_shot_topics_df.empty: + # If you have already created revised zero shot topics, concat to the current + existing_topic_summary_df = pd.concat( + [existing_topic_summary_df, zero_shot_topics_df] + ) + + existing_topic_summary_df["Number of responses"] = "" + # existing_topic_summary_df.fillna("", inplace=True) + existing_topic_summary_df["General topic"] = ( + existing_topic_summary_df["General topic"].str.replace( + "(?i)^Nan$", "", regex=True + ) + ) + existing_topic_summary_df["Subtopic"] = existing_topic_summary_df[ + "Subtopic" + ].str.replace("(?i)^Nan$", "", regex=True) + existing_topic_summary_df = ( + existing_topic_summary_df.drop_duplicates() + ) + + # If user has chosen to try to force zero shot topics, then the prompt is changed to ask the model not to deviate at all from submitted topic list. + keep_cols = [ + col + for col in ["General topic", "Subtopic", "Description"] + if col in existing_topic_summary_df.columns + and not existing_topic_summary_df[col] + .replace(r"^\s*$", pd.NA, regex=True) + .isna() + .all() + ] + + # Create topics table to be presented to LLM + topics_df_for_markdown = existing_topic_summary_df[ + keep_cols + ].drop_duplicates(keep_cols) + if ( + "General topic" in topics_df_for_markdown.columns + and "Subtopic" in topics_df_for_markdown.columns + ): + topics_df_for_markdown = topics_df_for_markdown.sort_values( + ["General topic", "Subtopic"] + ) + + if "Description" in existing_topic_summary_df: + if existing_topic_summary_df["Description"].isnull().all(): + existing_topic_summary_df.drop( + "Description", axis=1, inplace=True + ) + + if produce_structured_summary_radio == "Yes": + if "General topic" in topics_df_for_markdown.columns: + topics_df_for_markdown.rename( + columns={"General topic": "Main heading"}, + inplace=True, + errors="ignore", + ) + if "Subtopic" in topics_df_for_markdown.columns: + topics_df_for_markdown.rename( + columns={"Subtopic": "Subheading"}, + inplace=True, + errors="ignore", + ) + + # Remove duplicate General topic and subtopic names, prioritising topics where a general topic is provided + if ( + "General topic" in topics_df_for_markdown.columns + and "Subtopic" in topics_df_for_markdown.columns + ): + topics_df_for_markdown = topics_df_for_markdown.sort_values( + ["General topic", "Subtopic"], ascending=[False, True] + ) + topics_df_for_markdown = topics_df_for_markdown.drop_duplicates( + ["General topic", "Subtopic"], keep="first" + ) + topics_df_for_markdown = topics_df_for_markdown.sort_values( + ["General topic", "Subtopic"], ascending=[True, True] + ) + elif "Subtopic" in topics_df_for_markdown.columns: + topics_df_for_markdown = topics_df_for_markdown.sort_values( + ["Subtopic"], ascending=[True] + ) + topics_df_for_markdown = topics_df_for_markdown.drop_duplicates( + ["Subtopic"], keep="first" + ) + topics_df_for_markdown = topics_df_for_markdown.sort_values( + ["Subtopic"], ascending=[True] + ) + elif ( + "Main heading" in topics_df_for_markdown.columns + and "Subheading" in topics_df_for_markdown.columns + ): + topics_df_for_markdown = topics_df_for_markdown.sort_values( + ["Main heading", "Subheading"], ascending=[True, True] + ) + topics_df_for_markdown = topics_df_for_markdown.drop_duplicates( + ["Main heading", "Subheading"], keep="first" + ) + topics_df_for_markdown = topics_df_for_markdown.sort_values( + ["Main heading", "Subheading"], ascending=[True, True] + ) + + print("Number of topics:", topics_df_for_markdown.shape[0]) + + # Clean topic names before converting to markdown to ensure consistent formatting + if "General topic" in topics_df_for_markdown.columns: + topics_df_for_markdown["General topic"] = ( + topics_df_for_markdown["General topic"].apply( + normalize_topic_name_for_llm + ) + ) + if "Subtopic" in topics_df_for_markdown.columns: + topics_df_for_markdown["Subtopic"] = topics_df_for_markdown[ + "Subtopic" + ].apply(normalize_topic_name_for_llm) + if "Main heading" in topics_df_for_markdown.columns: + topics_df_for_markdown["Main heading"] = topics_df_for_markdown[ + "Main heading" + ].apply(normalize_topic_name_for_llm) + if "Subheading" in topics_df_for_markdown.columns: + topics_df_for_markdown["Subheading"] = topics_df_for_markdown[ + "Subheading" + ].apply(normalize_topic_name_for_llm) + + unique_topics_markdown = topics_df_for_markdown.to_markdown( + index=False + ) + unique_topics_markdown = normalise_string(unique_topics_markdown) + + if force_zero_shot_radio == "Yes": + topic_assignment_prompt = force_existing_topics_prompt + else: + topic_assignment_prompt = allow_new_topics_prompt + + # Should the outputs force only one single topic assignment per response? + if force_single_topic_radio != "Yes": + force_single_topic_prompt = "" + else: + topic_assignment_prompt = ( + topic_assignment_prompt.replace( + "Assign topics", "Assign a topic" + ) + .replace("assign Subtopics", "assign a Subtopic") + .replace("Subtopics", "Subtopic") + .replace("Topics", "Topic") + .replace("topics", "a topic") + ) + + # Format the summary prompt with the response table and topics + if produce_structured_summary_radio != "Yes": + formatted_summary_prompt = add_existing_topics_prompt.format( + validate_prompt_prefix="", + response_table=response_table_prompt, + topics=unique_topics_markdown, + topic_assignment=topic_assignment_prompt, + force_single_topic=force_single_topic_prompt, + sentiment_choices=sentiment_prompt, + response_reference_format=response_reference_format, + add_existing_topics_summary_format=additional_instructions_summary_format, + previous_table_introduction="", + previous_table="", + validate_prompt_suffix="", + ) + else: + formatted_summary_prompt = structured_summary_prompt.format( + response_table=response_table_prompt, + topics=unique_topics_markdown, + summary_format=additional_instructions_summary_format, + ) + + batch_file_path_details = f"{file_name_clean}_batch_{latest_batch_completed + 1}_size_{batch_size}_col_{in_column_cleaned}" + + # Use the helper function to process the batch + ( + new_topic_df, + new_reference_df, + new_topic_summary_df, + is_error, + current_prompt_content_logged, + current_summary_content_logged, + current_conversation_content_logged, + current_metadata_content_logged, + topic_table_out_path, + reference_table_out_path, + topic_summary_df_out_path, + ) = process_batch_with_llm( + is_first_batch=False, + formatted_system_prompt=formatted_system_prompt, + formatted_prompt=formatted_summary_prompt, + batch_file_path_details=batch_file_path_details, + model_source=model_source, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + max_tokens=max_tokens, + azure_api_key_textbox=azure_api_key_textbox, + azure_endpoint_textbox=azure_endpoint_textbox, + reasoning_suffix=reasoning_suffix, + local_model=local_model, + tokenizer=tokenizer, + bedrock_runtime=bedrock_runtime, + reported_batch_no=reported_batch_no, + response_text="", + whole_conversation=list(), + all_metadata_content=list(), + start_row=start_row, + end_row=end_row, + model_choice_clean=model_choice_clean, + log_files_output_paths=log_files_output_paths, + existing_reference_df=existing_reference_df, + existing_topic_summary_df=existing_topic_summary_df, + batch_size=batch_size, + batch_basic_response_df=batch_basic_response_df, + group_name=group_name, + produce_structured_summary_radio=produce_structured_summary_radio, + output_folder=output_folder, + output_debug_files=output_debug_files, + task_type=task_type, + assistant_prefill=add_existing_topics_assistant_prefill, + api_url=api_url, + ) + + print("Completed batch processing") + + all_prompts_content.append(current_prompt_content_logged) + all_responses_content.append(current_summary_content_logged) + all_conversation_content.append(current_conversation_content_logged) + all_metadata_content.append(current_metadata_content_logged) + all_groups_content.append(group_name) + all_batches_content.append(f"{reported_batch_no}:") + all_model_choice_content.append(model_choice_clean_short) + all_validated_content.append("No") + all_task_type_content.append(task_type) + all_file_names_content.append(original_full_file_name) + + ## Response ID table mapping response numbers to topics + if output_debug_files == "True": + new_reference_df.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv( + reference_table_out_path, index=None, encoding="utf-8-sig" + ) + out_file_paths.append(reference_table_out_path) + + ## Unique topic list + new_topic_summary_df = pd.concat( + [new_topic_summary_df, existing_topic_summary_df] + ).drop_duplicates("Subtopic") + + new_topic_summary_df["Group"] = group_name + + if output_debug_files == "True": + new_topic_summary_df.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv( + topic_summary_df_out_path, index=None, encoding="utf-8-sig" + ) + out_file_paths.append(topic_summary_df_out_path) + + # Outputs for markdown table output + unique_table_df_display_table = new_topic_summary_df.apply( + lambda col: col.map( + lambda x: wrap_text(x, max_text_length=max_text_length) + ) + ) + + if produce_structured_summary_radio == "Yes": + unique_table_df_display_table = unique_table_df_display_table[ + ["General topic", "Subtopic", "Summary"] + ] + unique_table_df_display_table.rename( + columns={ + "General topic": "Main heading", + "Subtopic": "Subheading", + }, + inplace=True, + ) + else: + unique_table_df_display_table = unique_table_df_display_table[ + [ + "General topic", + "Subtopic", + "Sentiment", + "Number of responses", + "Summary", + ] + ] + + unique_table_df_display_table_markdown = ( + unique_table_df_display_table.to_markdown(index=False) + ) + + all_metadata_content_str = " ".join(all_metadata_content) + + out_file_paths = [ + col for col in out_file_paths if str(reported_batch_no) in col + ] + log_files_output_paths = [ + col for col in out_file_paths if str(reported_batch_no) in col + ] + # If this is the first batch, run this + else: + formatted_system_prompt = initial_table_system_prompt.format( + consultation_context=context_textbox, column_name=chosen_cols + ) + + # Format the summary prompt with the response table and topics + if produce_structured_summary_radio != "Yes": + formatted_initial_table_prompt = initial_table_prompt.format( + validate_prompt_prefix="", + response_table=response_table_prompt, + sentiment_choices=sentiment_prompt, + response_reference_format=response_reference_format, + add_existing_topics_summary_format=additional_instructions_summary_format, + previous_table_introduction="", + previous_table="", + validate_prompt_suffix="", + ) + else: + unique_topics_markdown = ( + "No suggested headings for this summary" + ) + formatted_initial_table_prompt = ( + structured_summary_prompt.format( + response_table=response_table_prompt, + topics=unique_topics_markdown, + ) + ) + + batch_file_path_details = f"{file_name_clean}_batch_{latest_batch_completed + 1}_size_{batch_size}_col_{in_column_cleaned}" + + # Use the helper function to process the batch + ( + new_topic_df, + new_reference_df, + new_topic_summary_df, + is_error, + current_prompt_content_logged, + current_summary_content_logged, + current_conversation_content_logged, + current_metadata_content_logged, + topic_table_out_path, + reference_table_out_path, + topic_summary_df_out_path, + ) = process_batch_with_llm( + is_first_batch=True, + formatted_system_prompt=formatted_system_prompt, + formatted_prompt=formatted_initial_table_prompt, + batch_file_path_details=batch_file_path_details, + model_source=model_source, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + max_tokens=max_tokens, + azure_api_key_textbox=azure_api_key_textbox, + azure_endpoint_textbox=azure_endpoint_textbox, + reasoning_suffix=reasoning_suffix, + local_model=local_model, + tokenizer=tokenizer, + bedrock_runtime=bedrock_runtime, + reported_batch_no=reported_batch_no, + response_text="", + whole_conversation=list(), + all_metadata_content=list(), + start_row=start_row, + end_row=end_row, + model_choice_clean=model_choice_clean, + log_files_output_paths=log_files_output_paths, + existing_reference_df=existing_reference_df, + existing_topic_summary_df=existing_topic_summary_df, + batch_size=batch_size, + batch_basic_response_df=batch_basic_response_df, + group_name=group_name, + produce_structured_summary_radio=produce_structured_summary_radio, + output_folder=output_folder, + output_debug_files=output_debug_files, + task_type=task_type, + assistant_prefill=initial_table_assistant_prefill, + api_url=api_url, + ) + + all_prompts_content.append(current_prompt_content_logged) + all_responses_content.append(current_summary_content_logged) + all_conversation_content.append(current_conversation_content_logged) + all_metadata_content.append(current_metadata_content_logged) + all_groups_content.append(group_name) + all_batches_content.append(f"{reported_batch_no}:") + all_model_choice_content.append(model_choice_clean_short) + all_validated_content.append("No") + all_task_type_content.append(task_type) + all_file_names_content.append(original_full_file_name) + + if output_debug_files == "True": + + # Output reference table + new_reference_df.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv( + reference_table_out_path, index=None, encoding="utf-8-sig" + ) + out_file_paths.append(reference_table_out_path) + + ## Unique topic list + + new_topic_summary_df = pd.concat( + [new_topic_summary_df, existing_topic_summary_df] + ).drop_duplicates("Subtopic") + + new_topic_summary_df["Group"] = group_name + + if output_debug_files == "True": + new_topic_summary_df.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv( + topic_summary_df_out_path, index=None, encoding="utf-8-sig" + ) + out_file_paths.append(topic_summary_df_out_path) + + all_metadata_content.append(all_metadata_content_str) + all_metadata_content_str = ". ".join(all_metadata_content) + + else: + print( + "Current batch of responses contains no text, moving onto next. Batch number:", + str(latest_batch_completed + 1), + ". Start row:", + start_row, + ". End row:", + end_row, + ) + + # Increase latest file completed count unless we are over the last batch number, then go back around + num_batches = int(num_batches) + latest_batch_completed = int(latest_batch_completed) + if latest_batch_completed <= num_batches: + latest_batch_completed += 1 + + toc = time.perf_counter() + final_time = toc - tic + + if final_time > max_time_for_loop: + print("Max time reached, breaking loop.") + topics_loop.close() + tqdm._instances.clear() + break + + # Overwrite 'existing' elements to add new tables + existing_reference_df = new_reference_df.dropna(how="all") + # Rebuild topic_summary_df from accumulated reference_df to ensure consistency + # This properly aggregates response counts and handles all dimensions (General topic, Subtopic, Sentiment, Group) + existing_topic_summary_df = create_topic_summary_df_from_reference_table( + existing_reference_df, sentiment_checkbox=sentiment_checkbox + ) + existing_topics_table = new_topic_df.dropna(how="all") + + # Clean General topic and Subtopic columns using the same process as zero-shot topics + for col_name in ["General topic", "Subtopic"]: + for df in [ + existing_reference_df, + existing_topic_summary_df, + existing_topics_table, + ]: + if col_name in df.columns and not df[col_name].isnull().all(): + df[col_name] = df[col_name].apply(normalize_topic_name_for_llm) + + # Deduplicate topics after each batch if enabled and conditions are met + if ( + ENABLE_BATCH_DEDUPLICATION + and force_zero_shot_radio == "No" + and produce_structured_summary_radio == "No" + and not existing_reference_df.empty + and not existing_topic_summary_df.empty + and latest_batch_completed % BATCH_DEDUPLICATION_EVERY_N_BATCHES == 0 + ): + print(f"Deduplicating topics after batch {latest_batch_completed}...") + # Create temporary file names for deduplication + reference_table_file_name = ( + f"{file_name_clean}_batch_{latest_batch_completed}_reference" + ) + unique_topics_table_file_name = ( + f"{file_name_clean}_batch_{latest_batch_completed}_unique_topics" + ) + + # Prepare file_data for deduplication if available + in_data_files_for_dedup = in_data_file + extract_num_batches = None + extract_data_file_names_textbox = None + if not file_data.empty and chosen_cols: + # Pass file_data as DataFrame and use available num_batches and file_name + in_data_files_for_dedup = file_data + extract_num_batches = num_batches + extract_data_file_names_textbox = file_name + + try: + original_topic_summary_df = existing_topic_summary_df.copy() + topics_before = existing_topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + # Check Group values before deduplication for better diagnostics + input_groups_before = ( + existing_topic_summary_df["Group"].nunique() + if "Group" in existing_topic_summary_df.columns + else 0 + ) + input_ref_groups_before = ( + existing_reference_df["Group"].nunique() + if "Group" in existing_reference_df.columns + else 0 + ) + # Call deduplicate_topics function + ( + existing_reference_df, + existing_topic_summary_df, + _, + _, + _, + ) = deduplicate_topics( + reference_df=existing_reference_df, + topic_summary_df=existing_topic_summary_df, + reference_table_file_name=reference_table_file_name, + unique_topics_table_file_name=unique_topics_table_file_name, + in_excel_sheets=( + ",".join(in_excel_sheets) + if isinstance(in_excel_sheets, list) + else in_excel_sheets + ), + merge_sentiment="No", + merge_general_topics="No", + score_threshold=95, + in_data_files=in_data_files_for_dedup, + chosen_cols=( + chosen_cols + if isinstance(chosen_cols, list) + else [chosen_cols] if chosen_cols else [] + ), + output_folder=output_folder, + deduplicate_topics="Yes", + output_files="False", + total_number_of_batches=extract_num_batches, + data_file_names_textbox=extract_data_file_names_textbox, + ) + topics_after = existing_topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + # Check Group values after deduplication + output_groups_after = ( + existing_topic_summary_df["Group"].nunique() + if "Group" in existing_topic_summary_df.columns + else 0 + ) + output_ref_groups_after = ( + existing_reference_df["Group"].nunique() + if "Group" in existing_reference_df.columns + else 0 + ) + if topics_after < topics_before: + print( + f"Topics deduplicated successfully after batch {latest_batch_completed}. " + f"Topics before: {topics_before}, Topics after: {topics_after}" + ) + if topics_after > topics_before: + print( + f"WARNING: Number of topics increased from {topics_before} to {topics_after}. " + f"This may be due to Group column differences or new topic combinations created during deduplication. " + f"Input topic_summary_df had {input_groups_before} unique Group values, " + f"Input reference_df had {input_ref_groups_before} unique Group values. " + f"Output topic_summary_df has {output_groups_after} unique Group values, " + f"Output reference_df has {output_ref_groups_after} unique Group values." + ) + original_topic_summary_df.to_csv( + f"{output_folder}{file_name_clean}_batch_{latest_batch_completed}_original_topic_summary_df.csv", + index=None, + encoding="utf-8-sig", + ) + existing_topic_summary_df.to_csv( + f"{output_folder}{file_name_clean}_batch_{latest_batch_completed}_deduplicated_topic_summary_df.csv", + index=None, + encoding="utf-8-sig", + ) + except Exception as e: + print( + f"Warning: Deduplication failed after batch {latest_batch_completed}: {str(e)}. Continuing with original topics." + ) + + # Check if number of topics exceeds max_topics_number and use LLM deduplication if needed + topics_after_dedup = existing_topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + if topics_after_dedup > max_topics_number: + print( + f"Number of topics ({topics_after_dedup}) exceeds max_topics_number ({max_topics_number}). " + f"Using LLM-based deduplication after batch {latest_batch_completed}..." + ) + try: + # Prepare file names for LLM deduplication + reference_table_file_name_llm = f"{file_name_clean}_batch_{latest_batch_completed}_reference_llm" + unique_topics_table_file_name_llm = f"{file_name_clean}_batch_{latest_batch_completed}_unique_topics_llm" + + # Prepare in_data_files - use file_data if available, otherwise use in_data_file + in_data_files_for_llm_dedup = ( + file_data if not file_data.empty else in_data_file + ) + + # Store topics count before LLM deduplication + topics_before_llm_dedup = ( + existing_topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + ) + + # Call deduplicate_topics_llm function + ( + existing_reference_df, + existing_topic_summary_df, + _, + _, + _, + _, + _, + _, + _, + ) = deduplicate_topics_llm( + reference_df=existing_reference_df, + topic_summary_df=existing_topic_summary_df, + reference_table_file_name=reference_table_file_name_llm, + unique_topics_table_file_name=unique_topics_table_file_name_llm, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + model_source=model_source, + local_model=( + local_model if "local_model" in locals() else None + ), + tokenizer=tokenizer if "tokenizer" in locals() else None, + assistant_model=None, # assistant_model not typically available in this scope + in_excel_sheets=( + ",".join(in_excel_sheets) + if isinstance(in_excel_sheets, list) + else in_excel_sheets + ), + merge_sentiment="No", + merge_general_topics="Yes", + in_data_files=in_data_files_for_llm_dedup, + chosen_cols=( + chosen_cols + if isinstance(chosen_cols, list) + else [chosen_cols] if chosen_cols else [] + ), + output_folder=output_folder, + candidate_topics=None, + azure_endpoint=azure_endpoint_textbox, + output_debug_files=output_debug_files, + api_url=api_url, + aws_access_key_textbox=aws_access_key_textbox, + aws_secret_key_textbox=aws_secret_key_textbox, + aws_region_textbox=aws_region_textbox, + azure_api_key_textbox=azure_api_key_textbox, + model_name_map=model_name_map, + output_files="False", + sentiment_checkbox=sentiment_checkbox, + ) + # Rebuild topic_summary_df from updated reference_df to ensure consistency + # This ensures the topic count reflects the actual merged state + if not existing_reference_df.empty: + existing_topic_summary_df = ( + create_topic_summary_df_from_reference_table( + existing_reference_df, + sentiment_checkbox=sentiment_checkbox, + ) + ) + + topics_after_llm_dedup = ( + existing_topic_summary_df.drop_duplicates( + subset=["General topic", "Subtopic"] + ).shape[0] + ) + print( + f"LLM-based deduplication completed successfully after batch {latest_batch_completed}. " + f"Number of topics before LLM deduplication: {topics_before_llm_dedup}, Number of topics after LLM deduplication: {topics_after_llm_dedup}" + ) + except Exception as e: + print( + f"Warning: LLM-based deduplication failed after batch {latest_batch_completed}: {str(e)}. " + f"Continuing with existing topics." + ) + + # The topic table that can be modified does not need the summary column + # Check if DataFrame is not empty and has Summary column before dropping + if ( + not existing_topic_summary_df.empty + and "Summary" in existing_topic_summary_df.columns + ): + modifiable_topic_summary_df = existing_topic_summary_df.drop( + "Summary", axis=1 + ) + else: + modifiable_topic_summary_df = existing_topic_summary_df.copy() + + out_time = f"{final_time:0.1f} seconds." + + out_message.append("All queries successfully completed in") + + final_message_out = "\n".join(out_message) + final_message_out = final_message_out + " " + out_time + + print(final_message_out) + + # If we have extracted topics from the last batch, return the input out_message and file list to the relevant components + if latest_batch_completed >= num_batches: + + group_combined_logged_content = [ + { + "prompt": prompt, + "response": summary, + "metadata": metadata, + "batch": batch, + "model_choice": model_choice, + "validated": validated, + "group": group, + "task_type": task_type, + "file_name": file_name, + } + for prompt, summary, metadata, batch, model_choice, validated, group, task_type, file_name in zip( + all_prompts_content, + all_responses_content, + all_metadata_content, + all_batches_content, + all_model_choice_content, + all_validated_content, + all_groups_content, + all_task_type_content, + all_file_names_content, + ) + ] + + # VALIDATION LOOP - Run validation if enabled + if ENABLE_VALIDATION == "True": + + # Use the standalone validation function + ( + existing_reference_df, + existing_topic_summary_df, + group_combined_logged_content, + validation_conversation_metadata_str, + ) = validate_topics( + file_data=file_data, + reference_df=existing_reference_df, + topic_summary_df=existing_topic_summary_df, + file_name=file_name, + chosen_cols=chosen_cols, + batch_size=batch_size, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + max_tokens=max_tokens, + azure_api_key_textbox=azure_api_key_textbox, + azure_endpoint_textbox=azure_endpoint_textbox, + reasoning_suffix=reasoning_suffix, + group_name=group_name, + produce_structured_summary_radio=produce_structured_summary_radio, + force_zero_shot_radio=force_zero_shot_radio, + force_single_topic_radio=force_single_topic_radio, + context_textbox=context_textbox, + additional_instructions_summary_format=additional_instructions_summary_format, + additional_validation_issues_provided=additional_validation_issues_provided, + output_folder=output_folder, + output_debug_files=output_debug_files, + original_full_file_name=original_full_file_name, + max_time_for_loop=max_time_for_loop, + sentiment_checkbox=sentiment_checkbox, + logged_content=group_combined_logged_content, + api_url=api_url, + max_topics_number=max_topics_number, + ) + + # Add validation conversation metadata to the main conversation metadata + if validation_conversation_metadata_str: + all_metadata_content_str = ( + all_metadata_content_str + + ". " + + validation_conversation_metadata_str + ) + + print("Last batch reached, returning batch:", str(latest_batch_completed)) + + join_file_paths = list() + + toc = time.perf_counter() + final_time = (toc - tic) + time_taken + out_time = f"Everything finished in {round(final_time,1)} seconds." + print(out_time) + + print("All batches completed. Exporting outputs.") + + all_groups_logged_content = ( + all_groups_logged_content + group_combined_logged_content + ) + + # file_path_details = create_batch_file_path_details(file_name, in_column=chosen_cols) + + # Create a pivoted reference table + existing_reference_df_pivot = convert_reference_table_to_pivot_table( + existing_reference_df + ) + + # Save the new DataFrame to CSV + reference_table_out_pivot_path = ( + output_folder + + file_name_clean + + "_final_reference_table_pivot_" + + model_choice_clean_short + + "_temp_" + + str(temperature) + + ".csv" + ) + reference_table_out_path = ( + output_folder + + file_name_clean + + "_final_reference_table_" + + model_choice_clean_short + + "_temp_" + + str(temperature) + + ".csv" + ) + topic_summary_df_out_path = ( + output_folder + + file_name_clean + + "_final_unique_topics_" + + model_choice_clean_short + + "_temp_" + + str(temperature) + + ".csv" + ) + basic_response_data_out_path = ( + output_folder + + file_name_clean + + "_simplified_data_file_" + + model_choice_clean_short + + "_temp_" + + str(temperature) + + ".csv" + ) + + ## Response ID table mapping response numbers to topics + # Ensure Group is present for grouped pipelines and downstream concatenation + if "Group" not in existing_reference_df.columns: + existing_reference_df["Group"] = group_name + else: + existing_reference_df["Group"] = group_name + + existing_reference_df.to_csv( + reference_table_out_path, index=None, encoding="utf-8-sig" + ) + out_file_paths.append(reference_table_out_path) + join_file_paths.append(reference_table_out_path) + + # Create final unique topics table from reference table to ensure consistent numbers + final_out_topic_summary_df = create_topic_summary_df_from_reference_table( + existing_reference_df, sentiment_checkbox=sentiment_checkbox + ) + final_out_topic_summary_df["Group"] = group_name + + ## Unique topic list + final_out_topic_summary_df.to_csv( + topic_summary_df_out_path, index=None, encoding="utf-8-sig" + ) + out_file_paths.append(topic_summary_df_out_path) + + # Outputs for markdown table output + unique_table_df_display_table = final_out_topic_summary_df.apply( + lambda col: col.map(lambda x: wrap_text(x, max_text_length=max_text_length)) + ) + + if produce_structured_summary_radio == "Yes": + unique_table_df_display_table = unique_table_df_display_table[ + ["General topic", "Subtopic", "Summary"] + ] + unique_table_df_display_table.rename( + columns={"General topic": "Main heading", "Subtopic": "Subheading"}, + inplace=True, + ) + else: + # create_topic_summary_df_from_reference_table omits Sentiment when + # sentiment_checkbox is "Do not assess sentiment" + display_cols = ["General topic", "Subtopic"] + if ( + _assess_sentiment(sentiment_checkbox) + and "Sentiment" in unique_table_df_display_table.columns + ): + display_cols.append("Sentiment") + display_cols.extend(["Number of responses", "Summary"]) + display_cols = [ + c for c in display_cols if c in unique_table_df_display_table.columns + ] + unique_table_df_display_table = unique_table_df_display_table[display_cols] + + unique_table_df_display_table_markdown = ( + unique_table_df_display_table.to_markdown(index=False) + ) + + # Ensure that we are only returning the final results to outputs + out_file_paths = [x for x in out_file_paths if "_final_" in x] + + ## Response ID table mapping response numbers to topics + existing_reference_df_pivot["Group"] = group_name + existing_reference_df_pivot.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv(reference_table_out_pivot_path, index=None, encoding="utf-8-sig") + log_files_output_paths.append(reference_table_out_pivot_path) + + ## Create a dataframe for missing response references: + # Assuming existing_reference_df and file_data are already defined + # Simplify table to just responses column and the Response reference number + basic_response_data = get_basic_response_data(file_data, chosen_cols) + + # Save simplified file data to log outputs + pd.DataFrame(basic_response_data).drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv(basic_response_data_out_path, index=None, encoding="utf-8-sig") + log_files_output_paths.append(basic_response_data_out_path) + + # Note: missing_df creation moved to wrapper functions to handle grouped processing correctly + missing_df = pd.DataFrame() + + out_file_paths = list(set(out_file_paths)) + log_files_output_paths = list(set(log_files_output_paths)) + + final_out_file_paths = [ + file_path for file_path in out_file_paths if "final_" in file_path + ] + + # The topic table that can be modified does not need the summary column + # Check if DataFrame is not empty and has Summary column before dropping + if ( + not final_out_topic_summary_df.empty + and "Summary" in final_out_topic_summary_df.columns + ): + modifiable_topic_summary_df = final_out_topic_summary_df.drop( + "Summary", axis=1 + ) + else: + modifiable_topic_summary_df = final_out_topic_summary_df.copy() + + # Ensure Group column is present on the returned reference_df for wrapper aggregation + if "Group" not in existing_reference_df.columns: + existing_reference_df["Group"] = group_name + + return ( + unique_table_df_display_table_markdown, + existing_topics_table, + final_out_topic_summary_df, + existing_reference_df, + final_out_file_paths, + final_out_file_paths, + latest_batch_completed, + log_files_output_paths, + log_files_output_paths, + all_metadata_content_str, + final_time, + final_out_file_paths, + final_out_file_paths, + modifiable_topic_summary_df, + final_out_file_paths, + join_file_paths, + existing_reference_df_pivot, + missing_df, + all_groups_logged_content, + ) + + return ( + unique_table_df_display_table_markdown, + existing_topics_table, + existing_topic_summary_df, + existing_reference_df, + out_file_paths, + out_file_paths, + latest_batch_completed, + log_files_output_paths, + log_files_output_paths, + all_metadata_content_str, + final_time, + out_file_paths, + out_file_paths, + modifiable_topic_summary_df, + out_file_paths, + join_file_paths, + existing_reference_df_pivot, + missing_df, + all_groups_logged_content, + ) + + +@spaces.GPU(duration=MAX_SPACES_GPU_RUN_TIME) +def wrapper_extract_topics_per_column_value( + grouping_col: str, + in_data_file: Any, + file_data: pd.DataFrame, + initial_existing_topics_table: pd.DataFrame, + initial_existing_reference_df: pd.DataFrame, + initial_existing_topic_summary_df: pd.DataFrame, + initial_unique_table_df_display_table_markdown: str, + original_file_name: str, # Original file name, to be modified per segment + total_number_of_batches: int, + in_api_key: str, + temperature: float, + chosen_cols: List[str], + model_choice: str, + candidate_topics: gr.FileData = None, + initial_first_loop_state: bool = True, + initial_all_metadata_content_str: str = "", + initial_latest_batch_completed: int = 0, + initial_time_taken: float = 0, + initial_table_prompt: str = initial_table_prompt, + initial_table_system_prompt: str = initial_table_system_prompt, + add_existing_topics_system_prompt: str = add_existing_topics_system_prompt, + add_existing_topics_prompt: str = add_existing_topics_prompt, + number_of_prompts_used: int = 1, + batch_size: int = 50, # Crucial for calculating num_batches per segment + context_textbox: str = "", + sentiment_checkbox: str = "Negative, Neutral, or Positive", + force_zero_shot_radio: str = "No", + in_excel_sheets: List[str] = list(), + force_single_topic_radio: str = "No", + produce_structured_summary_radio: str = "No", + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + hf_api_key_textbox: str = "", + azure_api_key_textbox: str = "", + azure_endpoint_textbox: str = "", + output_folder: str = OUTPUT_FOLDER, + existing_logged_content: list = list(), + additional_instructions_summary_format: str = "", + additional_validation_issues_provided: str = "", + show_previous_table: str = "Yes", + api_url: str = None, + force_single_topic_prompt: str = force_single_topic_prompt, + max_tokens: int = max_tokens, + model_name_map: dict = model_name_map, + max_time_for_loop: int = max_time_for_loop, # This applies per call to extract_topics + reasoning_suffix: str = reasoning_suffix, + CHOSEN_LOCAL_MODEL_TYPE: str = CHOSEN_LOCAL_MODEL_TYPE, + output_debug_files: str = output_debug_files, + model: object = None, + tokenizer: object = None, + assistant_model: object = None, + max_rows: int = max_rows, + max_topics_number: int = MAXIMUM_ALLOWED_TOPICS, + progress=Progress(track_tqdm=True), # type: ignore +) -> Tuple: # Mimicking the return tuple structure of extract_topics + """ + A wrapper function that iterates through unique values in a specified grouping column + and calls the `extract_topics` function for each segment of the data. + It accumulates results from each call and returns a consolidated output. + + :param grouping_col: The name of the column to group the data by. + :param in_data_file: The input data file object (e.g., Gradio FileData). + :param file_data: The full DataFrame containing all data. + :param initial_existing_topics_table: Initial DataFrame of existing topics. + :param initial_existing_reference_df: Initial DataFrame mapping responses to topics. + :param initial_existing_topic_summary_df: Initial DataFrame summarizing topics. + :param initial_unique_table_df_display_table_markdown: Initial markdown string for topic display. + :param original_file_name: The original name of the input file. + :param total_number_of_batches: The total number of batches across all data. + :param in_api_key: API key for the chosen LLM. + :param temperature: Temperature setting for the LLM. + :param chosen_cols: List of columns from `file_data` to be processed. + :param model_choice: The chosen LLM model (e.g., "Gemini", "AWS Claude"). + :param candidate_topics: Optional Gradio FileData for candidate topics (zero-shot). + :param initial_first_loop_state: Boolean indicating if this is the very first loop iteration. + :param initial_all_metadata_content_str: Initial metadata string for the whole conversation. + :param initial_latest_batch_completed: The batch number completed in the previous run. + :param initial_time_taken: Initial time taken for processing. + :param initial_table_prompt: The initial prompt for table summarization. + :param initial_table_system_prompt: The initial system prompt for table summarization. + :param add_existing_topics_system_prompt: System prompt for adding existing topics. + :param add_existing_topics_prompt: Prompt for adding existing topics. + :param number_of_prompts_used: Number of prompts used in the LLM call. + :param batch_size: Number of rows to process in each batch for the LLM. + :param context_textbox: Additional context provided by the user. + :param sentiment_checkbox: Choice for sentiment assessment (e.g., "Negative, Neutral, or Positive"). + :param force_zero_shot_radio: Option to force responses into zero-shot topics. + :param in_excel_sheets: List of Excel sheet names if applicable. + :param force_single_topic_radio: Option to force a single topic per response. + :param produce_structured_summary_radio: Option to produce a structured summary. + :param aws_access_key_textbox: AWS access key for Bedrock. + :param aws_secret_key_textbox: AWS secret key for Bedrock. + :param hf_api_key_textbox: Hugging Face API key for local models. + :param azure_api_key_textbox: Azure/OpenAI API key for Azure/OpenAI AI Inference. + :param output_folder: The folder where output files will be saved. + :param existing_logged_content: A list of existing logged content. + :param force_single_topic_prompt: Prompt for forcing a single topic. + :param additional_instructions_summary_format: Initial instructions to guide the format for the initial summary of the topics. + :param additional_validation_issues_provided: Additional validation issues provided by the user. + :param show_previous_table: Whether to show the previous table ("Yes" or "No"). + :param max_tokens: Maximum tokens for LLM generation. + :param model_name_map: Dictionary mapping model names to their properties. + :param max_time_for_loop: Maximum time allowed for the processing loop. + :param reasoning_suffix: Suffix to append for reasoning. + :param CHOSEN_LOCAL_MODEL_TYPE: Type of local model chosen. + :param output_debug_files: Whether to output debug files ("True" or "False"). + :param model: Model object for local inference. + :param tokenizer: Tokenizer object for local inference. + :param assistant_model: Assistant model object for local inference. + :param max_rows: The maximum number of rows to process. + :param progress: Gradio Progress object for tracking progress. + :return: A tuple containing consolidated results, mimicking the return structure of `extract_topics`. + """ + + # Ensure custom model_choice is registered in model_name_map + ensure_model_in_map(model_choice, model_name_map) + + acc_input_tokens = 0 + acc_output_tokens = 0 + acc_number_of_calls = 0 + out_message = list() + + # Logged content + all_groups_logged_content = existing_logged_content + + # If you have a file input but no file data it hasn't yet been loaded. Load it here. + if file_data.empty: + print("No data table found, loading from file") + try: + ( + in_colnames_drop, + in_excel_sheets, + file_name, + join_colnames, + join_colnames_drop, + ) = put_columns_in_df(in_data_file) + file_data, file_name, num_batches = load_in_data_file( + in_data_file, chosen_cols, batch_size_default, in_excel_sheets + ) + except Exception as e: + out_message = "Could not load in data file due to: " + str(e) + print(out_message) + raise Exception(out_message) + + if file_data.shape[0] > max_rows: + out_message = ( + "Your data has more than " + + str(max_rows) + + " rows, which has been set as the maximum in the application configuration." + ) + print(out_message) + raise Exception(out_message) + + if grouping_col is None: + print("No grouping column found") + file_data["group_col"] = "All" + grouping_col = "group_col" + + if grouping_col not in file_data.columns: + raise ValueError(f"Selected column '{grouping_col}' not found in file_data.") + + unique_values = file_data[grouping_col].unique() + + if len(unique_values) > MAX_GROUPS: + print( + f"Warning: More than {MAX_GROUPS} unique values found in '{grouping_col}'. Processing only the first {MAX_GROUPS}." + ) + unique_values = unique_values[:MAX_GROUPS] + + # Initialise accumulators for results across all unique values + # DataFrames are built upon iteratively + acc_topics_table = initial_existing_topics_table.copy() + acc_reference_df = initial_existing_reference_df.copy() + acc_topic_summary_df = initial_existing_topic_summary_df.copy() + acc_reference_df_pivot = pd.DataFrame() + acc_missing_df = pd.DataFrame() + + # Lists are extended + acc_out_file_paths = list() + acc_log_files_output_paths = list() + acc_join_file_paths = ( + list() + ) # join_file_paths seems to be overwritten, so maybe last one or extend? Let's extend. + + # Single value outputs - typically the last one is most relevant, or sum for time + acc_markdown_output = initial_unique_table_df_display_table_markdown + acc_latest_batch_completed = ( + initial_latest_batch_completed # From the last segment processed + ) + acc_all_metadata_content = initial_all_metadata_content_str + acc_total_time_taken = float(initial_time_taken) + acc_gradio_df = gr.Dataframe(value=pd.DataFrame()) # type: ignore # Placeholder for the last Gradio DF + acc_logged_content = list() + + wrapper_first_loop = initial_first_loop_state + + if len(unique_values) == 1: + # If only one unique value, no need for progress bar, iterate directly + loop_object = unique_values + else: + # If multiple unique values, use tqdm progress bar + loop_object = progress.tqdm( + unique_values, desc="Analysing group", unit="groups" + ) + + for i, group_value in enumerate(loop_object): + print( + f"\nProcessing group: {grouping_col} = {group_value} ({i+1}/{len(unique_values)})" + ) + + filtered_file_data = file_data.copy() + + filtered_file_data = filtered_file_data[ + filtered_file_data[grouping_col] == group_value + ] + + if filtered_file_data.empty: + print(f"No data for {grouping_col} = {group_value}. Skipping.") + continue + + # Calculate num_batches for this specific segment + current_num_batches = (len(filtered_file_data) + batch_size - 1) // batch_size + + # Modify file_name to be unique for this segment's outputs + # _grp_{clean_column_name(grouping_col, max_length=15)} + segment_file_name = f"{clean_column_name(original_file_name, max_length=15)}_{clean_column_name(str(group_value), max_length=15).replace(' ','_')}" + + # Determine first_loop_state for this call to extract_topics + # It's True only if this is the very first segment *and* the wrapper was told it's the first loop. + # For subsequent segments, it's False, as we are building on accumulated DFs. + current_first_loop_state = wrapper_first_loop if i == 0 else False + + # latest_batch_completed for extract_topics should be 0 for each new segment, + # as it processes the new filtered_file_data from its beginning. + # However, if it's the very first call, respect initial_latest_batch_completed. + current_latest_batch_completed = ( + initial_latest_batch_completed if i == 0 and wrapper_first_loop else 0 + ) + + # Call extract_topics for the current segment + try: + ( + seg_markdown, + seg_topics_table, + seg_topic_summary_df, + seg_reference_df, + seg_out_files1, + _seg_out_files2, # Often same as 1 + seg_batch_completed, # Specific to this segment's run + seg_log_files1, + _seg_log_files2, # Often same as 1 + seg_conversation_metadata, + seg_time_taken, + _seg_out_files3, # Often same as 1 + _seg_out_files4, # Often same as 1 + seg_gradio_df, + _seg_out_files5, # Often same as 1 + seg_join_files, + seg_reference_df_pivot, + seg_missing_df, + seg_logged_content, + ) = extract_topics( + in_data_file=in_data_file, + file_data=filtered_file_data, + existing_topics_table=pd.DataFrame(), # acc_topics_table.copy(), # Pass the accumulated table + existing_reference_df=pd.DataFrame(), # acc_reference_df.copy(), # Pass the accumulated table + existing_topic_summary_df=pd.DataFrame(), # acc_topic_summary_df.copy(), # Pass the accumulated table + unique_table_df_display_table_markdown="", # extract_topics will generate this + file_name=segment_file_name, + num_batches=current_num_batches, + in_api_key=in_api_key, + temperature=temperature, + chosen_cols=chosen_cols, + model_choice=model_choice, + candidate_topics=candidate_topics, + latest_batch_completed=current_latest_batch_completed, # Reset for each new segment's internal batching + out_message=list(), # Fresh for each call + out_file_paths=list(), # Fresh for each call + log_files_output_paths=list(), # Fresh for each call + first_loop_state=current_first_loop_state, # True only for the very first iteration of wrapper + all_metadata_content_str="", # Fresh for each call + initial_table_prompt=initial_table_prompt, + initial_table_system_prompt=initial_table_system_prompt, + add_existing_topics_system_prompt=add_existing_topics_system_prompt, + add_existing_topics_prompt=add_existing_topics_prompt, + number_of_prompts_used=number_of_prompts_used, + batch_size=batch_size, + context_textbox=context_textbox, + time_taken=0, # Time taken for this specific call, wrapper sums it. + sentiment_checkbox=sentiment_checkbox, + force_zero_shot_radio=force_zero_shot_radio, + in_excel_sheets=in_excel_sheets, + force_single_topic_radio=force_single_topic_radio, + output_folder=output_folder, + force_single_topic_prompt=force_single_topic_prompt, + group_name=group_value, + produce_structured_summary_radio=produce_structured_summary_radio, + aws_access_key_textbox=aws_access_key_textbox, + aws_secret_key_textbox=aws_secret_key_textbox, + aws_region_textbox=aws_region_textbox, + hf_api_key_textbox=hf_api_key_textbox, + azure_api_key_textbox=azure_api_key_textbox, + azure_endpoint_textbox=azure_endpoint_textbox, + max_tokens=max_tokens, + model_name_map=model_name_map, + max_time_for_loop=max_time_for_loop, + CHOSEN_LOCAL_MODEL_TYPE=CHOSEN_LOCAL_MODEL_TYPE, + output_debug_files=output_debug_files, + reasoning_suffix=reasoning_suffix, + model=model, + tokenizer=tokenizer, + assistant_model=assistant_model, + max_rows=max_rows, + existing_logged_content=all_groups_logged_content, + original_full_file_name=original_file_name, + additional_instructions_summary_format=additional_instructions_summary_format, + additional_validation_issues_provided=additional_validation_issues_provided, + api_url=api_url, + max_topics_number=max_topics_number, + progress=progress, + ) + + # Aggregate results + # The DFs returned by extract_topics are already cumulative for *its own run*. + # We now make them cumulative for the *wrapper's run*. + # Reset indices before concatenation to avoid reindexing errors + if not acc_reference_df.empty: + acc_reference_df = acc_reference_df.reset_index(drop=True) + if not seg_reference_df.empty: + seg_reference_df = seg_reference_df.reset_index(drop=True) + acc_reference_df = pd.concat( + [acc_reference_df, seg_reference_df], ignore_index=True + ) + + if not acc_topic_summary_df.empty: + acc_topic_summary_df = acc_topic_summary_df.reset_index(drop=True) + if not seg_topic_summary_df.empty: + seg_topic_summary_df = seg_topic_summary_df.reset_index(drop=True) + acc_topic_summary_df = pd.concat( + [acc_topic_summary_df, seg_topic_summary_df], ignore_index=True + ) + + if not acc_reference_df_pivot.empty: + acc_reference_df_pivot = acc_reference_df_pivot.reset_index(drop=True) + if not seg_reference_df_pivot.empty: + seg_reference_df_pivot = seg_reference_df_pivot.reset_index(drop=True) + acc_reference_df_pivot = pd.concat( + [acc_reference_df_pivot, seg_reference_df_pivot], ignore_index=True + ) + + if not acc_missing_df.empty: + acc_missing_df = acc_missing_df.reset_index(drop=True) + if not seg_missing_df.empty: + seg_missing_df = seg_missing_df.reset_index(drop=True) + acc_missing_df = pd.concat( + [acc_missing_df, seg_missing_df], ignore_index=True + ) + + # For lists, extend. Use set to remove duplicates if paths might be re-added. + acc_out_file_paths.extend( + f for f in seg_out_files1 if f not in acc_out_file_paths + ) + acc_log_files_output_paths.extend( + f for f in seg_log_files1 if f not in acc_log_files_output_paths + ) + acc_join_file_paths.extend( + f for f in seg_join_files if f not in acc_join_file_paths + ) + + acc_markdown_output = seg_markdown # Keep the latest markdown + acc_latest_batch_completed = seg_batch_completed # Keep latest batch count + acc_all_metadata_content += ( + ("\n---\n" if acc_all_metadata_content else "") + + f"Segment {grouping_col}={group_value}:\n" + + seg_conversation_metadata + ) + acc_total_time_taken += float(seg_time_taken) + acc_gradio_df = seg_gradio_df # Keep the latest Gradio DF + acc_logged_content.extend(seg_logged_content) + + print( + f"Group {grouping_col} = {group_value} processed. Time: {seg_time_taken:.2f}s" + ) + + except Exception as e: + print(f"Error processing segment {grouping_col} = {group_value}: {e}") + # Optionally, decide if you want to continue with other segments or stop + # For now, it will continue + continue + + overall_file_name = clean_column_name(original_file_name, max_length=20) + model_choice_clean = model_name_map[model_choice]["short_name"] + model_choice_clean_short = clean_column_name( + model_choice_clean, max_length=20, front_characters=False + ) + column_clean = clean_column_name(chosen_cols, max_length=20) + + # Need to join "Topic number" onto acc_reference_df + # If any blanks, there is an issue somewhere, drop and redo + if "Topic number" in acc_reference_df.columns: + try: + # Convert to numeric first to avoid dtype issues with object arrays + topic_num_series = pd.to_numeric( + acc_reference_df["Topic number"], errors="coerce" + ) + if topic_num_series.isnull().any(): + acc_reference_df = acc_reference_df.drop("Topic number", axis=1) + except (TypeError, ValueError): + # If conversion fails, drop the column to be safe + acc_reference_df = acc_reference_df.drop("Topic number", axis=1) + + if "Topic number" not in acc_reference_df.columns: + if "Topic number" in acc_topic_summary_df.columns: + if "General topic" in acc_topic_summary_df.columns: + on_cols, summary_cols = _merge_topic_number_keys( + sentiment_checkbox, acc_reference_df, acc_topic_summary_df + ) + if on_cols and summary_cols: + acc_reference_df = acc_reference_df.merge( + acc_topic_summary_df[summary_cols], + on=on_cols, + how="left", + ) + # Sort output dataframes + acc_reference_df["Response ID"] = ( + acc_reference_df["Response ID"].astype(float).astype(int) + ) + acc_reference_df["Start row of group"] = acc_reference_df[ + "Start row of group" + ].astype(int) + grp_ref_sort_cols = [ + "Group", + "Start row of group", + "Response ID", + "General topic", + "Subtopic", + ] + if ( + _assess_sentiment(sentiment_checkbox) + and "Sentiment" in acc_reference_df.columns + ): + grp_ref_sort_cols.append("Sentiment") + acc_reference_df.sort_values(grp_ref_sort_cols, inplace=True) + elif "Main heading" in acc_topic_summary_df.columns: + acc_reference_df = acc_reference_df.merge( + acc_topic_summary_df[ + ["Main heading", "Subheading", "Topic number"] + ], + on=["Main heading", "Subheading"], + how="left", + ) + # Sort output dataframes + acc_reference_df["Response ID"] = ( + acc_reference_df["Response ID"].astype(float).astype(int) + ) + acc_reference_df["Start row of group"] = acc_reference_df[ + "Start row of group" + ].astype(int) + acc_reference_df.sort_values( + [ + "Group", + "Start row of group", + "Response ID", + "Main heading", + "Subheading", + "Topic number", + ], + inplace=True, + ) + + if "General topic" in acc_topic_summary_df.columns: + acc_topic_summary_df["Number of responses"] = acc_topic_summary_df[ + "Number of responses" + ].astype(int) + grp_topic_sort_cols = [ + "Group", + "Number of responses", + "General topic", + "Subtopic", + ] + grp_topic_sort_asc = [True, False, True, True] + if ( + _assess_sentiment(sentiment_checkbox) + and "Sentiment" in acc_topic_summary_df.columns + ): + grp_topic_sort_cols.append("Sentiment") + grp_topic_sort_asc.append(True) + acc_topic_summary_df.sort_values( + grp_topic_sort_cols, + ascending=grp_topic_sort_asc, + inplace=True, + ) + elif "Main heading" in acc_topic_summary_df.columns: + acc_topic_summary_df["Number of responses"] = acc_topic_summary_df[ + "Number of responses" + ].astype(int) + acc_topic_summary_df.sort_values( + [ + "Group", + "Number of responses", + "Main heading", + "Subheading", + "Topic number", + ], + ascending=[True, False, True, True, True], + inplace=True, + ) + + if "Group" in acc_reference_df.columns: + # Create missing references dataframe using consolidated data from all groups + # This ensures we correctly identify missing references across all groups + # Get all basic_response_data from all groups + all_basic_response_data = list() + for logged_item in acc_logged_content: + if "basic_response_data" in logged_item: + all_basic_response_data.extend(logged_item["basic_response_data"]) + + if all_basic_response_data: + all_basic_response_df = pd.DataFrame(all_basic_response_data) + acc_missing_df = create_missing_references_df( + all_basic_response_df, acc_reference_df + ) + else: + # Fallback: if no logged content, create empty missing_df + acc_missing_df = pd.DataFrame( + columns=["Missing Response ID", "Response Character Count"] + ) + + acc_reference_df_path = ( + output_folder + + overall_file_name + + "_col_" + + column_clean + + "_all_final_reference_table_" + + model_choice_clean_short + + ".csv" + ) + acc_topic_summary_df_path = ( + output_folder + + overall_file_name + + "_col_" + + column_clean + + "_all_final_unique_topics_" + + model_choice_clean_short + + ".csv" + ) + acc_reference_df_pivot_path = ( + output_folder + + overall_file_name + + "_col_" + + column_clean + + "_all_final_reference_pivot_" + + model_choice_clean_short + + ".csv" + ) + acc_missing_df_path = ( + output_folder + + overall_file_name + + "_col_" + + column_clean + + "_all_missing_df_" + + model_choice_clean_short + + ".csv" + ) + + acc_reference_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + acc_reference_df_path, index=None, encoding="utf-8-sig" + ) + acc_topic_summary_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + acc_topic_summary_df_path, index=None, encoding="utf-8-sig" + ) + acc_reference_df_pivot.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + acc_reference_df_pivot_path, index=None, encoding="utf-8-sig" + ) + acc_missing_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + acc_missing_df_path, index=None, encoding="utf-8-sig" + ) + + acc_log_files_output_paths.append(acc_missing_df_path) + + # Remove the existing output file list and replace with the updated concatenated outputs + substring_list_to_remove = [ + "_final_reference_table_pivot_", + "_final_reference_table_", + "_final_unique_topics_", + ] + acc_out_file_paths = [ + x + for x in acc_out_file_paths + if not any(sub in x for sub in substring_list_to_remove) + ] + + acc_out_file_paths.extend([acc_reference_df_path, acc_topic_summary_df_path]) + + # Outputs for markdown table output + unique_table_df_display_table = acc_topic_summary_df.apply( + lambda col: col.map(lambda x: wrap_text(x, max_text_length=max_text_length)) + ) + if produce_structured_summary_radio == "Yes": + unique_table_df_display_table = unique_table_df_display_table[ + ["General topic", "Subtopic", "Summary", "Group"] + ] + unique_table_df_display_table.rename( + columns={"General topic": "Main heading", "Subtopic": "Subheading"}, + inplace=True, + ) + acc_markdown_output = unique_table_df_display_table.to_markdown(index=False) + else: + display_cols = [ + "General topic", + "Subtopic", + "Number of responses", + "Summary", + "Group", + ] + if ( + _assess_sentiment(sentiment_checkbox) + and "Sentiment" in unique_table_df_display_table.columns + ): + display_cols.insert(2, "Sentiment") + + acc_markdown_output = unique_table_df_display_table[ + display_cols + ].to_markdown(index=False) + + acc_input_tokens, acc_output_tokens, acc_number_of_calls = ( + calculate_tokens_from_metadata( + acc_all_metadata_content, model_choice, model_name_map + ) + ) + + out_message = "\n".join(out_message) + out_message = ( + out_message + + " " + + f"Topic extraction finished processing all groups. Total time: {acc_total_time_taken:.2f}s" + ) + print(out_message) + + out_logged_content_df_path = ( + output_folder + + overall_file_name + + "_col_" + + column_clean + + "_logs_" + + model_choice_clean_short + + ".json" + ) + + with open( + out_logged_content_df_path, "w", encoding="utf-8-sig", errors="replace" + ) as f: + f.write(json.dumps(acc_logged_content)) + + acc_log_files_output_paths.append(out_logged_content_df_path) + + # The return signature should match extract_topics. + # The aggregated lists will be returned in the multiple slots. + return ( + acc_markdown_output, + acc_topics_table, + acc_topic_summary_df, + acc_reference_df, + acc_out_file_paths, # Slot 1 for out_file_paths + acc_out_file_paths, # Slot 2 for out_file_paths + acc_latest_batch_completed, # From the last successfully processed segment + acc_log_files_output_paths, # Slot 1 for log_files_output_paths + acc_log_files_output_paths, # Slot 2 for log_files_output_paths + acc_all_metadata_content, + acc_total_time_taken, + acc_out_file_paths, # Slot 3 + acc_out_file_paths, # Slot 4 + acc_gradio_df, # Last Gradio DF + acc_out_file_paths, # Slot 5 + acc_join_file_paths, + acc_missing_df, + acc_input_tokens, + acc_output_tokens, + acc_number_of_calls, + out_message, + acc_logged_content, + ) + + +def discover_topics_from_sample( + grouping_col: str, + in_data_file: Any, + file_data: pd.DataFrame, + original_file_name: str, + total_number_of_batches: int, + in_api_key: str, + temperature: float, + chosen_cols: List[str], + model_choice: str, + candidate_topics: gr.FileData = None, + sample_fraction_percent: float = 20, + random_seed: int = LLM_SEED, + context_textbox: str = "", + sentiment_checkbox: str = "Negative, Neutral, or Positive", + force_zero_shot_radio: str = "No", + in_excel_sheets: List[str] = list(), + force_single_topic_radio: str = "No", + produce_structured_summary_radio: str = "No", + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + hf_api_key_textbox: str = "", + azure_api_key_textbox: str = "", + azure_endpoint_textbox: str = "", + output_folder: str = OUTPUT_FOLDER, + initial_table_prompt: str = "", + initial_table_system_prompt: str = "", + add_existing_topics_system_prompt: str = "", + add_existing_topics_prompt: str = "", + number_of_prompts_used: int = 1, + batch_size: int = 5, + additional_instructions_summary_format: str = "", + additional_validation_issues_provided: str = "", + show_previous_table: str = "Yes", + api_url: str = None, + force_single_topic_prompt: str = "", + max_tokens: int = LLM_MAX_NEW_TOKENS, + model_name_map: dict = model_name_map, + max_time_for_loop: int = max_time_for_loop, + reasoning_suffix: str = REASONING_SUFFIX, + CHOSEN_LOCAL_MODEL_TYPE: str = CHOSEN_LOCAL_MODEL_TYPE, + output_debug_files: str = output_debug_files, + model: object = None, + tokenizer: object = None, + assistant_model: object = None, + max_rows: int = max_rows, + max_topics_number: int = MAXIMUM_ALLOWED_TOPICS, + progress=Progress(track_tqdm=True), +) -> Tuple: + """ + Run topic extraction on a random subsample and write a suggested-topics CSV + for reuse in a full-data analysis. + """ + ensure_model_in_map(model_choice, model_name_map) + + warnings: list[str] = [] + if produce_structured_summary_radio == "Yes": + raise ValueError( + "Topic discovery cannot run with structured summaries enabled. " + "Set 'Ask the model to produce structured summaries' to No." + ) + + if candidate_topics is not None: + has_candidate = False + if isinstance(candidate_topics, list): + has_candidate = bool(candidate_topics) + elif isinstance(candidate_topics, str): + has_candidate = bool(candidate_topics.strip()) + elif getattr(candidate_topics, "name", None): + has_candidate = True + if has_candidate: + warnings.append( + "Uploaded candidate topics were ignored for this discovery run." + ) + + try: + sample_fraction_percent = float(sample_fraction_percent) + except (TypeError, ValueError) as exc: + raise ValueError( + f"sample_fraction_percent must be a number between 1 and 100, got {sample_fraction_percent!r}" + ) from exc + + if sample_fraction_percent <= 0 or sample_fraction_percent > 100: + raise ValueError( + f"sample_fraction_percent must be between 1 and 100, got {sample_fraction_percent}" + ) + + sample_fraction = sample_fraction_percent / 100.0 + random_seed = int(random_seed) + + if file_data.empty: + ( + in_colnames_drop, + in_excel_sheets, + file_name, + join_colnames, + join_colnames_drop, + ) = put_columns_in_df(in_data_file) + file_data, file_name, total_number_of_batches = load_in_data_file( + in_data_file, chosen_cols, batch_size, in_excel_sheets + ) + else: + file_name = original_file_name + + if file_data.empty: + raise ValueError("No data available for topic discovery.") + + original_row_count = len(file_data) + group_col_for_sampling = grouping_col if grouping_col else None + + sampled_df, sample_metadata = subsample_responses_for_topic_discovery( + file_data=file_data, + sample_fraction=sample_fraction, + random_seed=random_seed, + group_col=group_col_for_sampling, + ) + + manifest_df = sampled_df.copy() + sampled_df = sampled_df.drop( + columns=["_discovery_original_row_index"], errors="ignore" + ) + + wrapper_results = wrapper_extract_topics_per_column_value( + grouping_col=grouping_col, + in_data_file=in_data_file, + file_data=sampled_df, + initial_existing_topics_table=pd.DataFrame(), + initial_existing_reference_df=pd.DataFrame(), + initial_existing_topic_summary_df=pd.DataFrame(), + initial_unique_table_df_display_table_markdown="", + original_file_name=file_name, + total_number_of_batches=max( + 1, (len(sampled_df) + batch_size - 1) // batch_size + ), + in_api_key=in_api_key, + temperature=temperature, + chosen_cols=chosen_cols, + model_choice=model_choice, + candidate_topics=None, + initial_first_loop_state=True, + initial_all_metadata_content_str="", + initial_latest_batch_completed=0, + initial_time_taken=0, + initial_table_prompt=initial_table_prompt, + initial_table_system_prompt=initial_table_system_prompt, + add_existing_topics_system_prompt=add_existing_topics_system_prompt, + add_existing_topics_prompt=add_existing_topics_prompt, + number_of_prompts_used=number_of_prompts_used, + batch_size=batch_size, + context_textbox=context_textbox, + sentiment_checkbox=sentiment_checkbox, + force_zero_shot_radio=force_zero_shot_radio, + in_excel_sheets=in_excel_sheets, + force_single_topic_radio=force_single_topic_radio, + produce_structured_summary_radio="No", + aws_access_key_textbox=aws_access_key_textbox, + aws_secret_key_textbox=aws_secret_key_textbox, + aws_region_textbox=aws_region_textbox, + hf_api_key_textbox=hf_api_key_textbox, + azure_api_key_textbox=azure_api_key_textbox, + azure_endpoint_textbox=azure_endpoint_textbox, + output_folder=output_folder, + existing_logged_content=list(), + additional_instructions_summary_format=additional_instructions_summary_format, + additional_validation_issues_provided=additional_validation_issues_provided, + show_previous_table=show_previous_table, + api_url=api_url, + force_single_topic_prompt=force_single_topic_prompt, + max_tokens=max_tokens, + model_name_map=model_name_map, + max_time_for_loop=max_time_for_loop, + reasoning_suffix=reasoning_suffix, + CHOSEN_LOCAL_MODEL_TYPE=CHOSEN_LOCAL_MODEL_TYPE, + output_debug_files=output_debug_files, + model=model, + tokenizer=tokenizer, + assistant_model=assistant_model, + max_rows=max_rows, + max_topics_number=max_topics_number, + progress=progress, + ) + + acc_topic_summary_df = wrapper_results[2] + chosen_col_str = ( + chosen_cols[0] + if isinstance(chosen_cols, list) and chosen_cols + else str(chosen_cols) if chosen_cols else "" + ) + file_name_cleaned = clean_column_name( + file_name, max_length=20, front_characters=True + ) + in_column_cleaned = clean_column_name(chosen_col_str, max_length=20) + model_choice_clean_short = clean_column_name( + model_name_map[model_choice]["short_name"], + max_length=20, + front_characters=False, + ) + pct_label = int(round(sample_fraction_percent)) + + discovery_base = ( + f"{file_name_cleaned}_col_{in_column_cleaned}_{model_choice_clean_short}" + f"_discovery_sample{pct_label}" + ) + discovery_topics_csv_path = output_folder + discovery_base + "_suggested_topics.csv" + discovery_manifest_csv_path = output_folder + discovery_base + "_manifest.csv" + + discovery_csv_path = write_candidate_topics_csv( + acc_topic_summary_df, discovery_topics_csv_path + ) + if not discovery_csv_path: + raise ValueError( + "Topic discovery completed but no General topic / Subtopic pairs were found." + ) + + write_topic_discovery_manifest_csv( + manifest_df, + discovery_manifest_csv_path, + group_col=group_col_for_sampling, + ) + + discovery_message = ( + f"Topic discovery complete: {sample_metadata['sampled_rows']} of " + f"{original_row_count} responses sampled ({pct_label}%, seed {random_seed}). " + "Upload the suggested topics CSV and run a full extraction with suggested topics." + ) + if warnings: + discovery_message = " ".join(warnings) + " " + discovery_message + + out_message = wrapper_results[-2] + if isinstance(out_message, str): + combined_message = discovery_message + " " + out_message + else: + combined_message = discovery_message + + acc_out_file_paths = list(wrapper_results[4]) + for extra_path in [discovery_csv_path, discovery_manifest_csv_path]: + if extra_path not in acc_out_file_paths: + acc_out_file_paths.append(extra_path) + + return ( + wrapper_results[0], + wrapper_results[1], + wrapper_results[2], + wrapper_results[3], + acc_out_file_paths, + acc_out_file_paths, + wrapper_results[6], + wrapper_results[7], + wrapper_results[8], + wrapper_results[9], + wrapper_results[10], + acc_out_file_paths, + acc_out_file_paths, + wrapper_results[13], + acc_out_file_paths, + wrapper_results[15], + wrapper_results[16], + wrapper_results[17], + wrapper_results[18], + wrapper_results[19], + combined_message, + wrapper_results[21], + discovery_csv_path, + ) + + +@spaces.GPU(duration=MAX_SPACES_GPU_RUN_TIME) +def discover_topics_from_sample_wrapper( + grouping_col: str, + in_data_file: Any, + file_data: pd.DataFrame, + original_file_name: str, + total_number_of_batches: int, + in_api_key: str, + temperature: float, + chosen_cols: List[str], + model_choice: str, + candidate_topics: gr.FileData = None, + sample_fraction_percent: float = 20, + random_seed: int = LLM_SEED, + context_textbox: str = "", + sentiment_checkbox: str = "Negative, Neutral, or Positive", + force_zero_shot_radio: str = "No", + in_excel_sheets: List[str] = list(), + force_single_topic_radio: str = "No", + produce_structured_summary_radio: str = "No", + aws_access_key_textbox: str = "", + aws_secret_key_textbox: str = "", + aws_region_textbox: str = "", + hf_api_key_textbox: str = "", + azure_api_key_textbox: str = "", + azure_endpoint_textbox: str = "", + output_folder: str = OUTPUT_FOLDER, + initial_table_prompt: str = "", + initial_table_system_prompt: str = "", + add_existing_topics_system_prompt: str = "", + add_existing_topics_prompt: str = "", + number_of_prompts_used: int = 1, + batch_size: int = 5, + additional_instructions_summary_format: str = "", + additional_validation_issues_provided: str = "", + show_previous_table: str = "Yes", + api_url: str = None, + max_topics_number: int = MAXIMUM_ALLOWED_TOPICS, + progress=Progress(track_tqdm=True), +) -> Tuple: + """Gradio wrapper for discover_topics_from_sample with app-compatible outputs.""" + results = discover_topics_from_sample( + grouping_col=grouping_col, + in_data_file=in_data_file, + file_data=file_data, + original_file_name=original_file_name, + total_number_of_batches=total_number_of_batches, + in_api_key=in_api_key, + temperature=temperature, + chosen_cols=chosen_cols, + model_choice=model_choice, + candidate_topics=candidate_topics, + sample_fraction_percent=sample_fraction_percent, + random_seed=random_seed, + context_textbox=context_textbox, + sentiment_checkbox=sentiment_checkbox, + force_zero_shot_radio=force_zero_shot_radio, + in_excel_sheets=in_excel_sheets, + force_single_topic_radio=force_single_topic_radio, + produce_structured_summary_radio=produce_structured_summary_radio, + aws_access_key_textbox=aws_access_key_textbox, + aws_secret_key_textbox=aws_secret_key_textbox, + aws_region_textbox=aws_region_textbox, + hf_api_key_textbox=hf_api_key_textbox, + azure_api_key_textbox=azure_api_key_textbox, + azure_endpoint_textbox=azure_endpoint_textbox, + output_folder=output_folder, + initial_table_prompt=initial_table_prompt, + initial_table_system_prompt=initial_table_system_prompt, + add_existing_topics_system_prompt=add_existing_topics_system_prompt, + add_existing_topics_prompt=add_existing_topics_prompt, + number_of_prompts_used=number_of_prompts_used, + batch_size=batch_size, + additional_instructions_summary_format=additional_instructions_summary_format, + additional_validation_issues_provided=additional_validation_issues_provided, + show_previous_table=show_previous_table, + api_url=api_url, + max_topics_number=max_topics_number, + progress=progress, + ) + + discovery_csv_path = results[-1] + extract_outputs = results[:-1] + + return ( + *extract_outputs, + discovery_csv_path, + ) + + +def join_modified_topic_names_to_ref_table( + modified_topic_summary_df: pd.DataFrame, + original_topic_summary_df: pd.DataFrame, + reference_df: pd.DataFrame, +): + """ + Take a unique topic table that has been modified by the user, and apply the topic name changes to the long-form reference table. + """ + + # Drop rows where Number of responses is either NA or null + modified_topic_summary_df = modified_topic_summary_df[ + ~modified_topic_summary_df["Number of responses"].isnull() + ] + modified_topic_summary_df.drop_duplicates( + ["General topic", "Subtopic", "Sentiment", "Topic number"], inplace=True + ) + + # First, join the modified topics to the original topics dataframe based on index to have the modified names alongside the original names + original_topic_summary_df_m = original_topic_summary_df.merge( + modified_topic_summary_df[ + ["General topic", "Subtopic", "Sentiment", "Topic number"] + ], + on="Topic number", + how="left", + suffixes=("", "_mod"), + ) + + original_topic_summary_df_m.drop_duplicates( + ["General topic", "Subtopic", "Sentiment", "Topic number"], inplace=True + ) + + # Then, join these new topic names onto the reference_df, merge based on the original names + modified_reference_df = reference_df.merge( + original_topic_summary_df_m[ + ["Topic number", "General Topic_mod", "Subtopic_mod", "Sentiment_mod"] + ], + on=["Topic number"], + how="left", + ) + + modified_reference_df.drop( + ["General topic", "Subtopic", "Sentiment"], + axis=1, + inplace=True, + errors="ignore", + ) + + modified_reference_df.rename( + columns={ + "General Topic_mod": "General topic", + "Subtopic_mod": "Subtopic", + "Sentiment_mod": "Sentiment", + }, + inplace=True, + ) + + modified_reference_df.drop( + ["General Topic_mod", "Subtopic_mod", "Sentiment_mod"], + inplace=True, + errors="ignore", + ) + + # modified_reference_df.drop_duplicates(["Response ID", "General topic", "Subtopic", "Sentiment"], inplace=True) + + modified_reference_df.sort_values( + [ + "Start row of group", + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + ], + inplace=True, + ) + + modified_reference_df = modified_reference_df.loc[ + :, + [ + "Response ID", + "General topic", + "Subtopic", + "Sentiment", + "Summary", + "Start row of group", + "Topic number", + ], + ] + + # Drop rows where Response ID is either NA or null + modified_reference_df = modified_reference_df[ + ~modified_reference_df["Response ID"].isnull() + ] + + return modified_reference_df + + +# MODIFY EXISTING TABLE +def modify_existing_output_tables( + original_topic_summary_df: pd.DataFrame, + modifiable_topic_summary_df: pd.DataFrame, + reference_df: pd.DataFrame, + text_output_file_list_state: List[str], + output_folder: str = OUTPUT_FOLDER, +) -> Tuple: + """ + Take a unique_topics table that has been modified, apply these new topic names to the long-form reference_df, and save both tables to file. + """ + + # Ensure text_output_file_list_state is a flat list + if any(isinstance(i, list) for i in text_output_file_list_state): + text_output_file_list_state = [ + item for sublist in text_output_file_list_state for item in sublist + ] # Flatten list + + # Extract file paths safely + reference_files = [x for x in text_output_file_list_state if "reference" in x] + unique_files = [x for x in text_output_file_list_state if "unique" in x] + + # Ensure files exist before accessing + reference_file_path = ( + os.path.basename(reference_files[0]) if reference_files else None + ) + unique_table_file_path = os.path.basename(unique_files[0]) if unique_files else None + + output_file_list = list() + + if reference_file_path and unique_table_file_path: + + reference_df = join_modified_topic_names_to_ref_table( + modifiable_topic_summary_df, original_topic_summary_df, reference_df + ) + + ## Response ID table mapping response numbers to topics + reference_table_file_name = reference_file_path.replace(".csv", "_mod") + new_reference_df_file_path = output_folder + reference_table_file_name + ".csv" + reference_df.drop(["1", "2", "3"], axis=1, errors="ignore").to_csv( + new_reference_df_file_path, index=None, encoding="utf-8-sig" + ) + output_file_list.append(new_reference_df_file_path) + + # Drop rows where Number of responses is NA or null + modifiable_topic_summary_df = modifiable_topic_summary_df[ + ~modifiable_topic_summary_df["Number of responses"].isnull() + ] + + # Convert 'Number of responses' to numeric (forcing errors to NaN if conversion fails) + modifiable_topic_summary_df["Number of responses"] = pd.to_numeric( + modifiable_topic_summary_df["Number of responses"], errors="coerce" + ) + + # Drop any rows where conversion failed (original non-numeric values) + modifiable_topic_summary_df.dropna(subset=["Number of responses"], inplace=True) + + # Sort values + modifiable_topic_summary_df.sort_values( + ["Number of responses"], ascending=False, inplace=True + ) + + unique_table_file_name = unique_table_file_path.replace(".csv", "_mod") + modified_unique_table_file_path = ( + output_folder + unique_table_file_name + ".csv" + ) + modifiable_topic_summary_df.drop( + ["1", "2", "3"], axis=1, errors="ignore" + ).to_csv(modified_unique_table_file_path, index=None, encoding="utf-8-sig") + output_file_list.append(modified_unique_table_file_path) + + else: + output_file_list = text_output_file_list_state + reference_table_file_name = reference_file_path + unique_table_file_name = unique_table_file_path + raise Exception("Response ID and unique topic tables not found.") + + # Outputs for markdown table output + unique_table_df_revised_display = modifiable_topic_summary_df.apply( + lambda col: col.map(lambda x: wrap_text(x, max_text_length=max_text_length)) + ) + deduplicated_unique_table_markdown = unique_table_df_revised_display.to_markdown( + index=False + ) + + return ( + modifiable_topic_summary_df, + reference_df, + output_file_list, + output_file_list, + output_file_list, + output_file_list, + reference_table_file_name, + unique_table_file_name, + deduplicated_unique_table_markdown, + ) + + +@spaces.GPU(duration=MAX_SPACES_GPU_RUN_TIME) +def all_in_one_pipeline( + grouping_col: str, + in_data_files: List[str], + file_data: pd.DataFrame, + existing_topics_table: pd.DataFrame, + existing_reference_df: pd.DataFrame, + existing_topic_summary_df: pd.DataFrame, + unique_table_df_display_table_markdown: str, + original_file_name: str, + total_number_of_batches: int, + in_api_key: str, + temperature: float, + chosen_cols: List[str], + model_choice: str, + candidate_topics: gr.FileData, + first_loop_state: bool, + conversation_metadata_text: str, + latest_batch_completed: int, + time_taken_so_far: float, + initial_table_prompt_text: str, + initial_table_system_prompt_text: str, + add_existing_topics_system_prompt_text: str, + add_existing_topics_prompt_text: str, + number_of_prompts_used: int, + batch_size: int, + context_text: str, + sentiment_choice: str, + force_zero_shot_choice: str, + in_excel_sheets: List[str], + force_single_topic_choice: str, + produce_structures_summary_choice: str, + aws_access_key_text: str, + aws_secret_key_text: str, + aws_region_text: str, + hf_api_key_text: str, + azure_api_key_text: str, + azure_endpoint_text: str, + output_folder: str = OUTPUT_FOLDER, + merge_sentiment: str = "No", + merge_general_topics: str = "Yes", + score_threshold: int = 90, + summarise_format: str = "", + random_seed: int = 42, + log_files_output_list_state: List[str] = list(), + model_name_map_state: dict = model_name_map, + usage_logs_location: str = "", + existing_logged_content: list = list(), + additional_instructions_summary_format: str = "", + additional_validation_issues_provided: str = "", + show_previous_table: str = "Yes", + sample_reference_table_checkbox: bool = True, + api_url: str = None, + output_debug_files: str = output_debug_files, + model: object = None, + tokenizer: object = None, + assistant_model: object = None, + max_rows: int = max_rows, + max_topics_number: int = MAXIMUM_ALLOWED_TOPICS, + progress=Progress(track_tqdm=True), +): + """ + Orchestrates the full All-in-one flow: extract → deduplicate → summarise → overall summary → Excel export. + + Args: + grouping_col (str): The column used for grouping data. + in_data_files (List[str]): List of input data file paths. + file_data (pd.DataFrame): The input data as a pandas DataFrame. + existing_topics_table (pd.DataFrame): DataFrame of existing topics. + existing_reference_df (pd.DataFrame): DataFrame of existing reference data. + existing_topic_summary_df (pd.DataFrame): DataFrame of existing topic summaries. + unique_table_df_display_table_markdown (str): Markdown string for displaying unique topics. + original_file_name (str): The original name of the input file. + total_number_of_batches (int): Total number of batches for processing. + in_api_key (str): API key for the LLM. + temperature (float): Temperature setting for the LLM. + chosen_cols (List[str]): List of columns chosen for analysis. + model_choice (str): The chosen LLM model. + candidate_topics (gr.FileData): Gradio file data for candidate topics. + first_loop_state (bool): State indicating if it's the first loop. + conversation_metadata_text (str): Text containing conversation metadata. + latest_batch_completed (int): The latest batch number completed. + time_taken_so_far (float): Cumulative time taken so far. + initial_table_prompt_text (str): Initial prompt text for table generation. + initial_table_system_prompt_text (str): Initial system prompt text for table generation. + add_existing_topics_system_prompt_text (str): System prompt for adding existing topics. + add_existing_topics_prompt_text (str): Prompt for adding existing topics. + number_of_prompts_used (int): Number of prompts used in sequence. + batch_size (int): Size of each processing batch. + context_text (str): Additional context for the LLM. + sentiment_choice (str): Choice for sentiment analysis (e.g., "Yes", "No"). + force_zero_shot_choice (str): Choice to force zero-shot prompting. + in_excel_sheets (List[str]): List of sheet names in the input Excel file. + force_single_topic_choice (str): Choice to force single topic extraction. + produce_structures_summary_choice (str): Choice to produce structured summaries. + aws_access_key_text (str): AWS access key. + aws_secret_key_text (str): AWS secret key. + hf_api_key_text (str): Hugging Face API key. + azure_api_key_text (str): Azure/OpenAI API key. + output_folder (str, optional): Folder to save output files. Defaults to OUTPUT_FOLDER. + merge_sentiment (str, optional): Whether to merge sentiment. Defaults to "No". + merge_general_topics (str, optional): Whether to merge general topics. Defaults to "Yes". + score_threshold (int, optional): Score threshold for topic matching. Defaults to 90. + summarise_format (str, optional): Format for summarization. Defaults to "". + random_seed (int, optional): Random seed for reproducibility. Defaults to 42. + log_files_output_list_state (List[str], optional): List of log file paths. Defaults to list(). + model_name_map_state (dict, optional): Mapping of model names. Defaults to model_name_map. + usage_logs_location (str, optional): Location for usage logs. Defaults to "". + existing_logged_content (list, optional): Existing logged content. Defaults to list(). + additional_instructions_summary_format (str, optional): Summary format for adding existing topics. Defaults to "". + additional_validation_issues_provided (str, optional): Additional validation issues provided by the user. Defaults to "". + show_previous_table (str, optional): Whether to show the previous table ("Yes" or "No"). Defaults to "Yes". + sample_reference_table_checkbox (bool, optional): Whether to sample summaries before creating revised summaries. + api_url (str, optional): API URL for inference-server models. Defaults to None. + output_debug_files (str, optional): Whether to output debug files. Defaults to "False". + model (object, optional): Loaded local model object. Defaults to None. + tokenizer (object, optional): Loaded local tokenizer object. Defaults to None. + assistant_model (object, optional): Loaded local assistant model object. Defaults to None. + max_rows (int, optional): Maximum number of rows to process. Defaults to max_rows. + progress (Progress, optional): Gradio Progress object for tracking. Defaults to Progress(track_tqdm=True). + + Returns: + A tuple matching the UI components updated during the original chained flow. + """ + + # Ensure custom model_choice is registered in model_name_map_state + ensure_model_in_map(model_choice, model_name_map_state) + + # Load local model if it's not already loaded + if ( + (model_name_map_state[model_choice]["source"] == "Local") + & (RUN_LOCAL_MODEL == "1") + & (not model) + ): + model = get_model() + tokenizer = get_tokenizer() + assistant_model = get_assistant_model() + + total_input_tokens = 0 + total_output_tokens = 0 + total_number_of_calls = 0 + total_time_taken = 0 + out_message = list() + out_logged_content = list() + + print( + "Analysing file: ", + original_file_name, + "column(s): ", + chosen_cols, + "with model: ", + model_choice, + ) + + model_source = model_name_map_state[model_choice]["source"] + + # 1) Extract topics (group-aware) + ( + display_markdown, + out_topics_table, + out_topic_summary_df, + out_reference_df, + out_file_paths_1, + _out_file_paths_dup, + out_latest_batch_completed, + out_log_files, + _out_log_files_dup, + out_conversation_metadata, + out_time_taken, + out_file_paths_2, + _out_file_paths_3, + out_gradio_df, + out_file_paths_4, + out_join_files, + out_missing_df, + out_input_tokens, + out_output_tokens, + out_number_of_calls, + out_message_text, + out_logged_content, + ) = wrapper_extract_topics_per_column_value( + grouping_col=grouping_col, + in_data_file=in_data_files, + file_data=file_data, + initial_existing_topics_table=existing_topics_table, + initial_existing_reference_df=existing_reference_df, + initial_existing_topic_summary_df=existing_topic_summary_df, + initial_unique_table_df_display_table_markdown=unique_table_df_display_table_markdown, + original_file_name=original_file_name, + total_number_of_batches=total_number_of_batches, + in_api_key=in_api_key, + temperature=temperature, + chosen_cols=chosen_cols, + model_choice=model_choice, + candidate_topics=candidate_topics, + initial_first_loop_state=first_loop_state, + initial_all_metadata_content_str=conversation_metadata_text, + initial_latest_batch_completed=latest_batch_completed, + initial_time_taken=time_taken_so_far, + initial_table_prompt=initial_table_prompt_text, + initial_table_system_prompt=initial_table_system_prompt_text, + add_existing_topics_system_prompt=add_existing_topics_system_prompt_text, + add_existing_topics_prompt=add_existing_topics_prompt_text, + number_of_prompts_used=number_of_prompts_used, + batch_size=batch_size, + context_textbox=context_text, + sentiment_checkbox=sentiment_choice, + force_zero_shot_radio=force_zero_shot_choice, + in_excel_sheets=in_excel_sheets, + force_single_topic_radio=force_single_topic_choice, + produce_structured_summary_radio=produce_structures_summary_choice, + aws_access_key_textbox=aws_access_key_text, + aws_secret_key_textbox=aws_secret_key_text, + aws_region_textbox=aws_region_text, + hf_api_key_textbox=hf_api_key_text, + azure_api_key_textbox=azure_api_key_text, + azure_endpoint_textbox=azure_endpoint_text, + output_folder=output_folder, + existing_logged_content=existing_logged_content, + model_name_map=model_name_map_state, + output_debug_files=output_debug_files, + model=model, + tokenizer=tokenizer, + assistant_model=assistant_model, + max_rows=max_rows, + additional_instructions_summary_format=additional_instructions_summary_format, + additional_validation_issues_provided=additional_validation_issues_provided, + show_previous_table=show_previous_table, + api_url=api_url, + max_topics_number=max_topics_number, + ) + + total_input_tokens += out_input_tokens + total_output_tokens += out_output_tokens + total_number_of_calls += out_number_of_calls + total_time_taken += out_time_taken + out_message.append(out_message_text) + + # Prepare outputs after extraction, matching wrapper outputs + topic_extraction_output_files = out_file_paths_1 + text_output_file_list_state = out_file_paths_1 + log_files_output_list_state = out_log_files + + # If producing structured summaries, return the outputs after extraction + if produce_structures_summary_choice == "Yes": + + # Write logged content to file + column_clean = clean_column_name(chosen_cols, max_length=20) + model_choice_clean = model_name_map[model_choice]["short_name"] + model_choice_clean_short = clean_column_name( + model_choice_clean, max_length=20, front_characters=False + ) + + out_logged_content_df_path = ( + output_folder + + original_file_name + + "_col_" + + column_clean + + "_logs_" + + model_choice_clean_short + + ".json" + ) + + with open( + out_logged_content_df_path, "w", encoding="utf-8-sig", errors="replace" + ) as f: + f.write(json.dumps(out_logged_content)) + + log_files_output_list_state.append(out_logged_content_df_path) + out_log_files.append(out_logged_content_df_path) + + # Map to the UI outputs list expected by the new single-call wiring + return ( + display_markdown, + out_topics_table, + out_topic_summary_df, + out_reference_df, + topic_extraction_output_files, + text_output_file_list_state, + out_latest_batch_completed, + out_log_files, + log_files_output_list_state, + out_conversation_metadata, + total_time_taken, + out_file_paths_1, + list(), # summarisation_input_files is not available yet + out_gradio_df, + list(), # modification_input_files placeholder + out_join_files, + out_missing_df, + total_input_tokens, + total_output_tokens, + total_number_of_calls, + out_message[0], + pd.DataFrame(), # summary_reference_table_sample_state is not available yet + "", # summarised_references_markdown is not available yet + out_topic_summary_df, + out_reference_df, + list(), # summary_output_files is not available yet + list(), # summarised_outputs_list is not available yet + 0, # latest_summary_completed_num is not available yet + list(), # overall_summarisation_input_files is not available yet + list(), # overall_summary_output_files is not available yet + "", # overall_summarised_output_markdown is not available yet + pd.DataFrame(), # summarised_output_df is not available yet + out_logged_content, + ) + + # 2) Deduplication + print("Deduplicating topic names with fuzzy matching") + ( + ref_df_loaded, + unique_df_loaded, + latest_batch_completed_no_loop, + deduplication_input_files_status, + working_data_file_name_textbox, + unique_topics_table_file_name_textbox, + ) = load_in_previous_data_files(out_file_paths_1) + + ( + ref_df_after_dedup, + unique_df_after_dedup, + summarisation_input_files, + log_files_output_dedup, + summarised_output_markdown, + ) = deduplicate_topics( + reference_df=ref_df_loaded if not ref_df_loaded.empty else out_reference_df, + topic_summary_df=( + unique_df_loaded if not unique_df_loaded.empty else out_topic_summary_df + ), + reference_table_file_name=working_data_file_name_textbox, + unique_topics_table_file_name=unique_topics_table_file_name_textbox, + in_excel_sheets=in_excel_sheets, + merge_sentiment=merge_sentiment, + merge_general_topics=merge_general_topics, + score_threshold=score_threshold, + in_data_files=in_data_files, + chosen_cols=chosen_cols, + output_folder=output_folder, + sentiment_checkbox=sentiment_choice, + ) + + # LLM-based deduplication if enabled + if force_zero_shot_choice == "No" and ALL_IN_ONE_USE_LLM_DEDUP: + # Set up model source and bedrock runtime if needed + + print("Deduplicating topic names with LLM") + + # Call LLM deduplication + ( + ref_df_after_dedup, + unique_df_after_dedup, + summarisation_input_files, + log_files_output_dedup, + summarised_output_markdown, + dedup_input_tokens, + dedup_output_tokens, + dedup_number_of_calls, + dedup_estimated_time_taken, + ) = deduplicate_topics_llm( + reference_df=ref_df_after_dedup, + topic_summary_df=unique_df_after_dedup, + reference_table_file_name=working_data_file_name_textbox, + unique_topics_table_file_name=unique_topics_table_file_name_textbox, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + model_source=model_source, + local_model=model, + tokenizer=tokenizer, + assistant_model=assistant_model, + in_excel_sheets=in_excel_sheets, + merge_sentiment=merge_sentiment, + merge_general_topics=merge_general_topics, + in_data_files=in_data_files, + chosen_cols=chosen_cols, + output_folder=output_folder, + candidate_topics=candidate_topics, + azure_endpoint=azure_endpoint_text, + output_debug_files=output_debug_files, + api_url=api_url, + aws_access_key_textbox=aws_access_key_text, + aws_secret_key_textbox=aws_secret_key_text, + aws_region_textbox=aws_region_text, + azure_api_key_textbox=azure_api_key_text, + model_name_map=model_name_map_state, + sentiment_checkbox=sentiment_choice, + ) + + # Update token counts and time taken + total_input_tokens += dedup_input_tokens + total_output_tokens += dedup_output_tokens + total_number_of_calls += dedup_number_of_calls + total_time_taken += dedup_estimated_time_taken + out_message.append( + f"LLM deduplication completed. Total time: {dedup_estimated_time_taken:.2f}s" + ) + + # 3) Summarisation + ( + ref_df_loaded_2, + unique_df_loaded_2, + _latest_batch_completed_no_loop_2, + _deduplication_input_files_status_2, + _working_name_2, + _unique_name_2, + ) = load_in_previous_data_files(summarisation_input_files) + + ( + summary_reference_table_sample_state, + master_unique_topics_df_revised_summaries_state, + master_reference_df_revised_summaries_state, + summary_output_files, + summarised_outputs_list, + latest_summary_completed_num, + conversation_metadata_text_updated, + display_markdown_updated, + log_files_output_after_sum, + overall_summarisation_input_files, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + output_messages_textbox, + out_logged_content, + ) = wrapper_summarise_output_topics_per_group( + grouping_col=grouping_col, + sampled_reference_table_df=ref_df_after_dedup, + topic_summary_df=unique_df_after_dedup, + reference_table_df=ref_df_after_dedup, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + reference_data_file_name=working_data_file_name_textbox, + summarised_outputs=list(), + latest_summary_completed=0, + out_metadata_str=out_conversation_metadata, + in_data_files=in_data_files, + in_excel_sheets=in_excel_sheets, + chosen_cols=chosen_cols, + log_output_files=log_files_output_list_state, + summarise_format_radio=summarise_format, + output_folder=output_folder, + context_textbox=context_text, + aws_access_key_textbox=aws_access_key_text, + aws_secret_key_textbox=aws_secret_key_text, + aws_region_textbox=aws_region_text, + model_name_map=model_name_map_state, + hf_api_key_textbox=hf_api_key_text, + azure_endpoint_textbox=azure_endpoint_text, + existing_logged_content=out_logged_content, + sample_reference_table=sample_reference_table_checkbox, + no_of_sampled_summaries=100, + random_seed=random_seed, + output_debug_files=output_debug_files, + api_url=api_url, + additional_summary_instructions_provided=additional_instructions_summary_format, + local_model=model, + tokenizer=tokenizer, + assistant_model=assistant_model, + ) + + # Generate summarised_references_markdown from the sampled reference table + summarised_references_markdown = summary_reference_table_sample_state.to_markdown( + index=False + ) + + total_input_tokens += input_tokens_num + total_output_tokens += output_tokens_num + total_number_of_calls += number_of_calls_num + total_time_taken += estimated_time_taken_number + out_message.append(output_messages_textbox) + + # 4) Overall summary + ( + _ref_df_loaded_3, + _unique_df_loaded_3, + _latest_batch_completed_no_loop_3, + _deduplication_input_files_status_3, + _working_name_3, + _unique_name_3, + ) = load_in_previous_data_files(overall_summarisation_input_files) + + ( + overall_summary_output_files, + overall_summarised_output_markdown, + summarised_output_df, + conversation_metadata_textbox, + input_tokens_num, + output_tokens_num, + number_of_calls_num, + estimated_time_taken_number, + output_messages_textbox, + out_logged_content, + ) = overall_summary( + topic_summary_df=master_unique_topics_df_revised_summaries_state, + model_choice=model_choice, + in_api_key=in_api_key, + temperature=temperature, + reference_data_file_name=working_data_file_name_textbox, + output_folder=output_folder, + chosen_cols=chosen_cols, + context_textbox=context_text, + aws_access_key_textbox=aws_access_key_text, + aws_secret_key_textbox=aws_secret_key_text, + aws_region_textbox=aws_region_text, + model_name_map=model_name_map_state, + hf_api_key_textbox=hf_api_key_text, + azure_endpoint_textbox=azure_endpoint_text, + local_model=model, + tokenizer=tokenizer, + assistant_model=assistant_model, + existing_logged_content=out_logged_content, + output_debug_files=output_debug_files, + api_url=api_url, + ) + + total_input_tokens += input_tokens_num + total_output_tokens += output_tokens_num + total_number_of_calls += number_of_calls_num + total_time_taken += estimated_time_taken_number + out_message.append(output_messages_textbox) + + out_message = "\n".join(out_message) + out_message = ( + out_message + "\n" + f"Overall time for all processes: {total_time_taken:.2f}s" + ) + print(out_message) + + # Write logged content to file + column_clean = clean_column_name(chosen_cols, max_length=20) + model_choice_clean = model_name_map[model_choice]["short_name"] + model_choice_clean_short = clean_column_name( + model_choice_clean, max_length=20, front_characters=False + ) + + out_logged_content_df_path = ( + output_folder + + original_file_name + + "_col_" + + column_clean + + "_logs_" + + model_choice_clean_short + + ".json" + ) + + with open( + out_logged_content_df_path, "w", encoding="utf-8-sig", errors="replace" + ) as f: + f.write(json.dumps(out_logged_content)) + + log_files_output_list_state.append(out_logged_content_df_path) + log_files_output_after_sum.append(out_logged_content_df_path) + + # Map to the UI outputs list expected by the new single-call wiring + # Use the original markdown with renamed columns if produce_structured_summary_radio is "Yes" + final_display_markdown = ( + display_markdown_updated if display_markdown_updated else display_markdown + ) + if produce_structures_summary_choice == "Yes": + final_display_markdown = unique_table_df_display_table_markdown + + return ( + final_display_markdown, + out_topics_table, + unique_df_after_dedup, + ref_df_after_dedup, + topic_extraction_output_files, + text_output_file_list_state, + out_latest_batch_completed, + log_files_output_after_sum if log_files_output_after_sum else out_log_files, + log_files_output_list_state, + ( + conversation_metadata_text_updated + if conversation_metadata_text_updated + else out_conversation_metadata + ), + total_time_taken, + out_file_paths_1, + summarisation_input_files, + out_gradio_df, + list(), # modification_input_files placeholder + out_join_files, + out_missing_df, + total_input_tokens, + total_output_tokens, + total_number_of_calls, + out_message, + summary_reference_table_sample_state, + summarised_references_markdown, + master_unique_topics_df_revised_summaries_state, + master_reference_df_revised_summaries_state, + summary_output_files, + summarised_outputs_list, + latest_summary_completed_num, + overall_summarisation_input_files, + overall_summary_output_files, + overall_summarised_output_markdown, + summarised_output_df, + out_logged_content, + ) + + +# When LLM deduplication button pressed, deduplicate data using LLM +def deduplicate_topics_llm_wrapper( + reference_df, + topic_summary_df, + reference_table_file_name, + unique_topics_table_file_name, + model_choice, + in_api_key, + temperature, + in_excel_sheets, + merge_sentiment, + merge_general_topics, + in_data_files, + chosen_cols, + output_folder, + candidate_topics=None, + azure_endpoint="", + api_url=None, + aws_access_key_textbox="", + aws_secret_key_textbox="", + aws_region_textbox="", + azure_api_key_textbox="", + sentiment_checkbox="Negative or Positive", +): + # Ensure custom model_choice is registered in model_name_map + ensure_model_in_map(model_choice) + model_source = model_name_map[model_choice]["source"] + return deduplicate_topics_llm( + reference_df, + topic_summary_df, + reference_table_file_name, + unique_topics_table_file_name, + model_choice, + in_api_key, + temperature, + model_source, + None, + None, + None, + None, + in_excel_sheets, + merge_sentiment, + merge_general_topics, + in_data_files, + chosen_cols, + output_folder, + candidate_topics, + azure_endpoint, + OUTPUT_DEBUG_FILES, + api_url, + aws_access_key_textbox, + aws_secret_key_textbox, + aws_region_textbox, + azure_api_key_textbox, + sentiment_checkbox=sentiment_checkbox, + ) diff --git a/tools/llm_funcs.py b/tools/llm_funcs.py new file mode 100644 index 0000000000000000000000000000000000000000..f518a9a02a8899458a1f07510605376cc237c21f --- /dev/null +++ b/tools/llm_funcs.py @@ -0,0 +1,2108 @@ +import json +import os +import re +import time +from typing import List, Tuple + +import boto3 +import pandas as pd +import requests + +# Import mock patches if in test mode +if os.environ.get("USE_MOCK_LLM") == "1" or os.environ.get("TEST_MODE") == "1": + try: + # Try to import and apply mock patches + import sys + + # Add project root to sys.path so we can import test.mock_llm_calls + project_root = os.path.dirname(os.path.dirname(__file__)) + if project_root not in sys.path: + sys.path.insert(0, project_root) + try: + from test.mock_llm_calls import apply_mock_patches + + apply_mock_patches() + except ImportError: + # If mock module not found, continue without mocking + pass + except Exception: + # If anything fails, continue without mocking + pass +from google import genai as ai +from google.genai import types +from gradio import Progress +from huggingface_hub import hf_hub_download +from openai import OpenAI +from tqdm import tqdm + +model_type = None # global variable setup +full_text = ( + "" # Define dummy source text (full text) just to enable highlight function to load +) + +# Global variables for model and tokenizer +_model = None +_tokenizer = None +_assistant_model = None + +from tools.config import ( + ASSISTANT_MODEL, + BATCH_SIZE_DEFAULT, + CHOSEN_LOCAL_MODEL_TYPE, + COMPILE_MODE, + COMPILE_TRANSFORMERS, + DEDUPLICATION_THRESHOLD, + HF_TOKEN, + INFERENCE_SERVER_DISABLE_THINKING, + INFERENCE_SERVER_READ_TIMEOUT_SECONDS, + INT8_WITH_OFFLOAD_TO_CPU, + K_QUANT_LEVEL, + LLM_BATCH_SIZE, + LLM_CONTEXT_LENGTH, + LLM_LAST_N_TOKENS, + LLM_MAX_GPU_LAYERS, + LLM_MAX_NEW_TOKENS, + LLM_MIN_P, + LLM_REPETITION_PENALTY, + LLM_RESET, + LLM_SAMPLE, + LLM_SEED, + LLM_STOP_STRINGS, + LLM_STREAM, + LLM_TEMPERATURE, + LLM_THREADS, + LLM_TOP_K, + LLM_TOP_P, + LOAD_LOCAL_MODEL_AT_START, + LOCAL_MODEL_FILE, + LOCAL_MODEL_FOLDER, + LOCAL_REPO_ID, + MAX_COMMENT_CHARS, + MAX_TIME_FOR_LOOP, + MODEL_DTYPE, + MULTIMODAL_PROMPT_FORMAT, + NUM_PRED_TOKENS, + NUMBER_OF_RETRY_ATTEMPTS, + RUN_LOCAL_MODEL, + SPECULATIVE_DECODING, + TIMEOUT_WAIT, + USE_BITSANDBYTES, + USE_LLAMA_CPP, + USE_LLAMA_SWAP, + V_QUANT_LEVEL, +) +from tools.helper_functions import _get_env_list + +if SPECULATIVE_DECODING == "True": + SPECULATIVE_DECODING = True +else: + SPECULATIVE_DECODING = False + + +if isinstance(NUM_PRED_TOKENS, str): + NUM_PRED_TOKENS = int(NUM_PRED_TOKENS) +if isinstance(LLM_MAX_GPU_LAYERS, str): + LLM_MAX_GPU_LAYERS = int(LLM_MAX_GPU_LAYERS) +if isinstance(LLM_THREADS, str): + LLM_THREADS = int(LLM_THREADS) + +if LLM_RESET == "True": + reset = True +else: + reset = False + +if LLM_STREAM == "True": + stream = True +else: + stream = False + +if LLM_SAMPLE == "True": + sample = True +else: + sample = False + +if LLM_STOP_STRINGS: + LLM_STOP_STRINGS = _get_env_list(LLM_STOP_STRINGS, strip_strings=False) + +max_tokens = LLM_MAX_NEW_TOKENS +timeout_wait = TIMEOUT_WAIT +inference_server_read_timeout_seconds = INFERENCE_SERVER_READ_TIMEOUT_SECONDS +number_of_api_retry_attempts = NUMBER_OF_RETRY_ATTEMPTS +max_time_for_loop = MAX_TIME_FOR_LOOP +batch_size_default = BATCH_SIZE_DEFAULT +deduplication_threshold = DEDUPLICATION_THRESHOLD +max_comment_character_length = MAX_COMMENT_CHARS + +temperature = LLM_TEMPERATURE +top_k = LLM_TOP_K +top_p = LLM_TOP_P +min_p = LLM_MIN_P +repetition_penalty = LLM_REPETITION_PENALTY +last_n_tokens = LLM_LAST_N_TOKENS +LLM_MAX_NEW_TOKENS: int = LLM_MAX_NEW_TOKENS +seed: int = LLM_SEED +reset: bool = reset +stream: bool = stream +batch_size: int = LLM_BATCH_SIZE +context_length: int = LLM_CONTEXT_LENGTH +sample = LLM_SAMPLE +stop_strings = LLM_STOP_STRINGS +speculative_decoding = SPECULATIVE_DECODING +if LLM_MAX_GPU_LAYERS != 0: + gpu_layers = int(LLM_MAX_GPU_LAYERS) + torch_device = "cuda" +else: + gpu_layers = 0 + torch_device = "cpu" + +if not LLM_THREADS: + threads = 1 +else: + threads = LLM_THREADS + + +class llama_cpp_init_config_gpu: + def __init__( + self, + last_n_tokens=last_n_tokens, + seed=seed, + n_threads=threads, + n_batch=batch_size, + n_ctx=context_length, + n_gpu_layers=gpu_layers, + reset=reset, + ): + + self.last_n_tokens = last_n_tokens + self.seed = seed + self.n_threads = n_threads + self.n_batch = n_batch + self.n_ctx = n_ctx + self.n_gpu_layers = n_gpu_layers + self.reset = reset + # self.stop: list[str] = field(default_factory=lambda: [stop_string]) + + def update_gpu(self, new_value): + self.n_gpu_layers = new_value + + def update_context(self, new_value): + self.n_ctx = new_value + + +class llama_cpp_init_config_cpu(llama_cpp_init_config_gpu): + def __init__(self): + super().__init__() + self.n_gpu_layers = gpu_layers + self.n_ctx = context_length + + +gpu_config = llama_cpp_init_config_gpu() +cpu_config = llama_cpp_init_config_cpu() + + +class LlamaCPPGenerationConfig: + def __init__( + self, + temperature=temperature, + top_k=top_k, + min_p=min_p, + top_p=top_p, + repeat_penalty=repetition_penalty, + seed=seed, + stream=stream, + max_tokens=LLM_MAX_NEW_TOKENS, + reset=reset, + ): + self.temperature = temperature + self.top_k = top_k + self.top_p = top_p + self.repeat_penalty = repeat_penalty + self.seed = seed + self.max_tokens = max_tokens + self.stream = stream + self.reset = reset + + def update_temp(self, new_value): + self.temperature = new_value + + +# ResponseObject class for AWS Bedrock calls +class ResponseObject: + def __init__(self, text, usage_metadata): + self.text = text + self.usage_metadata = usage_metadata + + +### +# LOCAL MODEL FUNCTIONS +### + + +def get_model_path( + repo_id=LOCAL_REPO_ID, + model_filename=LOCAL_MODEL_FILE, + model_dir=LOCAL_MODEL_FOLDER, + hf_token=HF_TOKEN, +): + # Construct the expected local path + local_path = os.path.join(model_dir, model_filename) + + print("local path for model load:", local_path) + + try: + if os.path.exists(local_path): + print(f"Model already exists at: {local_path}") + + return local_path + else: + if hf_token: + print("Downloading model from Hugging Face Hub with HF token") + downloaded_model_path = hf_hub_download( + repo_id=repo_id, token=hf_token, filename=model_filename + ) + + return downloaded_model_path + else: + print( + "No HF token found, downloading model from Hugging Face Hub without token" + ) + downloaded_model_path = hf_hub_download( + repo_id=repo_id, filename=model_filename + ) + + return downloaded_model_path + + except Exception as e: + print("Error loading model:", e) + raise Warning("Error loading model:", e) + + +def load_model( + local_model_type: str = CHOSEN_LOCAL_MODEL_TYPE, + gpu_layers: int = gpu_layers, + max_context_length: int = context_length, + gpu_config: llama_cpp_init_config_gpu = gpu_config, + cpu_config: llama_cpp_init_config_cpu = cpu_config, + torch_device: str = torch_device, + repo_id=LOCAL_REPO_ID, + model_filename=LOCAL_MODEL_FILE, + model_dir=LOCAL_MODEL_FOLDER, + compile_mode=COMPILE_MODE, + model_dtype=MODEL_DTYPE, + hf_token=HF_TOKEN, + speculative_decoding=speculative_decoding, + model=None, + tokenizer=None, + assistant_model=None, +): + """ + Load in a model from Hugging Face hub via the transformers package, or using llama_cpp_python by downloading a GGUF file from Huggingface Hub. + + Args: + local_model_type (str): The type of local model to load (e.g., "llama-cpp"). + gpu_layers (int): The number of GPU layers to offload to the GPU. + max_context_length (int): The maximum context length for the model. + gpu_config (llama_cpp_init_config_gpu): Configuration object for GPU-specific Llama.cpp parameters. + cpu_config (llama_cpp_init_config_cpu): Configuration object for CPU-specific Llama.cpp parameters. + torch_device (str): The device to load the model on ("cuda" for GPU, "cpu" for CPU). + repo_id (str): The Hugging Face repository ID where the model is located. + model_filename (str): The specific filename of the model to download from the repository. + model_dir (str): The local directory where the model will be stored or downloaded. + compile_mode (str): The compilation mode to use for the model. + model_dtype (str): The data type to use for the model. + hf_token (str): The Hugging Face token to use for the model. + speculative_decoding (bool): Whether to use speculative decoding. + model (Llama/transformers model): The model to load. + tokenizer (list/transformers tokenizer): The tokenizer to load. + assistant_model (transformers model): The assistant model for speculative decoding. + Returns: + tuple: A tuple containing: + - model (Llama/transformers model): The loaded Llama.cpp/transformers model instance. + - tokenizer (list/transformers tokenizer): An empty list (tokenizer is not used with Llama.cpp directly in this setup), or a transformers tokenizer. + - assistant_model (transformers model): The assistant model for speculative decoding (if speculative_decoding is True). + """ + + if model: + return model, tokenizer, assistant_model + + print("Loading model:", local_model_type) + + # Verify the device and cuda settings + # Check if CUDA is enabled + + import torch + + torch.cuda.empty_cache() + print("Is CUDA enabled? ", torch.cuda.is_available()) + print("Is a CUDA device available on this computer?", torch.backends.cudnn.enabled) + if torch.cuda.is_available(): + torch_device = "cuda" + gpu_layers = int(LLM_MAX_GPU_LAYERS) + print("CUDA version:", torch.version.cuda) + # try: + # os.system("nvidia-smi") + # except Exception as e: + # print("Could not print nvidia-smi settings due to:", e) + else: + torch_device = "cpu" + gpu_layers = 0 + + print("Running on device:", torch_device) + print("GPU layers assigned to cuda:", gpu_layers) + + if not LLM_THREADS: + threads = torch.get_num_threads() + else: + threads = LLM_THREADS + print("CPU threads:", threads) + + # GPU mode + if torch_device == "cuda": + torch.cuda.empty_cache() + gpu_config.update_gpu(gpu_layers) + gpu_config.update_context(max_context_length) + + if USE_LLAMA_CPP == "True": + from llama_cpp import Llama + from llama_cpp.llama_speculative import LlamaPromptLookupDecoding + + model_path = get_model_path( + repo_id=repo_id, model_filename=model_filename, model_dir=model_dir + ) + + try: + print("GPU load variables:", vars(gpu_config)) + if speculative_decoding: + model = Llama( + model_path=model_path, + type_k=K_QUANT_LEVEL, + type_v=V_QUANT_LEVEL, + flash_attn=True, + draft_model=LlamaPromptLookupDecoding( + num_pred_tokens=NUM_PRED_TOKENS + ), + **vars(gpu_config), + ) + else: + model = Llama( + model_path=model_path, + type_k=K_QUANT_LEVEL, + type_v=V_QUANT_LEVEL, + flash_attn=True, + **vars(gpu_config), + ) + + except Exception as e: + print("GPU load failed due to:", e, "Loading model in CPU mode") + # If fails, go to CPU mode + model = Llama(model_path=model_path, **vars(cpu_config)) + + else: + from transformers import ( + AutoModelForCausalLM, + BitsAndBytesConfig, + ) + from unsloth import FastLanguageModel + + print("Loading model from transformers") + # Use the official model ID for Gemma 3 4B + model_id = ( + repo_id.split("https://huggingface.co/")[-1] + if "https://huggingface.co/" in repo_id + else repo_id + ) + # 1. Set Data Type (dtype) + # For H200/Hopper: 'bfloat16' + # For RTX 3060/Ampere: 'float16' + dtype_str = model_dtype # os.environ.get("MODEL_DTYPE", "bfloat16").lower() + if dtype_str == "bfloat16": + torch_dtype = torch.bfloat16 + elif dtype_str == "float16": + torch_dtype = torch.float16 + else: + torch_dtype = torch.float32 # A safe fallback + + # 2. Set Compilation Mode + # 'max-autotune' is great for both but can be slow initially. + # 'reduce-overhead' is a faster alternative for compiling. + + print("--- System Configuration ---") + print(f"Using model id: {model_id}") + print(f"Using dtype: {torch_dtype}") + print(f"Using compile mode: {compile_mode}") + print(f"Using bitsandbytes: {USE_BITSANDBYTES}") + print("--------------------------\n") + + # --- Load Tokenizer and Model --- + + try: + + # Load Tokenizer and Model + # tokenizer = AutoTokenizer.from_pretrained(model_id) + + if USE_BITSANDBYTES == "True": + + if INT8_WITH_OFFLOAD_TO_CPU == "True": + # This will be very slow. Requires at least 4GB of VRAM and 32GB of RAM + print( + "Using bitsandbytes for quantisation to 8 bits, with offloading to CPU" + ) + max_memory = {0: "4GB", "cpu": "32GB"} + BitsAndBytesConfig( + load_in_8bit=True, + max_memory=max_memory, + llm_int8_enable_fp32_cpu_offload=True, # Note: if bitsandbytes has to offload to CPU, inference will be slow + ) + else: + # For Gemma 4B, requires at least 6GB of VRAM + print("Using bitsandbytes for quantisation to 4 bits") + BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", # Use the modern NF4 quantisation for better performance + bnb_4bit_compute_dtype=torch_dtype, + bnb_4bit_use_double_quant=True, # Optional: uses a second quantisation step to save even more memory + ) + + # print("Loading model with bitsandbytes quantisation config:", quantisation_config) + + model, tokenizer = FastLanguageModel.from_pretrained( + model_id, + max_seq_length=max_context_length, + dtype=torch_dtype, + device_map="auto", + load_in_4bit=True, + # quantization_config=quantisation_config, # Not actually used in Unsloth + token=hf_token, + ) + + FastLanguageModel.for_inference(model) + else: + print("Loading model without bitsandbytes quantisation") + model, tokenizer = FastLanguageModel.from_pretrained( + model_id, + max_seq_length=max_context_length, + dtype=torch_dtype, + device_map="auto", + token=hf_token, + ) + + FastLanguageModel.for_inference(model) + + if not tokenizer.pad_token: + tokenizer.pad_token = tokenizer.eos_token + + except Exception as e: + print("Error loading model with bitsandbytes quantisation config:", e) + raise Warning( + "Error loading model with bitsandbytes quantisation config:", e + ) + + # Compile the Model with the selected mode 🚀 + if COMPILE_TRANSFORMERS == "True": + try: + model = torch.compile(model, mode=compile_mode, fullgraph=True) + except Exception as e: + print(f"Could not compile model: {e}. Running in eager mode.") + + print( + "Loading with", + gpu_config.n_gpu_layers, + "model layers sent to GPU and a maximum context length of", + gpu_config.n_ctx, + ) + + # CPU mode + else: + if USE_LLAMA_CPP == "False": + raise Warning( + "Using transformers model in CPU mode is not supported. Please change your config variable USE_LLAMA_CPP to True if you want to do CPU inference." + ) + + model_path = get_model_path( + repo_id=repo_id, model_filename=model_filename, model_dir=model_dir + ) + + # gpu_config.update_gpu(gpu_layers) + cpu_config.update_gpu(gpu_layers) + + # Update context length according to slider + # gpu_config.update_context(max_context_length) + cpu_config.update_context(max_context_length) + + # if speculative_decoding: + # model = Llama( + # model_path=model_path, + # draft_model=LlamaPromptLookupDecoding(num_pred_tokens=NUM_PRED_TOKENS), + # **vars(cpu_config), + # ) + # else: + # model = Llama(model_path=model_path, **vars(cpu_config)) + + # print( + # "Loading with", + # cpu_config.n_gpu_layers, + # "model layers sent to GPU and a maximum context length of", + # cpu_config.n_ctx, + # ) + + model = None + print("Llama_cpp_python model not currently supported") + + print("Finished loading model:", local_model_type) + print("GPU layers assigned to cuda:", gpu_layers) + + # Load assistant model for speculative decoding if enabled + if speculative_decoding and USE_LLAMA_CPP == "False" and torch_device == "cuda": + print("Loading assistant model for speculative decoding:", ASSISTANT_MODEL) + try: + from transformers import AutoModelForCausalLM + + # Load the assistant model with the same configuration as the main model + assistant_model = AutoModelForCausalLM.from_pretrained( + ASSISTANT_MODEL, dtype=torch_dtype, device_map="auto", token=hf_token + ) + + # assistant_model.config._name_or_path = model.config._name_or_path + + # Compile the assistant model if compilation is enabled + if COMPILE_TRANSFORMERS == "True": + try: + assistant_model = torch.compile( + assistant_model, mode=compile_mode, fullgraph=True + ) + except Exception as e: + print( + f"Could not compile assistant model: {e}. Running in eager mode." + ) + + print("Successfully loaded assistant model for speculative decoding") + + except Exception as e: + print(f"Error loading assistant model: {e}") + assistant_model = None + else: + assistant_model = None + + return model, tokenizer, assistant_model + + +def get_model(): + """Get the globally loaded model. Load it if not already loaded.""" + global _model, _tokenizer, _assistant_model + if _model is None: + _model, _tokenizer, _assistant_model = load_model( + local_model_type=CHOSEN_LOCAL_MODEL_TYPE, + gpu_layers=gpu_layers, + max_context_length=context_length, + gpu_config=gpu_config, + cpu_config=cpu_config, + torch_device=torch_device, + repo_id=LOCAL_REPO_ID, + model_filename=LOCAL_MODEL_FILE, + model_dir=LOCAL_MODEL_FOLDER, + compile_mode=COMPILE_MODE, + model_dtype=MODEL_DTYPE, + hf_token=HF_TOKEN, + model=_model, + tokenizer=_tokenizer, + assistant_model=_assistant_model, + ) + return _model + + +def get_tokenizer(): + """Get the globally loaded tokenizer. Load it if not already loaded.""" + global _model, _tokenizer, _assistant_model + if _tokenizer is None: + _model, _tokenizer, _assistant_model = load_model( + local_model_type=CHOSEN_LOCAL_MODEL_TYPE, + gpu_layers=gpu_layers, + max_context_length=context_length, + gpu_config=gpu_config, + cpu_config=cpu_config, + torch_device=torch_device, + repo_id=LOCAL_REPO_ID, + model_filename=LOCAL_MODEL_FILE, + model_dir=LOCAL_MODEL_FOLDER, + compile_mode=COMPILE_MODE, + model_dtype=MODEL_DTYPE, + hf_token=HF_TOKEN, + model=_model, + tokenizer=_tokenizer, + assistant_model=_assistant_model, + ) + return _tokenizer + + +def get_assistant_model(): + """Get the globally loaded assistant model. Load it if not already loaded.""" + global _model, _tokenizer, _assistant_model + if _assistant_model is None: + _model, _tokenizer, _assistant_model = load_model( + local_model_type=CHOSEN_LOCAL_MODEL_TYPE, + gpu_layers=gpu_layers, + max_context_length=context_length, + gpu_config=gpu_config, + cpu_config=cpu_config, + torch_device=torch_device, + repo_id=LOCAL_REPO_ID, + model_filename=LOCAL_MODEL_FILE, + model_dir=LOCAL_MODEL_FOLDER, + compile_mode=COMPILE_MODE, + model_dtype=MODEL_DTYPE, + hf_token=HF_TOKEN, + model=_model, + tokenizer=_tokenizer, + assistant_model=_assistant_model, + ) + return _assistant_model + + +def set_model(model, tokenizer, assistant_model=None): + """Set the global model, tokenizer, and assistant model.""" + global _model, _tokenizer, _assistant_model + _model = model + _tokenizer = tokenizer + _assistant_model = assistant_model + + +# Initialize model at startup if configured +if LOAD_LOCAL_MODEL_AT_START == "True" and RUN_LOCAL_MODEL == "1": + get_model() # This will trigger loading + + +def call_llama_cpp_model(formatted_string: str, gen_config: str, model=None): + """ + Calls your generation model with parameters from the LlamaCPPGenerationConfig object. + + Args: + formatted_string (str): The formatted input text for the model. + gen_config (LlamaCPPGenerationConfig): An object containing generation parameters. + model: Optional model instance. If None, will use the globally loaded model. + """ + if model is None: + model = get_model() + + if model is None: + raise ValueError( + "No model available. Either pass a model parameter or ensure LOAD_LOCAL_MODEL_AT_START is True." + ) + + # Extracting parameters from the gen_config object + temperature = gen_config.temperature + top_k = gen_config.top_k + top_p = gen_config.top_p + repeat_penalty = gen_config.repeat_penalty + seed = gen_config.seed + max_tokens = gen_config.max_tokens + stream = gen_config.stream + + # Now you can call your model directly, passing the parameters: + output = model( + formatted_string, + temperature=temperature, + top_k=top_k, + top_p=top_p, + repeat_penalty=repeat_penalty, + seed=seed, + max_tokens=max_tokens, + stream=stream, # , + # stop=["<|eot_id|>", "\n\n"] + ) + + return output + + +def call_llama_cpp_chatmodel( + formatted_string: str, + system_prompt: str, + gen_config: LlamaCPPGenerationConfig, + model=None, +): + """ + Calls your Llama.cpp chat model with a formatted user message and system prompt, + using generation parameters from the LlamaCPPGenerationConfig object. + + Args: + formatted_string (str): The formatted input text for the user's message. + system_prompt (str): The system-level instructions for the model. + gen_config (LlamaCPPGenerationConfig): An object containing generation parameters. + model: Optional model instance. If None, will use the globally loaded model. + """ + if model is None: + model = get_model() + + if model is None: + raise ValueError( + "No model available. Either pass a model parameter or ensure LOAD_LOCAL_MODEL_AT_START is True." + ) + + # Extracting parameters from the gen_config object + temperature = gen_config.temperature + top_k = gen_config.top_k + top_p = gen_config.top_p + repeat_penalty = gen_config.repeat_penalty + seed = gen_config.seed + max_tokens = gen_config.max_tokens + stream = gen_config.stream + reset = gen_config.reset + + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": formatted_string}, + ] + + input_tokens = len( + model.tokenize( + (system_prompt + "\n" + formatted_string).encode("utf-8"), special=True + ) + ) + + if stream: + final_tokens = list() + output_tokens = 0 + for chunk in model.create_chat_completion( + messages=messages, + temperature=temperature, + top_k=top_k, + top_p=top_p, + repeat_penalty=repeat_penalty, + seed=seed, + max_tokens=max_tokens, + stream=True, + stop=stop_strings, + ): + delta = chunk["choices"][0].get("delta", {}) + token = delta.get("content") or chunk["choices"][0].get("text") or "" + if token: + print(token, end="", flush=True) + final_tokens.append(token) + output_tokens += 1 + print() # newline after stream finishes + + text = "".join(final_tokens) + + if reset: + model.reset() + + return { + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": text}, + } + ], + # Provide a usage object so downstream code can read it + "usage": { + "prompt_tokens": input_tokens, # unknown during streaming + "completion_tokens": output_tokens, # unknown during streaming + "total_tokens": input_tokens + + output_tokens, # unknown during streaming + }, + } + + else: + response = model.create_chat_completion( + messages=messages, + temperature=temperature, + top_k=top_k, + top_p=top_p, + repeat_penalty=repeat_penalty, + seed=seed, + max_tokens=max_tokens, + stream=False, + stop=stop_strings, + ) + + if reset: + model.reset() + + return response + + +def call_inference_server_api( + formatted_string: str, + system_prompt: str, + gen_config: LlamaCPPGenerationConfig, + api_url: str = "http://localhost:8080", + model_name: str = None, + use_llama_swap: bool = USE_LLAMA_SWAP, + disable_thinking: bool = INFERENCE_SERVER_DISABLE_THINKING, + read_timeout_seconds: int = None, +): + """ + Calls a inference-server API endpoint with a formatted user message and system prompt, + using generation parameters from the LlamaCPPGenerationConfig object. + + This function provides the same interface as call_llama_cpp_chatmodel but calls + a remote inference-server instance instead of a local model. + + Args: + formatted_string (str): The formatted input text for the user's message. + system_prompt (str): The system-level instructions for the model. + gen_config (LlamaCPPGenerationConfig): An object containing generation parameters. + api_url (str): The base URL of the inference-server API (default: "http://localhost:8080"). + model_name (str): Optional model name to use. If None, uses the default model. + use_llama_swap (bool): Whether to use llama-swap for the model. + disable_thinking (bool): When True, adds chat_template_kwargs={"enable_thinking": False} + for vLLM Qwen3/Qwen3.5-style chat templates. Defaults to INFERENCE_SERVER_DISABLE_THINKING. + Returns: + dict: Response in the same format as call_llama_cpp_chatmodel + + Example: + # Create generation config + gen_config = LlamaCPPGenerationConfig(temperature=0.7, max_tokens=100) + + # Call the API + response = call_inference_server_api( + formatted_string="Hello, how are you?", + system_prompt="You are a helpful assistant.", + gen_config=gen_config, + api_url="http://localhost:8080" + ) + + # Extract the response text + response_text = response['choices'][0]['message']['content'] + + Integration Example: + # To use inference-server instead of local model: + # 1. Set model_source to "inference-server" + # 2. Provide api_url parameter + # 3. Call your existing functions as normal + + responses, conversation_history, whole_conversation, whole_conversation_metadata, response_text = call_llm_with_markdown_table_checks( + batch_prompts=["Your prompt here"], + system_prompt="Your system prompt", + conversation_history=[], + whole_conversation=[], + whole_conversation_metadata=[], + client=None, # Not used for inference-server + client_config=None, # Not used for inference-server + model_choice="your-model-name", # Model name on the server + temperature=0.7, + reported_batch_no=1, + local_model=None, # Not used for inference-server + tokenizer=None, # Not used for inference-server + bedrock_runtime=None, # Not used for inference-server + model_source="inference-server", + MAX_OUTPUT_VALIDATION_ATTEMPTS=3, + api_url="http://localhost:8080" + ) + """ + # Extract parameters from the gen_config object + temperature = gen_config.temperature + top_k = gen_config.top_k + top_p = gen_config.top_p + repeat_penalty = gen_config.repeat_penalty + seed = gen_config.seed + max_tokens = gen_config.max_tokens + stream = gen_config.stream + + # Prepare the request payload + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": formatted_string}, + ] + + payload = { + "messages": messages, + "temperature": temperature, + "top_k": top_k, + "top_p": top_p, + "repeat_penalty": repeat_penalty, + "seed": seed, + "max_tokens": max_tokens, + "stream": stream, + "stop": LLM_STOP_STRINGS if LLM_STOP_STRINGS else [], + } + # Add model name if specified and use llama-swap + if model_name and use_llama_swap: + payload["model"] = model_name + + if disable_thinking: + payload["chat_template_kwargs"] = {"enable_thinking": False} + + # Determine the endpoint based on streaming preference + if stream: + endpoint = f"{api_url}/v1/chat/completions" + else: + endpoint = f"{api_url}/v1/chat/completions" + + try: + timeout_seconds = ( + int(read_timeout_seconds) + if read_timeout_seconds is not None + else int(inference_server_read_timeout_seconds) + ) + if stream: + # Handle streaming response + response = requests.post( + endpoint, + json=payload, + headers={"Content-Type": "application/json"}, + stream=True, + timeout=timeout_seconds, + ) + response.raise_for_status() + + final_tokens = [] + output_tokens = 0 + + for line in response.iter_lines(): + if line: + line = line.decode("utf-8") + if line.startswith("data: "): + data = line[6:] # Remove 'data: ' prefix + if data.strip() == "[DONE]": + break + try: + chunk = json.loads(data) + if "choices" in chunk and len(chunk["choices"]) > 0: + delta = chunk["choices"][0].get("delta", {}) + token = delta.get("content", "") + if token: + print(token, end="", flush=True) + final_tokens.append(token) + output_tokens += 1 + except json.JSONDecodeError: + continue + + print() # newline after stream finishes + + text = "".join(final_tokens) + + # Estimate input tokens (rough approximation) + input_tokens = len((system_prompt + "\n" + formatted_string).split()) + + return { + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": text}, + } + ], + "usage": { + "prompt_tokens": input_tokens, + "completion_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + } + else: + # Handle non-streaming response + response = requests.post( + endpoint, + json=payload, + headers={"Content-Type": "application/json"}, + timeout=timeout_seconds, + ) + response.raise_for_status() + + result = response.json() + + # Ensure the response has the expected format + if "choices" not in result: + raise ValueError("Invalid response format from inference-server") + + return result + + except requests.exceptions.HTTPError as e: + # Preserve server error details (llama.cpp includes useful context-limit info in the body) + detail = "" + try: + if e.response is not None: + detail = e.response.text + except Exception: + detail = "" + raise ConnectionError( + f"Failed to connect to inference-server at {api_url}: {str(e)}" + + (f" | server_response={detail}" if detail else "") + ) + except requests.exceptions.RequestException as e: + raise ConnectionError( + f"Failed to connect to inference-server at {api_url}: {str(e)}" + ) + except json.JSONDecodeError as e: + raise ValueError(f"Invalid JSON response from inference-server: {str(e)}") + except Exception as e: + raise RuntimeError(f"Error calling inference-server API: {str(e)}") + + +### +# LLM FUNCTIONS +### + + +def construct_gemini_generative_model( + in_api_key: str, + temperature: float, + model_choice: str, + system_prompt: str, + max_tokens: int, + random_seed=seed, +) -> Tuple[object, dict]: + """ + Constructs a GenerativeModel for Gemini API calls. + ... + """ + # Construct a GenerativeModel + try: + if in_api_key: + # print("Getting API key from textbox") + api_key = in_api_key + client = ai.Client(api_key=api_key) + elif "GOOGLE_API_KEY" in os.environ: + # print("Searching for API key in environmental variables") + api_key = os.environ["GOOGLE_API_KEY"] + client = ai.Client(api_key=api_key) + else: + print("No Gemini API key found") + raise Warning("No Gemini API key found.") + except Exception as e: + print("Error constructing Gemini generative model:", e) + raise Warning("Error constructing Gemini generative model:", e) + + config = types.GenerateContentConfig( + temperature=temperature, max_output_tokens=max_tokens, seed=random_seed + ) + + return client, config + + +def construct_azure_client(in_api_key: str, endpoint: str) -> Tuple[object, dict]: + """ + Constructs an OpenAI client for Azure/OpenAI AI Inference. + """ + try: + key = None + if in_api_key: + key = in_api_key + elif os.environ.get("AZURE_OPENAI_API_KEY"): + key = os.environ["AZURE_OPENAI_API_KEY"] + if not key: + raise Warning("No Azure/OpenAI API key found.") + + if not endpoint: + endpoint = os.environ.get("AZURE_OPENAI_INFERENCE_ENDPOINT", "") + if not endpoint: + # Assume using OpenAI API + client = OpenAI( + api_key=key, + ) + else: + # Use the provided endpoint + client = OpenAI( + api_key=key, + base_url=f"{endpoint}", + ) + + return client, dict() + except Exception as e: + print("Error constructing Azure/OpenAI client:", e) + raise + + +def call_aws_bedrock( + prompt: str, + system_prompt: str, + temperature: float, + max_tokens: int, + model_choice: str, + bedrock_runtime: boto3.Session.client, + assistant_prefill: str = "", +) -> ResponseObject: + """ + This function sends a request to AWS Claude with the following parameters: + - prompt: The user's input prompt to be processed by the model. + - system_prompt: A system-defined prompt that provides context or instructions for the model. + - temperature: A value that controls the randomness of the model's output, with higher values resulting in more diverse responses. + - max_tokens: The maximum number of tokens (words or characters) in the model's response. + - model_choice: The specific model to use for processing the request. + - bedrock_runtime: The client object for boto3 Bedrock runtime + - assistant_prefill: A string indicating the text that the response should start with. + + The function constructs the request configuration, invokes the model, extracts the response text, and returns a ResponseObject containing the text and metadata. + """ + + inference_config = { + "maxTokens": max_tokens, + "topP": 0.999, + "temperature": temperature, + } + + # Using an assistant prefill only works for Anthropic models. + if assistant_prefill and "anthropic" in model_choice: + assistant_prefill_added = True + messages = [ + { + "role": "user", + "content": [ + {"text": prompt}, + ], + }, + { + "role": "assistant", + # Pre-filling with '|' + "content": [{"text": assistant_prefill}], + }, + ] + else: + assistant_prefill_added = False + messages = [ + { + "role": "user", + "content": [ + {"text": prompt}, + ], + } + ] + + system_prompt_list = [{"text": system_prompt}] + + # The converse API call. + api_response = bedrock_runtime.converse( + modelId=model_choice, + messages=messages, + system=system_prompt_list, + inferenceConfig=inference_config, + ) + + output_message = api_response["output"]["message"] + + if "reasoningContent" in output_message["content"][0]: + # Extract the reasoning text + output_message["content"][0]["reasoningContent"]["reasoningText"]["text"] + + # Extract the output text + if assistant_prefill_added: + text = assistant_prefill + output_message["content"][1]["text"] + else: + text = output_message["content"][1]["text"] + else: + if assistant_prefill_added: + text = assistant_prefill + output_message["content"][0]["text"] + else: + text = output_message["content"][0]["text"] + + # The usage statistics are neatly provided in the 'usage' key. + usage = api_response["usage"] + + # The full API response metadata is in 'ResponseMetadata' if you still need it. + api_response["ResponseMetadata"] + + # Create ResponseObject with the cleanly extracted data. + response = ResponseObject(text=text, usage_metadata=usage) + + return response + + +def call_transformers_model( + prompt: str, + system_prompt: str, + gen_config: LlamaCPPGenerationConfig, + model=None, + tokenizer=None, + assistant_model=None, + speculative_decoding=speculative_decoding, +): + """ + This function sends a request to a transformers model (through Unsloth) with the given prompt, system prompt, and generation configuration. + """ + from transformers import TextStreamer + + if model is None: + model = get_model() + if tokenizer is None: + tokenizer = get_tokenizer() + if assistant_model is None and speculative_decoding: + assistant_model = get_assistant_model() + + if model is None or tokenizer is None: + raise ValueError( + "No model or tokenizer available. Either pass them as parameters or ensure LOAD_LOCAL_MODEL_AT_START is True." + ) + + # 1. Define the conversation as a list of dictionaries + # Note: The multimodal format [{"type": "text", "text": text}] is only needed for actual multimodal models + # with images/videos. For text-only content, even multimodal models expect plain strings. + + # Check if system_prompt is meaningful (not empty/None) + has_system_prompt = system_prompt and str(system_prompt).strip() + + # Always use string format for text-only content, regardless of MULTIMODAL_PROMPT_FORMAT setting + # MULTIMODAL_PROMPT_FORMAT should only be used when you actually have multimodal inputs (images, etc.) + if MULTIMODAL_PROMPT_FORMAT == "True": + conversation = [] + if has_system_prompt: + conversation.append( + { + "role": "system", + "content": [{"type": "text", "text": str(system_prompt)}], + } + ) + conversation.append( + {"role": "user", "content": [{"type": "text", "text": str(prompt)}]} + ) + else: + conversation = [] + if has_system_prompt: + conversation.append({"role": "system", "content": str(system_prompt)}) + conversation.append({"role": "user", "content": str(prompt)}) + + # 2. Apply the chat template + # Get the device from the model (handles both single device and device_map="auto" cases) + try: + model_device = next(model.parameters()).device + except (StopIteration, AttributeError): + # Fallback to cuda if we can't determine device + model_device = "cuda" + + try: + # Try applying chat template with system prompt (if present) + input_ids = tokenizer.apply_chat_template( + conversation, + add_generation_prompt=True, + tokenize=True, + return_tensors="pt", + ).to(model_device) + except (TypeError, KeyError, IndexError, ValueError) as e: + # If chat template fails, try without system prompt (some models don't support it) + if has_system_prompt: + print( + f"Chat template failed with system prompt ({e}), trying without system prompt..." + ) + # Try again with only user prompt + user_only_conversation = [{"role": "user", "content": str(prompt)}] + try: + input_ids = tokenizer.apply_chat_template( + user_only_conversation, + add_generation_prompt=True, + tokenize=True, + return_tensors="pt", + ).to(model_device) + print("Successfully applied chat template without system prompt") + except Exception as e2: + print( + f"Chat template failed without system prompt ({e2}), using manual tokenization" + ) + # Combine system and user prompts manually as fallback + full_prompt = ( + f"{system_prompt}\n\n{prompt}" if has_system_prompt else prompt + ) + # Tokenize manually with special tokens + encoded = tokenizer( + full_prompt, return_tensors="pt", add_special_tokens=True + ) + if encoded is None: + raise ValueError( + "Tokenizer returned None - tokenizer may not be properly initialized" + ) + if not hasattr(encoded, "input_ids") or encoded.input_ids is None: + raise ValueError("Tokenizer output does not contain input_ids") + input_ids = encoded.input_ids.to(model_device) + else: + # No system prompt, but chat template still failed - use manual tokenization + print(f"Chat template failed ({e}), using manual tokenization") + full_prompt = str(prompt) + encoded = tokenizer( + full_prompt, return_tensors="pt", add_special_tokens=True + ) + if encoded is None: + raise ValueError( + "Tokenizer returned None - tokenizer may not be properly initialized" + ) + if not hasattr(encoded, "input_ids") or encoded.input_ids is None: + raise ValueError("Tokenizer output does not contain input_ids") + input_ids = encoded.input_ids.to(model_device) + except Exception as e: + print("Error applying chat template:", e) + import traceback + + traceback.print_exc() + raise + + # Map LlamaCPP parameters to transformers parameters + generation_kwargs = { + "max_new_tokens": gen_config.max_tokens, + "temperature": gen_config.temperature, + "top_p": gen_config.top_p, + "top_k": gen_config.top_k, + "do_sample": True, + #'pad_token_id': tokenizer.eos_token_id + } + + if gen_config.stream: + streamer = TextStreamer(tokenizer, skip_prompt=True) + else: + streamer = None + + # Remove parameters that don't exist in transformers + if hasattr(gen_config, "repeat_penalty"): + generation_kwargs["repetition_penalty"] = gen_config.repeat_penalty + + # --- Timed Inference Test --- + print("\nStarting model inference...") + start_time = time.time() + + # Use speculative decoding if assistant model is available + try: + if speculative_decoding and assistant_model is not None: + # print("Using speculative decoding with assistant model") + outputs = model.generate( + input_ids, + assistant_model=assistant_model, + **generation_kwargs, + streamer=streamer, + ) + else: + # print("Generating without speculative decoding") + outputs = model.generate(input_ids, **generation_kwargs, streamer=streamer) + except Exception as e: + error_msg = str(e) + # Check if this is a CUDA compilation error + if ( + "sm_120" in error_msg + or "LLVM ERROR" in error_msg + or "Cannot select" in error_msg + ): + print("\n" + "=" * 80) + print("CUDA COMPILATION ERROR DETECTED") + print("=" * 80) + print( + "\nThe error is caused by torch.compile() trying to compile CUDA kernels" + ) + print( + "with incompatible settings. This is a known issue with certain CUDA/PyTorch" + ) + print("combinations.\n") + print( + "SOLUTION: Disable model compilation by setting COMPILE_TRANSFORMERS=False" + ) + print("in your config file (config/app_config.env).") + print( + "\nThe model will still work without compilation, just slightly slower." + ) + print("=" * 80 + "\n") + raise RuntimeError( + "CUDA compilation error detected. Please set COMPILE_TRANSFORMERS=False " + "in your config file to disable model compilation and avoid this error." + ) from e + else: + # Re-raise other errors as-is + raise + + end_time = time.time() + + # --- Decode and Display Results --- + new_tokens = outputs[0][input_ids.shape[-1] :] + assistant_reply = tokenizer.decode(new_tokens, skip_special_tokens=True) + + num_input_tokens = input_ids.shape[ + -1 + ] # This gets the sequence length (number of tokens) + num_generated_tokens = len(new_tokens) + duration = end_time - start_time + tokens_per_second = num_generated_tokens / duration + + print("\n--- Performance ---") + print(f"Time taken: {duration:.2f} seconds") + print(f"Generated tokens: {num_generated_tokens}") + print(f"Tokens per second: {tokens_per_second:.2f}") + + return assistant_reply, num_input_tokens, num_generated_tokens + + +# Function to send a request and update history +def send_request( + prompt: str, + conversation_history: List[dict], + client: ai.Client | OpenAI, + config: types.GenerateContentConfig, + model_choice: str, + system_prompt: str, + temperature: float, + bedrock_runtime: boto3.Session.client, + model_source: str, + local_model=list(), + tokenizer=None, + assistant_model=None, + assistant_prefill="", + progress=Progress(track_tqdm=True), + api_url: str = None, +) -> Tuple[str, List[dict]]: + """Sends a request to a language model and manages the conversation history. + + This function constructs the full prompt by appending the new user prompt to the conversation history, + generates a response from the model, and updates the conversation history with the new prompt and response. + It handles different model sources (Gemini, AWS, Local, inference-server) and includes retry logic for API calls. + + Args: + prompt (str): The user's input prompt to be sent to the model. + conversation_history (List[dict]): A list of dictionaries representing the ongoing conversation. + Each dictionary should have 'role' and 'parts' keys. + client (ai.Client): The API client object for the chosen model (e.g., Gemini `ai.Client`, or Azure/OpenAI `OpenAI`). + config (types.GenerateContentConfig): Configuration settings for content generation (e.g., Gemini `types.GenerateContentConfig`). + model_choice (str): The specific model identifier to use (e.g., "gemini-pro", "claude-v2"). + system_prompt (str): An optional system-level instruction or context for the model. + temperature (float): Controls the randomness of the model's output, with higher values leading to more diverse responses. + bedrock_runtime (boto3.Session.client): The boto3 Bedrock runtime client object for AWS models. + model_source (str): Indicates the source/provider of the model (e.g., "Gemini", "AWS", "Local", "inference-server"). + local_model (list, optional): A list containing the local model and its tokenizer (if `model_source` is "Local"). Defaults to []. + tokenizer (object, optional): The tokenizer object for local models. Defaults to None. + assistant_model (object, optional): An optional assistant model used for speculative decoding with local models. Defaults to None. + assistant_prefill (str, optional): A string to pre-fill the assistant's response, useful for certain models like Claude. Defaults to "". + progress (Progress, optional): A progress object for tracking the operation, typically from `tqdm`. Defaults to Progress(track_tqdm=True). + api_url (str, optional): The API URL for inference-server calls. Required when model_source is 'inference-server'. + + Returns: + Tuple[str, List[dict]]: A tuple containing the model's response text and the updated conversation history. + """ + # Constructing the full prompt from the conversation history + full_prompt = "Conversation history:\n" + num_transformer_input_tokens = 0 + num_transformer_generated_tokens = 0 + response_text = "" + + for entry in conversation_history: + role = entry[ + "role" + ].capitalize() # Assuming the history is stored with 'role' and 'parts' + message = " ".join(entry["parts"]) # Combining all parts of the message + full_prompt += f"{role}: {message}\n" + + # Adding the new user prompt + full_prompt += f"\nUser: {prompt}" + + # Clear any existing progress bars + tqdm._instances.clear() + + progress_bar = range(0, number_of_api_retry_attempts) + + # Generate the model's response + if "Gemini" in model_source: + + for i in progress_bar: + try: + print("Calling Gemini model, attempt", i + 1) + + response = client.models.generate_content( + model=model_choice, contents=full_prompt, config=config + ) + + # print("Successful call to Gemini model.") + break + except Exception as e: + # If fails, try again after X seconds in case there is a throttle limit + print( + "Call to Gemini model failed:", + e, + " Waiting for ", + str(timeout_wait), + "seconds and trying again.", + ) + + time.sleep(timeout_wait) + + if i == number_of_api_retry_attempts: + return ( + ResponseObject(text="", usage_metadata={"RequestId": "FAILED"}), + conversation_history, + response_text, + num_transformer_input_tokens, + num_transformer_generated_tokens, + ) + + elif "AWS" in model_source: + for i in progress_bar: + try: + print("Calling AWS Bedrock model, attempt", i + 1) + response = call_aws_bedrock( + prompt, + system_prompt, + temperature, + max_tokens, + model_choice, + bedrock_runtime=bedrock_runtime, + assistant_prefill=assistant_prefill, + ) + + # print("Successful call to Claude model.") + break + except Exception as e: + # If fails, try again after X seconds in case there is a throttle limit + print( + "Call to Bedrock model failed:", + e, + " Waiting for ", + str(timeout_wait), + "seconds and trying again.", + ) + time.sleep(timeout_wait) + + if i == number_of_api_retry_attempts: + return ( + ResponseObject(text="", usage_metadata={"RequestId": "FAILED"}), + conversation_history, + response_text, + num_transformer_input_tokens, + num_transformer_generated_tokens, + ) + elif "Azure/OpenAI" in model_source: + for i in progress_bar: + try: + print("Calling Azure/OpenAI inference model, attempt", i + 1) + + messages = [ + { + "role": "system", + "content": system_prompt, + }, + { + "role": "user", + "content": prompt, + }, + ] + + response_raw = client.chat.completions.create( + messages=messages, + model=model_choice, + temperature=temperature, + max_completion_tokens=max_tokens, + ) + + response_text = response_raw.choices[0].message.content + usage = getattr(response_raw, "usage", None) + input_tokens = 0 + output_tokens = 0 + if usage is not None: + input_tokens = getattr( + usage, "input_tokens", getattr(usage, "prompt_tokens", 0) + ) + output_tokens = getattr( + usage, "output_tokens", getattr(usage, "completion_tokens", 0) + ) + response = ResponseObject( + text=response_text, + usage_metadata={ + "inputTokens": input_tokens, + "outputTokens": output_tokens, + }, + ) + break + except Exception as e: + print( + "Call to Azure/OpenAI model failed:", + e, + " Waiting for ", + str(timeout_wait), + "seconds and trying again.", + ) + time.sleep(timeout_wait) + if i == number_of_api_retry_attempts: + return ( + ResponseObject(text="", usage_metadata={"RequestId": "FAILED"}), + conversation_history, + response_text, + num_transformer_input_tokens, + num_transformer_generated_tokens, + ) + elif "Local" in model_source: + # This is the local model + for i in progress_bar: + try: + print("Calling local model, attempt", i + 1) + + gen_config = LlamaCPPGenerationConfig() + gen_config.update_temp(temperature) + + if USE_LLAMA_CPP == "True": + response = call_llama_cpp_chatmodel( + prompt, system_prompt, gen_config, model=local_model + ) + + else: + ( + response, + num_transformer_input_tokens, + num_transformer_generated_tokens, + ) = call_transformers_model( + prompt, + system_prompt, + gen_config, + model=local_model, + tokenizer=tokenizer, + assistant_model=assistant_model, + ) + response_text = response + + break + except Exception as e: + # If fails, try again after X seconds in case there is a throttle limit + print( + "Call to local model failed:", + e, + " Waiting for ", + str(timeout_wait), + "seconds and trying again.", + ) + + time.sleep(timeout_wait) + + if i == number_of_api_retry_attempts: + return ( + ResponseObject(text="", usage_metadata={"RequestId": "FAILED"}), + conversation_history, + response_text, + num_transformer_input_tokens, + num_transformer_generated_tokens, + ) + elif "inference-server" in model_source: + # This is the inference-server API + for i in progress_bar: + try: + print("Calling inference-server API, attempt", i + 1) + + if api_url is None: + raise ValueError( + "api_url is required when model_source is 'inference-server'" + ) + + gen_config = LlamaCPPGenerationConfig() + gen_config.update_temp(temperature) + + response = call_inference_server_api( + prompt, + system_prompt, + gen_config, + api_url=api_url, + model_name=model_choice, + ) + + break + except Exception as e: + # If fails, try again after X seconds in case there is a throttle limit + print( + "Call to inference-server API failed:", + e, + " Waiting for ", + str(timeout_wait), + "seconds and trying again.", + ) + + time.sleep(timeout_wait) + + if i == number_of_api_retry_attempts: + return ( + ResponseObject(text="", usage_metadata={"RequestId": "FAILED"}), + conversation_history, + response_text, + num_transformer_input_tokens, + num_transformer_generated_tokens, + ) + else: + print("Model source not recognised") + return ( + ResponseObject(text="", usage_metadata={"RequestId": "FAILED"}), + conversation_history, + response_text, + num_transformer_input_tokens, + num_transformer_generated_tokens, + ) + + # Update the conversation history with the new prompt and response + conversation_history.append({"role": "user", "parts": [prompt]}) + + # Check if is a LLama.cpp model response or inference-server response + if isinstance(response, ResponseObject): + response_text = response.text + elif "choices" in response: # LLama.cpp model response or inference-server response + # Check for GPT-OSS thinking models (case-insensitive, handle both hyphen and underscore) + if "gpt-oss" in model_choice.lower() or "gpt_oss" in model_choice.lower(): + content = response["choices"][0]["message"]["content"] + # Split on the final channel marker to extract only the final output (not thinking tokens) + parts = content.split("<|start|>assistant<|channel|>final<|message|>") + if len(parts) > 1: + response_text = parts[1] + # Following format may be from llama.cpp inference-server response + elif len(parts) == 1: + parts = content.split("<|end|>") + if len(parts) > 1: + response_text = parts[1] + else: + print( + "Warning: Could not find final channel marker in GPT-OSS response. Using full content." + ) + response_text = content + else: + # Fallback: if marker not found, use the full content (may include thinking tokens) + print( + "Warning: Could not find final channel marker in GPT-OSS response. Using full content." + ) + response_text = content + else: + response_text = response["choices"][0]["message"]["content"] + elif model_source == "Gemini": + response_text = response.text + else: # Assume transformers model response + # Check for GPT-OSS thinking models (case-insensitive, handle both hyphen and underscore) + if "gpt-oss" in model_choice.lower() or "gpt_oss" in model_choice.lower(): + # Split on the final channel marker to extract only the final output (not thinking tokens) + parts = response.split("<|start|>assistant<|channel|>final<|message|>") + if len(parts) > 1: + response_text = parts[1] + else: + # Fallback: if marker not found, use the full content (may include thinking tokens) + print( + "Warning: Could not find final channel marker in GPT-OSS response. Using full content." + ) + response_text = response + else: + response_text = response + + # Strip <|end|> tags (used by GPT-OSS thinking models to mark end of thinking) + response_text = re.sub(r"<\|end\|>", "", response_text) + + # Replace multiple spaces with single space + response_text = re.sub(r" {2,}", " ", response_text) + response_text = response_text.strip() + + conversation_history.append({"role": "assistant", "parts": [response_text]}) + + return ( + response, + conversation_history, + response_text, + num_transformer_input_tokens, + num_transformer_generated_tokens, + ) + + +def process_requests( + prompts: List[str], + system_prompt: str, + conversation_history: List[dict], + whole_conversation: List[str], + whole_conversation_metadata: List[str], + client: ai.Client | OpenAI, + config: types.GenerateContentConfig, + model_choice: str, + temperature: float, + bedrock_runtime: boto3.Session.client, + model_source: str, + batch_no: int = 1, + local_model=list(), + tokenizer=None, + assistant_model=None, + master: bool = False, + assistant_prefill="", + api_url: str = None, +) -> Tuple[List[ResponseObject], List[dict], List[str], List[str]]: + """ + Processes a list of prompts by sending them to the model, appending the responses to the conversation history, and updating the whole conversation and metadata. + + Args: + prompts (List[str]): A list of prompts to be processed. + system_prompt (str): The system prompt. + conversation_history (List[dict]): The history of the conversation. + whole_conversation (List[str]): The complete conversation including prompts and responses. + whole_conversation_metadata (List[str]): Metadata about the whole conversation. + client (object): The client to use for processing the prompts, from either Gemini or OpenAI client. + config (dict): Configuration for the model. + model_choice (str): The choice of model to use. + temperature (float): The temperature parameter for the model. + model_source (str): Source of the model, whether local, AWS, Gemini, or inference-server + batch_no (int): Batch number of the large language model request. + local_model: Local gguf model (if loaded) + master (bool): Is this request for the master table. + assistant_prefill (str, optional): Is there a prefill for the assistant response. Currently only working for AWS model calls + bedrock_runtime: The client object for boto3 Bedrock runtime + api_url (str, optional): The API URL for inference-server calls. Required when model_source is 'inference-server'. + + Returns: + Tuple[List[ResponseObject], List[dict], List[str], List[str]]: A tuple containing the list of responses, the updated conversation history, the updated whole conversation, and the updated whole conversation metadata. + """ + responses = list() + + # Clear any existing progress bars + tqdm._instances.clear() + + for prompt in prompts: + + ( + response, + conversation_history, + response_text, + num_transformer_input_tokens, + num_transformer_generated_tokens, + ) = send_request( + prompt, + conversation_history, + client=client, + config=config, + model_choice=model_choice, + system_prompt=system_prompt, + temperature=temperature, + local_model=local_model, + tokenizer=tokenizer, + assistant_model=assistant_model, + assistant_prefill=assistant_prefill, + bedrock_runtime=bedrock_runtime, + model_source=model_source, + api_url=api_url, + ) + + responses.append(response) + whole_conversation.append(system_prompt) + whole_conversation.append(prompt) + whole_conversation.append(response_text) + + whole_conversation_metadata.append(f"Batch {batch_no}:") + + try: + if "AWS" in model_source: + output_tokens = response.usage_metadata.get("outputTokens", 0) + input_tokens = response.usage_metadata.get("inputTokens", 0) + + elif "Gemini" in model_source: + output_tokens = response.usage_metadata.candidates_token_count + input_tokens = response.usage_metadata.prompt_token_count + + elif "Azure/OpenAI" in model_source: + input_tokens = response.usage_metadata.get("inputTokens", 0) + output_tokens = response.usage_metadata.get("outputTokens", 0) + + elif "Local" in model_source: + if USE_LLAMA_CPP == "True": + output_tokens = response["usage"].get("completion_tokens", 0) + input_tokens = response["usage"].get("prompt_tokens", 0) + + if USE_LLAMA_CPP == "False": + input_tokens = num_transformer_input_tokens + output_tokens = num_transformer_generated_tokens + + elif "inference-server" in model_source: + # inference-server returns the same format as llama-cpp + output_tokens = response["usage"].get("completion_tokens", 0) + input_tokens = response["usage"].get("prompt_tokens", 0) + + else: + input_tokens = 0 + output_tokens = 0 + + whole_conversation_metadata.append( + "input_tokens: " + + str(input_tokens) + + " output_tokens: " + + str(output_tokens) + ) + + except KeyError as e: + print(f"Key error: {e} - Check the structure of response.usage_metadata") + + return ( + responses, + conversation_history, + whole_conversation, + whole_conversation_metadata, + response_text, + ) + + +def call_llm_with_markdown_table_checks( + batch_prompts: List[str], + system_prompt: str, + conversation_history: List[dict], + whole_conversation: List[str], + whole_conversation_metadata: List[str], + client: ai.Client | OpenAI, + client_config: types.GenerateContentConfig, + model_choice: str, + temperature: float, + reported_batch_no: int, + local_model: object, + tokenizer: object, + bedrock_runtime: boto3.Session.client, + model_source: str, + MAX_OUTPUT_VALIDATION_ATTEMPTS: int, + assistant_prefill: str = "", + master: bool = False, + CHOSEN_LOCAL_MODEL_TYPE: str = CHOSEN_LOCAL_MODEL_TYPE, + random_seed: int = seed, + api_url: str = None, +) -> Tuple[List[ResponseObject], List[dict], List[str], List[str], str]: + """ + Call the large language model with checks for a valid markdown table. + + Parameters: + - batch_prompts (List[str]): A list of prompts to be processed. + - system_prompt (str): The system prompt. + - conversation_history (List[dict]): The history of the conversation. + - whole_conversation (List[str]): The complete conversation including prompts and responses. + - whole_conversation_metadata (List[str]): Metadata about the whole conversation. + - client (ai.Client | OpenAI): The client object for running Gemini or Azure/OpenAI API calls. + - client_config (types.GenerateContentConfig): Configuration for the model. + - model_choice (str): The choice of model to use. + - temperature (float): The temperature parameter for the model. + - reported_batch_no (int): The reported batch number. + - local_model (object): The local model to use. + - tokenizer (object): The tokenizer to use. + - bedrock_runtime (boto3.Session.client): The client object for boto3 Bedrock runtime. + - model_source (str): The source of the model, whether in AWS, Gemini, local, or inference-server. + - MAX_OUTPUT_VALIDATION_ATTEMPTS (int): The maximum number of attempts to validate the output. + - assistant_prefill (str, optional): The text to prefill the LLM response. Currently only working with AWS Claude calls. + - master (bool, optional): Boolean to determine whether this call is for the master output table. + - CHOSEN_LOCAL_MODEL_TYPE (str, optional): String to determine model type loaded. + - random_seed (int, optional): The random seed used for LLM generation. + - api_url (str, optional): The API URL for inference-server calls. Required when model_source is 'inference-server'. + + Returns: + - Tuple[List[ResponseObject], List[dict], List[str], List[str], str]: A tuple containing the list of responses, the updated conversation history, the updated whole conversation, the updated whole conversation metadata, and the response text. + """ + + call_temperature = temperature # This is correct now with the fixed parameter name + + # Update Gemini config with the new temperature settings + client_config = types.GenerateContentConfig( + temperature=call_temperature, max_output_tokens=max_tokens, seed=random_seed + ) + + for attempt in range(MAX_OUTPUT_VALIDATION_ATTEMPTS): + # Process requests to large language model + ( + responses, + conversation_history, + whole_conversation, + whole_conversation_metadata, + response_text, + ) = process_requests( + batch_prompts, + system_prompt, + conversation_history, + whole_conversation, + whole_conversation_metadata, + client, + client_config, + model_choice, + call_temperature, + bedrock_runtime, + model_source, + reported_batch_no, + local_model, + tokenizer=tokenizer, + master=master, + assistant_prefill=assistant_prefill, + api_url=api_url, + ) + + stripped_response = response_text.strip() + + # Check if response meets our criteria (length and contains table) OR is "No change" + if ( + len(stripped_response) > 120 and "|" in stripped_response + ) or stripped_response.lower().startswith("no change"): + if stripped_response.lower().startswith("no change"): + print(f"Attempt {attempt + 1} produced 'No change' response.") + else: + print(f"Attempt {attempt + 1} produced response with markdown table.") + break # Success - exit loop + + # Increase temperature for next attempt + call_temperature = max(1.0, temperature + (0.1 * (attempt + 1))) + print( + f"Attempt {attempt + 1} resulted in invalid table: {stripped_response}. " + f"Trying again with temperature: {call_temperature}" + ) + + else: # This runs if no break occurred (all attempts failed) + print( + f"Failed to get valid response after {MAX_OUTPUT_VALIDATION_ATTEMPTS} attempts" + ) + + return ( + responses, + conversation_history, + whole_conversation, + whole_conversation_metadata, + stripped_response, + ) + + +def create_missing_references_df( + basic_response_df: pd.DataFrame, existing_reference_df: pd.DataFrame +) -> pd.DataFrame: + """ + Identifies references in basic_response_df that are not present in existing_reference_df. + Returns a DataFrame with the missing references and the character count of their responses. + + Args: + basic_response_df (pd.DataFrame): DataFrame containing 'Response ID' and 'Response' columns. + existing_reference_df (pd.DataFrame): DataFrame containing 'Response ID' column. + + Returns: + pd.DataFrame: A DataFrame with 'Missing Response ID' and 'Response Character Count' columns. + 'Response Character Count' will be 0 for empty strings and NaN for actual missing data. + """ + # Ensure columns are treated as strings for robust comparison + existing_references_unique = ( + existing_reference_df["Response ID"].astype(str).unique() + ) + + # Step 1: Identify all rows from basic_response_df that correspond to missing references + # We want the entire row to access the 'Response' column later + missing_data_rows = basic_response_df[ + ~basic_response_df["Response ID"].astype(str).isin(existing_references_unique) + ].copy() # .copy() to avoid SettingWithCopyWarning + + # Step 2: Create the new DataFrame + # Populate the 'Missing Response ID' column directly + missing_df = pd.DataFrame({"Missing Response ID": missing_data_rows["Response ID"]}) + + # Step 3: Calculate and add 'Response Character Count' + # .str.len() works on Series of strings, handling empty strings (0) and NaN (NaN) + missing_df["Response Character Count"] = missing_data_rows["Response"].str.len() + + # Optional: Add the actual response text for easier debugging/inspection if needed + # missing_df['Response Text'] = missing_data_rows['Response'] + + # Reset index to have a clean, sequential index for the new DataFrame + missing_df = missing_df.reset_index(drop=True) + + return missing_df + + +def calculate_tokens_from_metadata( + metadata_string: str, model_choice: str, model_name_map: dict +): + """ + Calculate the number of input and output tokens for given queries based on metadata strings. + + Args: + metadata_string (str): A string containing all relevant metadata from the string. + model_choice (str): A string describing the model name + model_name_map (dict): A dictionary mapping model name to source + """ + + model_name_map[model_choice]["source"] + + # Regex to find the numbers following the keys in the "Query summary metadata" section + # This ensures we get the final, aggregated totals for the whole query. + input_regex = r"input_tokens: (\d+)" + output_regex = r"output_tokens: (\d+)" + + # re.findall returns a list of all matching strings (the captured groups). + input_token_strings = re.findall(input_regex, metadata_string) + output_token_strings = re.findall(output_regex, metadata_string) + + # Convert the lists of strings to lists of integers and sum them up + total_input_tokens = sum([int(token) for token in input_token_strings]) + total_output_tokens = sum([int(token) for token in output_token_strings]) + + number_of_calls = len(input_token_strings) + + print(f"Found {number_of_calls} LLM call entries in metadata.") + print("-" * 20) + print(f"Total Input Tokens: {total_input_tokens}") + print(f"Total Output Tokens: {total_output_tokens}") + + return total_input_tokens, total_output_tokens, number_of_calls diff --git a/tools/prompts.py b/tools/prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..1735ddfdaa990d217038acb3a91dbab912dd80b1 --- /dev/null +++ b/tools/prompts.py @@ -0,0 +1,268 @@ +### +# System prompt +### + +generic_system_prompt = """You are a researcher analysing responses from an open text dataset. You are analysing a single column from this dataset. Use British English spelling and grammar.""" + +system_prompt = """You are a researcher analysing responses from an open text dataset. You are analysing a single column from this dataset called '{column_name}'. {consultation_context} Use British English spelling and grammar.""" + +markdown_additional_prompt = """ You will be given a request for a markdown table. You must respond with ONLY the markdown table. Do not include any introduction, explanation, or concluding text.""" + +### +# Initial topic table prompt +### +initial_table_system_prompt = system_prompt + markdown_additional_prompt + +initial_table_assistant_prefill = "|" + +default_response_reference_format = "In the next column named 'Response ID', list each specific Response reference number that is relevant to the Subtopic, separated by commas. Do not write any other text in this column." + +initial_table_prompt = """{validate_prompt_prefix}Your task is to create one new markdown table based on open text responses in the reponse table below. +In the first column named 'General topic', identify general topics relevant to responses. Create as many general topics as you can. +In the second column named 'Subtopic', list subtopics relevant to responses. Make the subtopics as specific as possible and make sure they cover every issue mentioned. The subtopic should never be empty. +{sentiment_choices}{response_reference_format} +In the final column named 'Summary', write a summary of the subtopic based on relevant responses - highlight specific issues that appear. Do not mention specific response numbers in the summary. {add_existing_topics_summary_format} +Do not add any other columns. Do not add any other text to your response. Only mention topics that are relevant to at least one response. + +{response_table} + +New table:{previous_table_introduction}{previous_table}{validate_prompt_suffix}""" + +### +# Adding existing topics to consultation responses +### + +add_existing_topics_system_prompt = system_prompt + markdown_additional_prompt + +add_existing_topics_assistant_prefill = "|" + +force_existing_topics_prompt = """Create a new markdown table. In the first column named 'Placeholder', write 'Not assessed'. In the second column named 'Subtopics', assign Topics from the above table to Responses. Assign topics only if they are very relevant to the text of the Response. The assigned Subtopics should be chosen from the topics table above, exactly as written. Do not add any new topics, or modify existing topic names.""" + +allow_new_topics_prompt = """Create a new markdown table. In the first column named 'General topic', and the second column named 'Subtopic', assign General Topics and Subtopics to Responses. Assign topics from the Topics table above only if they are very relevant to the text of the Response. Fill in the General topic and Subtopic for the Topic if they do not already exist. If you find a new topic that does not exist in the Topics table, add a new row to the new table. Make the General topic and Subtopic as specific as possible. The subtopic should never be blank or empty.""" + +force_single_topic_prompt = """ Assign each response to one single topic only.""" + +add_existing_topics_prompt = """{validate_prompt_prefix}Your task is to create one new markdown table, assigning responses from the Response table below to topics. +{topic_assignment}{force_single_topic} +{sentiment_choices}{response_reference_format} +In the final column named 'Summary', write a summary of the Subtopic based on relevant responses - highlight specific issues that appear. Do not mention specific response numbers in the summary. {add_existing_topics_summary_format} +Do not add any other columns. Do not add any other text to your response. Only mention topics that are relevant to at least one response. + +Choose from among the following topic names to assign to the responses, only if they are directly relevant to responses from the response table below: + +{topics} + +{response_table} + +New table:{previous_table_introduction}{previous_table}{validate_prompt_suffix}""" + +### +# VALIDATION PROMPTS +### +# These are prompts used to validate previous LLM outputs, and create corrected versions of the outputs if errors are found. +validation_system_prompt = system_prompt + +validation_prompt_prefix_default = """The following instructions were previously provided to create an output table:\n'""" + +previous_table_introduction_default = ( + """'\n\nThe following output table was created based on the above instructions:\n""" +) + +validation_prompt_suffix_default = """\n\nBased on the above information, you need to create a corrected version of the output table. Examples of issues to correct include: + +- Remove rows where responses are not relevant to the assigned topic, or where responses are not relevant to any topic. +- Remove rows where a topic is not assigned to any specific response. +- If the current topic assignment does not cover all information in a response, assign responses to relevant topics from the suggested topics table, or create a new topic if necessary. +- Correct any false information in the summary column, which is a summary of the relevant response text. +{additional_validation_issues} +- Any other obvious errors that you can identify. + +With the above issues in mind, create a new, corrected version of the markdown table below. If there are no issues to correct, write simply "No change". Return only the corrected table without additional text, or 'no change' alone.""" + +validation_prompt_suffix_struct_summary_default = """\n\nBased on the above information, you need to create a corrected version of the output table. Examples of issues to correct include: + +- Any misspellings in the Main heading or Subheading columns +- Correct any false information in the summary column, which is a summary of the relevant response text. +{additional_validation_issues} +- Any other obvious errors that you can identify. + +With the above issues in mind, create a new, corrected version of the markdown table below. If there are no issues to correct, write simply "No change". Return only the corrected table without additional text, or 'no change' alone.""" + +### +# SENTIMENT CHOICES +### + +negative_neutral_positive_sentiment_prompt = ( + "write the sentiment of the Subtopic: Negative, Neutral, or Positive" +) +negative_or_positive_sentiment_prompt = ( + "write the sentiment of the Subtopic: Negative or Positive" +) +do_not_assess_sentiment_prompt = "write the text 'Not assessed'" # Not used anymore. Instead, the column is filled in automatically with 'Not assessed' +default_sentiment_prompt = ( + "write the sentiment of the Subtopic: Negative, Neutral, or Positive" +) + +### +# STRUCTURED SUMMARY PROMPT +### + +structured_summary_prompt = """Your task is to write a structured summary for open text responses. + +Create a new markdown table based on the response table below with the headings 'Main heading', 'Subheading' and 'Summary'. + +For each of the responses in the Response table, you will create a row for each summary associated with each of the Main headings and Subheadings from the Headings table. If there is no Headings table, created your own headings. In the first and second columns, write a Main heading and Subheading from the Headings table. Then in Summary, write a detailed and comprehensive summary that covers all information relevant to the Main heading and Subheading on the same row. +{summary_format} + +Do not add any other columns. Do not add any other text to your response. + +{response_table} + +Headings to structure the summary are in the following table: +{topics} + +New table:""" + +### +# SUMMARISE TOPICS PROMPT +### + +summary_assistant_prefill = "" + +summarise_topic_descriptions_system_prompt = system_prompt + +summarise_topic_descriptions_prompt = """Your task is to make a consolidated summary of the text below. {summary_format} +Return only the summary and no other text. Do not mention specific response numbers in the summary. + +Text to summarise: +{summaries} + +Summary:""" + +single_para_summary_format_prompt = "Return a concise summary up to one paragraph long that summarises only the most important themes from the original text" + +two_para_summary_format_prompt = "Return a summary up to two paragraphs long that includes as much detail as possible from the original text" + +### +# OVERALL SUMMARY PROMPTS +### + +summarise_everything_system_prompt = system_prompt + +summarise_everything_prompt = """Below is a table that gives an overview of the main topics from a dataset of open text responses along with a summary of each topic, and the number of responses that mentioned each topic: + +{topic_summary_table} + +Your task is to summarise the above table. {summary_format} Return only the summary and no other text. + +Summary:""" + +comprehensive_summary_format_prompt = "Return a comprehensive summary that covers all the important topics and themes described in the table. Structure the summary with General Topics as headings, with significant Subtopics described in bullet points below them in order of relative significance. Do not explicitly mention the Sentiment, Number of responses, or Group values. Do not use the words 'General topic' or 'Subtopic' directly in the summary. Format the output for Excel display using: **bold text** for main headings, • bullet points for sub-items, and line breaks between sections. Avoid markdown symbols like # or ##." + +comprehensive_summary_format_prompt_by_group = "Return a comprehensive summary that covers all the important topics and themes described in the table. Structure the summary with General Topics as headings, with significant Subtopics described in bullet points below them in order of relative significance. Do not explicitly mention the Sentiment, Number of responses, or Group values. Do not use the words 'General topic' or 'Subtopic' directly in the summary. Compare and contrast differences between the topics and themes from each Group. Format the output for Excel display using: **bold text** for main headings, • bullet points for sub-items, and line breaks between sections. Avoid markdown symbols like # or ##." + +# Alternative Excel formatting options +excel_rich_text_format_prompt = "Return a comprehensive summary that covers all the important topics and themes described in the table. Structure the summary with General Topics as headings, with significant Subtopics described in bullet points below them in order of relative significance. Do not explicitly mention the Sentiment, Number of responses, or Group values. Do not use the words 'General topic' or 'Subtopic' directly in the summary. Format for Excel using: BOLD for main headings, bullet points (•) for sub-items, and line breaks between sections. Use simple text formatting that Excel can interpret." + +excel_plain_text_format_prompt = "Return a comprehensive summary that covers all the important topics and themes described in the table. Structure the summary with General Topics as headings, with significant Subtopics described in bullet points below them in order of relative significance. Do not explicitly mention the Sentiment, Number of responses, or Group values. Do not use the words 'General topic' or 'Subtopic' directly in the summary. Format as plain text with clear structure: use ALL CAPS for main headings, bullet points (•) for sub-items, and line breaks between sections. Avoid any special formatting symbols." + +### +# LLM-BASED TOPIC DEDUPLICATION PROMPTS +### + +llm_deduplication_system_prompt = """You are an expert at analysing and consolidating topic categories. Your task is to identify semantically similar topics that should be merged together, even if they use different wording or synonyms.""" + +llm_deduplication_prompt = """You are given a table of topics with their General topics, Subtopics{sentiment_text}. Your task is to identify topics that are semantically similar and CONSOLIDATE multiple different topics into the SAME merged topic names. The goal is to REDUCE the total number of unique topics by merging similar ones together. Only merge topics that are almost identical in terms of meaning - if in doubt, do not merge. The user has specified that there should be a maximum of {max_number_of_topics} topics, so if the current number of topics is greater than this, merge topics until the number of topics is less than or equal to {max_number_of_topics}. If no duplicates are found, return an empty table with the column headings suggested below. + +IMPORTANT: You must CONSOLIDATE multiple different topics into the same merged topic names. For example, if you have: +- "Transportation issues" +- "Public transport problems" +- "Bus and train delays" + +These should ALL be merged to the SAME merged topic (e.g., "Transportation | Public Transport Issues"), not just renamed individually. The goal is to reduce the number of unique topics by grouping similar ones together. + +Analyse the following topics table and identify groups of topics that describe essentially the same concept but may use different words or phrases. When you find similar topics, merge ALL of them to the SAME merged topic name. For example: +- Multiple topics like "Housing costs", "Rent prices", and "Accommodation expenses" should all merge to the SAME merged topic (e.g., "Housing | Costs") +- Multiple topics like "Environmental concerns" and "Green issues" should merge to the SAME merged topic (e.g., "Environment | Concerns") + +Create a markdown table with the following columns: +1. 'Original General topic' - The current general topic name +2. 'Original Subtopic' - The current subtopic name{sentiment_columns} +3. 'Merged General topic' - The consolidated general topic name (use the most descriptive, and ensure multiple similar topics map to the SAME merged name) +4. 'Merged Subtopic' - The consolidated subtopic name (use the most descriptive, and ensure multiple similar topics map to the SAME merged name){merged_sentiment_columns} +5. 'Merge Reason' - Brief explanation of why these topics should be merged + +CRITICAL: Multiple different original topics that are similar should map to the SAME merged topic combination. Do not just rename topics individually - consolidate them. Only include rows where topics should actually be merged. If a topic has no semantic duplicates, do not include it in the table. Produce only a markdown table in the format described above. Do not add any other text to your response. + +Topics to analyse: +{topics_table} + +Merged topics table:""" + +llm_deduplication_prompt_with_candidates = """You are given a table of topics with their General topics, Subtopics{sentiment_text}. Your task is to identify topics that are semantically similar and CONSOLIDATE multiple different topics into the SAME merged topic names. The goal is to REDUCE the total number of unique topics by merging similar ones together, even if they use different wording. + +Additionally, you have been provided with a list of candidate topics that represent preferred topic categories. When merging topics, prioritise fitting similar topics into these existing candidate categories rather than creating new ones. Only merge topics that are almost identical in terms of meaning - if in doubt, do not merge. The user has specified that there should be a maximum of {max_number_of_topics} topics, so if the current number of topics is greater than this, merge topics until the number of topics is less than or equal to {max_number_of_topics}. + +IMPORTANT: You must CONSOLIDATE multiple different topics into the same merged topic names. For example, if you have: +- "Transportation issues" +- "Public transport problems" +- "Bus and train delays" + +These should ALL be merged to the SAME merged topic (preferably matching a candidate topic if available), not just renamed individually. The goal is to reduce the number of unique topics by grouping similar ones together. + +Analyse the following topics table and identify groups of topics that describe essentially the same concept but may use different words or phrases. When you find similar topics, merge ALL of them to the SAME merged topic name. For example: +- Multiple topics like "Housing costs", "Rent prices", and "Accommodation expenses" should all merge to the SAME merged topic (preferably matching a candidate topic) +- Multiple topics like "Environmental concerns" and "Green issues" should merge to the SAME merged topic (preferably matching a candidate topic) + +When merging topics, consider the candidate topics provided below and try to map similar topics to these preferred categories when possible. Multiple similar original topics should map to the SAME candidate topic. + +Create a markdown table with the following columns: +1. 'Original General topic' - The current general topic name +2. 'Original Subtopic' - The current subtopic name{sentiment_columns} +3. 'Merged General topic' - The consolidated general topic name (prefer candidate topics when similar, and ensure multiple similar topics map to the SAME merged name) +4. 'Merged Subtopic' - The consolidated subtopic name (prefer candidate topics when similar, and ensure multiple similar topics map to the SAME merged name){merged_sentiment_columns} +5. 'Merge Reason' - Brief explanation of why these topics should be merged + +CRITICAL: Multiple different original topics that are similar should map to the SAME merged topic combination. Do not just rename topics individually - consolidate them. Only include rows where topics should actually be merged. If a topic has no semantic duplicates, do not include it in the table. Produce only a markdown table in the format described above. Do not add any other text to your response. + +Topics to analyse: +{topics_table} + +Candidate topics to consider for mapping: +{candidate_topics_table} + +Merged topics table:""" + +### +# VERIFY EXISTING DESCRIPTIONS/TITLES - Currently not used +### + +verify_assistant_prefill = "|" + +verify_titles_system_prompt = system_prompt + +verify_titles_prompt = """Response numbers alongside the Response text and assigned descriptions are shown in the table below: +{response_table} + +The criteria for a suitable description for these responses is that they should be readable, concise, and fully encapsulate the main subject of the response. + +Create a markdown table with four columns. +The first column is 'Response ID', and should contain just the response number under consideration. +The second column is 'Is this a suitable description', answer the question with 'Yes' or 'No', with no other text. +The third column is 'Explanation', give a short explanation for your response in the second column. +The fourth column is 'Alternative description', suggest an alternative description for the response that meet the criteria stated above. +Do not add any other text to your response. + +Output markdown table:""" + + +## The following didn't work well in testing and so is not currently used + +create_general_topics_system_prompt = system_prompt + +create_general_topics_prompt = """Subtopics known to be relevant to this dataset are shown in the following Topics table: +{topics} + +Your task is to create a General topic name for each Subtopic. The new Topics table should have the columns 'General topic' and 'Subtopic' only. Write a 'General topic' text label relevant to the Subtopic next to it in the new table. The text label should describe the general theme of the Subtopic. Do not add any other text, thoughts, or notes to your response. + +New Topics table:""" diff --git a/tools/view_logs.py b/tools/view_logs.py new file mode 100644 index 0000000000000000000000000000000000000000..213e18291168fde422e57958e25df852e1ef9e8a --- /dev/null +++ b/tools/view_logs.py @@ -0,0 +1,382 @@ +""" +Functions for viewing and filtering JSON log files containing LLM prompts and responses. +""" + +import json +from typing import Dict, List, Optional, Tuple + +import gradio as gr + + +def load_log_file_handler(log_file): + """Handle log file upload and initialize dropdowns.""" + if log_file is None: + return ( + [], # log_data_state + gr.Dropdown(choices=[]), # log_batch_dropdown + gr.Dropdown(choices=[]), # log_task_type_dropdown + gr.Dropdown(choices=[]), # log_group_dropdown + gr.Dropdown(choices=[]), # log_model_choice_dropdown + gr.Dropdown(choices=[]), # log_validated_dropdown + "### Prompt\n\nNo file uploaded.", # log_prompt_markdown + "### Response\n\nNo file uploaded.", # log_response_markdown + ) + + file_path = log_file.name if hasattr(log_file, "name") else log_file + log_data = load_log_file(file_path) + + if not log_data: + return ( + [], # log_data_state + gr.Dropdown(choices=[]), # log_batch_dropdown + gr.Dropdown(choices=[]), # log_task_type_dropdown + gr.Dropdown(choices=[]), # log_group_dropdown + gr.Dropdown(choices=[]), # log_model_choice_dropdown + gr.Dropdown(choices=[]), # log_validated_dropdown + "### Prompt\n\nError: Could not load log file or file is empty.", # log_prompt_markdown + "### Response\n\nError: Could not load log file or file is empty.", # log_response_markdown + ) + + # Extract unique values for all filters + batches, task_types, groups, model_choices, validated_values = ( + extract_unique_filter_values(log_data) + ) + + # Format choices as strings for dropdown, add "All" option + batch_choices = ["All"] + [str(b) for b in batches] if batches else ["All"] + task_type_choices = ["All"] + task_types if task_types else ["All"] + group_choices = ["All"] + groups if groups else ["All"] + model_choice_choices = ["All"] + model_choices if model_choices else ["All"] + validated_choices = ["All"] + validated_values if validated_values else ["All"] + + # Get default prompt and response (show all if no filters) + prompt, response = get_prompt_and_response(log_data) + + # Format for markdown display + prompt_display = f"### Prompt\n\n{prompt}" + response_display = f"### Response\n\n{response}" + + return ( + log_data, # log_data_state + gr.Dropdown(choices=batch_choices, value="All"), # log_batch_dropdown + gr.Dropdown(choices=task_type_choices, value="All"), # log_task_type_dropdown + gr.Dropdown(choices=group_choices, value="All"), # log_group_dropdown + gr.Dropdown( + choices=model_choice_choices, value="All" + ), # log_model_choice_dropdown + gr.Dropdown(choices=validated_choices, value="All"), # log_validated_dropdown + prompt_display, # log_prompt_markdown + response_display, # log_response_markdown + ) + + +def filter_log_display(log_data, batch_str, task_type, group, model_choice, validated): + """Filter and display log entries based on all filter criteria.""" + if not log_data: + return ( + "### Prompt\n\nNo log data available.", + "### Response\n\nNo log data available.", + ) + + # Convert batch string to int, handle "All" option + batch = None + if batch_str and batch_str != "All": + try: + batch = int(batch_str) + except (ValueError, TypeError): + batch = None + + # Handle "All" option for all filters + task_type_filter = None if task_type == "All" else task_type + group_filter = None if group == "All" else group + model_choice_filter = None if model_choice == "All" else model_choice + validated_filter = None if validated == "All" else validated + + # Get filtered prompt and response + prompt, response = get_prompt_and_response( + log_data, + batch, + task_type_filter, + group_filter, + model_choice_filter, + validated_filter, + ) + + # Format for markdown display + prompt_display = f"### Prompt\n\n{prompt}" + response_display = f"### Response\n\n{response}" + + return prompt_display, response_display + + +# Handle batch and task type changes +def update_log_display_on_filter( + log_data, batch_str, task_type, group, model_choice, validated +): + """Update display when any filter changes.""" + return filter_log_display( + log_data, batch_str, task_type, group, model_choice, validated + ) + + +def load_log_file(file_path: str) -> List[Dict]: + """ + Load a JSON log file containing LLM prompts and responses. + + Args: + file_path: Path to the JSON file + + Returns: + List of dictionaries containing log entries + """ + try: + with open(file_path, "r", encoding="utf-8-sig", errors="replace") as f: + data = json.load(f) + return data if isinstance(data, list) else [] + except (FileNotFoundError, json.JSONDecodeError, Exception) as e: + print(f"Error loading log file: {e}") + return [] + + +def parse_batch_number(batch: str) -> int: + """ + Parse batch number from string format. + Handles formats like "1", "1:", "2", "2:", etc. + + Args: + batch: Batch string (e.g., "1", "1:", "2") + + Returns: + Integer batch number + """ + if not batch: + return 0 + # Remove colon if present and strip whitespace + batch_clean = batch.strip().rstrip(":") + try: + return int(batch_clean) + except (ValueError, TypeError): + return 0 + + +def extract_unique_batches_and_task_types( + log_data: List[Dict], +) -> Tuple[List[int], List[str]]: + """ + Extract unique batch numbers and task types from log data. + + Args: + log_data: List of log entry dictionaries + + Returns: + Tuple of (sorted list of unique batch numbers, sorted list of unique task types) + """ + batches = set() + task_types = set() + + for entry in log_data: + if "batch" in entry: + batch_num = parse_batch_number(str(entry["batch"])) + if batch_num > 0: + batches.add(batch_num) + if "task_type" in entry: + task_type = entry["task_type"] + if task_type: + task_types.add(task_type) + + return sorted(list(batches)), sorted(list(task_types)) + + +def extract_unique_filter_values( + log_data: List[Dict], +) -> Tuple[List[int], List[str], List[str], List[str], List[str]]: + """ + Extract unique values for all filter fields from log data. + + Args: + log_data: List of log entry dictionaries + + Returns: + Tuple of (batches, task_types, groups, model_choices, validated_values) + """ + batches = set() + task_types = set() + groups = set() + model_choices = set() + validated_values = set() + + for entry in log_data: + if "batch" in entry: + batch_num = parse_batch_number(str(entry["batch"])) + if batch_num > 0: + batches.add(batch_num) + if "task_type" in entry: + task_type = entry["task_type"] + if task_type: + task_types.add(task_type) + if "group" in entry: + group = entry["group"] + if group: + groups.add(str(group)) + if "model_choice" in entry: + model_choice = entry["model_choice"] + if model_choice: + model_choices.add(str(model_choice)) + if "validated" in entry: + validated = entry["validated"] + if validated: + validated_values.add(str(validated)) + + return ( + sorted(list(batches)), + sorted(list(task_types)), + sorted(list(groups)), + sorted(list(model_choices)), + sorted(list(validated_values)), + ) + + +def filter_log_entries( + log_data: List[Dict], + batch: Optional[int] = None, + task_type: Optional[str] = None, + group: Optional[str] = None, + model_choice: Optional[str] = None, + validated: Optional[str] = None, +) -> List[Dict]: + """ + Filter log entries by batch number, task type, group, model_choice, and/or validated. + + Args: + log_data: List of log entry dictionaries + batch: Optional batch number to filter by + task_type: Optional task type to filter by + group: Optional group to filter by + model_choice: Optional model choice to filter by + validated: Optional validated value to filter by + + Returns: + Filtered list of log entries + """ + filtered = [] + + for entry in log_data: + match_batch = True + match_task_type = True + match_group = True + match_model_choice = True + match_validated = True + + if batch is not None: + entry_batch = parse_batch_number(str(entry.get("batch", "0"))) + match_batch = entry_batch == batch + + if task_type is not None: + entry_task_type = entry.get("task_type", "") + match_task_type = entry_task_type == task_type + + if group is not None: + entry_group = str(entry.get("group", "")) + match_group = entry_group == group + + if model_choice is not None: + entry_model_choice = str(entry.get("model_choice", "")) + match_model_choice = entry_model_choice == model_choice + + if validated is not None: + entry_validated = str(entry.get("validated", "")) + match_validated = entry_validated == validated + + if ( + match_batch + and match_task_type + and match_group + and match_model_choice + and match_validated + ): + filtered.append(entry) + + return filtered + + +def get_prompt_and_response( + log_data: List[Dict], + batch: Optional[int] = None, + task_type: Optional[str] = None, + group: Optional[str] = None, + model_choice: Optional[str] = None, + validated: Optional[str] = None, +) -> Tuple[str, str]: + """ + Get prompt and response text for filtered log entries. + If multiple entries match, concatenate them. + + Args: + log_data: List of log entry dictionaries + batch: Optional batch number to filter by + task_type: Optional task type to filter by + group: Optional group to filter by + model_choice: Optional model choice to filter by + validated: Optional validated value to filter by + + Returns: + Tuple of (prompt_text, response_text) + """ + filtered = filter_log_entries( + log_data, batch, task_type, group, model_choice, validated + ) + + if not filtered: + return "No entries found for the selected batch and task type.", "" + + # If multiple entries, combine them with separators + prompts = [] + responses = [] + + for idx, entry in enumerate(filtered, 1): + prompt = entry.get("prompt", "") + response = entry.get("response", "") + + # Add entry number if multiple entries + entry_header = f"**Entry {idx}**" if len(filtered) > 1 else "" + + if prompt: + if entry_header: + prompts.append(f"{entry_header}\n\n{prompt}") + else: + prompts.append(prompt) + if response: + if entry_header: + responses.append(f"{entry_header}\n\n{response}") + else: + responses.append(response) + + prompt_text = "\n\n---\n\n".join(prompts) if prompts else "No prompt found." + response_text = "\n\n---\n\n".join(responses) if responses else "No response found." + + return prompt_text, response_text + + +def load_and_initialize_log_viewer( + file_path: str, +) -> Tuple[List[int], List[str], str, str]: + """ + Load log file and initialize viewer with default values. + + Args: + file_path: Path to the JSON log file + + Returns: + Tuple of (batch_numbers, task_types, default_prompt, default_response) + """ + log_data = load_log_file(file_path) + + if not log_data: + return [], [], "No log data found in file.", "" + + batches, task_types = extract_unique_batches_and_task_types(log_data) + + # Get default prompt and response (first entry or all if no filters) + prompt, response = get_prompt_and_response(log_data) + + return batches, task_types, prompt, response diff --git a/windows_install_llama-cpp-python.txt b/windows_install_llama-cpp-python.txt new file mode 100644 index 0000000000000000000000000000000000000000..ad7024d73335c0ab8ed1bf0a5a3bd697a1524445 --- /dev/null +++ b/windows_install_llama-cpp-python.txt @@ -0,0 +1,111 @@ +--- + +#How to build llama-cpp-python on Windows: Step-by-Step Guide + +First, you need to set up a proper C++ development environment. + +# Step 1: Install the C++ Compiler +Scroll down the page past the main programs to "Tools for Visual Studio" and download the "Build Tools for Visual Studio". This is a standalone installer that gives you the C++ compiler and libraries without installing the full Visual Studio IDE. + +Run the installer. In the "Workloads" tab, check the box for "Desktop development with C++". + +MSVC v143 +C++ ATL +C++ Profiling tools +C++ CMake tools for Windows +C++ MFC +C++ Modules +Windows 10 SDK (10.0.20348.0) + +Proceed with the installation. + +Need to use 'x64 Native Tools Command Prompt for VS 2022' to install the below. Run as administrator + +# Step 2: Install CMake +Go to the CMake download page: https://cmake.org/download + +Download the latest Windows installer (e.g., cmake-x.xx.x-windows-x86_64.msi). + +Run the installer. Crucially, when prompted, select the option to "Add CMake to the system PATH for all users" or "for the current user." This allows you to run cmake from any command prompt. + + +# Step 3: (FOR CPU INFERENCE ONLY) Download and Place OpenBLAS +This is often the trickiest part. + +Go to the OpenBLAS releases on GitHub. + +Find a recent release and download the pre-compiled version for Windows. It will typically be a file named something like OpenBLAS-0.3.21-x64.zip (the version number will change). Make sure you get the 64-bit (x64) version if you are using 64-bit Python. + +Create a folder somewhere easily accessible, for example, C:\libs\. + +Extract the contents of the OpenBLAS zip file into that folder. Your final directory structure should look something like this: + +C:\libs\OpenBLAS\ +├── bin\ +├── include\ +└── lib\ + +## 3.b. Install Chocolatey +https://chocolatey.org/install + +Step 1: Install Chocolatey (if you don't already have it) +Open PowerShell as an Administrator. (Right-click the Start Menu -> "Windows PowerShell (Admin)" or "Terminal (Admin)"). + +Run the following command to install Chocolatey. It's a single, long line: + +Set-ExecutionPolicy Bypass -Scope Process -Force; [System.Net.ServicePointManager]::SecurityProtocol = [System.Net.ServicePointManager]::SecurityProtocol -bor 3072; iex ((New-Object System.Net.WebClient).DownloadString('https://community.chocolatey.org/install.ps1')) + +Once it's done, close the Administrator PowerShell window. + +Step 2: Install pkg-config-lite using Chocolatey +IMPORTANT: Open a NEW command prompt or PowerShell window (as a regular user is fine). This is necessary so it recognises the new choco command. + +Run the following command in console to install a lightweight version of pkg-config: + +choco install pkgconfiglite + +Approve the installation by typing Y or A if prompted. + +# Step 4: Run the Installation Command +Now you have all the pieces. The final step is to run the command in a terminal that is aware of your new build environment. + +Open the "Developer Command Prompt for VS" from your Start Menu. This is important! This special command prompt automatically configures all the necessary paths for the C++ compiler. + +## For CPU + +set PKG_CONFIG_PATH=C:\\OpenBLAS\lib\pkgconfig # Set this in environment variables + +pip install llama-cpp-python==0.3.16 --force-reinstall --verbose --no-cache-dir -Ccmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS;-DBLAS_INCLUDE_DIRS=C://OpenBLAS/include;-DBLAS_LIBRARIES=C://OpenBLAS/lib/libopenblas.lib" + +pip install llama-cpp-python==0.3.16 --verbose --no-cache-dir -Ccmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS;-DBLAS_INCLUDE_DIRS=C:/Users/s_cas/libs/OpenBLAS/include;-DBLAS_LIBRARIES=C:/Users/s_cas/OpenBLAS/lib/libopenblas.lib";-DPKG_CONFIG_PATH=C:/users/s_cas/openblas/lib/pkgconfig" + +or to make a wheel: + +pip install llama-cpp-python==0.3.16 --wheel-dir dist --verbose --no-cache-dir -Ccmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS;-DBLAS_INCLUDE_DIRS=C://OpenBLAS/include;-DBLAS_LIBRARIES=C://OpenBLAS/lib/libopenblas.lib" + +pip wheel llama-cpp-python==0.3.16 --wheel-dir dist --verbose --no-cache-dir -Ccmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS;-DBLAS_INCLUDE_DIRS=C:/Users//libs/OpenBLAS/include;-DBLAS_LIBRARIES=C:/Users//libs/OpenBLAS/lib/libopenblas.lib" + + + +## With Cuda (NVIDIA GPUs only) + +Make sure that the have the CUDA 12.4 toolkit for windows installed: https://developer.nvidia.com/cuda-12-4-0-download-archive + +### Make sure you are using the x64 version of Developer command tools for the below, e.g. 'x64 Native Tools Command Prompt for VS 2022' ### + +Use NVIDIA GPU (cuBLAS): If you have an NVIDIA GPU, using cuBLAS is often easier because the CUDA Toolkit installer handles most of the setup. + +Install the NVIDIA CUDA Toolkit. + +Run the install command specifying cuBLAS (for faster inference): + +pip install llama-cpp-python==0.3.16 --force-reinstall --verbose -C cmake.args="-DGGML_CUDA=on -DGGML_CUBLAS=on" + +If you want to create a new wheel to help with future installs, you can run: + +cd first to a folder that you have edit access for + +pip wheel llama-cpp-python==0.3.16 --wheel-dir dist --verbose -C cmake.args="-DGGML_CUDA=on -DGGML_CUBLAS=on" + + +