Merge pull request #1441 from arc53/google-llm

application/llm/google_ai.py

@@ -0,0 +1,48 @@
+from application.llm.base import BaseLLM
+
+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

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