Classes/Owiki_Class.py

import os
import google.generativeai as genai
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.vectorstores import FAISS
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain.prompts import PromptTemplate
import json
import re
from Classes.Helper_Class import DB_Retriever
from typing import Optional




class OWiki:
    def __init__(self,**kwargs):
        temperature = kwargs['temperature']
        self.summary = kwargs['summary_length']
        model = kwargs["model"]
        self.db_loc = kwargs["db_loc"]
        self.api_key = kwargs["api_key"]
        os.environ["GOOGLE_API_KEY"]  = self.api_key
        genai.configure(api_key=self.api_key)
        self.llm = ChatGoogleGenerativeAI(model=model,
                            temperature=temperature)
        self.model_embedding = kwargs['model_embeddings']
    

    def get_summary_template(self):
        prompt = """Generate a summary for the following conversational data in less than {summary} lines.\nText:\n{text}\n\nSummary:"""
        prompt_template = PromptTemplate(template = prompt,input_variables=['summary','text'])
        return prompt_template
    
    def create_sql_prompt_template(self,schemas):
        prompt =  """Write an SQL query for the following questions whose schemas are as follows.\nSQL Schema:"""
        for table_name,table_schema in schemas.items():
            prompt+= f"Table Name: {table_name} Schema:"
            for key,value in table_schema.items():
                prompt+= f"{key} {value}"
        prompt+=  """\n\nQuestion:{question}\n\nAnswer:"""
        # print("Prompt",prompt,'\n\n')
        prompt_template = PromptTemplate(template = prompt,input_variables=['question'])
        # print("Prompt template",prompt_template)
        return prompt_template

    def create_prompt_for_OIC_bot(self):
        template = """You are expert OIC(Oracle Integration Cloud) Bot. You will be given a task to solve as best you can.
        
        Here are the rules you should always follow to solve your task:
        1. Always provide a Question Explanation with **Question Explanation:** Heading and Potential Solution with **Potential Solution:** Headings.
        2. Take care of not being too long typically not exceeding 5000 tokens.
        3. Response must contain all possible **Error Scenarios:** if applicable along with a **Summary:** Heading containing breif summary of the task solution at the end.
        4. If you don't know the answer or if it is not in the context, please answer as **I am not trained on this topics due to limited resources.**

        Here are a few examples:
        Question: How to configure a OFSC adapter in OIC.
        
        Now Begin! 
        Context:
        {context}
        
        Question: {question}
        """
        prompt = PromptTemplate.from_template(template)
        return prompt
    
    def create_sql_agent(self,question,schemas):
        prompt_template = self.create_sql_prompt_template(schemas)
        chain = prompt_template | self.llm | StrOutputParser()
        response = chain.invoke({"question":question})
        response = self.format_llm_response(response)
        return response
    
    def generate_summary(self,text):
        prompt_template = self.get_summary_template()
        chain = prompt_template | self.llm | StrOutputParser()
        response = chain.invoke({"text":text,"summary":self.summary})
        return response

    def format_llm_response(self,text):
        bold_pattern = r"\*\*(.*?)\*\*"
        italic_pattern = r"\*(.*?)\*"
        code_pattern = r"```(.*?)```"
        text = text.replace('\n', '<br>')
        formatted_text = re.sub(code_pattern,"<pre><code>\\1</code></pre>",text)
        formatted_text = re.sub(bold_pattern, "<b>\\1</b>", formatted_text)
        formatted_text = re.sub(italic_pattern, "<i>\\1</i>", formatted_text)
        return formatted_text

    def search_from_db(self, query : str, chat_history : Optional[str] ) -> str :
        db = DB_Retriever(self.db_loc,self.model_embedding)
        retriever = db.retrieve(query)
        prompt = self.create_prompt_for_OIC_bot()
        chat_history = self.generate_summary(chat_history)
        retrieval_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | self.llm
        | StrOutputParser()
        )
        response = retrieval_chain.invoke(query)
        response = self.format_llm_response(response)
        return response