Added support for Gemini AI model

Author: haseeb-heaven

.env

@@ -0,0 +1,3 @@
+PALM_API_KEY="XXXXXXXXXX"
+GEMINI_API_KEY="XXXXXXXXXX"
+OPENAI_API_KEY="XXXXXXXXXX"

libs/geminiai.py

@@ -0,0 +1,149 @@
+
+import traceback
+import google.generativeai as genai
+from dotenv import load_dotenv
+from libs.logger import logger
+import streamlit as st
+
+class GeminiAI:
+ def __init__(self, api_key, model="gemini-pro", temperature=0.1, max_output_tokens=2048,mode="balanced"):
+ self.api_key = api_key
+ self.model = model
+ self.temperature = temperature
+ self.max_output_tokens = max_output_tokens
+ self.mode = mode
+ self.top_k = 20
+ self.top_p = 0.85
+ self._configure()
+
+ # Dynamically construct guidelines based on session state
+ self.guidelines_list = []
+
+ if st.session_state["coding_guidelines"]["modular_code"]:
+ self.guidelines_list.append("- Ensure the method is modular in its approach.")
+ if st.session_state["coding_guidelines"]["exception_handling"]:
+ self.guidelines_list.append("- Integrate robust exception handling.")
+ if st.session_state["coding_guidelines"]["error_handling"]:
+ self.guidelines_list.append("- Add error handling to each module.")
+ if st.session_state["coding_guidelines"]["efficient_code"]:
+ self.guidelines_list.append("- Optimize the code to ensure it runs efficiently.")
+ if st.session_state["coding_guidelines"]["robust_code"]:
+ self.guidelines_list.append("- Ensure the code is robust against potential issues.")
+ if st.session_state["coding_guidelines"]["naming_conventions"]:
+ self.guidelines_list.append("- Follow standard naming conventions.")
+
+ # Convert the list to a string
+ self.guidelines = "\n".join(self.guidelines_list)
+ 
+ 
+ def _configure(self):
+ try:
+ logger.info("Configuring Gemini AI Pro...")
+ genai.configure(api_key=self.api_key)
+ self.generation_config = {
+ "temperature": self.temperature,
+ "top_p": self.top_p,
+ "top_k": self.top_k,
+ "max_output_tokens": self.max_output_tokens
+ }
+ self.model = genai.GenerativeModel(model_name=self.model,generation_config=self.generation_config)
+ logger.info("Gemini AI Pro configured successfully.")
+ except Exception as exception:
+ logger.error(f"Error configuring Gemini AI Pro: {str(exception)}")
+ traceback.print_exc()
+
+ def _extract_code(self, code):
+ """
+ Extracts the code from the provided string.
+ If the string contains '```', it extracts the code between them.
+ Otherwise, it returns the original string.
+ """
+ try:
+ if '```' in code:
+ start = code.find('```') + len('```\n')
+ end = code.find('```', start)
+ # Skip the first line after ```
+ start = code.find('\n', start) + 1
+ extracted_code = code[start:end]
+ logger.info("Code extracted successfully.")
+ return extracted_code
+ else:
+ logger.info("No special characters found in the code. Returning the original code.")
+ return code
+ except Exception as exception:
+ logger.error(f"Error occurred while extracting code: {exception}")
+ return None
+ 
+ def generate_code(self, code_prompt,code_language):
+ """
+ Function to generate text using the Gemini API.
+ """
+ try:
+ # Define top_k and top_p based on the mode
+ if self.mode == "precise":
+ top_k = 40
+ top_p = 0.95
+ self.temprature = 0
+ elif self.mode == "balanced":
+ top_k = 20
+ top_p = 0.85
+ self.temprature = 0.3
+ elif self.mode == "creative":
+ top_k = 10
+ top_p = 0.75
+ self.temprature = 1
+ else:
+ raise ValueError("Invalid mode. Choose from 'precise', 'balanced', 'creative'.")
+
+ logger.info(f"Generating code with mode: {self.mode}, top_k: {top_k}, top_p: {top_p}")
+
+ 
+ # check for valid prompt and language
+ if not code_prompt or len(code_prompt) == 0:
+ st.toast("Error in code generation: Please enter a valid prompt.", icon="❌")
+ logger.error("Error in code generation: Please enter a valid prompt.")
+ return
+ 
+ logger.info(f"Generating code for prompt: {code_prompt} in language: {code_language}")
+ if code_prompt and len(code_prompt) > 0 and code_language and len(code_language) > 0:
+ logger.info(f"Generating code for prompt: {code_prompt} in language: {code_language}")
+ 
+ # Plain and Simple Coding Task Prompt
+ prompt = f"""
+ Task: You're an experienced developer. Your mission is to create a program for {code_prompt} in {code_language} that takes {st.session_state.code_input} as input.
+
+ Your goal is clear: Craft a solution that showcases your expertise as a coder and problem solver.
+
+ Ensure that the program's output contains only the code you've written, with no extraneous information.
+
+ Show your skills and solve this challenge with confidence!
+ 
+ And follow the proper coding guidelines and dont add comment unless instructed to do so.
+ {self.guidelines}
+ """
+ 
+ gemini_completion = self.model.generate_content(prompt)
+ logger.info("Text generation completed successfully.")
+ 
+ code = None
+ if gemini_completion:
+ # extract the code from the gemini completion
+ code = gemini_completion.text
+ logger.info(f"GeminiAI coder is initialized.")
+ logger.info(f"Generated code: {code[:100]}...")
+ 
+ if gemini_completion:
+ # Extracted code from the gemini completion
+ extracted_code = self._extract_code(code)
+ 
+ # Check if the code or extracted code is not empty or null
+ if not code or not extracted_code:
+ raise Exception("Error: Generated code or extracted code is empty or null.")
+ 
+ return extracted_code
+ else:
+ raise Exception("Error in code generation: Please enter a valid code.")
+ 
+ except Exception as exception:
+ st.toast(f"Error in code generation: {exception}", icon="❌")
+ logger.error(f"Error in code generation: {traceback.format_exc()}")

libs/general_utils.py

@@ -401,24 +401,47 @@ def calculate_code_generation_cost(self,string,price=0.0005):
  total_cost = cost * number_of_words
 
  # Return the cost, cost per whole string and total cost
- return cost, cost_per_whole_string, total_cost
+ #return cost, cost_per_whole_string, total_cost
+
+ # Return the total cost
+ return total_cost
 
  def codey_generation_cost(self,string):
  codey_price = 0.0005
  return self.calculate_code_generation_cost(string,codey_price)
 
  def gpt_3_generation_cost(self,string):
- chatgpt_price = 0.0002
- return self.calculate_code_generation_cost(string,chatgpt_price)
+ model_price = 0.0002
+ return self.calculate_code_generation_cost(string,model_price)
 
  def gpt_3_5_turbo_generation_costself(self,string):
- chatgpt_price = 0.0080
- return self.calculate_code_generation_cost(string,chatgpt_price)
+ model_price = 0.0080
+ return self.calculate_code_generation_cost(string,model_price)
 
  def gpt_4_generation_cost(self,string):
- chatgpt_price = 0.06
- return self.calculate_code_generation_cost(string,chatgpt_price)
+ model_price = 0.06
+ return self.calculate_code_generation_cost(string,model_price)
 
  def gpt_text_davinci_generation_cost(self,string):
- chatgpt_price = 0.0060
- return self.calculate_code_generation_cost(string,chatgpt_price)
+ model_price = 0.0060
+ return self.calculate_code_generation_cost(string,model_price)
+ 
+ def palm_text_bison_generation_cost(self,string):
+ model_price = 0.00025
+ return self.calculate_code_generation_cost(string,model_price)
+ 
+ def palm_chat_bison_generation_cost(self,string):
+ model_price = 0.00025
+ return self.calculate_code_generation_cost(string,model_price)
+
+ def palm_embedding_gecko_generation_cost(self,string):
+ model_price = 0.0002
+ return self.calculate_code_generation_cost(string,model_price)
+ 
+ def gemini_pro_generation_cost(self,string):
+ model_price = 0.00025 
+ return self.calculate_code_generation_cost(string,model_price)
+ 
+ def gemini_pro_vision_generation_cost(self,string):
+ model_price = 0.00025 
+ return self.calculate_code_generation_cost(string,model_price)

libs/palm_coder.py renamed to libs/palmai.py

@@ -10,10 +10,6 @@
 import streamlit as st
 
 
-# Set up logging
-logging.basicConfig(filename='palm-coder.log', filemode='w', format='%(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
 class PalmAI:
  def __init__(self,api_key, model="text-bison-001", temperature=0.3, max_output_tokens=2048, mode="balanced"):
  """
@@ -87,8 +83,8 @@ def _extract_code(self, code):
  else:
  logger.info("No special characters found in the code. Returning the original code.")
  return code
- except Exception as e:
- logger.error(f"Error occurred while extracting code: {e}")
+ except Exception as exception:
+ logger.error(f"Error occurred while extracting code: {exception}")
  return None
 
  def _install_package(self, package_name):
@@ -164,19 +160,6 @@ def generate_code(self, code_prompt,code_language):
  And follow the proper coding guidelines and dont add comment unless instructed to do so.
  {self.guidelines}
  """
-
- # If graph were requested.
- if 'graph' in code_prompt.lower():
- prompt += "\n" + "using Python use Matplotlib save the graph in file called 'graph.png'"
-
- # if Chart were requested
- if 'chart' in code_prompt.lower() or 'plot' in code_prompt.lower():
- prompt += "\n" + "using Python use Plotly save the chart in file called 'chart.png'"
-
- # if Table were requested
- if 'table' in code_prompt.lower():
- prompt += "\n" + "using Python use Pandas save the table in file called 'table.md'"
- 
 
  palm_completion = palm.generate_text(
  model=self.model,
@@ -210,8 +193,8 @@ def generate_code(self, code_prompt,code_language):
  else:
  raise Exception("Error in code generation: Please enter a valid code.")
 
- except Exception as e:
- st.toast(f"Error in code generation: {e}", icon="❌")
+ except Exception as exception:
+ st.toast(f"Error in code generation: {exception}", icon="❌")
  logger.error(f"Error in code generation: {traceback.format_exc()}")
 
  def fix_generated_code(self, code, code_language, fix_instructions=""):
@@ -308,6 +291,6 @@ def fix_generated_code(self, code, code_language, fix_instructions=""):
  else:
  st.toast("Error in code fixing: Please enter a valid code and language.", icon="❌")
  logger.error("Error in code fixing: Please enter a valid code and language.")
- except Exception as e:
- st.toast(f"Error in code fixing: {e}", icon="❌")
+ except Exception as exception:
+ st.toast(f"Error in code fixing: {exception}", icon="❌")
  logger.error(f"Error in code fixing: {traceback.format_exc()}")