A small project to demonstarte usage of LangChain with Google Cloud's VertexAI and Gemini LLM.
This streamlit application accepts a code or a function as an input and generates the following outputs: a detailed description of what the code does, an alternative way to write the code using best practices, and some unit tests to check the code’s functionality.
- Google Cloud Account with Vertex AI API enabled
- for running locally - ensure your local auth is done/setup; https://cloud.google.com/sdk/gcloud/reference/auth
- for running on gcp (cloudrun or gke or computeengine) - ensure the account (serviceaccount) has permissions for VertexAI
- Python 3
- The required dependencies are in the requirements.txt
- LangChain
- Clone the repository.
- Create a python venv https://docs.python.org/3/library/venv.html
- Install dependencies
pip install -r requirements.txt
- Start the streamlit app, should open a brower window; if not open manually (the output of below command provides the url)
streamlit run main.py
It is wonderful to see how frameworks like LangChain make it easier to work with LLMs, they provide a simplified api which allows to interact with LLMs. Refer to the LangChain docs for more information.
Create an instance of LLM, in this case GCP's VertexAI's gemini-pro
from langchain.llms import VertexAI
llm = VertexAI(model_name="gemini-pro", max_output_tokens=1024)
Create prompt templates
prompt_explain = PromptTemplate.from_template("Explain what the below code does. {entered_code}")
prompt_alternate = PromptTemplate.from_template("Can you rewrite the below function using python best practices? {entered_code}")
prompt_unittests = PromptTemplate.from_template("Create unit test code/cases for the below function. {entered_code}")
Execute the query against the Gemini-pro LLM
- Get the entered code
entered_code = st.text_area(label="Enter your code here", height=200)
- Execute query, by providing the code to prompt template and invoking LLM
codeExplanation=llm(prompt_explain.format(entered_code=entered_code))
Refer to the code in
main.py
- Enter a python code/function in the given textarea and press submit
- Look for Tabs "Explain", "Alternate", "Unit Tests"