milanchndr/Supply-Chain-AI-Agent
AI Supply Chain Agent
Project Overview
The AI Supply Chain Agent is a sophisticated backend application designed to provide intelligent insights and perform actions related to supply chain management. It leverages a multi-tool Langchain agent capable of querying structured databases (PostgreSQL via Supabase), searching through unstructured PDF documents (company policies, guidelines), and accessing external information via web search. The agent implements Role-Based Access Control (RBAC) to ensure users only access data and functionalities appropriate for their roles and regions.
The system is built with Flask for the backend API, SQLAlchemy for database interactions, and custom Langchain components interacting with Amazon Bedrock models (via a Lambda proxy) for LLM and embedding capabilities.
Features
- Multi-Tool Agent: Utilizes a Langchain ReAct agent with the following tools:
- DocumentPolicySearch: Answers questions based on a corpus of PDF documents using a RetrievalQA chain (FAISS vector store and Bedrock embeddings).
- SupplyChainDatabaseQuery: Converts natural language questions into SQL queries to fetch data from a PostgreSQL database, respecting user roles and regions.
- ExternalWebSearch: Uses DuckDuckGo to find answers to general knowledge or current event questions.
- Role-Based Access Control (RBAC):
- User authentication via JWT.
- Application-level role and region checks.
- Database-level Row-Level Security (RLS) to filter data based on user's region and role (requires RLS policies to be set up in Supabase).
- Bedrock Integration (via Lambda):
- Utilizes Amazon Bedrock's Claude (e.g., Sonnet 3.5) for LLM tasks (SQL generation, agent reasoning).
- Uses Amazon Bedrock's Titan Multimodal Embeddings (e.g.,
amazon-embedding-v2) for document embeddings. - Interactions with Bedrock are proxied through a configurable AWS Lambda function.
- Document Processing:
- Extracts text from PDF documents.
- Chunks text and generates embeddings.
- Stores embeddings in a FAISS vector store for efficient similarity search.
- Database Management:
- Script (
load_db.py) to initialize the PostgreSQL database, create tables (supply_chain,users,audit_logs), and load sample data from a CSV. - API Endpoints:
/login: User authentication./query: Main endpoint to interact with the AI agent./health: Health check for the application.- Audit Logging: Records user queries and agent interactions in an
audit_logstable.
Project Structure
.
├── backend/
│ ├── .env # Environment variables (GITIGNORED!)
│ ├── agent_tools.py # Custom Langchain tools
│ ├── app.py # Flask application, API endpoints
│ ├── config.py # Application configuration, Bedrock IDs, DATABASE_URL
│ ├── logger_config.py # Logging setup
│ ├── requirements.txt # Python dependencies
│ ├── utils/ # Utility modules
│ │ ├── __init__.py
│ │ ├── agent_handler.py # Agent creation and prompt logic
│ │ ├── bedrock_utils.py # Bedrock LLM and Embedding wrappers (Lambda interaction)
│ │ ├── db_utils.py # Database schema introspection
│ │ ├── embed_documents.py# Script to process PDFs and create embeddings
│ │ ├── langchain_setup.py# Initializes QA chain and LLM instances
│ │ ├── load_db.py # Script to load data into the database
│ │ └── text_to_sql_utils.py # Text-to-SQL generation logic
│ ├── documents/ # Store PDF documents here for embedding
│ │ └── DataCoSupplyChainDataset.csv # Sample CSV for database loading
│ ├── embeddings/ # Stores FAISS index (GITIGNORED or managed via artifact store)
│ └── (other_python_files...)
├── .gitignore
└── README.mdSetup and Installation
Prerequisites
- Python 3.9+
- PostgreSQL database (Supabase project recommended)
- Access to Amazon Bedrock (or a Lambda function proxying Bedrock)
- Git
1. Clone the Repository
git clone <your-repository-url>
cd <your-repository-name>2. Set Up Environment Variables
Navigate to the backend directory and create a .env file:
cd backend
cp .env.example .env # If you have an example file, otherwise create .env manuallyEdit the .env file with your specific configurations:
# Bedrock Configuration
BEDROCK_API_KEY="your_bedrock_or_lambda_api_key"
LAMBDA_API_URL="your_aws_lambda_url_for_bedrock_proxy"
# BEDROCK_LLM_MODEL_ID="claude-3.5-sonnet" # Defined in config.py, can be overridden here if needed
# BEDROCK_EMBEDDING_MODEL_ID="amazon-embedding-v2" # Defined in config.py
# Database Configuration
SUPABASE_URL="postgresql://user:password@host:port/database" # Your Supabase DB connection string
DATABASE_URL=${SUPABASE_URL} # Can be the same or a different pooler URL
# JWT Configuration
JWT_SECRET_KEY="your_super_secret_jwt_key" # Change this to a strong random string
# Optional: Path to CSV for DB loading (relative to backend directory)
# CSV_PATH_FOR_DB_LOAD="../documents/DataCoSupplyChainDataset.csv" # Default is in config.py
# Optional: Logging Level
# LOG_LEVEL="INFO" # (DEBUG, INFO, WARNING, ERROR)Important: Ensure .env is listed in your .gitignore file to prevent committing secrets.
3. Set Up Python Virtual Environment and Install Dependencies
From the backend directory:
python -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate
pip install -r requirements.txt4. Set Up Supabase Database
- Ensure your Supabase project is running and you have the connection string.
- RLS Policies (Crucial for RBAC on Database): You need to set up Row-Level Security policies on your
supply_chaintable in Supabase. Example policies are discussed in the project and usually involve checkingcurrent_setting('request.jwt.claims.app_metadata.role', true)andcurrent_setting('request.jwt.claims.app_metadata.region', true). Refer to the Supabase RLS documentation and the RLS policy examples developed during this project. This typically involves:
-- Enable RLS on the table
ALTER TABLE supply_chain ENABLE ROW LEVEL SECURITY;
ALTER TABLE supply_chain FORCE ROW LEVEL SECURITY; -- Recommended
-- Helper functions for RLS (create these in your DB)
CREATE OR REPLACE FUNCTION public.get_jwt_claim_role() RETURNS TEXT ... STABLE SECURITY DEFINER;
CREATE OR REPLACE FUNCTION public.get_jwt_claim_region() RETURNS TEXT ... STABLE SECURITY DEFINER;
GRANT EXECUTE ON FUNCTION public.get_jwt_claim_role() TO authenticated;
GRANT EXECUTE ON FUNCTION public.get_jwt_claim_region() TO authenticated;
-- Example RLS Policy
CREATE POLICY "User region-based access on supply_chain"
ON supply_chain FOR SELECT USING (
(public.get_jwt_claim_role() = LOWER('Global Operations Manager')) OR
(LOWER(order_country) = public.get_jwt_claim_region())
);Ensure the helper functions and the policy are correctly implemented in your Supabase SQL editor.
- Indexes: Create necessary indexes for performance. Some are created by
load_db.py, but you can add more via the SQL editor.
CREATE INDEX IF NOT EXISTS idx_supply_chain_lower_order_country ON supply_chain (LOWER(order_country));
-- Add other relevant indexes on frequently queried columns5. Load Initial Data and Create Users
From the backend/utils directory (or adjust path if running from backend):
# Make sure your virtual environment is active
# Navigate to the directory containing load_db.py if not already in backend/utils/
# cd backend/utils (if load_db.py is there)
# python load_db.py
# Or if load_db.py is in backend/ and you run from project root:
# python backend/load_db.pyThis script will:
- Create the
supply_chain,users, andaudit_logstables. - Load data from
DataCoSupplyChainDataset.csvintosupply_chain. - Create a sample user (e.g.,
planner_india/mypassword). You can add more users directly in your Supabaseuserstable or modify the script.
6. Generate Document Embeddings
Place your PDF documents into the backend/documents/ folder. Then, run the embedding script from the backend/utils directory (or adjust path):
# Make sure your virtual environment is active
# Navigate to the directory containing embed_documents.py
# cd backend/utils
# python embed_documents.py
# Or if embed_documents.py is in backend/ and you run from project root:
# python backend/embed_documents.pyThis will create a backend/embeddings/ folder containing the FAISS index. Ensure this path aligns with what langchain_setup.py expects.
Running the Application
From the backend directory (ensure your virtual environment is active):
flask run --host=0.0.0.0 --port=5000
# Or, if app.py has app.run(debug=True...):
# python app.pyThe application will be available at http://localhost:5000.
API Endpoints
- `POST /login`:
- Authenticates a user.
- Request Body:
{"username": "your_username", "password": "your_password"} - Response:
{"access_token": "your_jwt_token"} - `POST /query`:
- Sends a question to the AI agent.
- Headers:
Authorization: Bearer <your_jwt_token> - Request Body:
{"question": "Your natural language question here"} - Response:
{"answer": "Agent's answer", "type": "agent_multi_tool"} - `GET /health`:
- Checks the health and initialization status of backend components.
Usage Example (e.g., using curl or Postman)
- Login:
curl -X POST http://localhost:5000/login \
-H "Content-Type: application/json" \
-d '{"username": "planner_india", "password": "mypassword"}'(Save the returned `access_token`)
- Query the Agent:
curl -X POST http://localhost:5000/query \
-H "Content-Type: application/json" \
-H "Authorization: Bearer <PASTE_YOUR_ACCESS_TOKEN_HERE>" \
-d '{"question": "What is the average sales per customer in India?"}'Key Configuration Points
- `backend/config.py`: Defines Bedrock model IDs, role permissions (
ROLES_PERMISSIONS), and default paths. - `backend/.env`: Stores all secrets and environment-specific URLs.
- RLS Policies in Supabase: Essential for database security and multi-tenancy/regional access.
- Lambda Proxy for Bedrock: The
LAMBDA_API_URLmust point to a working Lambda function that can invoke Bedrock models using the providedBEDROCK_API_KEY.
Troubleshooting
- "Function does not exist" errors (PostgreSQL): Ensure helper functions for RLS (
get_jwt_claim_role,get_jwt_claim_region) are created in the correct schema (usuallypublic) and that theauthenticatedrole (or other relevant roles) haveEXECUTEpermissions on them. RLS policies should use schema-qualified function names (e.g.,public.get_jwt_claim_role()). - RLS Not Filtering Data:
- Verify RLS is enabled on the table (
ALTER TABLE ... ENABLE ROW LEVEL SECURITY;). - Consider
ALTER TABLE ... FORCE ROW LEVEL SECURITY;. - Double-check the logic in your RLS policy and the values returned by helper functions within a simulated user session.
- Ensure data consistency between
users.regionandsupply_chain.order_country(or a similar column). - Timeouts ("Gateway Timeout", "Upstream Timeout"):
- Often due to RLS policies being inefficient on large tables. Optimize RLS using
STABLEhelper functions and appropriate indexes (e.g., onLOWER(order_country)). - Consider increasing statement timeouts in Supabase connection pooler settings for debugging, but aim to optimize queries/RLS for production.
- `current_setting('request.jwt.claims...', true)` returning NULL: This was a major debug point. Ensure the JWT is correctly set in the session before RLS policies try to read it. The order of
SET ROLEandset_configin your application's database interaction matters. Using explicit JSONB functions (::jsonb -> 'path' ->> 'key') in RLS policies oncurrent_setting('request.jwt.claims', true)is more robust than relying oncurrent_setting('request.jwt.claims.path.key', true). - Embedding Failures: Check
BEDROCK_EMBEDDING_MODEL_ID,LAMBDA_API_URL, andBEDROCK_API_KEY. Ensure PDF documents are in thebackend/documentsfolder.
