#!/usr/bin/env python
# coding: utf-8

# ### Set up functions and routes

# In[1]:


from datetime import datetime
from zoneinfo import ZoneInfo


def get_time(timezone: str) -> str:
    """Finds the current time in a specific timezone.

    :param timezone: The timezone to find the current time in, should
        be a valid timezone from the IANA Time Zone Database like
        "America/New_York" or "Europe/London".
    :type timezone: str
    :return: The current time in the specified timezone."""
    now = datetime.now(ZoneInfo(timezone))
    print(f"Invoked `get_time` function with timezone: `{timezone}`")
    return now.strftime("%H:%M")


def get_news(category: str, country: str) -> str:
    """Useful to get the news in a specific country"""
    print(
        f"Invoked: `get_news` function with category: `{category}` "
        f"and country: `{country}`"
    )
    return "Results from dummy news API"


# In[2]:


get_time("America/New_York")


# Define the LLM

# In[3]:


from semantic_router.llms import OpenAILLM

llm = OpenAILLM()


# Now generate a dynamic routing config for each function

# In[4]:


from semantic_router import Route

functions = [get_time, get_news]

routes = []

for function in functions:
    route = Route.from_dynamic_route(llm=llm, entities=function)
    routes.append(route)


# In[5]:


# You can manually add or remove routes

get_weather_route = Route(
    name="get_weather",
    utterances=[
        "what is the weather in SF",
        "what is the current temperature in London?",
        "tomorrow's weather in Paris?",
    ],
    function_schemas=None,
)
routes.append(get_weather_route)


# Add routes to the layer config

# In[6]:


from semantic_router.layer import LayerConfig

layer_config = LayerConfig(routes=routes)
layer_config.to_dict()


# In[7]:


# Get a route by name
layer_config.get("get_time")


# In[8]:


# Remove a route by name
layer_config.remove("get_weather")
layer_config.to_dict()


# Save config to a file (.json or .yaml)

# In[9]:


layer_config.to_file("output/layer_config.json")


# ### Define routing layer

# Load config from local file

# In[10]:


from semantic_router.layer import LayerConfig

layer_config = LayerConfig.from_file("output/layer_config.json")


# Initialize routing layer

# In[11]:


import os
from getpass import getpass
from semantic_router import RouteLayer

# https://dashboard.cohere.com/
os.environ["COHERE_API_KEY"] = os.getenv("COHERE_API_KEY") or getpass(
    "Enter Cohere API Key: "
)

layer = RouteLayer.from_config(config=layer_config)


# Do a function call with functions as tool

# In[12]:


layer("What is the time in Stockholm?")


# Define function execution method

# In[17]:


from semantic_router.schema import RouteChoice
from semantic_router.schema import Message

llm = OpenAILLM()


def route_and_execute(query, functions, layer):
    route_choice: RouteChoice = layer(query)

    for function in functions:
        if function.__name__ == route_choice.name:
            if route_choice.function_call:
                return function(**route_choice.function_call)

    # If no function is found, use the LLM for general queries
    msgs = [Message(role="user", content=query)]
    return llm(msgs)


queries = [
    "What is the time in Stockholm?",
    "What are the tech news in the US?",
    "The capital of France?",
]

for query in queries:
    print(f"Query: {query}")
    print(route_and_execute(query, functions, layer))


# In[ ]: