from together import Together
from openai import OpenAI
import os
import json

class AI_LLM:
    def __init__(self):
        self.llm_version = os.getenv("LLM_VERSION")
        if self.llm_version == "OPENAI":
            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
            self.model = "gpt-4.1-mini-2025-04-14"
        else:
            self.client = OpenAI(api_key=os.getenv("TOGETHER_API_KEY"), base_url="https://api.together.xyz/v1")
            self.model = "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
        print(f"Using model: {self.model})")

    def talk_to_LLM(self, messages, functions_list, functions_module, stage_key_values):
        function_output = { "text": []}
        completion = self.client.chat.completions.create(
        model = self.model,
        messages = messages,
        tools = functions_list,
        ).choices[0].message

        if completion.tool_calls:
            messages.append(completion)
            for tool_call in completion.tool_calls:
                function_name = tool_call.function.name
                function_args = json.loads(tool_call.function.arguments)
                function = getattr(functions_module, function_name, None)
                if function:
                    function_output = function(**function_args, stage_key_values=stage_key_values)
                    for msg in function_output["text"]:
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tool_call.id,
                            "content": msg,
                        })
                else:
                    print(f"Function {function_name} not found in module {functions_module.__name__}")

        else:
            messages.append({
                "role": "assistant",
                "content": completion.content,
            })
            function_output["text"].append(completion.content)
        
        return function_output
    
    def get_structured_output(self, messages, structured_output_schema, functions_module, stage_key_values):
        function_output = {"text": []}
        completion = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            response_format=structured_output_schema
        ).choices[0].message

        function = getattr(functions_module, "process_structured_output", None)
        if function:
            function_output = function(completion.content, stage_key_values=stage_key_values)
        else:
            print(f"Function 'process_structured_output' not found in module {functions_module.__name__}")
            messages.append({
                "role": "assistant",
                "content": completion.content,
            })
            function_output["text"].append(completion.content)

        return function_output