test_mistral.py

from langchain_mistralai import ChatMistralAI

model_name = "devstral-small-2507"
tools = [{
    "description": "Retrieve and display the current account balance, payment due dates, and payment method information for the customer's telecom account.",
    "properties": {
      "account_number": {
        "description": "The customer's unique account number or identifier.",
        "type": "string",
        "title": "Account_Number"
      },
      "include_payment_history": {
        "description": "Whether to include recent payment transaction history in the response.",
        "type": "boolean",
        "title": "Include_Payment_History"
      },
      "payment_history_months": {
        "description": "Number of months of payment history to retrieve if payment history is included.",
        "type": "string",
        "title": "Payment_History_Months",
        "enum": [
          "1",
          "3",
          "6",
          "12"
        ]
      }
    },
    "required": [
      "account_number"
    ],
    "title": "check_account_balance",
    "type": "object"
  },
  {
    "description": "Retrieves detailed transaction history for a customer's bank account with filtering and sorting options to help customers track their spending and account activity.",
    "properties": {
      "account_number": {
        "description": "The bank account number for which to retrieve transaction history.",
        "type": "string",
        "title": "Account_Number"
      },
      "start_date": {
        "description": "The start date for the transaction history search in YYYY-MM-DD format.",
        "type": "string",
        "title": "Start_Date"
      },
      "end_date": {
        "description": "The end date for the transaction history search in YYYY-MM-DD format.",
        "type": "string",
        "title": "End_Date"
      },
      "transaction_type": {
        "description": "Filter transactions by specific type.",
        "type": "string",
        "title": "Transaction_Type",
        "enum": [
          "all",
          "deposit",
          "withdrawal",
          "transfer",
          "payment",
          "fee",
          "interest"
        ]
      },
      "limit": {
        "description": "Maximum number of transactions to return, defaults to 50 if not specified.",
        "type": "integer",
        "title": "Limit"
      },
      "sort_order": {
        "description": "Sort order for the transaction results.",
        "type": "string",
        "title": "Sort_Order",
        "enum": [
          "newest_first",
          "oldest_first",
          "amount_high_to_low",
          "amount_low_to_high"
        ]
      }
    },
    "required": [
      "account_number",
      "start_date",
      "end_date"
    ],
    "title": "get_transaction_history",
    "type": "object"
  }
]

model_params = {"temperature": 0.0, "max_tokens": 4000}
llm = ChatMistralAI(model_name=model_name, **model_params).bind_tools(tools)

# call the llm with the tools
system_prompt = """You are a helpful assistant that can use the tools given to help the user. You can use multiple tools to answer the user's query."""
query = "What is the balance of my account 123 and 345. Also give transaction history for the last 1 week"

pprint(f"Query: {query}")
messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": query},
]

response = llm.invoke(messages)

print("\nModel Name:", model_name)
print("\nResponse Content:")
pprint(response.content)
print("\nTool Calls:")
pprint(response.tool_calls)
print("\nUsage Metadata:")
pprint(response.usage_metadata)
input_tokens, output_tokens = LLMHandler.get_token_usage_info(response)
print(f"\nToken Usage: input_tokens={input_tokens}, output_tokens={output_tokens}")