File Manager

Editing File: /home/ubuntu/codegamaai-test/tax_bot/main.py

from dotenv import load_dotenv

import openai
import sys, os
import logging
from llama_index import (
    SimpleDirectoryReader,
    LLMPredictor,
    GPTVectorStoreIndex,
    load_index_from_storage,
    set_global_service_context,
)
from llama_index.storage.storage_context import StorageContext
from llama_index.prompts import Prompt
from llama_index.indices.service_context import ServiceContext
from langchain.chat_models import ChatOpenAI
from llama_index.memory import ChatMemoryBuffer
from llama_index.llms import OpenAI
from src.constants import *
from src.chats import *
from src.utils import *
from src.calculate import *
from llama_index.tools.tool_spec.base import BaseToolSpec

load_dotenv()

openai.api_key = os.getenv("OPENAI_API_KEY")

#llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.2, model_name="gpt-3.5-turbo", max_tokens=1500))
service_context = ServiceContext.from_defaults(llm=OpenAI(model_name="gpt-4", temperature=0.2, max_tokens=1500 ))

directory_path = os.path.join(os.environ['DOCUMENTS_DIR'], "India")

class Llama:
    index = None

@staticmethod
    def init():
        Llama.index = Llama.load_index(directory_path)

@staticmethod
    def load_index(directory_path):
        documents = SimpleDirectoryReader(directory_path, filename_as_id=True).load_data()
        print(f"Loaded documents with {len(documents)} pages")

try:
            storage_context = StorageContext.from_defaults(persist_dir=os.environ['STORAGE_DIR'])
            index = load_index_from_storage(storage_context)
            logging.info("Index loaded from storage.")
        except FileNotFoundError:
            logging.info("Index not found. Creating a new one...")
            index = GPTVectorStoreIndex.from_documents(documents)
            index.storage_context.persist()
            logging.info("New index created and persisted to storage.")

refreshed_docs = index.refresh_ref_docs(documents,
                                                update_kwargs={"delete_kwargs": {'delete_from_docstore': True}})
        print(refreshed_docs)
        print('Number of newly inserted/refreshed docs: ', sum(refreshed_docs))

index.storage_context.persist()
        logging.info("Index refreshed and persisted to storage.")

return index

# Load the index when the server starts
Llama.init()

user_sessions = {}

def is_tax_related(query):
    try:
        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "user", "content": query},
                {"role": "assistant", "content": "Is the above query about tax calculation? Return 'yes' if it is, otherwise return 'no'."}
            ]
        )
        print(response)
        assistant_response = response.choices[0].message.content.strip().lower()
        print("Assistant response:", assistant_response)
        return "yes" in assistant_response
    except Exception as e:
        print(f"Error in OpenAI call: {e}")
        return False

def data_querying(input_text, user_id):
    global user_sessions
    if user_id not in user_sessions:
        user_sessions[user_id] = {"interaction_count": 0, "conversation_history": [], "awaiting_tax_details": False, "detail_prompted": False}

user_session = user_sessions[user_id]
    user_session["interaction_count"] += 1
    print(f"Debug: User {user_id} interaction count: {user_session['interaction_count']}")
    user_session["conversation_history"].append(input_text)

if is_tax_related(input_text):
        user_session["detail_prompted"] = True
        user_session["awaiting_tax_details"] = True
    if user_session["awaiting_tax_details"] and user_session["interaction_count"] > 1:
        user_session["awaiting_tax_details"] = False
        response_message = """To further assist you with tax calculation, please provide the following details:
    1. **Basic Salary**: Your basic salary amount for the year.
    2. **Dearness Allowances**: Total dearness allowances received throughout the year.
    3. **HRA Received**: The HRA (House Rent Allowance) received from your company for the year.
    4. **Special Allowance**: Any special allowance received throughout the year.
    5. **PF Deduction**: Total Provident Fund deduction for the year.
    6. **Deduction under 80C**: Investments or expenses qualifying for 80C deductions (like ELSS, PPF, etc.) up to a maximum limit of 1.5 lakhs.
    7. **Deduction under 80D**: Medical insurance premium paid, up to a maximum limit of 25K.
    8. **NPS Deduction**: Contributions to the National Pension Scheme (NPS), up to a maximum limit of 2 lakhs.
        
    Providing these details will help us in accurately calculating your tax liability."""

return response_message

text_qa_template_str = (
        "Context information is below.\n"
        "---------------------\n"
        "{context_str}\n"
        "---------------------\n"
        "Using both the context information and also using your own knowledge, "
        "answer the question: {query_str}\n"
        "If the context isn't helpful, Please respond with: I don't understand please try to rephrase your question\n"
    )
    refine_template_str = (
        "The original question is as follows: {query_str}\n"
        "We have provided an existing answer: {existing_answer}\n"
        "We have the opportunity to refine the existing answer "
        "(only if needed) with some more context below.\n"
        "------------\n"
        "{context_msg}\n"
        "------------\n"
        "Using both the new context and your own knowledege, update or repeat the existing answer.\n"
    )
    
    
    create_json(user_id)
    # Get existing memory from user
    history = read_json(user_id)
    chat_history = convert_chat_message(history)
    if history:
        memory = ChatMemoryBuffer.from_defaults(token_limit=3900, chat_history=chat_history)
    else:
        memory = ChatMemoryBuffer.from_defaults(token_limit=3900)
    refine_template = Prompt(refine_template_str)
    text_qa_template = Prompt(text_qa_template_str)

system_prompt = """You are a Tax Helper Assistant, specialized in addressing and assisting with tax-related inquiries for the fiscal year 2023-24, including tax calculation guidance. 
    For general queries, provide clear explanations on tax matters and tax-saving strategies.
    Respond accordingly don't deviate from context, If you not able respond Please respond with: I don't understand please try to rephrase your question
    If User ask to proivde tax calculation for "50 Lakh" ("5000000") or above please respond with: `Please contact us directly at support@haive.tech for more detailed assistance. We can't provide tax calculation for 50 Lakh or above.`
    """

response = Llama.index.as_chat_engine(text_qa_template=text_qa_template, refine_template=refine_template,
                                          chat_mode="context", memory=memory, system_promt=system_prompt,
                                          ).chat(input_text)

user_sessions[user_id] = user_session

# Update User Memory
    save_response(response.response, user_id, input_text)

return response.response

def answer_question(question, user_id):
    user_id = "static_user_id"
    try:
        # response = data_querying(question)
        response = data_querying(question, user_id)

return response
    except Exception as e:
        return str(e)

Back to Directory File Manager