from litellm import completion 
import os
from dotenv import load_dotenv
load_dotenv()
os.environ["TOGETHERAI_API_KEY"] = "tgp_v1_l1Go7VAvVghKE-5RDZ4si4oJEXTynZx2qDvflNIixbU"

from litellm import completion 


# messages = [{"role": "user", "content": "Write me a poem about the blue sky"}]

tools = [
    {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "Perform basic arithmetic operations",
            "parameters": {
                "type": "object",
                "properties": {
                    "operation": {"type": "string", "description": "The operation to perform"},
                    "a": {"type": "number", "description": "First number"},
                    "b": {"type": "number", "description": "Second number"}
                },
                "required": ["operation", "a", "b"]
            }
        }
    }
]

messages = [
    {"role": "user", "content": "Calculate 15 * 23"}
]

response = completion(
    model="together_ai/openai/gpt-oss-120b", 
    messages=messages,
    tools=tools,
    allowed_openai_params=['tools', 'tool_choice'],
    tool_choice="auto"
)
print("Response:", response.choices[0].message)
if hasattr(response.choices[0].message, 'tool_calls') and response.choices[0].message.tool_calls:
    print("Tool calls:", response.choices[0].message.tool_calls)



# from together import Together

# client = Together()

# response = client.chat.completions.create(
#     model="openai/gpt-oss-120b",
#     # messages=[
#     #   {
#     #     "role": "user",
#     #     "content": "What are some fun things to do in New York?"
#     #   }
#     # ],
#     messages=messages,
#     tools=tools,
#     tool_choice="auto"
# )
# print(response.choices[0].message.content)