diff --git a/common_scaffold/prompts/prompt_builder.py b/common_scaffold/prompts/prompt_builder.py index fb83f9fa86..3eb9c57fbd 100644 --- a/common_scaffold/prompts/prompt_builder.py +++ b/common_scaffold/prompts/prompt_builder.py @@ -58,6 +58,7 @@ - query_db: {"tool": "query_db", "args": {"db_name": "some_db_name", "query": "SELECT * FROM some_table LIMIT some_limit;"}} # for SQL databases {"tool": "query_db", "args": {"db_name": "some_db_name", "query": "{\"collection\": \"some_collection\", \"filter\": {some_filter}, \"projection\": {some_projection}, \"limit\": some_limit}"}} # for MongoDB databases +{"tool": "query_db", "args": {"db_name": "some_db_name", "query": "{\"collection\": \"some_collection\", \"pipeline\": [some_aggregation_stages]}"}} # MongoDB aggregation - list_db: {"tool": "list_db", "args": {"db_name": "some_db_name"}} diff --git a/common_scaffold/tools/QueryDBTool.py b/common_scaffold/tools/QueryDBTool.py index b47e919441..1cd48368c7 100644 --- a/common_scaffold/tools/QueryDBTool.py +++ b/common_scaffold/tools/QueryDBTool.py @@ -46,7 +46,7 @@ def get_spec(self): }, "query": { "type": "string", - "description": "The query to execute. For SQL databases, provide a SQL query string. For MongoDB, provide a JSON string representing the query, including `collection` field, and optional `filter`, `projection`, and `limit` fields." + "description": "The query to execute. For SQL databases, provide a SQL query string. For MongoDB, provide a JSON string representing the query, including `collection` field, and optional `filter`, `projection`, and `limit` fields; or `collection` plus `pipeline` (a list of aggregation stages) to run an aggregation instead." } }, "required": ["db_name", "query"] diff --git a/common_scaffold/tools/db_utils/mongo_utils.py b/common_scaffold/tools/db_utils/mongo_utils.py index 4e7e0a6e48..2830ac03aa 100644 --- a/common_scaffold/tools/db_utils/mongo_utils.py +++ b/common_scaffold/tools/db_utils/mongo_utils.py @@ -1,4 +1,5 @@ import json +import os import pandas as pd from pymongo import MongoClient from common_scaffold.tools.db_utils import db_config @@ -17,6 +18,8 @@ # Reference: https://www.mongodb.com/docs/manual/administration/install-community/?linux-distribution=ubuntu&linux-package=default&operating-system=linux&search-linux=with-search-linux#start-mongodb-12 def load_db(dump_folder: str, db_name: str): + if os.getenv("DAB_PERSISTENT_DBS") and check_db_exists(db_name): + return if check_db_exists(db_name) == True: logging.getLogger(__name__).warning(f"Database '{db_name}' already exists. Cleaning up before loading dump.") clean_up(db_name) @@ -53,6 +56,8 @@ def check_db_exists(db_name): raise FatalError(f"MongoDB existence check error ({type(e).__name__}): {str(e)}") def clean_up(db_name): + if os.getenv("DAB_PERSISTENT_DBS"): + return logging.getLogger(__name__).debug(f"Cleaning up MongoDB database '{db_name}'") with MongoClient(db_config.MONGO_URI) as mongo_client: try: @@ -77,17 +82,21 @@ def check_args(db_client: dict, query: str): collection = query_json["collection"] filter = query_json.get("filter", {}) projection = query_json.get("projection", None) - limit = query_json.get("limit", 5) + limit = query_json.get("limit", None) # no silent cap: matches SQL, which returns every row + pipeline = query_json.get("pipeline", None) + if pipeline is not None and not isinstance(pipeline, list): + raise ValueError("Invalid Mongo query: `pipeline` must be a list of aggregation stages") return { "db_name": db_client['db_name'], "collection": collection, "filter": filter, "projection": projection, - "limit": limit + "limit": limit, + "pipeline": pipeline, } @staticmethod - def exec(db_name, collection, filter, projection, limit): + def exec(db_name, collection, filter, projection, limit, pipeline=None): with MongoClient(db_config.MONGO_URI) as mongo_client: db = mongo_client[db_name] @@ -95,10 +104,13 @@ def exec(db_name, collection, filter, projection, limit): raise ValueError(f"Collection does not exist: {collection}") try: - cursor = db[collection].find(filter, projection) - if limit is not None: - cursor = cursor.limit(limit) - result = list(cursor) + if pipeline is not None: + result = list(db[collection].aggregate(pipeline)) + else: + cursor = db[collection].find(filter, projection) + if limit is not None: + cursor = cursor.limit(limit) + result = list(cursor) except Exception as e: raise ValueError(f"MongoDB query execution error ({type(e).__name__}): {str(e)}") diff --git a/common_scaffold/tools/db_utils/postgres_utils.py b/common_scaffold/tools/db_utils/postgres_utils.py index 8b49e6e07a..f4855682c5 100644 --- a/common_scaffold/tools/db_utils/postgres_utils.py +++ b/common_scaffold/tools/db_utils/postgres_utils.py @@ -22,6 +22,8 @@ def pg_ident(identifier: str) -> str: def load_db(sql_file: str, db_name: str): + if os.getenv("DAB_PERSISTENT_DBS") and check_db_exists(db_name): + return if check_db_exists(db_name): logging.getLogger(__name__).warning(f"Database '{db_name}' already exists. Cleaning up before loading dump.") clean_up(db_name) @@ -91,6 +93,8 @@ def check_db_exists(db_name): def clean_up(db_name): + if os.getenv("DAB_PERSISTENT_DBS"): + return logging.getLogger(__name__).debug(f"Cleaning up PostgreSQL database '{db_name}'") try: conn = psycopg2.connect( diff --git a/query_bank_fresh_direct/db_config.yaml b/query_bank_fresh_direct/db_config.yaml new file mode 100644 index 0000000000..ac2b8bbbb5 --- /dev/null +++ b/query_bank_fresh_direct/db_config.yaml @@ -0,0 +1,21 @@ +db_clients: + core_banking: + db_type: postgres + db_name: bank_cbs + sql_file: query_dataset/cbs.sql + cards: + db_type: sqlite + db_path: query_dataset/cards.db + lending: + db_type: duckdb + db_path: query_dataset/lending.duckdb + payments: + db_type: sqlite + db_path: query_dataset/payments.db + crm: + db_type: mongo + db_name: bank_crm + dump_folder: query_dataset/bank_crm_dump + insurance: + db_type: duckdb + db_path: query_dataset/insurance.duckdb diff --git a/query_bank_fresh_direct/db_description.txt b/query_bank_fresh_direct/db_description.txt new file mode 100644 index 0000000000..c2d089c321 --- /dev/null +++ b/query_bank_fresh_direct/db_description.txt @@ -0,0 +1,55 @@ +You are working with six databases to solve this query. They are the operational systems of one retail bank in India. Amounts are in Indian rupees (INR). They are the bank's live systems: every transaction is visible in them as soon as it is posted. + +Here are the descriptions of these six databases: + +1. core_banking + - This database is stored in a PostgreSQL database and is the bank's core banking system: customers, their deposit accounts, account transactions, cheques, and fixed and recurring deposits. + - This database consists of six tables: + - customers: Bank customers + - Fields: customer_id, name, age, gender, senior_citizen, kyc_status, monthly_income, comm_line1, comm_line2, comm_city, comm_state, comm_pincode, perm_line1, perm_line2, perm_city, perm_state, perm_pincode (comm_ = communication address, perm_ = permanent address) + - accounts: Savings and current accounts + - Fields: account_id, customer_id, account_type, status, ifsc, home_branch, nominee_name, nominee_relationship, open_date, available_balance, hold_amount, currency + - account_holders: Holders of each account, in order (holder_seq 1 is the primary holder) + - Fields: account_id, holder_seq, holder_name + - transactions: Account transactions + - Fields: reference_id, account_id, txn_date, description, amount, direction, channel, balance_after + - cheques: Cheque leaves issued on accounts + - Fields: account_id, cheque_number, status + - deposits: Fixed deposits (kind FD) and recurring deposits (kind RD) + - Fields: deposit_id, payout_account, kind, principal, rate, tenure_months, start_date, maturity_date, maturity_amount, payout_type, maturity_instruction, status, nominee_name, nominee_relationship, monthly_installment, installments_paid, total_deposited, source_account + +2. cards + - This database is stored in a SQLite database and is the bank's card management system. + - This database consists of two tables: + - cards: Debit and credit cards + - Fields: card_id, account_no, card_type, network, status, credit_limit, available_limit + - card_controls: Per-card usage controls, one row per channel (atm, online, pos) and region (domestic, international) + - Fields: card_id, account_no, channel, region, enabled, daily_limit + +3. lending + - This database is stored in a DuckDB database and is the bank's loan management system. + - This database consists of one table: + - loans: Home, personal and gold loans + - Fields: loan_id, borrower_ref, loan_type, outstanding, rate, emi, remaining_tenure_months, next_due_date, status + +4. payments + - This database is stored in a SQLite database and is the bank's payments hub. + - This database consists of one table: + - mandates: Standing payment mandates (UPI autopay, NACH and similar) that debit an account + - Fields: mandate_id, debit_account, payee, amount, frequency, next_debit_date, status, mandate_type + +5. crm + - This database is stored in a MongoDB database and is the bank's customer relationship management system. + - This database consists of three collections: + - customers: Customer relationship data + - Fields: _id, segment, profession, sector, employer, marital_status, email, phone + - service_requests: Service requests raised by customers + - Fields: _id, request_id, category, description, related_transaction_id, account_id, status, priority, created_at, eta + - offers: Product offers shown to customers + - Fields: _id, customer, offer_id, type, title, detail, indicative + +6. insurance + - This database is stored in a DuckDB database and is the feed of insurance policies the bank distributes for its partner insurers. + - This database consists of one table: + - policies: Insurance policies held by bank customers + - Fields: policy_id, policyholder_ref, policy_name, insurer, sum_assured, premium, frequency, status, next_premium_date, nominee_name, nominee_relationship diff --git a/query_bank_fresh_direct/db_description_withhint.txt b/query_bank_fresh_direct/db_description_withhint.txt new file mode 100644 index 0000000000..e69de29bb2 diff --git a/query_bank_fresh_direct/query1/ground_truth.csv b/query_bank_fresh_direct/query1/ground_truth.csv new file mode 100644 index 0000000000..78f789d4f3 --- /dev/null +++ b/query_bank_fresh_direct/query1/ground_truth.csv @@ -0,0 +1 @@ +1113900.00 diff --git a/query_bank_fresh_direct/query1/query.json b/query_bank_fresh_direct/query1/query.json new file mode 100644 index 0000000000..0c771ff880 --- /dev/null +++ b/query_bank_fresh_direct/query1/query.json @@ -0,0 +1 @@ +"Current time: 2026-07-01 08:13 IST. What is the available balance of account SB4516247616 right now?" \ No newline at end of file diff --git a/query_bank_fresh_direct/query1/validate.py b/query_bank_fresh_direct/query1/validate.py new file mode 100644 index 0000000000..1e656414c7 --- /dev/null +++ b/query_bank_fresh_direct/query1/validate.py @@ -0,0 +1,81 @@ +"""Bank freshness question 1 (control, 2026-07-01 08:13). Generated by make_fresh_datasets.py; do not edit.""" +import re + +SPEC = ['money', 1113900.0, 1.0] +DATETIME = re.compile('\\b\\d{4}-\\d{2}-\\d{2}(?:[ T]\\d{1,2}:\\d{2}(?::\\d{2})?)?\\b|\\b\\d{1,2}:\\d{2}\\b') + +import re + +NUMBER = re.compile(r"(?