from typing import Optional
from class_registry import ClassRegistry

from HeGFlow.llm.llm import LLM


class LLMRegistry:
    registry = ClassRegistry()

    @classmethod
    def register(cls, *args, **kwargs):
        return cls.registry.register(*args, **kwargs)
    
    @classmethod
    def keys(cls):
        return cls.registry.keys()

    @classmethod
    def get(cls, model_name: Optional[str] = None) -> LLM:
        if model_name is None or model_name=="":
            model_name = "gpt-4o"

        if 'llama' in model_name or 'qwen' in model_name:
            
            
            model = cls.registry.get('llama', model_name)
        elif 'deepseek' in model_name:
           
            model = cls.registry.get('deepseek', model_name)
        elif model_name == 'mock':
            model = cls.registry.get(model_name)
        else: 
           
            model = cls.registry.get('GPTChat', model_name)

        return model
