update

Author: CTY-git

patchwork/common/client/llm/aio.py

@@ -35,7 +35,6 @@ def __init__(self, *clients: LlmClient):
                 self.__clients.append(client)
             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]:
         if model_settings is None:
@@ -205,8 +204,6 @@ 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:
@@ -219,8 +216,9 @@ def create_aio_client(inputs) -> "AioLlmClient" | None:
             clients.append(client)
 
         google_key = inputs.get("google_api_key")
-        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)))
+        is_gcp = bool(client_args.get("is_gcp") or os.environ.get("GOOGLE_GENAI_USE_VERTEXAI") or False)
+        if google_key is not None or is_gcp:
+            client = GoogleLlmClient(api_key=google_key, is_gcp=is_gcp)
             clients.append(client)
 
         anthropic_key = inputs.get("anthropic_api_key")

patchwork/common/client/llm/anthropic.py

@@ -2,7 +2,7 @@
 
 import json
 import time
-from functools import cached_property, lru_cache
+from functools import cached_property
 from pathlib import Path
 
 from anthropic import Anthropic

patchwork/common/client/llm/google_.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
+import os
 import time
-from functools import lru_cache
+from functools import lru_cache, partial
 from pathlib import Path
 
 import magic
+import vertexai
 from google import genai
-from google.auth.credentials import Credentials
+from google.auth.exceptions import GoogleAuthError
 from google.genai import types
+from google.genai.errors import APIError
 from google.genai.types import (
     CountTokensConfig,
     File,
@@ -42,9 +45,11 @@
     Type,
     Union,
 )
+from vertexai.generative_models import GenerativeModel, SafetySetting
 
 from patchwork.common.client.llm.protocol import NOT_GIVEN, LlmClient, NotGiven
 from patchwork.common.client.llm.utils import json_schema_to_model
+from patchwork.logger import logger
 
 
 class GoogleLlmClient(LlmClient):
@@ -53,6 +58,28 @@ class GoogleLlmClient(LlmClient):
         dict(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_NONE"),
         dict(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="BLOCK_NONE"),
         dict(category="HARM_CATEGORY_HARASSMENT", threshold="BLOCK_NONE"),
+        dict(category="HARM_CATEGORY_CIVIC_INTEGRITY", threshold="BLOCK_NONE"),
+    ]
+    __VERTEX_SAFETY_SETTINGS = [
+        SafetySetting(
+            category=SafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+            threshold=SafetySetting.HarmBlockThreshold.OFF,
+        ),
+        SafetySetting(
+            category=SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+            threshold=SafetySetting.HarmBlockThreshold.OFF,
+        ),
+        SafetySetting(
+            category=SafetySetting.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
+            threshold=SafetySetting.HarmBlockThreshold.OFF,
+        ),
+        SafetySetting(
+            category=SafetySetting.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=SafetySetting.HarmBlockThreshold.OFF
+        ),
+        SafetySetting(
+            category=SafetySetting.HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY,
+            threshold=SafetySetting.HarmBlockThreshold.OFF,
+        ),
     ]
     __MODEL_PREFIX = "models/"
 
@@ -63,6 +90,12 @@ def __init__(self, api_key: Optional[str] = None, is_gcp: bool = False):
             self.client = genai.Client(api_key=api_key)
         else:
             self.client = genai.Client(api_key=api_key, vertexai=True)
+            location = os.environ.get("GOOGLE_CLOUD_LOCATION", "global")
+            vertexai.init(
+                project=os.environ.get("GOOGLE_CLOUD_PROJECT"),
+                location=location,
+                api_endpoint=f"{location}-aiplatform.googleapis.com",
+            )
 
     @lru_cache(maxsize=1)
     def __get_models_info(self) -> list[Model]:
@@ -173,6 +206,8 @@ def is_prompt_supported(
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
         file: Path | NotGiven = NOT_GIVEN,
     ) -> int:
+        if self.__is_gcp:
+            return 1
         system, contents = self.__openai_messages_to_google_messages(messages)
 
         file_ref = self.__upload(file)
@@ -188,7 +223,12 @@ def is_prompt_supported(
                 ),
             )
             token_count = token_response.total_tokens
+        except GoogleAuthError:
+            raise
+        except APIError:
+            raise
         except Exception as e:
+            logger.debug(f"Error during token count at GoogleLlmClient: {e}")
             return -1
         model_limit = self.__get_model_limits(model)
         return model_limit - token_count
@@ -255,15 +295,25 @@ def chat_completion(
         if file_ref is not None:
             contents.append(file_ref)
 
-        response = self.client.models.generate_content(
-            model=model,
-            contents=contents,
-            config=GenerateContentConfig(
-                system_instruction=system_content,
-                safety_settings=self.__SAFETY_SETTINGS,
-                **NotGiven.remove_not_given(generation_dict),
-            ),
-        )
+        if not self.__is_gcp:
+            generate_content_func = partial(
+                self.client.models.generate_content,
+                model=model,
+                config=GenerateContentConfig(
+                    system_instruction=system_content,
+                    safety_settings=self.__SAFETY_SETTINGS,
+                    **NotGiven.remove_not_given(generation_dict),
+                ),
+            )
+        else:
+            vertexai_model = GenerativeModel(model, system_instruction=system_content)
+            generate_content_func = partial(
+                vertexai_model.generate_content,
+                safety_settings=self.__VERTEX_SAFETY_SETTINGS,
+                generation_config=NotGiven.remove_not_given(generation_dict),
+            )
+
+        response = generate_content_func(contents=contents)
         return self.__google_response_to_openai_response(response, model)
 
     @staticmethod

patchwork/common/utils/step_typing.py

@@ -76,15 +76,15 @@ def validate_step_type_config_with_inputs(
                 or f"Missing required input data because {key_name} is set: {', '.join(missing_and_keys)}",
             )
 
-    or_keys = set(step_type_config.or_op)
-    if len(or_keys) > 0:
-        missing_or_keys = or_keys.difference(input_keys)
-        if not is_key_set and len(missing_or_keys) == len(or_keys):
-            return (
-                False,
-                step_type_config.msg
-                or f"Missing required input: At least one of {', '.join(sorted([key_name, *or_keys]))} has to be set",
-            )
+    # or_keys = set(step_type_config.or_op)
+    # if len(or_keys) > 0:
+    #     missing_or_keys = or_keys.difference(input_keys)
+    #     if not is_key_set and len(missing_or_keys) == len(or_keys):
+    #         return (
+    #             False,
+    #             step_type_config.msg
+    #             or f"Missing required input: At least one of {', '.join(sorted([key_name, *or_keys]))} has to be set",
+    #         )
 
     xor_keys = set(step_type_config.xor_op)
     if len(xor_keys) > 0: