import re
import json
from config import pc, index, EMBED_MODEL, hf_client, PROMPT, HF_MODEL, TRANSFER_PROMPT
from database import db_manager
from transfers import (
prepare_transfer,
confirm_transfer,
cancel_transfer,
get_pending_transfer,
get_account_balance,
get_sender_account,
)
MODEL_NAME = HF_MODEL
BASE_PROMPT = PROMPT or "You are a helpful banking customer service assistant."
TRANSFER_PROMPT = TRANSFER_PROMPT or (
"Current user telegram_id is {telegram_id}. "
"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 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. "
"If a tool result returns 'need_user_name': true, ask the user for their name and then call create_illusion_account with their name. "
"If a tool result indicates an illusion account was created (contains 'is_illusion': true or mentions testing), inform the user that an illusion account with 2000 YER balance was created for testing purposes. "
"If confirm_money_transfer returns a result with 'is_illusion': true, make sure to display the testing disclaimer message to the user."
)
def clean_ai_response(text: str) -> str:
if not text:
return ""
text = re.sub(r'.*?', '', text, flags=re.DOTALL)
text = re.sub(r'', '\n', text, flags=re.IGNORECASE)
text = re.sub(r']*>', '\n', text, flags=re.IGNORECASE)
text = re.sub(r'<[^>]+>', '', text)
text = re.sub(r'^\|.*\|\s*$', '', text, flags=re.MULTILINE)
text = re.sub(r'^[\s|:-]+$', '', text, flags=re.MULTILINE)
text = re.sub(r'^#{1,6}\s*', '', text, flags=re.MULTILINE)
text = re.sub(r'\n{3,}', '\n\n', text)
return text.strip()
async def search_bank_knowledge(query: str):
query_embedding = pc.inference.embed(
model=EMBED_MODEL,
inputs=[query],
parameters={"input_type": "query"}
)
search_results = index.query(
vector=query_embedding[0].values,
top_k=3,
include_metadata=True
)
return "\n".join([res.metadata['original_text'] for res in search_results.matches])
TOOLS = [
{
"type": "function",
"function": {
"name": "search_bank_knowledge",
"description": "Use this tool to search the official Hadhramout Bank profile for accurate information about services, organizational structure, capital, and policies.",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query (e.g., 'What is Hadhramout Bank capital?' or 'individual services')."
}
},
"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": []
}
}
},
{
"type": "function",
"function": {
"name": "create_illusion_account",
"description": "Use this tool when the user needs an illusion account created for testing purposes. Requires the user's name.",
"parameters": {
"type": "object",
"properties": {
"user_name": {
"type": "string",
"description": "The name of the user for the illusion account."
}
},
"required": ["user_name"]
}
}
}
]
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)
if tool_name == "create_illusion_account":
return json.dumps(get_sender_account(telegram_id, args["user_name"]), 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:
raw_history = await db_manager.get_conversation_history(telegram_id, limit=6)
raw_history.reverse()
for msg in raw_history:
if msg.get('message_text'):
role = "user" if msg['message_type'] == 'user' else "assistant"
conversation_history.append({"role": role, "content": msg['message_text']})
transfer_instructions = TRANSFER_PROMPT.format(telegram_id=telegram_id)
messages = [{"role": "system", "content": f"{BASE_PROMPT}\n\n{transfer_instructions}"}] + conversation_history + [{"role": "user", "content": user_query}]
import asyncio
loop = asyncio.get_event_loop()
def call_hf(msgs):
return hf_client.chat.completions.create(
model=MODEL_NAME,
messages=msgs,
tools=TOOLS,
tool_choice="auto",
temperature=0.1,
max_tokens=800
)
completion = await loop.run_in_executor(None, lambda: call_hf(messages))
response_message = completion.choices[0].message
for _ in range(4):
if not response_message.tool_calls:
break
messages.append(response_message)
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
return clean_ai_response(response_message.content if response_message.content else "")