pip install -U keras-hub
pip install yfinance ddgs psutil pytz

import os
import datetime
import psutil
import yfinance as yf
from ddgs import DDGS
import keras_hub
import re
import pytz
import warnings

warnings.filterwarnings("ignore")

os.environ["KERAS_BACKEND"] = "jax"

def web_search(query, max_results=5):
    """Search the web using DuckDuckGo.

    Args:
        query: Search query string
        max_results: Maximum number of results to return (default: 5)

    Returns:
        Dictionary with "results" key containing list of search results,
        each with "title", "description", and "link" keys.
        Returns {"error": message} on failure.
    """
    results = []
    try:
        with DDGS() as ddgs:
            for r in ddgs.text(query, max_results=max_results):
                results.append(
                    {
                        "title": r.get("title"),
                        "description": r.get("body"),
                        "link": r.get("href"),
                    }
                )
    except Exception as e:
        return {"error": f"Search failed: {str(e)}"}
    return {"results": results}

def get_stock_price(symbol):
    """Get real-time stock price from Yahoo Finance.

    Args:
        symbol: Stock ticker symbol (e.g., "AAPL", "GOOGL")

    Returns:
        Dictionary with "symbol", "price", and "currency" keys.
        Returns {"error": message} on failure.
    """
    try:
        t = yf.Ticker(symbol)
        info = t.fast_info
        if info.last_price is None:
            raise ValueError("Price not available")
        current_date = datetime.datetime.now().strftime("%Y-%m-%d")
        return {
            "symbol": symbol,
            "price": round(info.last_price, 2),
            "currency": info.currency,
            "date": current_date,
        }
    except Exception as e:
        return {"error": f"Failed to get stock price for '{symbol}'. Reason: {e}"}

def get_system_stats(**kwargs):
    """Get CPU and memory usage.

    Args:
        **kwargs: Ignore any additional arguments the model might hallucinate.
    """

    mem = psutil.virtual_memory()
    return {
        "cpu_percent": psutil.cpu_percent(interval=1),
        "mem_percent": mem.percent,
        "used_gb": round(mem.used / 1e9, 2),
        "total_gb": round(mem.total / 1e9, 2),
    }

TIMEZONE_MAP = {
    # Countries
    "india": "Asia/Kolkata",
    "usa": "America/New_York",
    "uk": "Europe/London",
    "japan": "Asia/Tokyo",
    "china": "Asia/Shanghai",
    "australia": "Australia/Sydney",
    "france": "Europe/Paris",
    "germany": "Europe/Berlin",
    "russia": "Europe/Moscow",
    "brazil": "America/Sao_Paulo",
    "canada": "America/Toronto",
    "mexico": "America/Mexico_City",
    "spain": "Europe/Madrid",
    "italy": "Europe/Rome",
    # US States & Major Cities
    "california": "America/Los_Angeles",
    "new york": "America/New_York",
    "texas": "America/Chicago",
    "florida": "America/New_York",
    "washington": "America/Los_Angeles",
    "oregon": "America/Los_Angeles",
    "nevada": "America/Los_Angeles",
    "arizona": "America/Phoenix",
    "colorado": "America/Denver",
    "utah": "America/Denver",
    "illinois": "America/Chicago",
    "michigan": "America/Detroit",
    "massachusetts": "America/New_York",
    "pennsylvania": "America/New_York",
    # Major World Cities
    "los angeles": "America/Los_Angeles",
    "san francisco": "America/Los_Angeles",
    "seattle": "America/Los_Angeles",
    "chicago": "America/Chicago",
    "boston": "America/New_York",
    "miami": "America/New_York",
    "london": "Europe/London",
    "paris": "Europe/Paris",
    "tokyo": "Asia/Tokyo",
    "beijing": "Asia/Shanghai",
    "shanghai": "Asia/Shanghai",
    "dubai": "Asia/Dubai",
    "singapore": "Asia/Singapore",
    "hong kong": "Asia/Hong_Kong",
    "seoul": "Asia/Seoul",
    "moscow": "Europe/Moscow",
    "sydney": "Australia/Sydney",
    "mumbai": "Asia/Kolkata",
    "delhi": "Asia/Kolkata",
    "berlin": "Europe/Berlin",
    "rome": "Europe/Rome",
    "madrid": "Europe/Madrid",
    "toronto": "America/Toronto",
    "vancouver": "America/Vancouver",
    "montreal": "America/Toronto",
    # Special
    "utc": "UTC",
    "gmt": "GMT",
}


def get_timezone_for_location(location):
    """Map location to timezone."""

    loc = location.lower().strip()
    return TIMEZONE_MAP.get(loc)

def get_current_time(time_location=None):
    """Get current time for a location."""

    if not time_location:
        time_location = "UTC"

    try:
        timezone = get_timezone_for_location(time_location)
        note = ""

        if not timezone:
            # Default to UTC if location not found
            timezone = "UTC"
            note = (
                f"Location '{time_location}' not found in database. Showing UTC time."
            )

        tz = pytz.timezone(timezone)
        now = datetime.datetime.now(tz)
        return {
            "location": time_location if not note else "UTC",
            "time": now.strftime("%H:%M:%S"),
            "date": now.strftime("%Y-%m-%d"),
            "day": now.strftime("%A"),
            "timezone": timezone,
            "note": note,
        }
    except Exception as e:
        return {"error": str(e)}

TOOL_DEFINITIONS = [
    {
        "type": "function",
        "function": {
            "name": "web_search",
            "description": "Search the web for information.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "The search query"}
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_stock_price",
            "description": "Get the current stock price for a given ticker symbol",
            "parameters": {
                "type": "object",
                "properties": {
                    "symbol": {
                        "type": "string",
                        "description": "Stock ticker symbol (e.g., AAPL, GOOGL, TSLA)",
                    }
                },
                "required": ["symbol"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_current_time",
            "description": "Get the current clock time for a specific city or country.",
            "parameters": {
                "type": "object",
                "properties": {
                    "time_location": {
                        "type": "string",
                        "description": "The city or country for which to get the time, e.g., 'Tokyo' or 'Japan'",
                    }
                },
                "required": ["time_location"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_system_stats",
            "description": "Get current CPU and memory usage statistics",
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
    },
]

# Tool registry
TOOLS = {
    "web_search": web_search,
    "get_stock_price": get_stock_price,
    "get_current_time": get_current_time,
    "get_system_stats": get_system_stats,
}

FUNCTION_CALL_PATTERN = re.compile(
    r"<start_function_call>call:(\w+)\{([^}]*)\}<end_function_call>"
)
ARGUMENT_PATTERN = re.compile(
    r'(\w+):\s*("(?:[^"\\]|\\.)*"|\'(?:[^\'\\]|\\.)*\'|[^,]+)'
)


def parse_function_call(output):
    """Parse function call from FunctionGemma model output.

    Extracts function name and arguments from the special FunctionGemma function calling
    format: <start_function_call>call:tool_name{arg1:value1}<end_function_call>

    Args:
        output: Raw model output string

    Returns:
        Dictionary with "name" (function name) and "arguments" (dict of args),

    """
    match = FUNCTION_CALL_PATTERN.search(output)

    if not match:
        return None

    tool_name = match.group(1)
    args_str = match.group(2).strip()

    # Parse arguments
    args = {}
    if args_str:
        # Matches key:"value" or key:'value' or key:value
        arg_pairs = ARGUMENT_PATTERN.findall(args_str)
        for key, value in arg_pairs:
            value = value.strip()

            # Remove <escape> tags
            value = value.replace("<escape>", "").replace("</escape>", "")

            # Remove quotes if present
            if value.startswith('"') and value.endswith('"'):
                value = value[1:-1]
            elif value.startswith("'") and value.endswith("'"):
                value = value[1:-1]

            args[key] = value

    return {"name": tool_name, "arguments": args}

def format_tool_declaration(tool):
    """Format a tool declaration for FunctionGemma function calling."""

    func = tool["function"]
    params = func.get("parameters", {})
    props = params.get("properties", {})
    required = params.get("required", [])

    # Build parameter string
    param_parts = []
    for name, details in props.items():
        param_str = f"{name}:{{description:<escape>{details.get('description', '')}<escape>,type:<escape>{details['type'].upper()}<escape>}}"
        param_parts.append(param_str)

    properties_str = ",".join(param_parts) if param_parts else ""
    required_str = ",".join(f"<escape>{r}<escape>" for r in required)

    # Build full declaration
    declaration = f"declaration:{func['name']}"
    declaration += f"{{description:<escape>{func['description']}<escape>"

    if properties_str or required_str:
        declaration += ",parameters:{"
        if properties_str:
            declaration += f"properties:{{{properties_str}}}"
        if required_str:
            if properties_str:
                declaration += ","
            declaration += f"required:[{required_str}]"
        declaration += ",type:<escape>OBJECT<escape>}"

    declaration += "}"

    return declaration

def format_prompt_with_tools(messages, tools):
    """Manually construct FunctionGemma function calling prompt.

    This function replicates the behavior of HuggingFace's apply_chat_template()
    for FunctionGemma native function calling format, since KerasHub doesn't yet
    provide this functionality.

    Args:
        messages: List of message dicts with "role" and "content" keys
        tools: List of tool definition dicts in OpenAI format

    Returns:
        Formatted prompt string with tool declarations and conversation history
    """
    prompt = "<bos>"

    # Add tool declarations
    if tools:
        prompt += "<start_of_turn>developer\n"
        for tool in tools:
            prompt += "<start_function_declaration>"
            prompt += format_tool_declaration(tool)
            prompt += "<end_function_declaration>"

        # Add few-shot examples to help the model distinguish between tools
        prompt += "\nHere are some examples of how to use the tools correctly:\n"
        prompt += "User: Who is Isaac Newton?\n"
        prompt += "Model: <start_function_call>call:web_search{query:Isaac Newton}<end_function_call>\n"
        prompt += "User: What is the time in USA?\n"
        prompt += "Model: <start_function_call>call:get_current_time{time_location:USA}<end_function_call>\n"
        prompt += "User: Stock price of Apple\n"
        prompt += "Model: <start_function_call>call:get_stock_price{symbol:AAPL}<end_function_call>\n"
        prompt += "User: Who is the PM of Japan?\n"
        prompt += "Model: <start_function_call>call:web_search{query:PM of Japan}<end_function_call>\n"
        prompt += "User: Check the weather in London\n"
        prompt += "Model: <start_function_call>call:web_search{query:weather in London}<end_function_call>\n"
        prompt += "User: Latest news about AI\n"
        prompt += "Model: <start_function_call>call:web_search{query:Latest news about AI}<end_function_call>\n"
        prompt += "User: System memory usage\n"
        prompt += (
            "Model: <start_function_call>call:get_system_stats{}<end_function_call>\n"
        )

        prompt += "<end_of_turn>\n"

    # Add conversation messages
    for message in messages:
        role = message["role"]
        if role == "assistant":
            role = "model"

        content = message.get("content", "")

        # Regular user/assistant messages
        prompt += f"<start_of_turn>{role}\n"
        prompt += content
        prompt += "<end_of_turn>\n"

    # Add generation prompt
    prompt += "<start_of_turn>model\n"

    return prompt

def format_tool_response(tool_name, result):
    """Format tool result for conversation."""
    if "error" in result:
        return f"Error: {result['error']}"

    # Handle Web Search Results (from web_search tool OR fallback from other tools)
    if "results" in result:
        results = result.get("results", [])
        if results:
            output = f"Found {len(results)} results:\n"
            for i, r in enumerate(results[:3], 1):
                output += f"\n{i}. {r['title']}\n   {r['description'][:150]}...\n   {r['link']}"
            return output
        return "No results found"

    if tool_name == "get_stock_price":
        return f"Stock {result['symbol']}: ${result['price']} {result['currency']} (as of {result['date']})"

    elif tool_name == "get_current_time":
        response = f"Time in {result['location']}: {result['time']} ({result['day']}, {result['date']})\nTimezone: {result['timezone']}"
        if result.get("note"):
            response += f"\nNote: {result['note']}"
        return response

    elif tool_name == "get_system_stats":
        return f"CPU: {result['cpu_percent']}% | Memory: {result['mem_percent']}% ({result['used_gb']}/{result['total_gb']}GB)"

    return str(result)

def load_model():
    """Load FunctionGemma 270M model."""
    print("Loading FunctionGemma 270M model...")
    model = keras_hub.models.Gemma3CausalLM.from_preset("function_gemma_instruct_270m")
    print("Model loaded!\n")
    return model

def chat(model):
    """Main chat loop with native function calling."""
    print("=" * 70)
    print(" " * 20 + "FUNCTION CALLING")
    print("=" * 70)
    print("\nCapabilities:")
    print("  🔍 Web Search        - e.g., 'Who is Newton?', 'Latest AI news'")
    print("  📈 Stock Prices      - e.g., 'Stock price of AAPL', 'TSLA price'")
    print(
        "  🕐 World Time        - e.g., 'Time in London', 'What time is it in Tokyo?'"
    )
    print("  💻 System Stats      - e.g., 'System memory', 'CPU usage'")
    print("\nType 'exit' to quit.\n")
    print("-" * 70)

    # Conversation history
    messages = []

    while True:
        try:
            user_input = input("\nYou: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n\nGoodbye! 👋")
            break

        if user_input.lower() in {"exit", "quit", "q"}:
            print("\nGoodbye! 👋")
            break

        if not user_input:
            continue

        # Add user message
        messages.append({"role": "user", "content": user_input})

        # Generate response with tools
        print("🤔 Thinking...", end="", flush=True)

        try:
            # Manually format prompt with tools (KerasHub doesn't have apply_chat_template)
            prompt = format_prompt_with_tools(messages, TOOL_DEFINITIONS)

            # Generate
            output = model.generate(prompt, max_length=2048)

            # Extract only the new response
            response = output[len(prompt) :].strip()
            if "<start_function_response>" in response:
                response = response.split("<start_function_response>")[0].strip()
            if "<start_of_turn>" in response:
                response = response.split("<start_of_turn>")[0].strip()
            # Remove end tokens
            response = response.replace("<end_of_turn>", "").strip()

            print(f"\r{' ' * 40}\r", end="")  # Clear status
            print(f"[DEBUG] Raw Model Response: {response}")

            # Check if model made a function call
            function_call = parse_function_call(response)

            if function_call:
                tool_name = function_call["name"]
                tool_args = function_call["arguments"]

                print(f"🔧 Calling {tool_name}...", end="", flush=True)

                # Execute tool
                if tool_name in TOOLS:
                    try:
                        result = TOOLS[tool_name](**tool_args)
                        formatted_result = format_tool_response(tool_name, result)

                        print(f"\r{' ' * 40}\r", end="")  # Clear status
                        print(f"\nAssistant: {formatted_result}")

                        # Clear conversation history after successful tool execution
                        # This prevents token overflow and keeps the model focused
                        messages = []

                    except Exception as e:
                        print(f"\r❌ Error executing {tool_name}: {e}")
                        messages.pop()  # Remove user message on error
                else:
                    print(f"\r❌ Unknown tool: {tool_name}")
                    messages.pop()
            else:
                # Regular text response (no function call)
                # Check if response is empty
                if not response or response == "...":
                    print(
                        f"\nAssistant: I'm not sure how to help with that. Please try rephrasing your question."
                    )
                    messages.pop()  # Remove problematic message
                else:
                    print(f"\nAssistant: {response}")
                    messages.append({"role": "assistant", "content": response})

                    # Keep only last 3 turns to prevent context overflow
                    if len(messages) > 6:  # 3 user + 3 assistant
                        messages = messages[-6:]

        except Exception as e:
            print(f"\r❌ Error: {e}")
            messages.pop()  # Remove user message on error
            continue

 if __name__ == "__main__":
     model = load_model()
     chat(model)