base.py
 1  from abc import ABC
 2  from typing import List, Optional, Dict
 3  from langchain_core.language_models.chat_models import BaseChatModel
 4  from ..capabilities import ModelCapability, CapableModel
 5  
 6  # Instead of multiple inheritance, use composition
 7  
 8  
 9  class LLMModel(BaseChatModel):
10      """Enhanced base model with capabilities"""
11  
12      def __init__(self, capabilities: set[ModelCapability]):
13          self._capabilities = capabilities
14          super().__init__()
15  
16      def get_capabilities(self) -> set[ModelCapability]:
17          return self._capabilities
18  
19  
20  class LLMProvider(ABC):
21      def create_model(self, name: str, model: str, **kwargs) -> LLMModel:
22          pass
23  
24      def list_models(self, capabilities: Optional[set[ModelCapability]] = None) -> Dict[str, set[ModelCapability]]:
25          pass
26  
27      @property
28      def name(self) -> str:
29          raise NotImplementedError
30  
31      @property
32      def capabilities(self) -> Dict[str, set[ModelCapability]]:
33          raise NotImplementedError