{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "956ff3a4",
   "metadata": {},
   "source": [
    "# Chat\n",
    "\n",
    "## Imports and Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "926fc482-dd03-40e1-8083-171aad3e2b26",
   "metadata": {
    "height": 217,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "\n",
    "# read local .env file (that should contain an \"OPENAI_API_KEY\")\n",
    "_ = load_dotenv(find_dotenv()) \n",
    "\n",
    "abscurdir = os.path.abspath(os.curdir)\n",
    "docsdir = os.path.join(os.path.dirname(abscurdir), 'docs')\n",
    "existing_dbname = 'chroma_20241124_132314'\n",
    "persist_directory = os.path.join(docsdir, existing_dbname)\n",
    "collection_name = 'MLbooks'\n",
    "# llm_name = 'ollama3.2.1b'\n",
    "llm_name = 'openai'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51ece06a",
   "metadata": {},
   "source": [
    "## Embeddings, Vectorstore and LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab80dc86",
   "metadata": {
    "height": 113,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain_community.embeddings import HuggingFaceEmbeddings\n",
    "lc_embeddings = HuggingFaceEmbeddings(model_name=\"all-MiniLM-L6-v2\")\n",
    "\n",
    "import chromadb\n",
    "from langchain_community.vectorstores import Chroma\n",
    "\n",
    "client = chromadb.PersistentClient(\n",
    "    path=persist_directory\n",
    "    )\n",
    "\n",
    "vectorstore = Chroma(\n",
    "    client=client,\n",
    "    collection_name=collection_name,\n",
    "    embedding_function=lc_embeddings,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef4228b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"instantiating llm model: {llm_name}\")\n",
    "if llm_name == 'ollama3.2.1b':\n",
    "    from langchain.llms import Ollama\n",
    "    llm = Ollama(model=\"llama3.2:1b\", temperature=0)\n",
    "elif llm_name == 'openai':\n",
    "    from langchain_openai import ChatOpenAI\n",
    "    llm = ChatOpenAI(model_name=\"gpt-4o-mini\", temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "025962c3",
   "metadata": {},
   "source": [
    "## QA Chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8aa6c11",
   "metadata": {
    "height": 351,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Build prompt\n",
    "from langchain.prompts import PromptTemplate\n",
    "template = \"\"\"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. Use three sentences maximum. Keep the answer as concise as possible. Always say \"thanks for asking!\" at the end of the answer. \n",
    "{context}\n",
    "Question: {question}\n",
    "Helpful Answer:\"\"\"\n",
    "QA_CHAIN_PROMPT = PromptTemplate(input_variables=[\"context\", \"question\"],template=template,)\n",
    "\n",
    "# Run chain\n",
    "from langchain.chains import RetrievalQA\n",
    "question = 'How can machine learning models help a business?'\n",
    "qa_chain = RetrievalQA.from_chain_type(\n",
    "    llm,\n",
    "    retriever=vectorstore.as_retriever(),\n",
    "    return_source_documents=True,\n",
    "    chain_type_kwargs={\"prompt\": QA_CHAIN_PROMPT})\n",
    "\n",
    "result = qa_chain({\"query\": question})\n",
    "result[\"result\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84f8d2e2",
   "metadata": {},
   "source": [
    "### Memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e99e754",
   "metadata": {
    "height": 98,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationBufferMemory\n",
    "memory = ConversationBufferMemory(\n",
    "    memory_key=\"chat_history\",\n",
    "    return_messages=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "244e1b7a",
   "metadata": {},
   "source": [
    "### ConversationalRetrievalChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23700502",
   "metadata": {
    "height": 132,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.chains import ConversationalRetrievalChain\n",
    "retriever = vectorstore.as_retriever(\n",
    "    # search_type=\"mmr\",\n",
    "    search_type=\"similarity\",\n",
    "    search_kwargs={\"k\": 3})\n",
    "crc = ConversationalRetrievalChain.from_llm(\n",
    "    llm,\n",
    "    chain_type='stuff',\n",
    "    retriever=retriever,\n",
    "    memory=memory\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e25b9fbb",
   "metadata": {
    "height": 47,
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = 'How does multiclass classification work?'\n",
    "result = crc({\"question\": question})\n",
    "result['answer']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b280672",
   "metadata": {
    "height": 47,
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = \"Look at the context that was earlier provided.\"\n",
    "result = crc({\"question\": question})\n",
    "result['answer']"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain_311",
   "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}