import json
from ollama import chat

fc_model_path = 'qwen2.5-0.5b-fc'

class FCModel:
    def __init__(self, fc_model_path):
        self._model_path = fc_model_path

    def get_chat(self, text):
        messages = [
            {
                "role": "system",
                "content": (
                    "你是一个人工智能助手。根据用户的输入返回相应的function内容，如果没有相应的function call指令则返回 None 。"
                )
            },
            {
                "role": "user",
                "content": text
            }
        ] 
        response = chat(
            model=self._model_path,
            messages=messages,
        )
        return response

    def func_response(self, text, func_map):
        response = self.get_chat(text)
        print("response:", response)
        content = response['message']['content']
        print("content:", content)
    
        try:
            content = json.loads(content)
            func_name = content.get('function', '').lower()
            if not func_name:
                print("No function name")
                return False
            if func_name not in func_map:
                print(f"function name {func_name} not in function map")
                return False
            args = content.get('arguments', {})
            print("start to execute function:", func_name)
            if not args:
                func_map[func_name]()
                return True
            else:
                func_map[func_name](**args)
                return True
        except Exception as e:
            return False

