gradio_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.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"
# set True to compute cosine similarity
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})


# relating to refer to Indian Penal Code(IPC), CrPC(Code of Criminal Procedure) for most cases and therefore laws
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.83.177.108:8080/v1',
                     openai_api_key='none',)

    conv_agent_executor = create_conversational_retrieval_agent(
        llm, tools, verbose=True,
    )
    return conv_agent_executor


class ChatWrapper:
    def __init__(self):
        self.lock = Lock()

    def __call__(
        self, inp: str, history: Optional[Tuple[str, str]], chain: Optional[ConversationChain]
    ):
        """Execute the chat functionality."""
        self.lock.acquire()
        try:
            history = history or []
            # Run chain and append input.
            # output = chain({'input': inp})
            output = 'this is an output'
            history.append((inp, output))
        except Exception as e:
            raise e
        finally:
            self.lock.release()
        return history, history


chat = ChatWrapper()

block = gr.Blocks(css=".gradio-container {background-color: red}")

with block:
    chatbot = gr.Chatbot()

    with gr.Row():
        message = gr.Textbox(
            label="What's your question?",
            placeholder="What's the answer to life, the universe, and everything?",
            lines=1,
        )
        submit = gr.Button(value="Send", variant="secondary").style(
            full_width=False)

    gr.Examples(
        examples=[
            "Hi! How's it going?",
            "What should I do tonight?",
            "Whats 2 + 2?",
        ],
        inputs=message,
    )

    state = gr.State()
    agent_state = gr.State()

    load_chain()

    submit.click(chat, inputs=[message,
                 state, agent_state], outputs=[chatbot, state])
    message.submit(chat, inputs=[
        message, state, agent_state], outputs=[chatbot, state])

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

    msg.submit(respond, [msg, chatbot], [msg, chatbot])

if __name__ == "__main__":
    demo.launch()