Create app.py

app.py

@@ -0,0 +1,120 @@
+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()