milanchndr/Supply-Chain-AI-Agent

1# backend/app.py2import sys3import os4sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '.')))5from flask import Flask, request, jsonify6from utils.langchain_setup import setup_qa_chain7from sqlalchemy import create_engine, text8from config import BEDROCK_API_KEY9 10app = Flask(__name__)11 12# Initialize QA chain for document queries13try:14    qa_chain = setup_qa_chain()15except Exception as e:16    print(f"Error setting up QA chain: {e}")17    qa_chain = None18 19# Initialize database connection20try:21    engine = create_engine('postgresql://postgres:123@localhost:5432/hackathon_db')22except Exception as e:23    print(f"Error setting up database: {e}")24    engine = None25 26def is_document_query(question):27    """Simple heuristic to classify query type."""28    doc_keywords = ["policy", "guidelines", "procedure", "write-off", "ethical"]29    return any(keyword in question.lower() for keyword in doc_keywords)30 31@app.route('/')32def home():33    return "AI Supply Chain Agent Backend"34 35@app.route('/query', methods=['POST'])36def query():37    """Handle document or database queries."""38    if not qa_chain or not engine:39        return jsonify({"error": "Backend not initialized"}), 50040 41    data = request.json42    question = data.get("question", "").strip()43    if not question:44        return jsonify({"error": "No question provided"}), 40045 46    try:47        if is_document_query(question):48            # Document-based query49            result = qa_chain.invoke({"query": question})50            answer = result["result"]51            sources = [{"content": doc.page_content, "source": doc.metadata["source"]} for doc in result["source_documents"]]52            return jsonify({"answer": answer, "sources": sources, "type": "document"})53        else:54            # Database query (adjusted for DataCo dataset)55            with engine.connect() as connection:56                query = text('SELECT SUM("Order Item Quantity") as total_inventory FROM supply_chain WHERE "Region" = :region')57                result = connection.execute(query, {"region": "Southwest"}).fetchone()58                answer = f"Total inventory in Southwest region: {result[0] if result else 0} units"59                return jsonify({"answer": answer, "sources": [], "type": "database"})60    except Exception as e:61        return jsonify({"error": f"Error processing query: {str(e)}"}), 50062 63if __name__ == "__main__":64    app.run(debug=True, host="0.0.0.0", port=5000)