Untitled

from ollama import chat

QUERY_ROUTER_MODEL="qwen3:14b"  # replace with your model
QUERY_ROUTER_ENABLE_THINKING=False # set to True to have the agent use reasoning (better results, but slower)

def query_omdb(query:str, year:int|None=None) -> str:
    """Search the Open Movie Database (OMDB) for a given query and return the movie details.

    Args:
        query: The search query string.
        year: The release year of the movie.

    Returns:
        The movie details as a string.
    """

    return "implement me"

def query_wikipedia(query:str) -> str:
    """Search Wikipedia for a given query and return the article contents.

    Args:
        query: The search query string.

    Returns:
        The contents of the Wikipedia article as a string.
    """

    return "implement me"

def query_google(query:str) -> str:
    """Perform a Google search for a given query and return the top results.

    Args:
        query: The search query string.

    Returns:
        The top Google search results as a string.
    """

    return "implement me"

def get_llm_tool_calls_from_query(query:str):
    response = chat(
        model=QUERY_ROUTER_MODEL,
        messages=[
            {
                "role": "system",
                "content": """
                You a query tool routing agent. You have tools to query OMDB, Wikipedia, and Google.
                You follow these strict rules:
                    - Google Searches are expensive, don't query Google unless necessary to answer the question, or the user specifically asks for a Google search. This includes:
                      - Questions about current events, current weather, or anything else that is updated too frequently to be in Wikipedia.
                    - Questions related to movies must query both OMDB and Wikipedia.
                    - All other general knowledge questions should query Wikipedia. Multiple Wikipedia calls are allowed if necessary to answer the question.
                    - If the user's query can be answered without external information, do not call any tools at all.
                    - If multiple tools are needed, call them all at once in a single response.
                    - Don't provide a response, just return the tool calls.

                You will probably need to rewrite or simplify the user's query to get the best results from each tool. Follow these guidelines:
                    - For Google, extract the main search terms from the user's query.
                    - For OMDB, extract the movie title and year (if available) from the user's query. If the year is not known, set the year to null.
                    - For movie queries to Wikipedia, use the format `title (year film)` as the query. If the year is not known, use the format `title`.
                    - For general knowledge questions to Wikipedia, extract the article title(s) from the user's query.
                      - Query only for article titles. If the query is about specific information, query just for the article title. The specific information can be found in the article itself.
                      - Wikipedia articles are usually titled after the subject/noun they're about. For example:
                        - The name of a person
                        - The name of a place
                        - The title of a book, movie, or other work/media
                        - The most common name for a theory, concept, or idea
                        - When there's ambiguity, an additional disambiguation term may be needed in parentheses. For example, "Mercury (element)" vs "Mercury (planet)".
                        - If you're not sure what the article name should be favor a shorter, broader name over a more specific one.
                      - You may need to query multiple articles to get the full answer - that's okay.
                """,
            },
            {
                "role": "user",
                "content": query,
            },
        ],
        tools=[query_omdb, query_wikipedia, query_google],
        think=QUERY_ROUTER_ENABLE_THINKING,
        options={"num_ctx": 8192},
    )

    # uncomment to see the agent's thought process
    print("Thinking:", response.message.thinking)
    return response.message.tool_calls

# a helper to print tool calls nicely
def print_tools(tool_calls):
    if tool_calls is None:
        print("No tools called")
        return
    print(", ".join([f"{tc.function.name}({tc.function.arguments})" for tc in tool_calls]))


# examples
example_queries = [
    "tell me about the 1995 movie rumplestilskin",
    "tell me about rumplestilskin (1995 film)",
    "Who directed Troll 2?",
    "What years did Napoleon Bonaparte rule France?",
    "Search google for the capital of France",
    "What's the population of Dallas, New York, and Los Angeles?",
    "What's the current temperature in New York City?",
    "What's 2 + 2?",
]

for query in example_queries:
    print(f"Query: {query}")
    tool_calls = get_llm_tool_calls_from_query(query)
    print("Tool calls:", end="")
    print_tools(tool_calls)
    print("-----")