model.py

import os

os.environ["GOOGLE_API_KEY"] = "AIzaSyAGoYnNPu__70AId7EJS7F_61i69Qmn-wM"
os.environ["OPENAI_API_TYPE"] = "azure"
# os.environ["OPENAI_API_VERSION"] = "2023-07-01-preview"
# # os.environ["OPENAI_API_KEY"] = "5b624f6b71884a488560a86b1fffbf42"
# os.environ["OPENAI_API_KEY"] = "9e337d6696ce4a22a9a1b901e2ebb5fb"


from embedder import CustomEmbeddings
from langchain.chat_models import AzureChatOpenAI, ChatOpenAI
from langchain.prompts.chat import (ChatPromptTemplate,
                                    HumanMessagePromptTemplate,
                                    SystemMessagePromptTemplate)
from langchain_google_genai import ChatGoogleGenerativeAI
from search import SimilaritySearch

embeddings = CustomEmbeddings(
    model="text-embedding-ada-002", 
    model_url="https://year-embedding-ada-002-aiservices-2136192926.openai.azure.com//openai/deployments/fresh-embedding-ada-002/embeddings?api-version=2023-10-01-preview",
    api_key="6eed3006cdd3445cb3f422a7358ce461"
)
vector_store = SimilaritySearch.load_from_disk(
    embedding_function=embeddings, 
    data_dir="./indexs/text-embedding-ada-002/"
    # data_dir="../indexs/basic-fno-text-embedding-ada-002/"
)

class Model:
    def __init__(self, model_name: str, **kwargs) -> None:
        self.model_name = model_name
        self.llm = self.load_llm(model_name=model_name, **kwargs)
    
    def load_llm(self, model_name: str, **kwargs):
        if self.model_name == "gemini-pro":
            self.retriever = vector_store.as_retriever(search_kwargs={"k": 2}, search_type="similarity")
            return ChatGoogleGenerativeAI(model=model_name, temperature=0, max_tokens=4096)
        elif self.model_name == "gpt-3.5-turbo":
            self.retriever = vector_store.as_retriever(search_kwargs={"k": 2}, search_type="similarity")
            return AzureChatOpenAI(
                deployment_name="latest-gpt-35-turbo-16k", 
                temperature=0, 
                max_tokens=4096, 
                # azure_endpoint="https://high-gpt4-32k-0613-aiservices336365459.openai.azure.com/",
                openai_api_key="9e337d6696ce4a22a9a1b901e2ebb5fb",
                # openai_api_base="https://jan-2024-gpt35-turbo16k-aiservices800630185.openai.azure.com/",
                openai_api_base = "https://fresh-gpt35-turbo-aiservices-2112150452.openai.azure.com/",
                openai_api_version="2023-07-01-preview"
            )
        elif self.model_name == "gpt4":
            self.retriever = vector_store.as_retriever(search_kwargs={"k": kwargs.get("k", 2)}, search_type="similarity")
            return AzureChatOpenAI(
                deployment_name="gpt-4-32k", 
                temperature=0, 
                max_tokens=4096, 
                # azure_endpoint="https://high-gpt4-32k-0613-aiservices336365459.openai.azure.com/",
                openai_api_key="e91a341abb2f4646ab7b0acd3b9d461e",
                openai_api_base="https://jan-2024-gpt4-ai-aiservices-1959882301.openai.azure.com/",
                openai_api_version="2023-07-01-preview"
            )

        self.retriever = vector_store.as_retriever(search_kwargs={"k": kwargs.get("k", 1)}, search_type="similarity")
        return ChatOpenAI(
            model=model_name,
            openai_api_key="EMPTY",
            openai_api_base="http://localhost:8000/v1",
            max_tokens=1024,
            temperature=0,
            model_kwargs={"stop": ["<|im_end|>", "Query:", "Question:"], "top_p": 0.95}
        )
        
        
    def run_qa_result(self, query: str):
        support_docs = self.retriever.get_relevant_documents(query)
        sources = list({d.metadata['source'] for d in support_docs})
        context = "\n\n".join([f"{i + 1}. {d.page_content}" for i, d in enumerate(support_docs)])
        return context, sources

    def return_prompt(self, system_prompt: str, query: str, context: str):
        
        # human_template = "Context:\n\n{context}\n\nQuery: {query}"
        # human_template = "E-Book:\n\n{context}\n\nQuestion: {query}"

        human_template = "\n\nContext:\n\n{context}\n\nQuestion: {query}"
        # human_template = "\n\nBook:\n\n{context}\n\nQuestion: {query}"

        messages = []
        if self.model_name in [
            "gemini-pro",
            "TheBloke/Mistral-7B-Instruct-v0.2-AWQ",
        ]:
            human_template = system_prompt + "\n\n" + human_template
            human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
            messages.append(human_message_prompt)
        else:
            system_message_prompt = SystemMessagePromptTemplate.from_template(system_prompt)
            human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
            messages.extend([system_message_prompt, human_message_prompt])

        chat_prompt = ChatPromptTemplate.from_messages(messages)
        return chat_prompt.format_prompt(context=context, query=query).to_messages()
        
    def run(self, system_prompt: str, query: str):
        context, sources = self.run_qa_result(query=query)
        chat_prompt = self.return_prompt(system_prompt=system_prompt, query=query, context=context)
        # text = "".join(resp.content for resp in self.llm.stream(chat_prompt))
        # text += "\nSources: \n" + "\n".join([f"{i + 1}. {d}" for i, d in enumerate(sources)])
        # return text, sources
        for resp in self.llm.stream(chat_prompt):
            yield resp.content.replace("$", "₹")
        
        yield sources
        # text = "".join(resp.content for resp in self.llm.stream(chat_prompt))
        # text += "\nSources: \n" + "\n".join([f"{i + 1}. {d}" for i, d in enumerate(sources)])
        # return text, sources
   
def get_sources(query):
    results = vector_store.similarity_search_with_relevance_scores(query, k=2)
    return [
        {
            "score": r[-1],
            "source": r[0].metadata['source']
        }
        for r in results
    ]     

if __name__ == "__main__":
    # model = Model(model_name="phi2")
    # model = Model(model_name="gpt-3.5-turbo")       
    # model = Model(model_name="gemini-pro")    
    # model = Model(model_name="TheBloke/zephyr-7B-beta-AWQ")
    # model = Model(model_name="TheBloke/neural-chat-7B-v3-3-AWQ")
    model = Model(model_name="TheBloke/Mistral-7B-Instruct-v0.2-AWQ")
    model = Model(model_name="gpt4")
    model = Model(model_name="gpt-3.5-turbo") 
    
    # query = "what is reliance?"
    # print("results: ", get_sources(query))
    
    # query = "explain FNO trading?"
    # print("results: ", get_sources(query))
      
#     query="What is FNO trading?"
#     query = "Describe ITM, ATM and OTM"
#     query = "give formula to calculate intrinsic value in Put and provide an example"
#     query = "what is the order of delta, theta, gamma and vega amongst options in a given options chain"
#     query = "Explain apple stock and nasdaq"
    
#     query = "generate a table with long and short in F&O instruments"
#     query = "how can we calculate intrinsic value and time value"
#     query = "give formula to calculate intrinsic value in Put"
    
    query = "explain exit from a put trade"
    # 
    # query = "what will be buying cost if I long tesla CE"
    
    
#     system_prompt="""Use the following pieces of context to answer the question in detail. Provide example only if it is in provided context and make sure to use them in rupees.""", 
    
#     system_prompt = """Use the following pieces of context to answer the question in detail. Provide example only if it is in context and make sure to use them in ₹.
# If you don't know the answer, just say 'Please rephrase the question I am unable to answer'"""

#     system_prompt = """Answer the question using the context. Provide examples only from the context and use only Rupees (₹) in examples. If you don't know the answer, just say 'Please rephrase the question I am unable to answer'"""
    
#     system_prompt = """Your task is to answer the question using the given context. 

# Follow the below rules while answering the question:
# - Only create example using the context
# - Use only Rupees '₹' to represent currency. 
# - If you don't know the answer, just say 'Please rephrase the question I am unable to answer'"""

#     system_prompt = """You are an Indian Stock Market Assistant. Your task is to answer the question using the given context. Only create example from the given context and don't use '$'."""
    
    # query = "what is reliance?"    
    # query = "what is python?"    
    query = "what is an apple stock and nasdq"
    query = "Generate a tabular format on playing long and short through options"
    query = "What is FNO Trading?"
    
    system_prompt = """Answer the question only from context. 
Provide examples only from the context. 
If you don't know the answer, just say 'Please rephrase the question I am unable to answer'"""
    
    system_prompt = "Answer the question only from the e-book. If it is not sufficient then respond as \"Unknown\""
    system_prompt = """Use the following pieces of book to answer the question at the end. \nIf you don't know the answer, please think rationally and answer from the book"""
    # system_prompt = """Answer the question using the context. Provide examples only from the context and use only Rupees (₹) in examples. If you don't know the answer, just say 'Please rephrase the question I am unable to answer'"""
    
    # system_prompt = """Answer the question from the context. Provide examples only from the context. If you don't know the answer, just say 'Please rephrase the question'"""
    # system_prompt = """Answer the question from the book. Provide examples only from the book. If you don't know the answer, just say 'Please rephrase the question'"""
    
    response = model.run(
        system_prompt=system_prompt,
        query=query
    )
    text = ""
    for resp in response:
        if isinstance(resp, list):
            sources = resp
            break
        text += resp
    
    text = text.split("Question")[0].strip("\n")
        
    print("text: ", text)
    open("./text.txt", "w").write(text)