ReadEmail Step

patchwork/common/client/llm/aio.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+from pathlib import Path
 
 from openai.types.chat import (
     ChatCompletion,
@@ -54,10 +55,11 @@ def is_prompt_supported(
         tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
+        file: Path | NotGiven = NOT_GIVEN,
     ) -> int:
         for client in self.__clients:
             if client.is_model_supported(model):
-                return client.is_prompt_supported(
+                inputs = dict(
                     messages=messages,
                     model=model,
                     frequency_penalty=frequency_penalty,
@@ -74,6 +76,9 @@ def is_prompt_supported(
                     top_logprobs=top_logprobs,
                     top_p=top_p,
                 )
+                if file is not NotGiven:
+                    inputs["file"] = file
+                return client.is_prompt_supported(**inputs)
         return -1
 
     def truncate_messages(
@@ -101,27 +106,31 @@ def chat_completion(
         tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
+        file: Path | NotGiven = NOT_GIVEN,
     ) -> ChatCompletion:
         for client in self.__clients:
             if client.is_model_supported(model):
                 logger.debug(f"Using {client.__class__.__name__} for model {model}")
-                return client.chat_completion(
-                    messages,
-                    model,
-                    frequency_penalty,
-                    logit_bias,
-                    logprobs,
-                    max_tokens,
-                    n,
-                    presence_penalty,
-                    response_format,
-                    stop,
-                    temperature,
-                    tools,
-                    tool_choice,
-                    top_logprobs,
-                    top_p,
+                inputs = dict(
+                    messages=messages,
+                    model=model,
+                    frequency_penalty=frequency_penalty,
+                    logit_bias=logit_bias,
+                    logprobs=logprobs,
+                    max_tokens=max_tokens,
+                    n=n,
+                    presence_penalty=presence_penalty,
+                    response_format=response_format,
+                    stop=stop,
+                    temperature=temperature,
+                    tools=tools,
+                    tool_choice=tool_choice,
+                    top_logprobs=top_logprobs,
+                    top_p=top_p,
                 )
+                if file is not NotGiven:
+                    inputs["file"] = file
+                return client.chat_completion(**inputs)
         client_names = [client.__class__.__name__ for client in self.__original_clients]
         raise ValueError(
             f"Model {model} is not supported by {client_names} clients. "

patchwork/common/client/llm/anthropic.py

@@ -3,6 +3,7 @@
 import json
 import time
 from functools import lru_cache
+from pathlib import Path
 
 from anthropic import Anthropic
 from anthropic.types import Message, MessageParam, TextBlockParam
@@ -224,6 +225,7 @@ def is_prompt_supported(
         tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
+        file: Path | NotGiven = NOT_GIVEN,
     ) -> int:
         model_limit = self.__get_model_limit(model)
         input_kwargs = self.__adapt_chat_completion_request(
@@ -273,6 +275,7 @@ def chat_completion(
         tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
+        file: Path | NotGiven = NOT_GIVEN,
     ) -> ChatCompletion:
         input_kwargs = self.__adapt_chat_completion_request(
             messages=messages,

patchwork/common/client/llm/google.py

@@ -1,17 +1,20 @@
 from __future__ import annotations
 
-import functools
 import time
-
-from google import generativeai
-from google.generativeai.types.content_types import (
-    add_object_type,
-    convert_to_nullable,
-    strip_titles,
-    unpack_defs,
+from functools import lru_cache
+from pathlib import Path
+
+import magic
+from google import genai
+from google.genai import types
+from google.genai.types import (
+    CountTokensConfig,
+    File,
+    GenerateContentConfig,
+    GenerateContentResponse,
+    Model,
+    Part,
 )
-from google.generativeai.types.generation_types import GenerateContentResponse
-from google.generativeai.types.model_types import Model
 from openai.types import CompletionUsage
 from openai.types.chat import (
     ChatCompletionMessage,
@@ -21,17 +24,13 @@
     completion_create_params,
 )
 from openai.types.chat.chat_completion import ChatCompletion, Choice
-from typing_extensions import Any, Dict, Iterable, List, Optional, Union
+from pydantic import BaseModel
+from typing_extensions import Any, Dict, Iterable, List, Optional, Type, Union
 
 from patchwork.common.client.llm.protocol import NOT_GIVEN, LlmClient, NotGiven
 from patchwork.common.client.llm.utils import json_schema_to_model
 
 
-@functools.lru_cache
-def _cached_list_model_from_google() -> list[Model]:
-    return list(generativeai.list_models())
-
-
 class GoogleLlmClient(LlmClient):
     __SAFETY_SETTINGS = [
         dict(category="HARM_CATEGORY_HATE_SPEECH", threshold="BLOCK_NONE"),
@@ -43,20 +42,54 @@ class GoogleLlmClient(LlmClient):
 
     def __init__(self, api_key: str):
         self.__api_key = api_key
-        generativeai.configure(api_key=api_key)
+        self.client = genai.Client(api_key=api_key)
+
+    @lru_cache(maxsize=1)
+    def __get_models_info(self) -> list[Model]:
+        return list(self.client.models.list())
 
     def __get_model_limits(self, model: str) -> int:
-        for model_info in _cached_list_model_from_google():
-            if model_info.name == f"{self.__MODEL_PREFIX}{model}":
+        for model_info in self.__get_models_info():
+            if model_info.name == f"{self.__MODEL_PREFIX}{model}" and model_info.input_token_limit is not None:
                 return model_info.input_token_limit
         return 1_000_000
 
+    @lru_cache
     def get_models(self) -> set[str]:
-        return {model.name.removeprefix(self.__MODEL_PREFIX) for model in _cached_list_model_from_google()}
+        return {model_info.name.removeprefix(self.__MODEL_PREFIX) for model_info in self.__get_models_info()}
 
     def is_model_supported(self, model: str) -> bool:
         return model in self.get_models()
 
+    def __upload(self, file: Path | NotGiven) -> Part | File | None:
+        if file is NotGiven:
+            return None
+
+        file_bytes = file.read_bytes()
+
+        try:
+            mime_type = magic.Magic(mime=True, uncompress=True).from_buffer(file_bytes)
+            return types.Part.from_bytes(data=file_bytes, mime_type=mime_type)
+        except Exception as e:
+            pass
+
+        cleaned_name = "".join([char for char in file.name.lower() if char.isalnum()])
+        try:
+            file_ref = self.client.files.get(name=cleaned_name)
+            if file_ref.error is None:
+                return file_ref
+        except Exception as e:
+            pass
+
+        try:
+            file_ref = self.client.files.upload(file=file, config=dict(name=cleaned_name))
+            if file_ref.error is None:
+                return file_ref
+        except Exception as e:
+            pass
+
+        return None
+
     def is_prompt_supported(
         self,
         messages: Iterable[ChatCompletionMessageParam],
@@ -74,11 +107,23 @@ def is_prompt_supported(
         tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
+        file: Path | NotGiven = NOT_GIVEN,
     ) -> int:
-        system, chat = self.__openai_messages_to_google_messages(messages)
-        gen_model = generativeai.GenerativeModel(model_name=model, system_instruction=system)
+        system, contents = self.__openai_messages_to_google_messages(messages)
+
+        file_ref = self.__upload(file)
+        if file_ref is not None:
+            contents.append(file_ref)
+
         try:
-            token_count = gen_model.count_tokens(chat).total_tokens
+            token_response = self.client.models.count_tokens(
+                model=model,
+                contents=contents,
+                config=CountTokensConfig(
+                    system_instruction=system,
+                ),
+            )
+            token_count = token_response.total_tokens
         except Exception as e:
             return -1
         model_limit = self.__get_model_limits(model)
@@ -122,6 +167,7 @@ def chat_completion(
         tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
+        file: Path | NotGiven = NOT_GIVEN,
     ) -> ChatCompletion:
         generation_dict = dict(
             stop_sequences=[stop] if isinstance(stop, str) else stop,
@@ -141,20 +187,25 @@ def chat_completion(
             )
 
         system_content, contents = self.__openai_messages_to_google_messages(messages)
+        file_ref = self.__upload(file)
+        if file_ref is not None:
+            contents.append(file_ref)
 
-        model_client = generativeai.GenerativeModel(
-            model_name=model,
-            safety_settings=self.__SAFETY_SETTINGS,
-            generation_config=NOT_GIVEN.remove_not_given(generation_dict),
-            system_instruction=system_content,
+        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),
+            ),
         )
-        response = model_client.generate_content(contents=contents)
         return self.__google_response_to_openai_response(response, model)
 
     @staticmethod
     def __google_response_to_openai_response(google_response: GenerateContentResponse, model: str) -> ChatCompletion:
         choices = []
-        for candidate in google_response.candidates:
+        for index, candidate in enumerate(google_response.candidates):
             # note that instead of system, from openai, its model, from google.
             parts = [part.text or part.inline_data for part in candidate.content.parts]
 
@@ -167,7 +218,7 @@ def __google_response_to_openai_response(google_response: GenerateContentRespons
 
             choice = Choice(
                 finish_reason=finish_reason_map.get(candidate.finish_reason, "stop"),
-                index=candidate.index,
+                index=index,
                 message=ChatCompletionMessage(
                     content="\n".join(parts),
                     role="assistant",
@@ -191,18 +242,9 @@ def __google_response_to_openai_response(google_response: GenerateContentRespons
         )
 
     @staticmethod
-    def json_schema_to_google_schema(json_schema: dict[str, Any] | None) -> dict[str, Any] | None:
+    def json_schema_to_google_schema(json_schema: dict[str, Any] | None) -> Type[BaseModel] | None:
         if json_schema is None:
             return None
 
         model = json_schema_to_model(json_schema)
-        parameters = model.model_json_schema()
-        defs = parameters.pop("$defs", {})
-
-        for name, value in defs.items():
-            unpack_defs(value, defs)
-        unpack_defs(parameters, defs)
-        convert_to_nullable(parameters)
-        add_object_type(parameters)
-        strip_titles(parameters)
-        return parameters
+        return model

patchwork/common/client/llm/openai_.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import functools
+from pathlib import Path
 
 import tiktoken
 from openai import OpenAI
@@ -82,6 +83,7 @@ def is_prompt_supported(
         tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
+        file: Path | NotGiven = NOT_GIVEN,
     ) -> int:
         # might not implement model endpoint
         if self.__is_not_openai_url():
@@ -125,6 +127,7 @@ def chat_completion(
         tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
+        file: Path | NotGiven = NOT_GIVEN,
     ) -> ChatCompletion:
         input_kwargs = dict(
             messages=messages,

patchwork/common/client/llm/protocol.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from pathlib import Path
+
 from openai.types.chat import (
     ChatCompletion,
     ChatCompletionMessageParam,
@@ -51,6 +53,7 @@ def is_prompt_supported(
         tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
+        file: Path | NotGiven = NOT_GIVEN,
     ) -> int:
         ...
 
@@ -135,5 +138,6 @@ def chat_completion(
         tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
         top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
         top_p: Optional[float] | NotGiven = NOT_GIVEN,
+        file: Path | NotGiven = NOT_GIVEN,
     ) -> ChatCompletion:
         ...