diff --git a/agent/base_agent/base_agent.py b/agent/base_agent/base_agent.py index 86392de..b63b2e9 100644 --- a/agent/base_agent/base_agent.py +++ b/agent/base_agent/base_agent.py @@ -29,6 +29,7 @@ from tools.deployment_config import ( log_api_key_warning, get_deployment_mode ) +from agent.context_injector import ContextInjector # Load environment variables load_dotenv() @@ -124,6 +125,9 @@ class BaseAgent: self.tools: Optional[List] = None self.model: Optional[ChatOpenAI] = None self.agent: Optional[Any] = None + + # Context injector for MCP tools + self.context_injector: Optional[ContextInjector] = None # Data paths self.data_path = os.path.join(self.base_log_path, self.signature) @@ -169,16 +173,27 @@ class BaseAgent: print("⚠️ OpenAI base URL not set, using default") try: - # Create MCP client - self.client = MultiServerMCPClient(self.mcp_config) + # Create context injector for injecting signature and today_date into tool calls + self.context_injector = ContextInjector( + signature=self.signature, + today_date=self.init_date # Will be updated per trading session + ) + + # Create MCP client with interceptor + self.client = MultiServerMCPClient( + self.mcp_config, + tool_interceptors=[self.context_injector] + ) # Get tools - self.tools = await self.client.get_tools() - if not self.tools: + raw_tools = await self.client.get_tools() + if not raw_tools: print("⚠️ Warning: No MCP tools loaded. MCP services may not be running.") print(f" MCP configuration: {self.mcp_config}") + self.tools = [] else: - print(f"✅ Loaded {len(self.tools)} MCP tools") + print(f"✅ Loaded {len(raw_tools)} MCP tools") + self.tools = raw_tools except Exception as e: raise RuntimeError( f"❌ Failed to initialize MCP client: {e}\n" @@ -336,6 +351,10 @@ Summary:""" """ print(f"📈 Starting trading session: {today_date}") + # Update context injector with current trading date + if self.context_injector: + self.context_injector.today_date = today_date + # Clear conversation history for new trading day self.clear_conversation_history() diff --git a/agent/context_injector.py b/agent/context_injector.py new file mode 100644 index 0000000..0184bf4 --- /dev/null +++ b/agent/context_injector.py @@ -0,0 +1,50 @@ +""" +Tool interceptor for injecting runtime context into MCP tool calls. + +This interceptor automatically injects `signature` and `today_date` parameters +into buy/sell tool calls to support concurrent multi-model simulations. +""" + +from typing import Any, Dict + + +class ContextInjector: + """ + Intercepts tool calls to inject runtime context (signature, today_date). + + Usage: + interceptor = ContextInjector(signature="gpt-5", today_date="2025-10-01") + client = MultiServerMCPClient(config, tool_interceptors=[interceptor]) + """ + + def __init__(self, signature: str, today_date: str): + """ + Initialize context injector. + + Args: + signature: Model signature to inject + today_date: Trading date to inject + """ + self.signature = signature + self.today_date = today_date + + def __call__(self, tool_name: str, tool_input: Dict[str, Any]) -> Dict[str, Any]: + """ + Intercept tool call and inject context parameters. + + Args: + tool_name: Name of the tool being called + tool_input: Original tool input parameters + + Returns: + Modified tool input with injected context + """ + # Only inject for trade tools (buy/sell) + if tool_name in ["buy", "sell"]: + # Inject signature and today_date if not already provided + if "signature" not in tool_input: + tool_input["signature"] = self.signature + if "today_date" not in tool_input: + tool_input["today_date"] = self.today_date + + return tool_input diff --git a/agent_tools/tool_trade.py b/agent_tools/tool_trade.py index 0055391..d766b30 100644 --- a/agent_tools/tool_trade.py +++ b/agent_tools/tool_trade.py @@ -13,41 +13,45 @@ mcp = FastMCP("TradeTools") @mcp.tool() -def buy(symbol: str, amount: int) -> Dict[str, Any]: +def buy(symbol: str, amount: int, signature: str = None, today_date: str = None) -> Dict[str, Any]: """ Buy stock function - + This function simulates stock buying operations, including the following steps: 1. Get current position and operation ID 2. Get stock opening price for the day 3. Validate buy conditions (sufficient cash) 4. Update position (increase stock quantity, decrease cash) 5. Record transaction to position.jsonl file - + Args: symbol: Stock symbol, such as "AAPL", "MSFT", etc. amount: Buy quantity, must be a positive integer, indicating how many shares to buy - + signature: Model signature (optional, will use config/env if not provided) + today_date: Trading date (optional, will use config/env if not provided) + Returns: Dict[str, Any]: - Success: Returns new position dictionary (containing stock quantity and cash balance) - Failure: Returns {"error": error message, ...} dictionary - + Raises: ValueError: Raised when SIGNATURE environment variable is not set - + Example: >>> result = buy("AAPL", 10) >>> print(result) # {"AAPL": 110, "MSFT": 5, "CASH": 5000.0, ...} """ # Step 1: Get environment variables and basic information - # Get signature (model name) from environment variable, used to determine data storage path - signature = get_config_value("SIGNATURE") + # Get signature (model name) from parameter or fallback to config/env if signature is None: - raise ValueError("SIGNATURE environment variable is not set") - - # Get current trading date from environment variable - today_date = get_config_value("TODAY_DATE") + signature = get_config_value("SIGNATURE") + if signature is None: + raise ValueError("SIGNATURE not provided and environment variable is not set") + + # Get current trading date from parameter or fallback to config/env + if today_date is None: + today_date = get_config_value("TODAY_DATE") # Step 2: Get current latest position and operation ID # get_latest_position returns two values: position dictionary and current maximum operation ID @@ -104,41 +108,45 @@ def buy(symbol: str, amount: int) -> Dict[str, Any]: return new_position @mcp.tool() -def sell(symbol: str, amount: int) -> Dict[str, Any]: +def sell(symbol: str, amount: int, signature: str = None, today_date: str = None) -> Dict[str, Any]: """ Sell stock function - + This function simulates stock selling operations, including the following steps: 1. Get current position and operation ID 2. Get stock opening price for the day 3. Validate sell conditions (position exists, sufficient quantity) 4. Update position (decrease stock quantity, increase cash) 5. Record transaction to position.jsonl file - + Args: symbol: Stock symbol, such as "AAPL", "MSFT", etc. amount: Sell quantity, must be a positive integer, indicating how many shares to sell - + signature: Model signature (optional, will use config/env if not provided) + today_date: Trading date (optional, will use config/env if not provided) + Returns: Dict[str, Any]: - Success: Returns new position dictionary (containing stock quantity and cash balance) - Failure: Returns {"error": error message, ...} dictionary - + Raises: ValueError: Raised when SIGNATURE environment variable is not set - + Example: >>> result = sell("AAPL", 10) >>> print(result) # {"AAPL": 90, "MSFT": 5, "CASH": 15000.0, ...} """ # Step 1: Get environment variables and basic information - # Get signature (model name) from environment variable, used to determine data storage path - signature = get_config_value("SIGNATURE") + # Get signature (model name) from parameter or fallback to config/env if signature is None: - raise ValueError("SIGNATURE environment variable is not set") - - # Get current trading date from environment variable - today_date = get_config_value("TODAY_DATE") + signature = get_config_value("SIGNATURE") + if signature is None: + raise ValueError("SIGNATURE not provided and environment variable is not set") + + # Get current trading date from parameter or fallback to config/env + if today_date is None: + today_date = get_config_value("TODAY_DATE") # Step 2: Get current latest position and operation ID # get_latest_position returns two values: position dictionary and current maximum operation ID