Coverage for packages/registry/src/langgate/registry/local.py: 54%

1"""LocalRegistryClient for direct registry access.""" 

2 

3from collections.abc import Sequence 

4from datetime import timedelta 

5from typing import Generic, cast, get_args 

6 

7from langgate.client.protocol import BaseRegistryClient, LLMInfoT 

8from langgate.core.logging import get_logger 

9from langgate.core.models import LLMInfo 

10from langgate.registry.models import ModelRegistry 

11 

12logger = get_logger(__name__) 

13 

14 

15class BaseLocalRegistryClient(BaseRegistryClient[LLMInfoT], Generic[LLMInfoT]): 

16 """ 

17 Base local registry client that calls ModelRegistry directly. 

18 

19 This client is used when you want to embed the registry in your application 

20 rather than connecting to a remote registry service. 

21 

22 Type Parameters: 

23 LLMInfoT: The LLMInfo-derived model class to use for responses 

24 """ 

25 

26 __orig_bases__: tuple 

27 model_info_cls: type[LLMInfoT] 

28 

29 def __init__(self, model_info_cls: type[LLMInfoT] | None = None): 

30 """Initialize the client with a ModelRegistry instance.""" 

31 # Cache refreshing is no-op for local registry clients. 

32 # Since this client is local, we don't need to refresh the cache. 

33 # TODO: Move caching to the base HTTP client class instead. 

34 cache_ttl = timedelta(days=365) 

35 super().__init__(cache_ttl=cache_ttl) 

36 self.registry = ModelRegistry() 

37 

38 # Set model_info_cls if provided, otherwise it is inferred from the class 

39 if model_info_cls is not None: 

40 self.model_info_cls = model_info_cls 

41 

42 logger.debug("initialized_base_local_registry_client") 

43 

44 def __init_subclass__(cls, **kwargs): 

45 """Set up model class when this class is subclassed.""" 

46 super().__init_subclass__(**kwargs) 

47 

48 if not hasattr(cls, "model_info_cls"): 

49 cls.model_info_cls = cls._get_model_info_class() 

50 

51 @classmethod 

52 def _get_model_info_class(cls) -> type[LLMInfoT]: 

53 """Extract the model class from generic type parameters.""" 

54 return get_args(cls.__orig_bases__[0])[0] 

55 

56 async def _fetch_model_info(self, model_id: str) -> LLMInfoT: 

57 """Get information about a model directly from registry. 

58 

59 Args: 

60 model_id: The ID of the model to get information for 

61 

62 Returns: 

63 Information about the requested model with the type expected by this client 

64 """ 

65 # Get the model info from the registry (always returns LLMInfo) 

66 info = self.registry.get_model_info(model_id) 

67 

68 # If model_info_cls is LLMInfo (not a subclass), we can return it as-is 

69 if self.model_info_cls is LLMInfo: 

70 return cast(LLMInfoT, info) 

71 

72 # Otherwise, validate against the subclass schema 

73 return self.model_info_cls.model_validate(info.model_dump()) 

74 

75 async def _fetch_all_models(self) -> Sequence[LLMInfoT]: 

76 """List all available models directly from registry. 

77 

78 Returns: 

79 A sequence of model information objects of the type expected by this client. 

80 """ 

81 models = self.registry.list_models() 

82 

83 # If model_info_cls is LLMInfo (not a subclass), we can return the list as-is 

84 if self.model_info_cls is LLMInfo: 

85 return cast(Sequence[LLMInfoT], models) 

86 

87 # Otherwise, we need to validate each model against the subclass schema 

88 return [ 

89 self.model_info_cls.model_validate(model.model_dump()) for model in models 

90 ] 

91 

92 

93class LocalRegistryClient(BaseLocalRegistryClient[LLMInfo]): 

94 """ 

95 Local registry client that calls ModelRegistry directly using the default LLMInfo schema. 

96 

97 This is implemented as a singleton for convenient access across an application. 

98 """ 

99 

100 _instance = None 

101 

102 def __new__(cls, *args, **kwargs): 

103 """Create or return the singleton instance.""" 

104 if cls._instance is None: 

105 cls._instance = super().__new__(cls) 

106 return cls._instance 

107 

108 def __init__(self): 

109 """Initialize the client with a ModelRegistry instance.""" 

110 if not hasattr(self, "_initialized"): 

111 super().__init__() 

112 self._initialized = True 

113 logger.debug("initialized_local_registry_client_singleton")