import ollama

OLLAMA_CLIENT = ollama.Client(host="172.16.2.96:11434")

def get_location() -> str:
    """
    Get the current geographic location.

    Returns:
    str: Current geographic location.
    """
    return "Shanghai"


def get_weather(location: str) -> str:
    """
    Get the weather conditions for a specific location.

    Args:
    location (str): Geographic location.

    Returns:
    str: Weather conditions.
    """
    return "Sunny, temperature 25°C."

available_functions = {
    'get_location': get_location,
    'get_weather': get_weather,
}

def chat_generate(options):
    tool_calls = []
    for part in OLLAMA_CLIENT.chat(**options):
        yield part.message.content  # Expecting incremental streaming here
        options["messages"].append(part.message)
        if part.message.tool_calls:
            for tool in part.message.tool_calls:
                if function_to_call := available_functions.get(tool.function.name):
                    output = function_to_call(**tool.function.arguments)
                    tool_calls.append({
                        'role': 'tool', 'content': str(output), 'name': tool.function.name
                    })

    for result in tool_calls:
        print(result)
        options["messages"].append(result)

    if len(tool_calls) > 0:
        for result in chat_generate(options):
            yield result

for result in chat_generate({
    "model": "qwq",
    "messages": [
        {"role": "user", "content": "How is the weather now?"},
    ],
    "stream": True,
    "tools": [get_location, get_weather]
}):
    print(result)