TradingAgents/tradingagents/agents/trader/paper_trading.py

386 lines
11 KiB
Python

"""
Paper Trading Integration - Webull API
Connects to Webull's paper trading environment to execute trades.
Requires Webull credentials and 2FA setup.
Installation:
pip install webull
Usage:
from tradingagents.agents.trader.paper_trading import PaperTrader
trader = PaperTrader(
email="your_email@example.com",
password="your_password",
is_paper=True,
)
# Buy
result = trader.place_buy_order("AAPL", 10, limit_price=150.0)
# Sell
result = trader.place_sell_order("AAPL", 5, limit_price=152.0)
# Get positions
positions = trader.get_positions()
"""
import os
import json
from typing import Optional, Dict, List
from datetime import datetime
import logging
logger = logging.getLogger(__name__)
class PaperTrader:
"""
Webull Paper Trading Interface
Handles buy/sell orders, position tracking, and account management
in Webull's paper trading environment.
"""
def __init__(
self,
email: Optional[str] = None,
password: Optional[str] = None,
is_paper: bool = True,
mfa_code: Optional[str] = None,
):
"""
Initialize Webull paper trader.
Args:
email: Webull account email
password: Webull account password
is_paper: Use paper trading (True) or live (False)
mfa_code: 2FA code if required
"""
self.email = email or os.getenv("WEBULL_EMAIL")
self.password = password or os.getenv("WEBULL_PASSWORD")
self.is_paper = is_paper
self.mfa_code = mfa_code
self.client = None
self.is_authenticated = False
self.positions: Dict[str, Dict] = {}
self.orders: List[Dict] = []
self.account_info: Dict = {}
if self.email and self.password:
self._initialize_client()
def _initialize_client(self):
"""Initialize Webull client"""
try:
if self.is_paper:
from webull import paper_webull as WebullClass
else:
from webull import webull as WebullClass
self.client = WebullClass()
logger.info("Webull client initialized (paper mode)")
except ImportError:
logger.error(
"Webull package not installed. "
"Run: pip install webull"
)
self.client = None
def login(self) -> bool:
"""
Login to Webull account.
Returns:
True if successful, False otherwise
"""
if not self.client or not self.email or not self.password:
logger.error("Client not initialized or credentials missing")
return False
try:
self.client.login(self.email, self.password)
self.is_authenticated = True
logger.info(f"Logged into Webull as {self.email}")
return True
except Exception as e:
logger.error(f"Failed to login to Webull: {e}")
return False
def get_trade_token(self, pin: str) -> bool:
"""
Get trade token (required for placing orders).
Args:
pin: Trading PIN from Webull account
Returns:
True if successful
"""
if not self.client or not self.is_authenticated:
logger.error("Not authenticated")
return False
try:
self.client.get_trade_token(pin)
logger.info("Trade token obtained")
return True
except Exception as e:
logger.error(f"Failed to get trade token: {e}")
return False
def place_buy_order(
self,
ticker: str,
quantity: int,
limit_price: Optional[float] = None,
order_type: str = "LMT", # LMT or MKT
) -> Optional[Dict]:
"""
Place a buy order.
Args:
ticker: Stock symbol
quantity: Number of shares
limit_price: Limit price (None for market order)
order_type: "LMT" for limit, "MKT" for market
Returns:
Order confirmation dict or None if failed
"""
if not self.client or not self.is_authenticated:
logger.error("Not authenticated")
return None
try:
order_result = self.client.place_order(
stock=ticker,
price=limit_price or 0,
qty=quantity,
orderType=order_type,
timeInForce="DAY", # Good for day
)
logger.info(f"Buy order placed: {ticker} x {quantity} @ ${limit_price}")
order_info = {
"action": "BUY",
"ticker": ticker,
"quantity": quantity,
"price": limit_price,
"order_type": order_type,
"timestamp": datetime.now().isoformat(),
"result": order_result,
}
self.orders.append(order_info)
return order_info
except Exception as e:
logger.error(f"Failed to place buy order: {e}")
return None
def place_sell_order(
self,
ticker: str,
quantity: int,
limit_price: Optional[float] = None,
order_type: str = "LMT",
) -> Optional[Dict]:
"""
Place a sell order.
Args:
ticker: Stock symbol
quantity: Number of shares
limit_price: Limit price (None for market order)
order_type: "LMT" for limit, "MKT" for market
Returns:
Order confirmation dict or None if failed
"""
if not self.client or not self.is_authenticated:
logger.error("Not authenticated")
return None
try:
order_result = self.client.place_order(
stock=ticker,
price=limit_price or 0,
qty=quantity,
orderType=order_type,
timeInForce="DAY",
)
logger.info(f"Sell order placed: {ticker} x {quantity} @ ${limit_price}")
order_info = {
"action": "SELL",
"ticker": ticker,
"quantity": quantity,
"price": limit_price,
"order_type": order_type,
"timestamp": datetime.now().isoformat(),
"result": order_result,
}
self.orders.append(order_info)
return order_info
except Exception as e:
logger.error(f"Failed to place sell order: {e}")
return None
def get_positions(self) -> Dict[str, Dict]:
"""
Get current positions from Webull.
Returns:
Dict of {ticker: position_info}
"""
if not self.client or not self.is_authenticated:
logger.error("Not authenticated")
return {}
try:
positions = self.client.get_positions()
self.positions = {p["ticker"]: p for p in positions}
return self.positions
except Exception as e:
logger.error(f"Failed to get positions: {e}")
return {}
def get_account_balance(self) -> Optional[Dict]:
"""
Get account balance and cash.
Returns:
Dict with balance info or None if failed
"""
if not self.client or not self.is_authenticated:
logger.error("Not authenticated")
return None
try:
account = self.client.get_account()
self.account_info = account
return {
"account_value": account.get("account_value"),
"cash": account.get("cash"),
"buying_power": account.get("buying_power"),
}
except Exception as e:
logger.error(f"Failed to get account balance: {e}")
return None
def get_orders(self) -> List[Dict]:
"""
Get current open orders.
Returns:
List of open orders
"""
if not self.client or not self.is_authenticated:
logger.error("Not authenticated")
return []
try:
orders = self.client.get_current_orders()
return orders
except Exception as e:
logger.error(f"Failed to get orders: {e}")
return []
def cancel_order(self, order_id: str) -> bool:
"""
Cancel an open order.
Args:
order_id: Order ID to cancel
Returns:
True if successful
"""
if not self.client or not self.is_authenticated:
logger.error("Not authenticated")
return False
try:
self.client.cancel_order(order_id)
logger.info(f"Order {order_id} cancelled")
return True
except Exception as e:
logger.error(f"Failed to cancel order: {e}")
return False
def cancel_all_orders(self) -> bool:
"""
Cancel all open orders.
Returns:
True if successful
"""
if not self.client or not self.is_authenticated:
logger.error("Not authenticated")
return False
try:
self.client.cancel_all_orders()
logger.info("All orders cancelled")
return True
except Exception as e:
logger.error(f"Failed to cancel all orders: {e}")
return False
def get_stock_quote(self, ticker: str) -> Optional[Dict]:
"""
Get current stock quote.
Args:
ticker: Stock symbol
Returns:
Quote dict with price info or None if failed
"""
if not self.client:
logger.error("Client not initialized")
return None
try:
quote = self.client.get_stock(ticker)
return {
"ticker": ticker,
"price": quote.get("price"),
"bid": quote.get("bid"),
"ask": quote.get("ask"),
"timestamp": datetime.now().isoformat(),
}
except Exception as e:
logger.error(f"Failed to get quote for {ticker}: {e}")
return None
def to_dict(self) -> Dict:
"""Serialize trader state"""
return {
"is_authenticated": self.is_authenticated,
"is_paper": self.is_paper,
"positions": self.positions,
"orders": self.orders,
"account_info": self.account_info,
}
def save_state(self, filepath: str):
"""Save trader state to JSON"""
with open(filepath, 'w') as f:
json.dump(self.to_dict(), f, indent=2)
@staticmethod
def demo_mode() -> "PaperTrader":
"""Create a demo paper trader (no auth required)"""
trader = PaperTrader(is_paper=True)
trader.is_authenticated = True # Simulate auth for demo
return trader