refinement

gemini.py

@@ -3,31 +3,39 @@ import os
 from dotenv import load_dotenv
 from typing import Optional
 
+generation_config=genai.types.GenerationConfig(
+        # Only one candidate for now.
+        #candidate_count=1,
+        #stop_sequences=['x'],
+        #max_output_tokens=1024,
+        temperature=1.0
+)
 class GeminiModel:
     """
-    This class is used to interact with the google llm models for text generation.
+    This class is used to interact with the Google LLM models for text generation.
 
     Args:
-        model: The name of the model to be used. Defaults to 'text-bison@002'.
+        model: The name of the model to be used. Defaults to 'gemini-pro'.
         max_output_tokens: The maximum number of tokens to generate. Defaults to 1024.
-        top_P: The probability of generating the next token. Defaults to 1.0.
+        top_p: The probability of generating the next token. Defaults to 1.0.
         temperature: The temperature of the model. Defaults to 0.0.
         top_k: The number of top tokens to consider. Defaults to 5.
     """
 
-    model=None
     def __init__(self,
                  model: Optional[str] = 'gemini-pro',
-                 max_output_tokens: Optional[int] = 1024,
-                 top_p: Optional[float] = 1.0,
-                 temperature: Optional[float] = 0.0,
-                 top_k: Optional[int] = 5
                  ):
-        load_dotenv() 
+        
+        
+        load_dotenv()
         genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
-        self.model=genai.GenerativeModel(model)
+        self.model = genai.GenerativeModel(model)
+        
+
+    def execute(self, prompt: str) -> str:
         
-    def execute(self,prompt):
-        print(prompt)
-        response=self.model.generate_content(prompt)
-        return response
+        try:
+            response = self.model.generate_content(prompt, generation_config=None)
+            return response.text
+        except Exception as e:
+            return f"An error occurred: {e}"

handler.py

@@ -15,5 +15,5 @@ def code_review(code,c_prompt=None):
     
     model=GeminiModel()    
     res= model.execute(prompt)
-    return res.text
+    return res