fix aio

Author: CTY-git

.github/workflows/test.yml

@@ -148,7 +148,7 @@ jobs:
           poetry run patchwork ResolveIssue --log debug \
           --patched_api_key=${{ secrets.PATCHED_API_KEY }} \
           --github_api_key=${{ secrets.SCM_GITHUB_KEY }} \
-          --issue_url=https://github.com/patched-codes/patchwork/issues/20 \
+          --issue_url=https://github.com/patched-codes/patchwork/issues/1039 \
           --disable_telemetry
 
   main-test:

patchwork/common/client/llm/google.py

@@ -16,6 +16,8 @@
 from openai.types.chat import (
     ChatCompletionMessage,
     ChatCompletionMessageParam,
+    ChatCompletionToolChoiceOptionParam,
+    ChatCompletionToolParam,
     completion_create_params,
 )
 from openai.types.chat.chat_completion import ChatCompletion, Choice
@@ -55,7 +57,24 @@ def get_models(self) -> set[str]:
     def is_model_supported(self, model: str) -> bool:
         return model in self.get_models()
 
-    def is_prompt_supported(self, messages: Iterable[ChatCompletionMessageParam], model: str) -> int:
+    def is_prompt_supported(
+            self,
+            messages: Iterable[ChatCompletionMessageParam],
+            model: str,
+            frequency_penalty: Optional[float] | NotGiven = NOT_GIVEN,
+            logit_bias: Optional[Dict[str, int]] | NotGiven = NOT_GIVEN,
+            logprobs: Optional[bool] | NotGiven = NOT_GIVEN,
+            max_tokens: Optional[int] | NotGiven = NOT_GIVEN,
+            n: Optional[int] | NotGiven = NOT_GIVEN,
+            presence_penalty: Optional[float] | NotGiven = NOT_GIVEN,
+            response_format: dict | completion_create_params.ResponseFormat | NotGiven = NOT_GIVEN,
+            stop: Union[Optional[str], List[str]] | NotGiven = NOT_GIVEN,
+            temperature: Optional[float] | NotGiven = NOT_GIVEN,
+            tools: Iterable[ChatCompletionToolParam] | NotGiven = NOT_GIVEN,
+            tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
+            top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
+            top_p: Optional[float] | NotGiven = NOT_GIVEN,
+    ) -> int:
         system, chat = self.__openai_messages_to_google_messages(messages)
         gen_model = generativeai.GenerativeModel(model_name=model, system_instruction=system)
         try:
@@ -96,9 +115,11 @@ def chat_completion(
         max_tokens: Optional[int] | NotGiven = NOT_GIVEN,
         n: Optional[int] | NotGiven = NOT_GIVEN,
         presence_penalty: Optional[float] | NotGiven = NOT_GIVEN,
-        response_format: completion_create_params.ResponseFormat | NotGiven = NOT_GIVEN,
+        response_format: str | completion_create_params.ResponseFormat | NotGiven = NOT_GIVEN,
         stop: Union[Optional[str], List[str]] | NotGiven = NOT_GIVEN,
         temperature: Optional[float] | NotGiven = NOT_GIVEN,
+        tools: Iterable[ChatCompletionToolParam] | NotGiven = NOT_GIVEN,
+        tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
     ) -> ChatCompletion:

patchwork/common/client/llm/openai_.py

@@ -63,7 +63,24 @@ def is_model_supported(self, model: str) -> bool:
     def __get_model_limits(self, model: str) -> int:
         return self.__MODEL_LIMITS.get(model, 128_000)
 
-    def is_prompt_supported(self, messages: Iterable[ChatCompletionMessageParam], model: str) -> int:
+    def is_prompt_supported(
+            self,
+            messages: Iterable[ChatCompletionMessageParam],
+            model: str,
+            frequency_penalty: Optional[float] | NotGiven = NOT_GIVEN,
+            logit_bias: Optional[Dict[str, int]] | NotGiven = NOT_GIVEN,
+            logprobs: Optional[bool] | NotGiven = NOT_GIVEN,
+            max_tokens: Optional[int] | NotGiven = NOT_GIVEN,
+            n: Optional[int] | NotGiven = NOT_GIVEN,
+            presence_penalty: Optional[float] | NotGiven = NOT_GIVEN,
+            response_format: dict | completion_create_params.ResponseFormat | NotGiven = NOT_GIVEN,
+            stop: Union[Optional[str], List[str]] | NotGiven = NOT_GIVEN,
+            temperature: Optional[float] | NotGiven = NOT_GIVEN,
+            tools: Iterable[ChatCompletionToolParam] | NotGiven = NOT_GIVEN,
+            tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
+            top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
+            top_p: Optional[float] | NotGiven = NOT_GIVEN,
+    ) -> int:
         # might not implement model endpoint
         if self.__is_not_openai_url():
             return 1