client.py•6.31 kB
import asyncio
import os
import json
from typing import Optional, List, Dict, Any
from contextlib import AsyncExitStack
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from dotenv import load_dotenv
load_dotenv()
class NtropyMCPClient:
"""Client for testing the Ntropy MCP server"""
def __init__(self):
"""Initialize the client"""
self.session: Optional[ClientSession] = None
self.exit_stack = AsyncExitStack()
async def connect_to_server(self):
"""Connect to the Ntropy MCP server using uvx"""
api_key = os.environ.get("NTROPY_API_KEY")
# Create server parameters
server_params = StdioServerParameters(
command="uvx",
args=["ntropy-mcp", "--api-key", api_key],
env=None
)
# Set up the connection
stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
read, write = stdio_transport
self.session = await self.exit_stack.enter_async_context(ClientSession(read, write))
# Initialize the connection
await self.session.initialize()
# List available tools
response = await self.session.list_tools()
tools = response.tools
print("\nAvailable tools:")
for tool in tools:
print(f"- {tool.name}: {tool.description}")
return tools
async def call_tool(self, tool_name: str, **kwargs):
"""Call a tool on the server
Args:
tool_name: Name of the tool to call
**kwargs: Arguments to pass to the tool
Returns:
The result of the tool call
"""
if not self.session:
raise ValueError("Not connected to server. Call connect_to_server first.")
print(f"\nCalling tool: {tool_name}")
print(f"Arguments: {json.dumps(kwargs, indent=2)}")
result = await self.session.call_tool(tool_name, kwargs)
print("\nResult:")
print(json.dumps(result.content, indent=2))
return result.content
async def check_connection(self):
"""Check connection to the Ntropy API"""
return await self.call_tool("check_connection")
async def create_account_holder(self, id: str, type: str, name: str):
"""Create an account holder"""
return await self.call_tool("create_account_holder", id=id, type=type, name=name)
async def enrich_transaction(self, id: str, description: str, date: str,
amount: float, entry_type: str, currency: str,
account_holder_id: str, country: str = None):
"""Enrich a transaction"""
args = {
"id": id,
"description": description,
"date": date,
"amount": amount,
"entry_type": entry_type,
"currency": currency,
"account_holder_id": account_holder_id
}
if country:
args["country"] = country
return await self.call_tool("enrich_transaction", **args)
async def list_transactions(self, account_holder_id: str, limit: int = 10, offset: int = 0):
"""List transactions for an account holder"""
return await self.call_tool("list_transactions",
account_holder_id=account_holder_id,
limit=limit,
offset=offset)
async def bulk_enrich_transactions(self, transactions: List[Dict[str, Any]]):
"""Bulk enrich multiple transactions"""
return await self.call_tool("bulk_enrich_transactions", transactions=transactions)
async def cleanup(self):
"""Clean up resources"""
await self.exit_stack.aclose()
async def run():
"""Run a demonstration of the Ntropy MCP client"""
client = NtropyMCPClient()
try:
# Connect to the server using uvx
await client.connect_to_server()
# Check connection to Ntropy API
print("\n1. Checking connection to Ntropy API...")
await client.check_connection()
# Create an account holder
print("\n2. Creating account holder...")
account_holder = await client.create_account_holder(
id="test_user_123",
type="individual",
name="Test User"
)
account_holder_id = account_holder.get("id", "test_user_123")
# Enrich a single transaction
print("\n3. Enriching a single transaction...")
await client.enrich_transaction(
id="tx_001",
description="AMAZON.COM*MK1AB6TE1",
date="2023-05-15",
amount=-29.99,
entry_type="debit",
currency="USD",
account_holder_id=account_holder_id,
country="US"
)
# Enrich a batch of transactions
print("\n4. Bulk enriching multiple transactions...")
transactions = [
{
"id": "tx_002",
"description": "NETFLIX.COM",
"date": "2023-05-16",
"amount": -13.99,
"entry_type": "debit",
"currency": "USD",
"account_holder_id": account_holder_id
},
{
"id": "tx_003",
"description": "Starbucks Coffee",
"date": "2023-05-17",
"amount": -5.65,
"entry_type": "debit",
"currency": "USD",
"account_holder_id": account_holder_id
}
]
await client.bulk_enrich_transactions(transactions)
# List transactions for the account holder
print("\n5. Listing transactions for account holder...")
await client.list_transactions(account_holder_id)
except Exception as e:
print(f"\nError: {str(e)}")
finally:
# Clean up resources
await client.cleanup()
if __name__ == "__main__":
asyncio.run(run())