From 0d036f6a6eb0221ad5fce9f85f868c51322a4492 Mon Sep 17 00:00:00 2001
From: "Khalid F. Ahmed" <khfaiz949@gmail.com>
Date: Tue, 10 Sep 2024 09:16:04 +0300
Subject: [PATCH] Adding Google Gemini

---
 requirements.txt |  1 +
 src/gpt.py       | 19 +++++++++++++++----
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index de21428..11127a8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,6 +17,7 @@ pdfminer.six==20221105
 inputimeout==1.0.4
 langchain-ollama==0.1.3
 langchain-anthropic==0.1.3
+langchain-google-genai==1.0.10
 jsonschema==4.23.0
 jsonschema-specifications==2023.12.1
 httpx~=0.27.2
diff --git a/src/gpt.py b/src/gpt.py
index e87c6f6..c82e02e 100644
--- a/src/gpt.py
+++ b/src/gpt.py
@@ -62,6 +62,16 @@ class OllamaModel(AIModel):
         return response
 
 
+class GeminiModel(AIModel):
+    def __init__(self, api_key:str, llm_model: str, llm_api_url: str):
+        from langchain_google_genai import ChatGoogleGenerativeAI
+        self.model = ChatGoogleGenerativeAI(model=llm_model, google_api_key=api_key)
+
+    def invoke(self, prompt: str) -> str:
+        response = self.model.invoke(prompt)
+        return response
+
+
 class AIAdapter:
     def __init__(self, config: dict, api_key: str):
         self.model = self._create_model(config, api_key)
@@ -79,6 +89,8 @@ class AIAdapter:
             return ClaudeModel(api_key, llm_model, llm_api_url)
         elif llm_model_type == "ollama":
             return OllamaModel(api_key, llm_model, llm_api_url)
+        elif llm_model_type == "gemini":
+            return GeminiModel(api_key, llm_model, llm_api_url)
         else:
             raise ValueError(f"Unsupported model type: {llm_model_type}")
 
@@ -88,7 +100,7 @@ class AIAdapter:
 
 class LLMLogger:
 
-    def __init__(self, llm: Union[OpenAIModel, OllamaModel, ClaudeModel]):
+    def __init__(self, llm: Union[OpenAIModel, OllamaModel, ClaudeModel, GeminiModel]):
 
         self.llm = llm
         logger.debug("LLMLogger successfully initialized with LLM: %s", llm)
@@ -203,7 +215,7 @@ class LLMLogger:
 
 class LoggerChatModel:
 
-    def __init__(self, llm: Union[OpenAIModel, OllamaModel, ClaudeModel]):
+    def __init__(self, llm: Union[OpenAIModel, OllamaModel, ClaudeModel, GeminiModel]):
         self.llm = llm
         logger.debug(
             "LoggerChatModel successfully initialized with LLM: %s", llm)
@@ -494,8 +506,7 @@ class GPTAnswerer:
         if resume_section is None:
             logger.error(
                 "Section '%s' not found in either resume or job_application_profile.", section_name)
-            raise ValueError(f"Section '{
-                             section_name}' not found in either resume or job_application_profile.")
+            raise ValueError(f"Section '{section_name}' not found in either resume or job_application_profile.")
         chain = chains.get(section_name)
         if chain is None:
             logger.error("Chain not defined for section '%s'", section_name)