#!/usr/bin/env python3

# dummy weather agent
# illustrating communication with a LLM

import json
import requests
from agent_tools import ALL_TOOLS, weather_tool

# configuration for an OpenAI-compatible API endpoint
API_URL = "https://litellm.s.studiumdigitale.uni-frankfurt.de/v1/chat/completions"
API_KEY = "XXX"
MODEL   = "openai-gpt-oss-120b"

# 'Bearer' is an authentification protocol
HEADERS = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json",
}

# ---------------------------
# 1. initialize messages with system and user prompts
# ---------------------------
prompt_system = "You are a helpful assistant. "\
                "Use the provided tools to answer weather queries."
prompt_user   = "Is it raining in Berlin?"
messages = [ { "role": "system", "content": prompt_system},
             { "role": "user"  , "content": prompt_user}, ]

# ---------------------------
# 2. prepare payload 
# ---------------------------
tools = [weather_tool,]      # list of tools

payload = {"model": MODEL, 
        "messages": messages,
           "tools": tools,
     'tool_choice': 'auto',
     'temperature': 0.2}

# ---------------------------
# 3. send prompts and tool definitions to the LLM
#    LLM may sent back incorrect format,
#    if format is correct, convert response string into JSON
# ---------------------------
try:
    resp = requests.post(API_URL, headers=HEADERS, json=payload)
    resp.raise_for_status()
except requests.exceptions.HTTPError as e:
    print(f"First request failed: {e}")
    print(f"First response body : {resp.text}")  # for details of the API raise
else:
    response = resp.json()

# ---------------------------
# 4. extract message
# ---------------------------
response_message = response["choices"][0]["message"]

# ---------------------------
# 5. check if the LLM decided to call a tool
# ---------------------------
tool_calls = response_message.get("tool_calls")

if tool_calls:
    messages.append(response_message)     # append to chat history

    for tool_call in tool_calls:
        function_name = tool_call["function"]["name"]
        function_args = json.loads(tool_call["function"]["arguments"])

        if function_name in ALL_TOOLS:    # executing tool
            tool_output = ALL_TOOLS[function_name](**function_args)

            messages.append(              # append result to chat history
                { "role"        : "tool",
                  "tool_call_id": tool_call["id"],
                  "content"     : tool_output,
                }          )
else:
    print("No tool call triggered.")
    print(reponse_message.get("content"))
    exit()

# ---------------------------
# 6. dump chat hitory to file
# ---------------------------

with open("chat_history.json", "w", encoding="utf-8") as f:
    json.dump(messages, f, indent=2, ensure_ascii=False)

# ---------------------------
# 7. send updated history back to the LLM 
#    for the final natural language answer
# ---------------------------

payload = {"model": MODEL, "messages": messages}
try:
    resp = requests.post(API_URL, headers=HEADERS, json=payload)
    resp.raise_for_status()
except requests.exceptions.HTTPError as e:
    print(f"Final request failed: {e}")
    print(f"Final response body : {resp.text}")  # for details of the API raise
else:
    response = resp.json()

LLM_answer = response["choices"][0]["message"]["content"]
print("\nuser prompt:\n", prompt_user)
print("\nLLM response:\n", LLM_answer)