From a4c95fd62b8e6e731b40b47a87ebb6f4255bfade Mon Sep 17 00:00:00 2001
From: Alex <a@tushynski.me>
Date: Fri, 15 Nov 2024 12:17:25 +0000
Subject: [PATCH 1/2] feat: add google ai

---
 application/llm/google_ai.py   | 52 ++++++++++++++++++++++++++++++++++
 application/llm/llm_creator.py |  4 ++-
 2 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 application/llm/google_ai.py

diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py
new file mode 100644
index 000000000..ea460345d
--- /dev/null
+++ b/application/llm/google_ai.py
@@ -0,0 +1,52 @@
+from application.llm.base import BaseLLM
+from application.core.settings import settings
+
+
+
+
+class GoogleLLM(BaseLLM):
+
+    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
+
+        super().__init__(*args, **kwargs)
+        self.api_key = api_key
+        self.user_api_key = user_api_key
+
+    def _clean_messages_google(self, messages):
+        return [
+            {
+                "role": "model" if message["role"] == "system" else message["role"],
+                "parts": [message["content"]],
+            }
+            for message in messages[1:]
+        ]
+
+    def _raw_gen(
+        self,
+        baseself,
+        model,
+        messages,
+        stream=False,
+        **kwargs
+    ):  
+        import google.generativeai as genai
+        genai.configure(api_key=self.api_key)
+        model = genai.GenerativeModel(model, system_instruction=messages[0]["content"])
+        response = model.generate_content(self._clean_messages_google(messages))
+        return response.text
+
+    def _raw_gen_stream(
+        self,
+        baseself,
+        model,
+        messages,
+        stream=True,
+        **kwargs
+    ):  
+        import google.generativeai as genai
+        genai.configure(api_key=self.api_key)
+        model = genai.GenerativeModel(model, system_instruction=messages[0]["content"])
+        response = model.generate_content(self._clean_messages_google(messages), stream=True)
+        for line in response:
+            if line.text is not None:
+                yield line.text
\ No newline at end of file
diff --git a/application/llm/llm_creator.py b/application/llm/llm_creator.py
index 6a19de107..f32089de8 100644
--- a/application/llm/llm_creator.py
+++ b/application/llm/llm_creator.py
@@ -6,6 +6,7 @@
 from application.llm.anthropic import AnthropicLLM
 from application.llm.docsgpt_provider import DocsGPTAPILLM
 from application.llm.premai import PremAILLM
+from application.llm.google_ai import GoogleLLM
 
 
 class LLMCreator:
@@ -18,7 +19,8 @@ class LLMCreator:
         "anthropic": AnthropicLLM,
         "docsgpt": DocsGPTAPILLM,
         "premai": PremAILLM,
-        "groq": GroqLLM
+        "groq": GroqLLM,
+        "google": GoogleLLM
     }
 
     @classmethod

From 144ab61e079cb37fa93baa4095d5ae569b66adbf Mon Sep 17 00:00:00 2001
From: Alex <a@tushynski.me>
Date: Fri, 15 Nov 2024 12:19:43 +0000
Subject: [PATCH 2/2] fix: ruff lint

---
 application/llm/google_ai.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py
index ea460345d..df252abfc 100644
--- a/application/llm/google_ai.py
+++ b/application/llm/google_ai.py
@@ -1,8 +1,4 @@
 from application.llm.base import BaseLLM
-from application.core.settings import settings
-
-
-
 
 class GoogleLLM(BaseLLM):