from groq import Groq

client = Groq()

# Define comprehensive tool set
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get the current weather in a given location",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The city and state, e.g. San Francisco, CA"
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                        "description": "The unit of temperature"
                    }
                },
                "required": ["location"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "calculate_math",
            "description": "Perform mathematical calculations",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "Mathematical expression to evaluate, e.g. '2 + 2' or 'sqrt(16)'"
                    }
                },
                "required": ["expression"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "search_web",
            "description": "Search the web for current information",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Search query"
                    },
                    "num_results": {
                        "type": "integer",
                        "description": "Number of results to return",
                        "minimum": 1,
                        "maximum": 10,
                        "default": 5
                    }
                },
                "required": ["query"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_time",
            "description": "Get the current time in a specific timezone",
            "parameters": {
                "type": "object",
                "properties": {
                    "timezone": {
                        "type": "string",
                        "description": "Timezone identifier, e.g. 'America/New_York' or 'UTC'"
                    }
                },
                "required": ["timezone"]
            }
        }
    }
]

def use_tools_with_caching():
    # First request - creates cache for all tool definitions
    first_request = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant with access to various tools. Use the appropriate tools to answer user questions accurately."
            },
            {
                "role": "user",
                "content": "What's the weather like in New York City?"
            }
        ],
        model="moonshotai/kimi-k2-instruct-0905",
        tools=tools
    )

    print("First request response:", first_request.choices[0].message)
    print("Usage:", first_request.usage)

    # Check if the model wants to use tools
    if first_request.choices[0].message.tool_calls:
        print("Tool calls requested:", first_request.choices[0].message.tool_calls)

    # Second request - tool definitions will be cached
    second_request = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant with access to various tools. Use the appropriate tools to answer user questions accurately."
            },
            {
                "role": "user",
                "content": "Can you calculate the square root of 144 and tell me what time it is in Tokyo?"
            }
        ],
        model="moonshotai/kimi-k2-instruct-0905",
        tools=tools
    )

    print("Second request response:", second_request.choices[0].message)
    print("Usage:", second_request.usage)

    if second_request.choices[0].message.tool_calls:
        print("Tool calls requested:", second_request.choices[0].message.tool_calls)

    # Third request - same tool definitions cached
    third_request = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant with access to various tools. Use the appropriate tools to answer user questions accurately."
            },
            {
                "role": "user",
                "content": "Search for recent news about artificial intelligence developments."
            }
        ],
        model="moonshotai/kimi-k2-instruct-0905",
        tools=tools
    )

    print("Third request response:", third_request.choices[0].message)
    print("Usage:", third_request.usage)

    if third_request.choices[0].message.tool_calls:
        print("Tool calls requested:", third_request.choices[0].message.tool_calls)

if __name__ == "__main__":
    use_tools_with_caching()

---

## Prompt Caching