app.py

import re
import traceback
from langchain import LLMChain, PromptTemplate
from langchain.llms import VertexAI
from libs.logger import logger
import streamlit as st
from google.oauth2 import service_account
from langchain.prompts import ChatPromptTemplate
import libs.general_utils

class VertexAILangChain:
    def __init__(self, project="", location="us-central1", model_name="code-bison", max_tokens=256, temperature:float=0.3, credentials_file_path=None):
        self.project = project
        self.location = location
        self.model_name = model_name
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.credentials_file_path = credentials_file_path
        self.vertexai_llm = None
        self.utils = libs.general_utils.GeneralUtils()

    def load_model(self, model_name, max_tokens, temperature):
        try:
            logger.info(f"Loading model... with project: {self.project} and location: {self.location}")
            # Set the GOOGLE_APPLICATION_CREDENTIALS environment variable
            credentials = service_account.Credentials.from_service_account_file(self.credentials_file_path)

            logger.info(f"Trying to set Vertex model with parameters: {model_name or self.model_name}, {max_tokens or self.max_tokens}, {temperature or self.temperature}, {self.location}")
            self.vertexai_llm = VertexAI(
                model_name=model_name or self.model_name,
                max_output_tokens=max_tokens or self.max_tokens,
                temperature=temperature or self.temperature,
                verbose=True,
                location=self.location,
                credentials=credentials,
            )
            logger.info("Vertex model loaded successfully.")
            return True
        except Exception as exception:
            logger.error(f"Error loading Vertex model: {str(exception)}")
            logger.error(traceback.format_exc())  # Add traceback details
            return False

    def generate_code(self, code_prompt, code_language):
        try:
            # Dynamically construct guidelines based on session state
            guidelines_list = []
            logger.info(f"Generating code with parameters: {code_prompt}, {code_language}")
            
            # Check for empty or null code prompt and code language
            if not code_prompt or len(code_prompt) == 0:
                logger.error("Code prompt is empty or null.")
                st.toast("Code prompt is empty or null.", icon="❌")
                return None
            
            if st.session_state["coding_guidelines"]["modular_code"]:
                logger.info("Modular code is enabled.")
                guidelines_list.append("- Ensure the method is modular in its approach.")
            if st.session_state["coding_guidelines"]["exception_handling"]:
                logger.info("Exception handling is enabled.")
                guidelines_list.append("- Integrate robust exception handling.")
            if st.session_state["coding_guidelines"]["error_handling"]:
                logger.info("Error handling is enabled.")
                guidelines_list.append("- Add error handling to each module.")
            if st.session_state["coding_guidelines"]["efficient_code"]:
                logger.info("Efficient code is enabled.")
                guidelines_list.append("- Optimize the code to ensure it runs efficiently.")
            if st.session_state["coding_guidelines"]["robust_code"]:
                logger.info("Robust code is enabled.")
                guidelines_list.append("- Ensure the code is robust against potential issues.")
            if st.session_state["coding_guidelines"]["naming_conventions"]:
                logger.info("Naming conventions is enabled.")
                guidelines_list.append("- Follow standard naming conventions.")
            
            logger.info("Guidelines: " + str(guidelines_list))

            # Convert the list to a string
            guidelines = "\n".join(guidelines_list)

            # Setting Prompt Template.
            input_section = f"Given the input for code: {st.session_state.code_input}" if st.session_state.code_input else "make sure the program doesn't ask for any input from the user"

            template = f"""
            Task: Design a program {{code_prompt}} in {{code_language}} with the following guidelines and
            make sure the output is printed on the screen.
            And make sure the output contains only the code and nothing else.
            {input_section}

            Guidelines:
            {guidelines}
            """
            
            prompt = PromptTemplate(template=template,input_variables=["code_prompt", "code_language"])
            formatted_prompt = prompt.format(code_prompt=code_prompt, code_language=code_language)
            logger.info(f"Formatted prompt: {formatted_prompt}")
            
            logger.info("Setting up LLMChain...")
            llm_chain = LLMChain(prompt=prompt, llm=self.vertexai_llm)
            logger.info("LLMChain setup successfully.")
            
            # Pass the required inputs as a dictionary to the chain
            logger.info("Running LLMChain...")
            response = llm_chain.run({"code_prompt": code_prompt, "code_language": code_language})
            if response or len(response) > 0:
                logger.info(f"Code generated successfully: {response}")
                
                # Extract text inside code block
                if response.startswith("```") or response.endswith("```"):
                    try:
                        generated_code = re.search('```(.*)```', response, re.DOTALL).group(1)
                    except AttributeError:
                        generated_code = response
                else:
                    st.toast(f"Error extracting code", icon="❌")
                    return response
                    
                if generated_code:
                    # Skip the language name in the first line.
                    response = generated_code.split("\n", 1)[1]
                    logger.info(f"Code generated successfully: {response}")
                else:
                    logger.error(f"Error generating code: {response}")
                    st.toast(f"Error generating code: {response}", icon="❌")
            return response
        except Exception as exception:
            stack_trace = traceback.format_exc()
            logger.error(f"Error generating code: {str(exception)} stack trace: {stack_trace}")
            st.toast(f"Error generating code: {str(exception)} stack trace: {stack_trace}", icon="❌")

    def generate_code_completion(self, code_prompt, code_language):
        try:
            if not code_prompt or len(code_prompt) == 0:
                logger.error("Code prompt is empty or null.")
                st.error("Code generateration cannot be performed as the code prompt is empty or null.")
                return None
            
            logger.info(f"Generating code completion with parameters: {code_prompt}, {code_language}")
            template = f"Complete the following {{code_language}} code: {{code_prompt}}"
            prompt_obj = PromptTemplate(template=template, input_variables=["code_language", "code_prompt"])
            
            max_tokens = st.session_state["vertexai"]["max_tokens"]
            temprature = st.session_state["vertexai"]["temperature"]
            
            # Check the maximum number of tokens of Gecko model i.e 65
            if max_tokens > 65:
                max_tokens = 65
                logger.info(f"Maximum number of tokens for Model Gecko can't exceed 65. Setting max_tokens to 65.")
                st.toast(f"Maximum number of tokens for Model Gecko can't exceed 65. Setting max_tokens to 65.", icon="⚠️")
                
            self.model_name = "code-gecko" # Define the code completion model name.
            self.llm = VertexAI(model_name=self.model_name,max_output_tokens=max_tokens, temperature=temprature)
            logger.info(f"Initialized VertexAI with model: {self.model_name}")
            llm_chain = LLMChain(prompt=prompt_obj, llm=self.llm)
            response = llm_chain.run({"code_prompt": code_prompt, "code_language": code_language})
            
            if response:
                logger.info(f"Code completion generated successfully: {response}")
                return response
            else:
                logger.warning("No response received from LLMChain.")
                return None
        except Exception as e:
            logger.error(f"Error generating code completion: {str(e)}")
            raise

    def set_temperature(self, temperature):
        self.temperature = temperature
        self.vertexai_llm.temperature = temperature
        # call load_model to reload the model with the new temperature and rest values should be same
        self.load_model(self.model_name, self.max_tokens, self.temperature)
        
    def set_max_tokens(self, max_tokens):
        self.max_tokens = max_tokens
        self.vertexai_llm.max_output_tokens = max_tokens
        # call load_model to reload the model with the new max_output_tokens and rest values should be same
        self.load_model(self.model_name, self.max_tokens, self.temperature)
        
    def set_model_name(self, model_name):
        self.model_name = model_name
        # call load_model to reload the model with the new model_name and rest values should be same
        self.load_model(self.model_name, self.max_tokens, self.temperature)