Merge pull request #9 from hrshdhgd/google

Implemented Google's `gemini-pro` model

README.md

@@ -2,12 +2,18 @@
 
 CoderGPT is a versatile command-line interface (CLI) designed to enhance coding workflows. It leverages the capabilities of Large Language Models (LLM) and Generative Pre-trained Transformers (GPT) to assist developers in various tasks such as commenting, optimizing, documenting, and testing their code. This tool integrates seamlessly with [langchain](https://github.com/langchain-ai/langchain), providing a powerful backend for code generation and modification.
 
+# Model Providers Implemented
+ - [x] OpenAI [`gpt-3.5-turbo`, `gpt-4`, `gpt-4-turbo-preview`(default)]
+ - [x] Google [`gemini-pro`]
+ - [ ] Anthropic [`Claude`] (coming soon!)
+
 ## Prerequisites
 
-Before you begin using CoderGPT, you must set the `OPENAI_API_KEY` environment variable on your machine. This key enables authentication with the OpenAI API, which is essential for the language model's operation.
+Before you begin using CoderGPT, you must set the `OPENAI_API_KEY` and `GOOGLE_API_KEY` environment variables on your machine. This key enables authentication with the OpenAI and Google APIs, which are essential for the language model's operation.
 
 ```sh
 export OPENAI_API_KEY='your-api-key-here'
+export GOOGLE_API_KEY='your-api-key-here'
 ```
 
 Ensure that you replace `your-api-key-here` with your actual OpenAI API key to enable the full functionality of CoderGPT.

docs/description.rst

@@ -10,16 +10,35 @@ CoderGPT is a command line interface for generating/modifying code. It allows de
 enhance code by commenting, optimizing, documenting and adding tests to their projects using 
 the power of LLM and GPT. This project is powered by `langchain <https://github.com/langchain-ai/langchain>`_.
 
+Model Providers Implemented
+---------------------------
+
+The following model providers have been implemented in CoderGPT:
+
+.. list-table::
+   :widths: 25 75
+   :header-rows: 1
+
+   * - Provider
+     - Status
+   * - OpenAI (``gpt-3.5-turbo``, ``gpt-4``, ``gpt-4-turbo-preview`` (default))
+     - ✓
+   * - Google (``gemini-pro``)
+     - ✓
+   * - Anthropic (``Claude``)
+     - Coming soon!
+
+
 .. note::
    **NOTE**
-   Before using CoderGPT, ensure that the environment variable ``OPENAI_API_KEY`` is set locally on your machine. This key is required for authentication with the OpenAI API which powers the underlying language model.
-
-   .. code-block:: sh
+   Before you begin using CoderGPT, you must set the ``OPENAI_API_KEY`` and ``GOOGLE_API_KEY`` environment variables on your machine. These keys are necessary for authentication with the respective APIs and are crucial for the operation of the language models.
 
-      export OPENAI_API_KEY='your-api-key-here'
+.. code-block:: sh
 
-   Replace ``your-api-key-here`` with your actual OpenAI API key. This step is crucial for the proper functioning of CoderGPT as it relies on the OpenAI API for generating and modifying code.
+   export OPENAI_API_KEY='your-api-key-here'
+   export GOOGLE_API_KEY='your-api-key-here'
 
+   Replace ``your-api-key-here`` with your actual OpenAI and Google API keys. This step is crucial for the proper functioning of CoderGPT as it relies on OpenAI and Google APIs for generating and modifying code.
 
 Installation
 ------------

pyproject.toml

@@ -14,6 +14,7 @@ langchain = "^0.1.6"
 poetry-dynamic-versioning = "^1.2.0"
 langchain-openai = "^0.0.5"
 tabulate = "^0.9.0"
+langchain-google-genai = "^0.0.8"
 
 [tool.poetry.group.dev.dependencies]
 pytest = {version = ">=7.1.2"}

src/codergpt/cli.py

@@ -27,6 +27,7 @@
 import click
 
 from codergpt import __version__
+from codergpt.constants import ALL_MODELS
 from codergpt.main import CoderGPT
 
 __all__ = [
@@ -36,27 +37,35 @@
 logger = logging.getLogger(__name__)
 
 path_argument = click.argument("path", type=click.Path(exists=True))
+model_option = click.option(
+    "-m",
+    "--model",
+    type=click.Choice(ALL_MODELS),
+    default="gpt-4",
+    help="Model to use for code generation.",
+)
 function_option = click.option("-f", "--function", help="Function name to explain or optimize.")
 class_option = click.option("-c", "--classname", help="Class name to explain or optimize.")
 overwrite_option = click.option(
     "--overwrite/--no-overwrite", is_flag=True, default=False, help="Overwrite the existing file."
 )
 output_option = click.option("-o", "--outfile", help="Output file to write to.")
 
-coder = CoderGPT()
-
 
 @click.group()
 @click.option("-v", "--verbose", count=True)
 @click.option("-q", "--quiet")
 @click.version_option(__version__)
-def main(verbose: int, quiet: bool):
+@model_option
+def main(verbose: int, quiet: bool, model: str):
     """
     CLI for CoderGPT.
 
     :param verbose: Verbosity while running.
     :param quiet: Boolean to be quiet or verbose.
     """
+    global coder
+    coder = CoderGPT(model=model)
     if verbose >= 2:
         logger.setLevel(level=logging.DEBUG)
     elif verbose == 1:

src/codergpt/constants.py

@@ -13,3 +13,13 @@
 GPT_3_5_TURBO = "gpt-3.5-turbo"
 GPT_4 = "gpt-4"
 GPT_4_TURBO = "gpt-4-turbo-preview"
+CLAUDE = "claude"
+GEMINI = "gemini-pro"
+
+ALL_MODELS = [
+    GPT_3_5_TURBO,
+    GPT_4,
+    GPT_4_TURBO,
+    CLAUDE,
+    GEMINI,
+]

src/codergpt/main.py

@@ -6,11 +6,12 @@
 
 import yaml
 from langchain_core.prompts import ChatPromptTemplate
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_openai import ChatOpenAI
 from tabulate import tabulate
 
 from codergpt.commenter.commenter import CodeCommenter
-from codergpt.constants import EXTENSION_MAP_FILE, GPT_4_TURBO, INSPECTION_HEADERS
+from codergpt.constants import EXTENSION_MAP_FILE, GEMINI, GPT_4_TURBO, INSPECTION_HEADERS
 from codergpt.documenter.documenter import CodeDocumenter
 from codergpt.explainer.explainer import CodeExplainer
 from codergpt.optimizer.optimizer import CodeOptimizer
@@ -22,7 +23,16 @@ class CoderGPT:
 
     def __init__(self, model: str = GPT_4_TURBO):
         """Initialize the CoderGPT class."""
-        self.llm = ChatOpenAI(openai_api_key=os.environ.get("OPENAI_API_KEY"), temperature=0.7, model=model)
+        if model is None or model.startswith("gpt-"):
+            self.llm = ChatOpenAI(openai_api_key=os.environ.get("OPENAI_API_KEY"), temperature=0.7, model=model)
+        # elif model == CLAUDE:
+        #     self.llm = ChatAnthropic()
+        #     print("Coming Soon!")
+        elif model == GEMINI:
+            self.llm = ChatGoogleGenerativeAI(model=model, convert_system_message_to_human=True)
+        else:
+            raise ValueError(f"The model {model} is not supported yet.")
+
         self.prompt = ChatPromptTemplate.from_messages(
             [("system", "You are world class software developer."), ("user", "{input}")]
         )