Data

Tools

Indexes

Applications

Experiments

Awesome

An open API service indexing awesome lists of open source software.

https://github.com/smit-parekh/proactive-supply-chain-disruption-agent

An AI Agent leveraging fine-tuned LLMs (Mistral-7B w/ PEFT) and LangGraph to proactively identify, assess, and suggest mitigations for supply chain disruptions, tailored to specific client needs (e.g., Shell, Pfizer). Includes MLOps integration using Vertex AI and MLflow.
https://github.com/smit-parekh/proactive-supply-chain-disruption-agent

ai-agent api fastapi fine-tuning forecasting google-cloud langchain langgraph llm lora machine-learning mlflow mlops peft python risk-management supply-chain vertex-ai

Last synced: 4 months ago
JSON representation

An AI Agent leveraging fine-tuned LLMs (Mistral-7B w/ PEFT) and LangGraph to proactively identify, assess, and suggest mitigations for supply chain disruptions, tailored to specific client needs (e.g., Shell, Pfizer). Includes MLOps integration using Vertex AI and MLflow.

Awesome Lists containing this project

README 

          
# AI Agent for Proactive Supply Chain Disruption Management



[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) 

[![Python Version](https://img.shields.io/badge/python-3.9%2B-blue.svg)](https://www.python.org/downloads/)



An AI Agent leveraging fine-tuned LLMs (Google/Gemma-3-1B-it w/ PEFT) and LangGraph to proactively identify, assess, and suggest mitigations for supply chain disruptions, tailored to specific client needs (e.g. from Petroleum and Pharma Industry). Includes MLOps integration using Vertex AI and MLflow for robust deployment and maintenance.



## Problem Statement



Managing complex global supply chains for diverse clients from Petroleum Industry (volatile energy markets) and Pharma Industry (temperature-sensitive pharma) involves constant risk of disruptions (e.g., port closures, extreme weather, customs delays, geopolitical events). Reactively managing these led to significant costs, SLA penalties, and client dissatisfaction for major global logistics providers. A system was needed to proactively identify potential disruptions, assess impact *specific* to client supply chains, and suggest mitigation actions *before* escalation.



## Solution Overview



This project implements a proactive AI agent built with **LangGraph** that orchestrates multiple components:

1.  **Event Ingestion & Filtering:** Monitors real-time data streams (news APIs like GDELT, weather feeds, internal alerts) and filters for potentially relevant supply chain events.

2.  **Context Gathering:** Retrieves client-specific data including network maps, contractual obligations (SLAs), product sensitivities (e.g., temperature constraints), and historical performance.

3.  **Client-Specific Risk Assessment:** Utilizes a Large Language Model (LLM), **fine-tuned using PEFT/LoRA on client-specific data**, to assess the potential impact of the event on that client's unique supply chain configuration and sensitivities.

4.  **Impact Forecasting:** Employs a hybrid approach combining a traditional ML model (XGBoost trained on historical data) to predict quantitative impact (e.g., probability of delay) and LLM reasoning to forecast qualitative impact (cost, operational issues).

5.  **RAG-Enhanced Mitigation Suggestions:** Generates context-aware mitigation actions using an LLM augmented with Retrieval-Augmented Generation (RAG). The RAG component retrieves relevant Standard Operating Procedures (SOPs) and successful past mitigation examples from a vector database.

6.  **Notification:** Dispatches tailored alerts to relevant internal teams (Operations, Client Managers) with the assessment, forecast, and suggested actions.



## Differentiation: Client-Specific Fine-Tuning



The core differentiator of this agent is the **client-specific fine-tuning** of the risk assessment LLM. Unlike generic alerting systems, this agent understands the unique vulnerabilities, contractual nuances, and product requirements of individual clients (demonstrated with Pharma vs. Petroleum Industry archetypes). This allows for highly contextualized risk scoring and mitigation suggestions, moving beyond a one-size-fits-all approach.



## Key Features



*   Proactive Disruption Identification from multiple sources.

*   **Client-Specific Impact Assessment** using fine-tuned LLMs (PEFT/LoRA on Google/Gemma-3-1B-it).

*   Hybrid Impact Forecasting (XGBoost for delay probability + LLM reasoning).

*   **RAG-Enhanced Mitigation** leveraging SOPs and historical data via Vector Search.

*   Stateful and Conditional Agent Workflow using **LangGraph**.

*   **End-to-End MLOps Pipeline** on Google Cloud Vertex AI (Pipelines, Training, Prediction, Monitoring).

*   Experiment Tracking & Model Registry integration with **MLflow**.

*   Containerized **FastAPI** application for potential integration.

*   Comprehensive testing suite (`pytest`).



## Tech Stack



*   **Orchestration:** Langchain, LangGraph

*   **LLMs:** Google/Gemma-3-1B-it (Base Model), Hugging Face Transformers, PEFT (LoRA)

*   **ML:** PyTorch (underlying Transformers), Scikit-learn, XGBoost

*   **Vector DB (RAG):** FAISS (for local demo), Vertex AI Vector Search (for cloud)

*   **Embeddings:** Sentence Transformers (Hugging Face)

*   **Cloud Platform:** Google Cloud Platform (GCP)

    *   **MLOps:** Vertex AI (Pipelines, Custom Training, Endpoints, Prediction, Model Registry, Experiments [via MLflow integration or native], Model Monitoring)

    *   **Data:** BigQuery, Cloud Storage, Pub/Sub

    *   **Compute/Serving:** Google Kubernetes Engine (GKE) or Cloud Run

*   **Experiment Tracking:** MLflow

*   **API:** FastAPI, Uvicorn

*   **Containerization:** Docker

*   **Testing:** Pytest, pytest-asyncio, HTTPX

*   **Language:** Python 3.9+



## Agent Graph Flow (LangGraph)



The agent operates as a state machine defined using LangGraph. The state (`AgentState`) tracks information about the event, affected clients, assessments, forecasts, and suggestions as it moves through the nodes. Edges determine the next step, often conditionally based on the current state (e.g., event relevance, presence of errors).



**Node Descriptions:**



1.  **`entry_node`**: Initializes state, assigns event ID, parses or fetches initial event data.

2.  **`filter_event`**: Checks event relevance using keywords/rules. Identifies potentially affected clients. *Edge:* Proceeds to `gather_context` if relevant AND clients identified, else END.

3.  **`gather_context`**: Fetches client-specific details (routes, SLAs) for relevant clients. *Edge:* Proceeds to `assess_risk` if context gathered for at least one client, else END.

4.  **`assess_risk`**: For each client, calls the fine-tuned LLM for specific risk assessment. *Edge:* Proceeds to `forecast_impact` (even with partial errors), END on critical failure.

5.  **`forecast_impact`**: For each client, calls XGBoost model for delay probability and LLM for qualitative impact/reasoning. *Edge:* Proceeds to `suggest_mitigation` (even with partial errors), END on critical failure.

6.  **`suggest_mitigation`**: For each client, performs RAG search on SOPs and calls LLM to generate mitigation steps. *Edge:* Proceeds to `dispatch_notification` (even with partial errors), END on critical failure.

7.  **`dispatch_notification`**: Formats analysis and simulates sending alerts.

8.  **`terminate_workflow` / `final_node`**: Logs final status (Completed/Failed) and performs cleanup, leading to LangGraph's `END`.



**Diagram** 



```mermaid

graph LR

    A[entry_node] --> B(filter_event);

    B -- Relevant --> C{gather_context};

    B -- Irrelevant --> Z(END);

    C -- Context Found --> D(assess_risk);

    C -- Context Error / No Clients --> Z;

    D -- Assessment Done --> E(forecast_impact);

    D -- Critical Assessment Error --> Z;

    E -- Forecast Done --> F(suggest_mitigation);

    E -- Critical Forecast Error --> Z;

    F -- Suggestions Done --> G(dispatch_notification);

    F -- Critical Suggestion Error --> Z;

    G --> H(terminate_workflow / final_node);

    H --> Z;



    style Z fill:#f9f,stroke:#333,stroke-width:2px;

```



## Repository Structure



```

proactive-supply-chain-disruption-agent/

├── .github/workflows/         # CI/CD workflows (e.g., lint, test, build)

├── .gitignore

├── LICENSE

├── README.md                  # This file

├── requirements.txt           # Python dependencies

├── setup.py                   # Setup script

├── setup.cfg                  # Build configuration

├── pyproject.toml             # Build system definition, tool configs

├── configs/                   # YAML configuration files (agent, api, ml)

├── data/                      # Sample/hypothetical data & schemas

│   ├── schemas/               # JSON schemas for data sources

│   ├── sample_*.csv/txt       # Small sample data files

│   └── fine-tuning/           # Sample raw/processed data for LLM tuning

├── deployment/                # Deployment artifacts (Dockerfile, GKE/Cloud Run YAMLs)

├── mlops/                     # MLOps pipelines and components (Vertex AI)

│   ├── vertex_pipelines/      # KFP pipeline definitions and components

│   └── monitoring/            # Monitoring config placeholders (e.g., dashboard JSON)

├── notebooks/                 # Jupyter notebooks for exploration and prototyping

├── src/                       # Source code

│   └── proactive_disruption_agent/ # Main Python package

│       ├── agent/             # LangGraph agent implementation (graph, nodes, state)

│       ├── api/               # FastAPI application (main, models, dependencies)

│       ├── data_ingestion/    # Scripts/modules for event/data ingestion

│       ├── fine_tuning/       # LLM fine-tuning scripts (preprocess, train)

│       ├── forecasting/       # XGBoost forecasting model (train, predict)

│       ├── llm_services/      # Clients for interacting with LLMs (Vertex AI, local)

│       ├── retrieval/         # RAG / Vector store interaction

│       └── utils/             # Utility functions (config, logging)

└── tests/                     # Unit and integration tests (pytest)

    ├── agent/

    ├── api/

    └── forecasting/

```



## Getting Started



### Prerequisites



*   Python 3.9+

*   [Google Cloud SDK](https://cloud.google.com/sdk/docs/install) initialized and authenticated (if using GCP resources).

*   Access to a GCP project with Vertex AI, BigQuery, Pub/Sub, Cloud Storage APIs enabled (for full MLOps/deployment).

*   [Docker](https://docs.docker.com/get-docker/) installed (for containerization).

*   [MLflow Tracking Server](https://mlflow.org/docs/latest/tracking.html#where-runs-are-recorded) (local, hosted, or Vertex AI Experiment) configured (optional, but recommended).

*   Access to a base LLM like Gemma3-1B (e.g., via Hugging Face Hub).

*   GPU (preferably NVIDIA with CUDA, >=16GB VRAM recommended) for local fine-tuning/inference experimentation.



### Installation



1.  **Clone the repository:**

    ```bash

    git clone https://github.com/Smit-Parekh/proactive-supply-chain-disruption-agent.git

    cd proactive-supply-chain-disruption-agent

    ```



2.  **Create and activate a virtual environment:**

    ```bash

    python -m venv venv

    source venv/bin/activate # On Windows use `venv\Scripts\activate`

    ```



3.  **Install dependencies:**

    ```bash

    pip install --upgrade pip

    pip install -r requirements.txt



    # Install the src package in editable mode

    pip install -e .

    ```

    *(Note: Some libraries like `torch`, `faiss`, `bitsandbytes` might require specific installation commands depending on your OS and CUDA version. Refer to their official documentation if the standard pip install fails.)*



### Configuration



1.  **Copy & Customize:** Examine the configuration files in `configs/`. You might copy `agent_config.yaml` to `agent_config_local.yaml` etc. for local overrides.

2.  **Fill Placeholders:** **Crucially, you MUST replace placeholders** in `configs/*.yaml` (e.g., `your-gcp-project-id`, `YOUR_*_ENDPOINT_ID`, `your-gcs-bucket-name`, `your_mlflow_or_vertex_experiment_uri`) with your actual values.

3.  **Secrets Management:** For sensitive values (API keys, service account keys), **DO NOT hardcode them in YAML files**. Use environment variables, Google Secret Manager, or other secure methods, especially for cloud deployments. The code should be adapted to read from these sources.

4.  **Authentication:**

    *   **GCP:** Use `gcloud auth application-default login` for local development or configure service accounts for deployments (Workload Identity for GKE/Cloud Run recommended). Set `GOOGLE_APPLICATION_CREDENTIALS` environment variable if using service account keys directly (keep key files secure and out of git).

    *   **MLflow:** Set `MLFLOW_TRACKING_URI` environment variable if needed.



### Data Setup



*   The `data/` folder contains **small, hypothetical sample files** (`.csv`, `.txt`) and schemas (`data/schemas/`).

*   The `data/fine-tuning/tuning_raw_data.csv` provides a *small sample* of the format needed for fine-tuning.

*   For local testing, the application reads from these files (configured via `configs/agent_config.yaml`).

*   For cloud deployment, data would typically reside in GCS, BigQuery, etc., and MLOps pipelines would reference those sources.



## Usage



### Running the Agent Locally (via Prototype Notebook)



The easiest way to test the agent logic locally, especially with a locally fine-tuned LLM adapter:



1.  Ensure you have run necessary preprocessing/training steps (e.g., Notebook 03 for forecast model, Notebook 02 for LLM adapter).

2.  Modify and run `notebooks/04_langgraph_agent_prototype.ipynb`. This notebook shows how to:

    *   Load local forecasting model/vector store.

    *   Load the locally fine-tuned LLM adapter.

    *   Patch the agent's LLM calls to use the local model.

    *   Execute the graph with sample input and inspect the state.



### Running the API Locally



1.  Ensure API configuration is appropriate (port usually set via command line).

2.  Ensure the agent dependencies (models, vector store) can be loaded based on `configs/*.yaml`. You might need to mock cloud LLM calls if not testing with a local adapter via patching.

3.  Run from the project root:

    ```bash

    uvicorn src.proactive_disruption_agent.api.main:app --reload --port 8000

    ```

4.  Access the API docs via Swagger UI at `http://localhost:8000/docs`.

5.  Send POST requests to `/analyze-event` and check status using GET `/analysis/{analysis_id}`.



### Running MLOps Pipelines (Vertex AI)



1.  **Build Container Images:** Build Docker images for training (`xgboost_training_container`, `llm_tuning_container`) containing the `src` code and dependencies. Push them to Google Container Registry (GCR) or Artifact Registry. Also build the image for the API service specified in `deployment/`.

2.  **Configure Pipeline Parameters:** Update `mlops/vertex_pipelines/pipeline.py` with your GCP project details, GCS paths, container image URIs, desired machine types, model names, etc.

3.  **Compile the Pipeline:**

    ```bash

    python -m mlops.vertex_pipelines.compile_pipeline --output mlops/vertex_pipelines/compiled_ml_pipeline.yaml

    ```

4.  **Submit the Pipeline Job** via the Google Cloud Console (Vertex AI > Pipelines) or using `gcloud`:

    ```bash

    gcloud ai pipelines jobs submit \

      --region= \

      --pipeline-spec=mlops/vertex_pipelines/compiled_ml_pipeline.yaml \

      --parameter-values=project_id="",pipeline_root="gs:///pipeline-root",incident_data_gcs_uri="gs://...", ... # Pass all required pipeline parameters

    ```



## Fine-Tuning (Client Specificity)



*   **Objective:** Adapt the base LLM (Google/Gemma-3-1B-it) to better understand client-specific context and generate relevant risk assessments or mitigation steps.

*   **Data:** Requires curated, high-quality instruction-following data. See `data/fine-tuning/tuning_raw_data.csv` for format (input = event + client context, output = desired assessment/mitigation).

*   **Process:**

    1.  **Preprocessing:** Use `src/fine_tuning/data_preprocessor.py` to convert raw data into JSONL format suitable for training.

    2.  **Training:** Use `src/fine_tuning/train.py` (or Notebook 02 for experiments). Leverages Hugging Face `transformers`, `peft` (LoRA), and `Trainer`. Can be run locally (GPU required) or orchestrated via the Vertex AI MLOps pipeline (`launch_llm_fine_tuning_job` component) for scalable cloud training.

    3.  **Output:** PEFT adapter weights saved to GCS (or locally).

*   **Configuration:** See `configs/ml_config.yaml` for fine-tuning parameters and `configs/agent_config.yaml` for referencing resulting models/endpoints. See Notebook 02 discussion in the project history/documentation for recommended hyperparameter starting points.



## MLOps Integration (Vertex AI & MLflow)



*   **Vertex AI Pipelines:** Orchestrate the ML workflow (data validation, training, evaluation, deployment). See `mlops/vertex_pipelines/`.

*   **Vertex AI Custom Training:** Execute containerized training scripts (`model.py`, `train.py`) scalably.

*   **Vertex AI Prediction Endpoints:** Host the trained XGBoost model and the fine-tuned LLM (base + adapter, using suitable serving containers like TGI/vLLM) for low-latency inference.

*   **Vertex AI Model Registry:** Store and version trained models (XGBoost and potentially LLM adapters/pointers).

*   **Vertex AI Experiments / MLflow:** Track hyperparameters, metrics, and artifacts for both XGBoost and LLM fine-tuning runs using MLflow (configure `MLFLOW_TRACKING_URI`).

*   **Vertex AI Model Monitoring:** Configure monitoring jobs to detect drift in predictions or skew in input features for the deployed XGBoost model (and potentially for LLM inputs/outputs if feasible).

*   **Automation:** Pipelines can be triggered manually, on a schedule (e.g., for retraining), or via Pub/Sub events (e.g., on new data arrival).



## API Details



*   Built with **FastAPI** (`src/proactive_disruption_agent/api/`).

*   Uses Pydantic models (`api/models.py`) for request/response validation.

*   Provides endpoints:

    *   `POST /analyze-event`: Asynchronously triggers the agent workflow for a given event. Returns an `analysis_id`.

    *   `GET /analysis/{analysis_id}`: Retrieves the status (PENDING, RUNNING, COMPLETED, FAILED) and results of a specific analysis.

*   Uses a background task runner for the potentially long-running agent execution.

*   Leverages FastAPI's dependency injection (`api/dependencies.py`) to manage the agent graph instance.

*   Uses a simple in-memory dictionary (`analysis_store`) for tracking job status (replace with Redis/database for production).



## Testing



*   Unit and integration tests are located in the `tests/` directory using `pytest`.

*   **`tests/agent/`**: Tests the LangGraph nodes and graph logic, heavily utilizing mocking (`unittest.mock`) for external dependencies (LLM, forecasting, RAG).

*   **`tests/forecasting/`**: Tests data preprocessing logic and the prediction function, potentially loading locally saved model artifacts.

*   **`tests/api/`**: Tests FastAPI endpoints using `TestClient`, mocking the agent execution or dependencies.

*   Run tests from the project root using:

    ```bash

    pytest tests/

    ```.



## Evaluation Metrics (Target / Illustrative Metrics)



*   Reduction in critical shipment delays: 18% (on pilot clients)

*   Response time improvement (proactive insight): Hours -> Minutes

*   Accuracy: LLM risk explanations verified by SMEs.

*   Qualitative feedback: Positive from pilot stakeholders (Ops VPs, Client Directors).



## Future Work



*   More sophisticated forecasting models (Deep Learning Time Series, ensemble methods).

*   Multi-agent systems (e.g., specialized agents for different disruption types or clients).

*   Reinforcement learning to optimize mitigation suggestions based on feedback/outcomes.

*   Integration with real-time shipment tracking data.

*   Full UI/Dashboard for visualizing risks, impacts, and suggested actions.

*   More robust error handling and state management (persistent job store).

*   Implementing LLM evaluation within the MLOps pipeline.



## License



This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.



## Acknowledgements



*   Langchain/LangGraph Team & Community

*   Hugging Face Community & Model Contributors (Google DeepMind)

*   MLflow Team

*   Google Cloud Vertex AI Team

*   GDELT Project and OpenWeatherMap