{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from langchain.vectorstores import Chroma\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.document_loaders import TextLoader\n",
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain.document_loaders import DirectoryLoader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings import HuggingFaceBgeEmbeddings\n",
    "\n",
    "model_name = \"BAAI/bge-base-en\"\n",
    "# set True to compute cosine similarity\n",
    "encode_kwargs = {'normalize_embeddings': True}\n",
    "\n",
    "model_norm = HuggingFaceBgeEmbeddings(\n",
    "    model_name=model_name,\n",
    "    model_kwargs={'device': 'cpu'},\n",
    "    encode_kwargs=encode_kwargs\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vectordb = FAISS.load_local('faissdb',embeddings=model_norm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = vectordb.as_retriever(search_type='similarity', search_kwargs={\"k\": 2})\n",
    "a = retriever.get_relevant_documents('Indian Penal Code 133')\n",
    "print([aa.metadata for aa in a])\n",
    "# a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "\n",
    "llm = ChatOpenAI(openai_api_base='http://20.124.240.6:8080/v1',\n",
    "           openai_api_key='none', callbacks=[StreamingStdOutCallbackHandler()], streaming=True,)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SIMPLE RETRIEVER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.agents.agent_toolkits import create_retriever_tool\n",
    "from langchain.agents.agent_toolkits import create_conversational_retrieval_agent\n",
    "from langchain.agents.openai_functions_agent.agent_token_buffer_memory import AgentTokenBufferMemory\n",
    "from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent\n",
    "from langchain.prompts import MessagesPlaceholder\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.tools import tool\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "memory = ConversationBufferMemory(\n",
    "    memory_key=\"chat_history\", return_messages=True)\n",
    "\n",
    "# This is needed for both the memory and the prompt\n",
    "\n",
    "# memory_key = \"history\"\n",
    "# memory = AgentTokenBufferMemory(memory_key=memory_key, llm=llm)\n",
    "\n",
    "\n",
    "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 consice and factually correct manner. 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.\n",
    "\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "Response:\"\"\"\n",
    "\n",
    "\n",
    "PROMPT = PromptTemplate(\n",
    "    template=prompt_template, input_variables=[\"context\", \"question\"]\n",
    ")\n",
    "\n",
    "\n",
    "class SearchInput(BaseModel):\n",
    "    query: str = Field(description=\"should be a search query in string format\")\n",
    "\n",
    "\n",
    "@tool('search', args_schema=SearchInput)\n",
    "def search(query: str) -> str:\n",
    "    \"\"\"Useful for retrieving documents related to Indian law.\"\"\"\n",
    "    retriever = vectordb.as_retriever(\n",
    "        search_type='similarity', search_kwargs={\"k\": 2})\n",
    "    res = retriever.get_relevant_documents(query)\n",
    "    print(res)\n",
    "    return res\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "tool = create_retriever_tool(\n",
    "    retriever,\n",
    "    \"search_legal_sections\",\n",
    "    \"Searches and returns documents regarding Indian law. Accept query as a string. For example: 'Section 298 of Indian Penal Code'.\",\n",
    ")\n",
    "tools = [tool]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "QA Chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_chain = RetrievalQA.from_chain_type(llm=llm,\n",
    "                                       chain_type_kwargs={\"prompt\": PROMPT},\n",
    "                                       retriever=retriever,\n",
    "                                       return_source_documents=False,\n",
    "                                       )\n",
    "\n",
    "conv_qa_chain = ConversationalRetrievalChain.from_llm(\n",
    "    llm, retriever, memory=memory, verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ReAct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import HuggingFaceTextGenInference\n",
    "\n",
    "llm_hg = HuggingFaceTextGenInference(\n",
    "    inference_server_url=\"http://20.83.177.108:8080/\",\n",
    "    max_new_tokens=2000,\n",
    "    # top_k=10,\n",
    "    # top_p=0.95,\n",
    "    # typical_p=0.95,\n",
    "    # temperature=0.6,\n",
    "    # repetition_penalty=1.1,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import initialize_agent, Tool\n",
    "from langchain.agents import AgentType\n",
    "from langchain.schema.messages import SystemMessage\n",
    "import langchain\n",
    "\n",
    "langchain.verbose = True\n",
    "\n",
    "\n",
    "\n",
    "agent_executor = initialize_agent(\n",
    "    tools, \n",
    "    llm, \n",
    "    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, \n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "agent_kwargs = {\n",
    "    \"extra_prompt_messages\": [MessagesPlaceholder(variable_name=\"memory\")],\n",
    "    \"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.\"),\n",
    "}\n",
    "\n",
    "conv_agent_executor = create_conversational_retrieval_agent(\n",
    "    llm, [search], verbose=False,  agent_kwargs=agent_kwargs, \n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# {'input': 'How is section 308 of Indian Penal Code different from section 299?'}\n",
    "conv_agent_executor(\n",
    "    {'input': 'Explain sections related to medical negligence.'}\n",
    "    )\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Flare"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We set this so we can see what exactly is going on\n",
    "from langchain.chains import FlareChain\n",
    "import langchain\n",
    "import os\n",
    "\n",
    "os.environ['OPENAI_API_KEY'] = 'none'\n",
    "os.environ['OPENAI_API_BASE'] = 'http://20.124.240.6:8080/v1'\n",
    "# os.environ['OPEN']\n",
    "\n",
    "langchain.verbose = True\n",
    "\n",
    "\n",
    "flare = FlareChain.from_llm(\n",
    "    llm,\n",
    "    retriever=retriever,\n",
    "    max_generation_len=164,\n",
    "    min_prob=0.3,\n",
    ")\n",
    "\n",
    "query = \"explain in great detail the difference between the langchain framework and baby agi\"\n",
    "flare.run(query)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plan and Execute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain_experimental.plan_and_execute import PlanAndExecute, load_agent_executor, load_chat_planner\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.utilities import SerpAPIWrapper\n",
    "from langchain.agents.tools import Tool\n",
    "from langchain.chains import LLMMathChain\n",
    "\n",
    "planner = load_chat_planner(llm)\n",
    "executor = load_agent_executor(llm, [search], verbose=True)\n",
    "plan_agent = PlanAndExecute(planner=planner, executor=executor, verbose=True)\n",
    "\n",
    "\n",
    "plan_agent.run('I bought a house in 2001 for 20 lakh rupees , i sold it in 2022 for 50 lakhs , what will be my profit?')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "DOCSTORE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import OpenAI\n",
    "from langchain.agents import initialize_agent, Tool\n",
    "from langchain.agents import AgentType\n",
    "from langchain.agents.react.base import DocstoreExplorer\n",
    "import langchain\n",
    "\n",
    "langchain.verbose = True\n",
    "\n",
    "\n",
    "docstore = DocstoreExplorer(vectordb)\n",
    "tools = [\n",
    "    Tool(\n",
    "        name=\"Search\",\n",
    "        func=docstore.search,\n",
    "        description=\"useful for when you need to ask with search\",\n",
    "    ),\n",
    "    Tool(\n",
    "        name=\"Lookup\",\n",
    "        func=docstore.lookup,\n",
    "        description=\"useful for when you need to ask with lookup\",\n",
    "    ),\n",
    "]\n",
    "\n",
    "\n",
    "react_docstore = initialize_agent(\n",
    "    tools, llm, agent=AgentType.REACT_DOCSTORE, verbose=True)\n",
    "\n",
    "react_docstore.run('hi')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from langchain.chat_models import ChatOpenAI\n",
    "# from langchain.document_loaders import TextLoader\n",
    "# from langchain.embeddings import OpenAIEmbeddings\n",
    "# from langchain.indexes import VectorstoreIndexCreator\n",
    "# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "\n",
    "# from langchain.schema import HumanMessage\n",
    "\n",
    "\n",
    "# llm_n = ChatOpenAI(openai_api_base='http://20.124.240.6:8080/v1',\n",
    "#                    openai_api_key='none', callbacks=[StreamingStdOutCallbackHandler()],streaming=True,)\n",
    "\n",
    "# questions = [\n",
    "#     \"Who is the speaker\",\n",
    "#     \"What did the president say about Ketanji Brown Jackson\",\n",
    "#     \"What are the threats to America\",\n",
    "#     \"Who are mentioned in the speech\",\n",
    "#     \"Who is the vice president\",\n",
    "#     \"How many projects were announced\",\n",
    "# ]\n",
    "\n",
    "\n",
    "# llm_n(\n",
    "#     [\n",
    "#         HumanMessage(\n",
    "#             content=\"What model are you?\"\n",
    "#         )\n",
    "#     ]\n",
    "# )\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}