data_aggregation/tool_use_generator.py

"""
Tool Use Generator Module

Generates training data for teaching LLMs to use tools and functions.
"""

import json
from typing import List, Dict, Any, Optional


class ToolUseGenerator:
    """Generate tool use training data for function calling."""

    def __init__(self, tools: Optional[List[Dict]] = None):
        """
        Initialize tool use generator.

        Args:
            tools: List of tool definitions
        """
        self.tools = tools or self._get_default_financial_tools()

    def _get_default_financial_tools(self) -> List[Dict]:
        """Get default financial tools."""
        return [
            {
                "name": "get_market_data",
                "description": "Get current market data for a stock",
                "parameters": {
                    "symbol": "string",
                    "data_type": "string (price, volume, market_cap)"
                },
                "returns": "dict"
            },
            {
                "name": "calculate_compound_interest",
                "description": "Calculate compound interest",
                "parameters": {
                    "principal": "float",
                    "rate": "float",
                    "time": "float",
                    "frequency": "int"
                },
                "returns": "float"
            },
            {
                "name": "get_portfolio_allocation",
                "description": "Get current portfolio allocation",
                "parameters": {
                    "user_id": "string"
                },
                "returns": "dict"
            },
            {
                "name": "calculate_retirement_needs",
                "description": "Calculate retirement savings needs",
                "parameters": {
                    "current_age": "int",
                    "retirement_age": "int",
                    "current_savings": "float",
                    "annual_contribution": "float"
                },
                "returns": "dict"
            }
        ]

    def generate_tool_use_examples(
        self,
        num_examples: int = 10,
        complexity: str = "single"
    ) -> List[Dict[str, Any]]:
        """
        Generate tool use training examples.

        Args:
            num_examples: Number of examples to generate
            complexity: "single" or "multi" (multi-step tool chains)

        Returns:
            List of tool use examples
        """
        examples = []

        for _ in range(num_examples):
            if complexity == "single":
                example = self._generate_single_tool_example()
            else:
                example = self._generate_multi_tool_example()

            if example:
                examples.append(example)

        return examples

    def _generate_single_tool_example(self) -> Dict[str, Any]:
        """Generate example with single tool call."""
        import random

        tool = random.choice(self.tools)

        # Create example based on tool
        if tool["name"] == "get_market_data":
            return {
                "instruction": "What is the current price of Apple stock?",
                "input": "",
                "output": "I'll check the current market data for Apple stock.",
                "tool_calls": [
                    {
                        "tool": "get_market_data",
                        "parameters": {
                            "symbol": "AAPL",
                            "data_type": "price"
                        }
                    }
                ]
            }
        elif tool["name"] == "calculate_compound_interest":
            return {
                "instruction": "If I invest $10,000 at 5% interest compounded quarterly, how much will I have in 10 years?",
                "input": "",
                "output": "I'll calculate that for you using compound interest formula.",
                "tool_calls": [
                    {
                        "tool": "calculate_compound_interest",
                        "parameters": {
                            "principal": 10000,
                            "rate": 0.05,
                            "time": 10,
                            "frequency": 4
                        }
                    }
                ]
            }

        # Generic example
        return {
            "instruction": f"Use the {tool['name']} tool",
            "input": "",
            "output": f"I'll use {tool['name']} to help with that.",
            "tool_calls": [{"tool": tool["name"], "parameters": {}}]
        }

    def _generate_multi_tool_example(self) -> Dict[str, Any]:
        """Generate example with multiple tool calls (chain)."""
        return {
            "instruction": "Should I rebalance my portfolio given current market conditions?",
            "input": "",
            "output": "Let me check your portfolio and current market data to provide a recommendation.",
            "tool_calls": [
                {
                    "tool": "get_portfolio_allocation",
                    "parameters": {"user_id": "user123"}
                },
                {
                    "tool": "get_market_data",
                    "parameters": {"symbol": "SPY", "data_type": "price"}
                }
            ],
            "reasoning": "First get the user's current portfolio allocation, then check market conditions, and finally provide rebalancing advice."
        }

    def add_custom_tool(self, tool: Dict[str, Any]) -> None:
        """
        Add a custom tool definition.

        Args:
            tool: Tool definition dict
        """
        required_keys = ["name", "description", "parameters", "returns"]
        if not all(key in tool for key in required_keys):
            raise ValueError(f"Tool must have keys: {required_keys}")

        self.tools.append(tool)

    def get_tools_schema(self) -> List[Dict]:
        """Get OpenAI-compatible tools schema."""
        schema = []

        for tool in self.tools:
            schema.append({
                "type": "function",
                "function": {
                    "name": tool["name"],
                    "description": tool["description"],
                    "parameters": {
                        "type": "object",
                        "properties": tool["parameters"],
                        "required": list(tool["parameters"].keys())
                    }
                }
            })

        return schema