update

Author: CTY-git

patchwork/common/client/llm/aio.py

@@ -31,9 +31,10 @@ def __init__(self, *clients: LlmClient):
         self.__supported_models = set()
         for client in clients:
             try:
-                self.__supported_models.update(client.get_models())
+                client.test()
                 self.__clients.append(client)
-            except Exception:
+            except Exception as e:
+                logger.error(f"{client.__class__.__name__} Failed with exception: {e}")
                 pass
 
     def __get_model(self, model_settings: ModelSettings | None) -> Optional[str]:
@@ -45,6 +46,9 @@ def __get_model(self, model_settings: ModelSettings | None) -> Optional[str]:
 
         return model_name
 
+    def test(self) -> None:
+        pass
+
     async def request(
         self,
         messages: list[ModelMessage],
@@ -94,9 +98,6 @@ def model_name(self) -> str:
     def system(self) -> str:
         return next(iter(self.__clients)).system
 
-    def get_models(self) -> set[str]:
-        return self.__supported_models
-
     def is_model_supported(self, model: str) -> bool:
         return any(client.is_model_supported(model) for client in self.__clients)
 
@@ -204,6 +205,8 @@ def create_aio_client(inputs) -> "AioLlmClient" | None:
         clients = []
 
         client_args = {key[len("client_") :]: value for key, value in inputs.items() if key.startswith("client_")}
+        if os.environ.get("GOOGLE_GENAI_USE_VERTEXAI") == "true":
+            client_args["is_gcp"] = True
 
         patched_key = inputs.get("patched_api_key")
         if patched_key is not None:
@@ -216,8 +219,8 @@ def create_aio_client(inputs) -> "AioLlmClient" | None:
             clients.append(client)
 
         google_key = inputs.get("google_api_key")
-        if google_key is not None:
-            client = GoogleLlmClient(google_key, **client_args)
+        if google_key is not None or "is_gcp" in client_args.keys():
+            client = GoogleLlmClient(api_key=google_key, is_gcp=bool(client_args.get("is_gcp", False)))
             clients.append(client)
 
         anthropic_key = inputs.get("anthropic_api_key")

patchwork/common/client/llm/anthropic.py

@@ -245,9 +245,8 @@ def __adapt_chat_completion_request(
 
         return NotGiven.remove_not_given(input_kwargs)
 
-    @lru_cache(maxsize=None)
-    def get_models(self) -> set[str]:
-        return self.__definitely_allowed_models.union(set(f"{self.__allowed_model_prefix}*"))
+    def test(self):
+        return
 
     def is_model_supported(self, model: str) -> bool:
         return model in self.__definitely_allowed_models or model.startswith(self.__allowed_model_prefix)

patchwork/common/client/llm/google_.py

@@ -56,17 +56,20 @@ class GoogleLlmClient(LlmClient):
     ]
     __MODEL_PREFIX = "models/"
 
-    def __init__(self, api_key: str, is_gcp: bool = False):
+    def __init__(self, api_key: Optional[str] = None, is_gcp: bool = False):
         self.__api_key = api_key
         self.__is_gcp = is_gcp
-        if not is_gcp:
+        if not self.__is_gcp:
             self.client = genai.Client(api_key=api_key)
         else:
-            self.client = genai.Client(api_key=api_key, vertexai=True, credentials=Credentials())
+            self.client = genai.Client(api_key=api_key, vertexai=True)
 
     @lru_cache(maxsize=1)
     def __get_models_info(self) -> list[Model]:
-        return list(self.client.models.list())
+        if not self.__is_gcp:
+            return list(self.client.models.list())
+        else:
+            return list()
 
     def __get_pydantic_model(self, model_settings: ModelSettings | None) -> PydanticAiModel:
         if model_settings is None:
@@ -112,12 +115,15 @@ def __get_model_limits(self, model: str) -> int:
                 return model_info.input_token_limit
         return 1_000_000
 
-    @lru_cache
-    def get_models(self) -> set[str]:
-        return {model_info.name.removeprefix(self.__MODEL_PREFIX) for model_info in self.__get_models_info()}
+    def test(self):
+        return
 
     def is_model_supported(self, model: str) -> bool:
-        return model in self.get_models()
+        if not self.__is_gcp:
+            model_names = {model_info.name.removeprefix(self.__MODEL_PREFIX) for model_info in self.__get_models_info()}
+            return model in model_names
+        else:
+            return True
 
     def __upload(self, file: Path | NotGiven) -> Part | File | None:
         if isinstance(file, NotGiven):

patchwork/common/client/llm/openai_.py

@@ -96,17 +96,18 @@ def __is_not_openai_url(self):
         # We mainly use this to skip using the model endpoints.
         return self.__base_url is not None and self.__base_url != "https://api.openai.com/v1"
 
-    def get_models(self) -> set[str]:
+    def test(self):
         if self.__is_not_openai_url():
-            return set()
+            return
 
-        return _cached_list_models_from_openai(self.__api_key)
+        _cached_list_models_from_openai(self.__api_key)
+        return
 
     def is_model_supported(self, model: str) -> bool:
         # might not implement model endpoint
         if self.__is_not_openai_url():
             return True
-        return model in self.get_models()
+        return model in _cached_list_models_from_openai(self.__api_key)
 
     def __get_model_limits(self, model: str) -> int:
         return self.__MODEL_LIMITS.get(model, 128_000)

patchwork/common/client/llm/protocol.py

@@ -33,7 +33,7 @@ def remove_not_given(obj: Any) -> Union[None, dict[Any, Any], list[Any], Any]:
 
 class LlmClient(Model):
     @abstractmethod
-    def get_models(self) -> set[str]:
+    def test(self) -> None:
         ...
 
     @abstractmethod