ai/backend/src/agent/tools/research_tools.py

"""Research and external data tools for trading analysis."""

from typing import Dict, Any, Optional
from langchain_core.tools import tool
from langchain_community.tools import (
    ArxivQueryRun,
    WikipediaQueryRun,
    DuckDuckGoSearchRun
)
from langchain_community.utilities import (
    ArxivAPIWrapper,
    WikipediaAPIWrapper,
    DuckDuckGoSearchAPIWrapper
)


@tool
def search_arxiv(query: str, max_results: int = 5) -> str:
    """Search arXiv for academic papers on quantitative finance, trading strategies, and machine learning.

    Use this to find research papers on topics like:
    - Market microstructure and order flow
    - Algorithmic trading strategies
    - Machine learning for finance
    - Time series forecasting
    - Risk management
    - Portfolio optimization

    Args:
        query: Search query (e.g., "machine learning algorithmic trading", "deep learning stock prediction")
        max_results: Maximum number of results to return (default: 5)

    Returns:
        Summary of papers including titles, authors, abstracts, and links

    Example:
        search_arxiv("reinforcement learning trading", max_results=3)
    """
    arxiv = ArxivQueryRun(api_wrapper=ArxivAPIWrapper(top_k_results=max_results))
    return arxiv.run(query)


@tool
def search_wikipedia(query: str) -> str:
    """Search Wikipedia for information on finance, trading, and economics concepts.

    Use this to get background information on:
    - Financial instruments and markets
    - Economic indicators
    - Trading terminology
    - Technical analysis concepts
    - Historical market events

    Args:
        query: Search query (e.g., "Black-Scholes model", "technical analysis", "options trading")

    Returns:
        Wikipedia article summary with key information

    Example:
        search_wikipedia("Bollinger Bands")
    """
    wikipedia = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
    return wikipedia.run(query)


@tool
def search_web(query: str, max_results: int = 5) -> str:
    """Search the web for current information on markets, news, and trading.

    Use this to find:
    - Latest market news and analysis
    - Company announcements and earnings
    - Economic events and indicators
    - Cryptocurrency updates
    - Exchange status and updates
    - Trading strategy discussions

    Args:
        query: Search query (e.g., "Bitcoin price news", "Fed interest rate decision")
        max_results: Maximum number of results to return (default: 5)

    Returns:
        Search results with titles, snippets, and links

    Example:
        search_web("Ethereum merge update", max_results=3)
    """
    # Lazy initialization to avoid hanging during import
    search = DuckDuckGoSearchRun(api_wrapper=DuckDuckGoSearchAPIWrapper())
    # Note: max_results parameter doesn't work properly with current wrapper
    return search.run(query)


@tool
def http_get(url: str, params: Optional[Dict[str, str]] = None) -> str:
    """Make HTTP GET request to fetch data from APIs or web pages.

    Use this to retrieve:
    - Exchange API data (if public endpoints)
    - Market data from external APIs
    - Documentation and specifications
    - News articles and blog posts
    - JSON/XML data from web services

    Args:
        url: The URL to fetch
        params: Optional query parameters as a dictionary

    Returns:
        Response text from the URL

    Raises:
        ValueError: If the request fails

    Example:
        http_get("https://api.coingecko.com/api/v3/simple/price",
                 params={"ids": "bitcoin", "vs_currencies": "usd"})
    """
    import requests

    try:
        response = requests.get(url, params=params, timeout=10)
        response.raise_for_status()
        return response.text
    except requests.RequestException as e:
        raise ValueError(f"HTTP GET request failed: {str(e)}")


@tool
def http_post(url: str, data: Dict[str, Any]) -> str:
    """Make HTTP POST request to send data to APIs.

    Use this to:
    - Submit data to external APIs
    - Trigger webhooks
    - Post analysis results
    - Interact with exchange APIs (if authenticated)

    Args:
        url: The URL to post to
        data: Dictionary of data to send in the request body

    Returns:
        Response text from the server

    Raises:
        ValueError: If the request fails

    Example:
        http_post("https://webhook.site/xxx", {"message": "Trade executed"})
    """
    import requests
    import json

    try:
        response = requests.post(url, json=data, timeout=10)
        response.raise_for_status()
        return response.text
    except requests.RequestException as e:
        raise ValueError(f"HTTP POST request failed: {str(e)}")


# Export tools list
RESEARCH_TOOLS = [
    search_arxiv,
    search_wikipedia,
    search_web,
    http_get,
    http_post
]