tool_registry.py

"""
Agent Tool Registry and Schema Definition Module.

This module serves as the central "switchboard" for the Financial Research &
Strategy Agent. It accomplishes two primary goals:
1. **Tool Specification**: It defines the JSON schemas (agent_tools) that
   describe the purpose and parameters of each tool to the LLM.
2. **Function Mapping**: It maintains a runtime dictionary (TOOL_MAP) that
   links the string-based tool names generated by the LLM to the actual
   executable Python functions.
"""

from tools.market import get_stock_prices
from tools.news import get_latest_news
from tools.internal_kb import query_internal_kb

# ==========================================
# 1. LLM Tool Definitions (JSON Schema)
# ==========================================

# This list is passed to the LLM during initialization (via .bind_tools()).
# It acts as the "API Documentation" that the LLM reads to understand
# what it can do and what data it needs to provide.
agent_tools = [
    {
        "type": "function",
        "function": {
            "name": "get_stock_prices",
            "description": (
                "Retrieves the real-time stock price and trailing P/E ratio "
                "for a given list of company tickers. Use this for valuation checks."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "tickers": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "A list of official stock ticker symbols, e.g., ['AAPL', 'RELIANCE.NS'].",
                    }
                },
                "required": ["tickers"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_latest_news",
            "description": (
                "Fetches the top 3 latest news headlines and summaries for a specific "
                "asset, company, or market query. Ideal for gathering current sentiment."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The search query for news retrieval (e.g., 'interest rate impact on tech').",
                    }
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "query_internal_kb",
            "description": (
                "Queries the firm's private vector database (ChromaDB) for risk tolerance, "
                "portfolio rebalancing guidelines, and compliance rules. Consult this "
                "before making any recommendation."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The specific policy or guideline to search for.",
                    }
                },
                "required": ["query"],
            },
        },
    },
]

# ==========================================
# 2. Runtime Tool Dispatcher (TOOL_MAP)
# ==========================================

# This dictionary is used by the 'execute_tools' node in the LangGraph workflow.
# When the LLM says "I want to run 'get_latest_news'", the orchestrator looks up
# that string here to find the actual Python function to call.
TOOL_MAP = {
    "get_stock_prices": get_stock_prices,
    "get_latest_news": get_latest_news,
    "query_internal_kb": query_internal_kb,
}