app.py

import os
import random
import time
from threading import Lock, Thread
from typing import Optional, Tuple

import gradio as gr
from langchain.agents import AgentType, Tool, initialize_agent
from langchain.agents.agent_toolkits import (
    create_conversational_retrieval_agent, create_retriever_tool)
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.chains import ConversationChain, RetrievalQA
from langchain.chat_models import ChatOpenAI
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.llms import HuggingFaceTextGenInference, OpenAI
from langchain.prompts import PromptTemplate
from langchain.schema.messages import SystemMessage
from langchain.tools import tool
from langchain.vectorstores import FAISS
from pydantic import BaseModel, Field


def reset_textbox():
    return gr.update(value='')


model_name = "BAAI/bge-base-en"

encode_kwargs = {'normalize_embeddings': True}

model_norm = HuggingFaceBgeEmbeddings(
    model_name=model_name,
    encode_kwargs=encode_kwargs
)


vectordb = FAISS.load_local('faissdb', embeddings=model_norm)
retriever = vectordb.as_retriever(
    search_type='similarity', search_kwargs={"k": 2})

prompt_template = """You are an expert legal assistant with extensive knowledge about Indian law. Your task is to respond to the given query in a factually correct and consise manner unless asked for a detailed explanation. Assume the query is asked by a common man unless explicitly specified otherwise, therefore no special acts or laws like ones for railway , army , police would apply to them. Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
{context}
Question: {question}
Response:"""


PROMPT = PromptTemplate(
    template=prompt_template, input_variables=["context", "question"]
)


class SearchInput(BaseModel):
    query: str = Field(description="should be a search query in string format")


@tool('search', args_schema=SearchInput)
def search(query: str) -> str:
    """Useful for retrieving documents related to Indian law."""
    retriever = vectordb.as_retriever(
        search_type='similarity', search_kwargs={"k": 2})
    res = retriever.get_relevant_documents(query)
    print(res)
    return res


def load_chain():
    # tool = create_retriever_tool(
    #     retriever,
    #     "search_legal_sections",
    #     "Searches and returns documents regarding Indian law. Accepts query as a string. For example: 'Section 298 of Indian Penal Code'."
    # )
    tools = [search]
    llm = ChatOpenAI(openai_api_base='http://20.124.240.6:8080/v1',
                     openai_api_key='none',)

    conv_agent_executor = create_conversational_retrieval_agent(
        llm, tools, verbose=False,
        system_message=SystemMessage(
            content="Your name is Votum, an expert legal assistant with extensive knowledge about Indian law. Your task is to respond to the given query in a factually correct and concise manner unless asked for a detailed explanation. Feel free to use any tools available to look up relevant information, only if necessary")
    )
    return conv_agent_executor


with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, chatbot])
    chain = load_chain()

    # def respond(message, chat_history):
    #     print('message is', message)
    #     bot_message = chain({'input': message})['output']
    #     chat_history.append((message, bot_message))
    #     time.sleep(2)
    #     return "", chat_history

    def user(user_message, history):
        return "", history + [[user_message, None]]

    def respond(history):
        print('message is', history[-1])
        bot_message = chain({'input': history[-1][0]})['output']

        if 'Final answer:' in bot_message:
            bot_message = bot_message.split('Final answer:')[-1]

        history[-1][1] = bot_message
        # for character in bot_message:
        #     history[-1][1] += character
        #     time.sleep(0.0)
        #     yield history
        return history

    clear.click(chain.memory.clear(),)
    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False
               ).then(respond, chatbot, chatbot)

if __name__ == "__main__":
    demo.queue(max_size=32).launch()