Secure account access and reorganize transfer module

ai_service.py

@@ -1,11 +1,18 @@
 import re
 import json
-import os
 from config import pc, index, EMBED_MODEL, hf_client, PROMPT, HF_MODEL
 from database import db_manager
+from transfers import (
+    prepare_transfer,
+    confirm_transfer,
+    cancel_transfer,
+    get_pending_transfer,
+    get_account_balance,
+)
 
 
 MODEL_NAME = HF_MODEL
+BASE_PROMPT = PROMPT or "You are a helpful banking customer service assistant."
 
 def clean_ai_response(text: str):
     if not text: return ""
@@ -50,9 +57,101 @@ TOOLS = [
                 "required": ["query"]
             }
         }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "check_account_balance",
+            "description": "Use this tool when the user wants to check their own account balance. The server identifies the current user from request context.",
+            "parameters": {
+                "type": "object",
+                "properties": {},
+                "required": []
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "prepare_money_transfer",
+            "description": "Use this tool when the user wants to transfer money. The server identifies the sender from request context, looks up the receiver by account serial ID, stores a pending transfer, and returns the receiver name for confirmation before any money is sent.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "receiver_serial_id": {
+                        "type": "string",
+                        "description": "The serial ID of the receiver account."
+                    },
+                    "amount": {
+                        "type": "number",
+                        "description": "Amount to transfer. Ask the user for it if missing."
+                    }
+                },
+                "required": ["receiver_serial_id"]
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "confirm_money_transfer",
+            "description": "Use this tool only after the user confirms that the receiver account is correct and the transfer should proceed.",
+            "parameters": {
+                "type": "object",
+                "properties": {},
+                "required": []
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "cancel_money_transfer",
+            "description": "Use this tool when the user says the receiver is wrong or wants to stop a pending money transfer.",
+            "parameters": {
+                "type": "object",
+                "properties": {},
+                "required": []
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "get_pending_money_transfer",
+            "description": "Use this tool to inspect the current pending transfer for the current user before asking for confirmation or when the user asks about the transfer details.",
+            "parameters": {
+                "type": "object",
+                "properties": {},
+                "required": []
+            }
+        }
     }
 ]
 
+
+async def run_tool(tool_name: str, args: dict, telegram_id: int):
+    if tool_name == "search_bank_knowledge":
+        return await search_bank_knowledge(args["query"])
+    if tool_name == "check_account_balance":
+        return json.dumps(get_account_balance(telegram_id), ensure_ascii=False)
+    if tool_name == "prepare_money_transfer":
+        return json.dumps(
+            prepare_transfer(
+                telegram_id=telegram_id,
+                receiver_serial_id=args["receiver_serial_id"],
+                amount=args.get("amount"),
+            ),
+            ensure_ascii=False,
+        )
+    if tool_name == "confirm_money_transfer":
+        return json.dumps(confirm_transfer(telegram_id), ensure_ascii=False)
+    if tool_name == "cancel_money_transfer":
+        return json.dumps(cancel_transfer(telegram_id), ensure_ascii=False)
+    if tool_name == "get_pending_money_transfer":
+        return json.dumps(get_pending_transfer(telegram_id), ensure_ascii=False)
+    return json.dumps({"success": False, "message": f"Unknown tool: {tool_name}"}, ensure_ascii=False)
+
 async def get_ai_response(user_query: str, telegram_id: int):
     conversation_history = []
     if db_manager:
@@ -63,7 +162,22 @@ async def get_ai_response(user_query: str, telegram_id: int):
                 role = "user" if msg['message_type'] == 'user' else "assistant"
                 conversation_history.append({"role": role, "content": msg['message_text']})
 
-    messages = [{"role": "system", "content": PROMPT}] + conversation_history + [{"role": "user", "content": user_query}]
+    transfer_instructions = (
+        f"Current user telegram_id is {telegram_id}. "
+        "The customer service system cannot create bank accounts. "
+        "The model must never choose, guess, extract, or override any telegram_id for tool calls. "
+        "Always act only for the current authenticated user from server-side request context. "
+        "If the user asks for another person's balance or provides another person's telegram ID, refuse and explain that you can only access the current user's own account. "
+        "If the sender does not already have a bank account, clearly tell them they must visit the bank to create one. "
+        "If the user asks for their balance, call check_account_balance. "
+        "For money transfers, first collect the receiver account serial ID and the amount if it is missing. "
+        "Then call prepare_money_transfer to fetch the receiver name and store the pending transfer. "
+        "Show the receiver name back to the user and ask for explicit confirmation. "
+        "Only call confirm_money_transfer after the user clearly agrees. "
+        "If the user rejects the receiver or wants to stop, call cancel_money_transfer. "
+        "Never claim a transfer is completed unless confirm_money_transfer returns success."
+    )
+    messages = [{"role": "system", "content": f"{BASE_PROMPT}\n\n{transfer_instructions}"}] + conversation_history + [{"role": "user", "content": user_query}]
     
     
     import asyncio
@@ -83,18 +197,19 @@ async def get_ai_response(user_query: str, telegram_id: int):
     completion = await loop.run_in_executor(None, lambda: call_hf(messages))
     response_message = completion.choices[0].message
 
-    # Handle tool call if model requests it
-    if response_message.tool_calls:
-        tool_call = response_message.tool_calls[0]
-        args = json.loads(tool_call.function.arguments)
-        tool_result = await search_bank_knowledge(args["query"])
+    for _ in range(4):
+        if not response_message.tool_calls:
+            break
 
         messages.append(response_message)
-        messages.append({
-            "role": "tool",
-            "tool_call_id": tool_call.id,
-            "content": tool_result
-        })
+        for tool_call in response_message.tool_calls:
+            args = json.loads(tool_call.function.arguments or "{}")
+            tool_result = await run_tool(tool_call.function.name, args, telegram_id)
+            messages.append({
+                "role": "tool",
+                "tool_call_id": tool_call.id,
+                "content": tool_result
+            })
 
         completion = await loop.run_in_executor(None, lambda: call_hf(messages))
         response_message = completion.choices[0].message
@@ -105,4 +220,4 @@ async def get_ai_response(user_query: str, telegram_id: int):
         db_manager.save_message(telegram_id, user_query, "user")
         db_manager.save_message(telegram_id, final_response, "assistant")
 
-    return final_response
+    return final_response

database_schema.sql

@@ -1,56 +0,0 @@
--- Conversation History Database Schema
--- SQLite Database Structure
-
--- Users table to store user information
-CREATE TABLE IF NOT EXISTS users (
-    id INTEGER PRIMARY KEY AUTOINCREMENT,
-    chat_id BIGINT UNIQUE NOT NULL,
-    username VARCHAR(255),
-    first_name VARCHAR(255),
-    last_name VARCHAR(255),
-    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-);
-
--- Messages table to store conversation history
-CREATE TABLE IF NOT EXISTS messages (
-    id INTEGER PRIMARY KEY AUTOINCREMENT,
-    chat_id BIGINT NOT NULL,
-    message_text TEXT NOT NULL,
-    message_type VARCHAR(20) NOT NULL CHECK (message_type IN ('user', 'assistant')),
-    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-    FOREIGN KEY (chat_id) REFERENCES users(chat_id) ON DELETE CASCADE
-);
-
--- Conversation sessions to group messages by session
-CREATE TABLE IF NOT EXISTS conversation_sessions (
-    id INTEGER PRIMARY KEY AUTOINCREMENT,
-    chat_id BIGINT NOT NULL,
-    session_start TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-    session_end TIMESTAMP NULL,
-    message_count INTEGER DEFAULT 0,
-    FOREIGN KEY (chat_id) REFERENCES users(chat_id) ON DELETE CASCADE
-);
-
--- Indexes for better performance
-CREATE INDEX IF NOT EXISTS idx_messages_chat_id_timestamp ON messages(chat_id, timestamp);
-CREATE INDEX IF NOT EXISTS idx_users_chat_id ON users(chat_id);
-CREATE INDEX IF NOT EXISTS idx_sessions_chat_id ON conversation_sessions(chat_id);
-
--- Trigger to update user's updated_at timestamp
-CREATE TRIGGER IF NOT EXISTS update_user_timestamp 
-    AFTER INSERT ON messages
-    FOR EACH ROW
-    BEGIN
-        UPDATE users SET updated_at = CURRENT_TIMESTAMP WHERE chat_id = NEW.chat_id;
-    END;
-
--- Trigger to update session message count
-CREATE TRIGGER IF NOT EXISTS update_session_count
-    AFTER INSERT ON messages
-    FOR EACH ROW
-    BEGIN
-        UPDATE conversation_sessions 
-        SET message_count = message_count + 1 
-        WHERE chat_id = NEW.chat_id AND session_end IS NULL;
-    END;

main.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI, Header, Request
+from schemas import AiTestRequest, WebhookData
 from security import validate_webhook_secret
-from telegram_handlers import telegram_webhook, WebhookData
+from telegram_handlers import telegram_webhook
 from utils import dns_test, test_ai_response
 
 app = FastAPI()
@@ -27,10 +28,11 @@ async def dns(
     validate_webhook_secret(request, x_telegram_bot_api_secret_token)
     return await dns_test()
 
-@app.get("/ai-test")
+@app.post("/ai-test")
 async def ai(
     request: Request,
+    data: AiTestRequest,
     x_telegram_bot_api_secret_token: str | None = Header(default=None),
 ):
     validate_webhook_secret(request, x_telegram_bot_api_secret_token)
-    return await test_ai_response()
+    return await test_ai_response(data.message, data.telegram_id)

rest.http

@@ -1,11 +1,10 @@
-GET https://codeboker-customer-service.hf.space/ai-test HTTP/1.1
+POST http://localhost:8000/ai-test HTTP/1.1
 Content-Type: application/json
+X-Telegram-Bot-Api-Secret-Token: {{$dotenv TELEGRAM_WEBHOOK_SECRET}}
 
 {
-    "message": {
-        "chat": {"id": 123},
-        "text": "هاي"
-    }
+    "message": "confirm then tell me how much that he has in his account",
+    "telegram_id": 12
 }
 ###
 GET https://codeboker-customer-service.hf.space/ai-test HTTP/1.1
@@ -13,9 +12,11 @@ GET https://codeboker-customer-service.hf.space/ai-test HTTP/1.1
 ###
 POST https://codeboker-customer-service.hf.space/webhook HTTP/1.1
 Content-Type: application/json
+X-Telegram-Bot-Api-Secret-Token: {{$dotenv TELEGRAM_WEBHOOK_SECRET}}
+
 {
     "message": {
         "chat": {"id": 123},
         "text": "السلام عليكم وؤحمة الله اريد ان استعلم عنلاالخدمات الالكترونية"
     }
-}
+}

schemas.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel
+
+
+class ChatInfo(BaseModel):
+    id: int
+    username: str = None
+    first_name: str = None
+    last_name: str = None
+
+
+class Message(BaseModel):
+    chat: ChatInfo
+    text: str = ""
+
+
+class WebhookData(BaseModel):
+    message: Message = None
+
+
+class AiTestRequest(BaseModel):
+    message: str
+    telegram_id: int = 12

telegram_handlers.py

@@ -1,9 +1,9 @@
-from pydantic import BaseModel
 import httpx
 import json
 from config import TELEGRAM_URL
 from ai_service import get_ai_response
 from database import db_manager
+from schemas import WebhookData
 
 TELEGRAM_IP = "149.154.167.220"
 MAX_TELEGRAM_MESSAGE_LENGTH = 4096
@@ -24,19 +24,6 @@ def _sanitize_telegram_text(text: str) -> str:
     return cleaned.strip()
 
 
-class ChatInfo(BaseModel):
-    id: int
-    username: str = None
-    first_name: str = None
-    last_name: str = None
-
-class Message(BaseModel):
-    chat: ChatInfo
-    text: str = ""
-
-class WebhookData(BaseModel):
-    message: Message = None
-
 async def telegram_webhook(data: WebhookData):
     try:
         if not data.message or not data.message.text:

transfers/__init__.py

@@ -0,0 +1,8 @@
+from .service import (
+    cancel_transfer,
+    confirm_transfer,
+    get_account_balance,
+    get_pending_transfer,
+    prepare_transfer,
+)
+

transfers/mock_bank_accounts.json

@@ -0,0 +1,68 @@
+{
+  "accounts": [
+    {
+      "telegram_id": 12,
+      "account_id": "ACC-00012",
+      "serial_id": "1001",
+      "name": "Test Sender",
+      "balance": 1900.0,
+      "currency": "YER"
+    },
+    {
+      "telegram_id": 991001,
+      "account_id": "ACC-991001",
+      "serial_id": "2001",
+      "name": "Ahmed Salem",
+      "balance": 1800.0,
+      "currency": "YER"
+    },
+    {
+      "telegram_id": 991002,
+      "account_id": "ACC-991002",
+      "serial_id": "2002",
+      "name": "Mona Ali",
+      "balance": 800.0,
+      "currency": "YER"
+    },
+    {
+      "telegram_id": 991003,
+      "account_id": "ACC-991003",
+      "serial_id": "2003",
+      "name": "Khaled Omar",
+      "balance": 300.0,
+      "currency": "YER"
+    }
+  ],
+  "transactions": [
+    {
+      "transaction_id": "TX-20260410193334671954",
+      "telegram_id": 12,
+      "sender_serial_id": "1001",
+      "receiver_serial_id": "2001",
+      "receiver_name": "Ahmed Salem",
+      "amount": 200.0,
+      "currency": "YER",
+      "created_at": "2026-04-10T19:33:34.671995"
+    },
+    {
+      "transaction_id": "TX-20260410193344247493",
+      "telegram_id": 12,
+      "sender_serial_id": "1001",
+      "receiver_serial_id": "2001",
+      "receiver_name": "Ahmed Salem",
+      "amount": 200.0,
+      "currency": "YER",
+      "created_at": "2026-04-10T19:33:44.247529"
+    },
+    {
+      "transaction_id": "TX-20260410193354495302",
+      "telegram_id": 12,
+      "sender_serial_id": "1001",
+      "receiver_serial_id": "2001",
+      "receiver_name": "Ahmed Salem",
+      "amount": 200.0,
+      "currency": "YER",
+      "created_at": "2026-04-10T19:33:54.495350"
+    }
+  ]
+}

transfers/pending_transfers.json

@@ -0,0 +1 @@
+{}

transfers/service.py

@@ -0,0 +1,257 @@
+import json
+from datetime import datetime
+from pathlib import Path
+from threading import Lock
+from typing import Any, Dict, Optional
+
+
+DATA_FILE = Path(__file__).with_name("mock_bank_accounts.json")
+PENDING_FILE = Path(__file__).with_name("pending_transfers.json")
+_LOCK = Lock()
+
+
+def _read_json(path: Path, default: Any) -> Any:
+    if not path.exists():
+        return default
+
+    with path.open("r", encoding="utf-8") as file:
+        return json.load(file)
+
+
+def _write_json(path: Path, data: Any) -> None:
+    with path.open("w", encoding="utf-8") as file:
+        json.dump(data, file, ensure_ascii=True, indent=2)
+
+
+def _load_bank_data() -> Dict[str, Any]:
+    return _read_json(DATA_FILE, {"accounts": [], "transactions": []})
+
+
+def _save_bank_data(data: Dict[str, Any]) -> None:
+    _write_json(DATA_FILE, data)
+
+
+def _load_pending_transfers() -> Dict[str, Any]:
+    return _read_json(PENDING_FILE, {})
+
+
+def _save_pending_transfers(data: Dict[str, Any]) -> None:
+    _write_json(PENDING_FILE, data)
+
+
+def _find_account_by_telegram_id(accounts: list[Dict[str, Any]], telegram_id: int) -> Optional[Dict[str, Any]]:
+    return next((account for account in accounts if account.get("telegram_id") == telegram_id), None)
+
+
+def _find_account_by_serial_id(accounts: list[Dict[str, Any]], serial_id: str) -> Optional[Dict[str, Any]]:
+    normalized_id = str(serial_id).strip()
+    return next((account for account in accounts if str(account.get("serial_id")) == normalized_id), None)
+
+
+def get_sender_account(telegram_id: int) -> Dict[str, Any]:
+    with _LOCK:
+        data = _load_bank_data()
+        sender = _find_account_by_telegram_id(data["accounts"], telegram_id)
+        if sender:
+            return {
+                "success": True,
+                "account": sender,
+            }
+
+        return {
+            "success": False,
+            "message": "You do not have an account in the bank system. Please visit the bank to create an account first.",
+        }
+
+
+def get_account_balance(telegram_id: int) -> Dict[str, Any]:
+    sender_result = get_sender_account(telegram_id)
+    if not sender_result["success"]:
+        return sender_result
+
+    account = sender_result["account"]
+    return {
+        "success": True,
+        "telegram_id": telegram_id,
+        "account_id": account["account_id"],
+        "serial_id": account["serial_id"],
+        "name": account["name"],
+        "balance": account.get("balance", 0.0),
+        "currency": account.get("currency", "YER"),
+    }
+
+
+def get_receiver_account_name(receiver_serial_id: str) -> Dict[str, Any]:
+    with _LOCK:
+        data = _load_bank_data()
+        receiver = _find_account_by_serial_id(data["accounts"], receiver_serial_id)
+
+        if not receiver:
+            return {
+                "success": False,
+                "message": f"No account was found for ID {receiver_serial_id}.",
+            }
+
+        return {
+            "success": True,
+            "serial_id": receiver["serial_id"],
+            "name": receiver["name"],
+            "account_id": receiver["account_id"],
+            "currency": receiver.get("currency", "YER"),
+        }
+
+
+def prepare_transfer(telegram_id: int, receiver_serial_id: str, amount: Optional[float] = None) -> Dict[str, Any]:
+    with _LOCK:
+        data = _load_bank_data()
+        sender = _find_account_by_telegram_id(data["accounts"], telegram_id)
+        if not sender:
+            return {
+                "success": False,
+                "message": "You do not have an account in the bank system. Please visit the bank to create an account first.",
+            }
+
+        receiver = _find_account_by_serial_id(data["accounts"], receiver_serial_id)
+        if not receiver:
+            return {
+                "success": False,
+                "message": f"No account was found for ID {receiver_serial_id}.",
+            }
+
+        if receiver["serial_id"] == sender["serial_id"]:
+            return {
+                "success": False,
+                "message": "You cannot transfer money to the same account.",
+            }
+
+        pending_transfers = _load_pending_transfers()
+        pending_payload = {
+            "telegram_id": telegram_id,
+            "sender_name": sender["name"],
+            "sender_serial_id": sender["serial_id"],
+            "receiver_name": receiver["name"],
+            "receiver_serial_id": receiver["serial_id"],
+            "receiver_account_id": receiver["account_id"],
+            "amount": amount,
+            "currency": sender.get("currency", "YER"),
+            "created_at": datetime.utcnow().isoformat(),
+        }
+        pending_transfers[str(telegram_id)] = pending_payload
+        _save_pending_transfers(pending_transfers)
+
+        return {
+            "success": True,
+            "pending_transfer": pending_payload,
+            "message": f"Receiver found: {receiver['name']}. Waiting for user confirmation.",
+        }
+
+
+def get_pending_transfer(telegram_id: int) -> Dict[str, Any]:
+    with _LOCK:
+        pending_transfers = _load_pending_transfers()
+        pending_transfer = pending_transfers.get(str(telegram_id))
+
+        if not pending_transfer:
+            return {
+                "success": False,
+                "message": "No pending transfer was found for this Telegram user.",
+            }
+
+        return {
+            "success": True,
+            "pending_transfer": pending_transfer,
+        }
+
+
+def confirm_transfer(telegram_id: int) -> Dict[str, Any]:
+    with _LOCK:
+        data = _load_bank_data()
+        pending_transfers = _load_pending_transfers()
+        pending_transfer = pending_transfers.get(str(telegram_id))
+
+        if not pending_transfer:
+            return {
+                "success": False,
+                "message": "There is no pending transfer to confirm.",
+            }
+
+        sender = _find_account_by_telegram_id(data["accounts"], telegram_id)
+        receiver = _find_account_by_serial_id(data["accounts"], pending_transfer["receiver_serial_id"])
+
+        if not sender or not receiver:
+            return {
+                "success": False,
+                "message": "The sender or receiver account could not be found.",
+            }
+
+        amount = pending_transfer.get("amount")
+        if amount is None:
+            return {
+                "success": False,
+                "message": "The transfer amount is missing. Ask the user for the amount before confirming.",
+            }
+
+        try:
+            amount_value = float(amount)
+        except (TypeError, ValueError):
+            return {
+                "success": False,
+                "message": "The transfer amount is invalid.",
+            }
+
+        if amount_value <= 0:
+            return {
+                "success": False,
+                "message": "The transfer amount must be greater than zero.",
+            }
+
+        if float(sender.get("balance", 0.0)) < amount_value:
+            return {
+                "success": False,
+                "message": f"Insufficient balance. Available balance is {sender.get('balance', 0.0):.2f} {sender.get('currency', 'YER')}.",
+            }
+
+        sender["balance"] = round(float(sender["balance"]) - amount_value, 2)
+        receiver["balance"] = round(float(receiver.get("balance", 0.0)) + amount_value, 2)
+
+        transaction = {
+            "transaction_id": f"TX-{datetime.utcnow().strftime('%Y%m%d%H%M%S%f')}",
+            "telegram_id": telegram_id,
+            "sender_serial_id": sender["serial_id"],
+            "receiver_serial_id": receiver["serial_id"],
+            "receiver_name": receiver["name"],
+            "amount": amount_value,
+            "currency": sender.get("currency", "YER"),
+            "created_at": datetime.utcnow().isoformat(),
+        }
+        data.setdefault("transactions", []).append(transaction)
+        _save_bank_data(data)
+
+        pending_transfers.pop(str(telegram_id), None)
+        _save_pending_transfers(pending_transfers)
+
+        return {
+            "success": True,
+            "transaction": transaction,
+            "sender_balance": sender["balance"],
+            "message": f"Transfer completed successfully to {receiver['name']}.",
+        }
+
+
+def cancel_transfer(telegram_id: int) -> Dict[str, Any]:
+    with _LOCK:
+        pending_transfers = _load_pending_transfers()
+        removed = pending_transfers.pop(str(telegram_id), None)
+        _save_pending_transfers(pending_transfers)
+
+        if not removed:
+            return {
+                "success": False,
+                "message": "There is no pending transfer to cancel.",
+            }
+
+        return {
+            "success": True,
+            "message": "The pending transfer has been canceled.",
+        }
+

utils.py

@@ -7,10 +7,14 @@ async def dns_test():
     except Exception as e:
         return {"error": str(e)}
 
-async def test_ai_response():
+async def test_ai_response(message: str, telegram_id: int):
     try:
         from ai_service import get_ai_response
-        response = await get_ai_response("عن ماذا كنت اسال", 12)
-        return {"response": response}
+        response = await get_ai_response(message, telegram_id)
+        return {
+            "message": message,
+            "telegram_id": telegram_id,
+            "response": response,
+        }
     except Exception as e:
         return {"error": str(e)}