PDF_RAG_with_LMstudio.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "aa2b9045-0485-4252-8348-680b126b591a",
   "metadata": {},
   "source": [
    "Scientific Publications as PDFs, Text Extraction, Storage SQL database, Chunks, Embedding and Vectorization, Retrieval, GPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a6f4f4c-7300-4cb9-8a96-31b4fec74b27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Not required, but can be usefull if a Cell need local proxy connection to Docker or LM Studio\n",
    "'''\n",
    "import os\n",
    "os.environ['http_proxy']=\"http://localhost:1238\"\n",
    "os.environ['https_proxy']=\"http://localhost:1238\"\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6f32765-9ce2-4a03-b023-c99450f64562",
   "metadata": {},
   "source": [
    "## Setting Up PDF Collection and DATABASE Paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0efea5a-4126-4a8a-bc78-b301549a83f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "cwd = os.getcwd()\n",
    "print(\"current working directory\",cwd)\n",
    "#pdf file directory\n",
    "pdf_collection_path='PDF_Pet'\n",
    "path= os.path.join(cwd, \"docs\", pdf_collection_path)\n",
    "#--------------------------------------------------------------------\n",
    "#database directory\n",
    "base_name='sample.db'\n",
    "database_name=os.path.join(cwd, \"docs\", pdf_collection_path, base_name)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1f44a23-c9cd-43c5-be3d-24fe4ed99cc5",
   "metadata": {},
   "source": [
    "## Read PDF and Extract TEXT, Stores in a SQLite3 database (db): GROBID based approach"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51c70f1c-43c1-4c48-a20a-fdfa32042b39",
   "metadata": {},
   "source": [
    "# Instructions for Running GROBID Docker\n",
    "1. Download the small GROBID docker image and copy `config.json` into the Jupyter notebook working directory.\n",
    "2. Run local Docker:\n",
    "    - Start Docker for Windows.\n",
    "    - Use CMD.\n",
    "3. Execute the following command to run the docker container:\n",
    "    ```bash\n",
    "    docker run --rm --init --ulimit core=0 -p 8070:8070 lfoppiano/grobid:0.8.0\n",
    "    ```\n",
    "4. Check the running docker with [http://localhost:8070](http://localhost:8070) and ensure this matches the one in your `config.json`.\n",
    "5. path to the PDF collections and database are setted above   \n",
    "6. Choose to recreate the SQL (y/n) or just reload the current one.\n",
    "7. Note that processing 70 PDFs can take over 10 minutes.\n",
    "8. Wait for the console message: \"Finished\".\n",
    "9. If you encounter an XML error or some \"localhost\" mix up, ensure you stay connected with the docker by avoiding proxy configurations elsewhere.\n",
    "10. todo: add new pdfs to database\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b7a4210-da75-415b-a13a-831f4a6714b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sqlite3\n",
    "import pandas as pd\n",
    "from grobid_client.grobid_client import GrobidClient\n",
    "import json\n",
    "import grobid_tei_xml\n",
    "\n",
    "\n",
    "\n",
    "class GrobidAuthor:\n",
    "    def __init__(self, full_name):\n",
    "        self.full_name = full_name\n",
    "\n",
    "class GrobidBiblio:\n",
    "    def __init__(self, index, authors, title, date, volume, pages, issue, journal, doi):\n",
    "        self.index = index\n",
    "        self.authors = authors\n",
    "        self.title = title\n",
    "        self.date = date\n",
    "        self.volume = volume\n",
    "        self.pages = pages\n",
    "        self.issue = issue\n",
    "        self.journal = journal\n",
    "        self.doi = doi\n",
    "\n",
    "def extract_bibliographic_details(biblios, print_choice=False):\n",
    "    for biblio in biblios:\n",
    "        if print_choice:\n",
    "            print(\"Index\", biblio.index, \"| Title:\", biblio.title)\n",
    "            print(\"Authors:\")\n",
    "            for author in biblio.authors:\n",
    "                print(\"-\", author.full_name)\n",
    "\n",
    "            print(\"Date:\", biblio.date)\n",
    "            print(\"Volume:\", biblio.volume)\n",
    "            print(\"Pages:\", biblio.pages)\n",
    "            # print(\"Issue:\", biblio.issue)\n",
    "            print(\"Journal:\", biblio.journal)\n",
    "            print(\"Doi:\", biblio.doi)\n",
    "            print()\n",
    "        else:\n",
    "            i=1# Add any other logic you want to perform when print_choice is False\n",
    "            pass\n",
    "    ref_list = '*'.join([\n",
    "        f\" * Index: {biblio.index} | Title: {biblio.title} | Authors: {', '.join([author.full_name for author in biblio.authors])} | Date: {biblio.date} | Volume: {biblio.volume} | Pages: {biblio.pages} | Journal: {biblio.journal} | Doi: {biblio.doi}\\|\"\n",
    "        for biblio in biblios\n",
    "    ])\n",
    "    return ref_list\n",
    "    \n",
    "    \n",
    "client = GrobidClient(config_path=\"config.json\")\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "print(\"\\n\\ndatabase_pathname\",database_name)\n",
    "\n",
    "\n",
    "\n",
    "service_name = \"processFulltextDocument\"\n",
    "\n",
    "# Check if the database file exists\n",
    "if os.path.exists(database_name):\n",
    "    while True:\n",
    "        choice = input(\"\\nThe database file already exists. Do you want to create a new one? (y/n): \")\n",
    "        if choice.lower() == 'y':\n",
    "            # Close the connection to the database if it is open\n",
    "            if 'conn' in locals():\n",
    "                conn.close()\n",
    "\n",
    "            # Delete the existing database file\n",
    "            conn = None  # Reset conn variable\n",
    "            os.remove(database_name)\n",
    "            break\n",
    "        elif choice.lower() == 'n':\n",
    "            break\n",
    "        else:\n",
    "            print(\"Invalid choice! Please enter 'y' or 'n'.\")\n",
    "\n",
    "# Connect to the SQLite database\n",
    "conn = sqlite3.connect(database_name)\n",
    "# Create the table if it doesn't exist\n",
    "cursor = conn.cursor()\n",
    "cursor.execute(\"SELECT name FROM sqlite_master WHERE type='table';\")\n",
    "result = cursor.fetchone()\n",
    "\n",
    "if result is None:\n",
    "    table_name = \"your_table_name\"  # Provide a table name of your choice\n",
    "    cursor.execute('''CREATE TABLE {table_name}\n",
    "                    (ID INTEGER PRIMARY KEY,\n",
    "                    Title TEXT,\n",
    "                    Authors TEXT,\n",
    "                    DOI TEXT,\n",
    "                    Citations INTEGER,\n",
    "                    Abstract TEXT,\n",
    "                    Body TEXT,\n",
    "                    Refs TEXT)'''.format(table_name=table_name))\n",
    "else:\n",
    "    table_name = result[0]\n",
    "    \n",
    "# Retrieve the maximum ID value from the table\n",
    "cursor.execute(\"SELECT MAX(ID) FROM {table_name}\".format(table_name=table_name))\n",
    "max_id = cursor.fetchone()[0]\n",
    "# Increment the ID value by one\n",
    "id_key = max_id + 1 if max_id is not None else 1\n",
    "\n",
    "\n",
    "if choice.lower() == 'y':\n",
    "    # Loop through the files in the directory\n",
    "    for filename in os.listdir(path):\n",
    "        if filename.endswith(\".pdf\"):\n",
    "            print(f\"filename: {filename}\")\n",
    "            file_path = os.path.join(path, filename)\n",
    "\n",
    "            pdf_file, status, text = client.process_pdf(service_name, \n",
    "                                     file_path, \n",
    "                                     generateIDs=True, \n",
    "                                     consolidate_header=True, \n",
    "                                     consolidate_citations=True, \n",
    "                                     include_raw_citations=True, \n",
    "                                     include_raw_affiliations=True, \n",
    "                                     tei_coordinates=True,                          \n",
    "                                     segment_sentences=True)\n",
    "\n",
    "\n",
    "            grobid_biblios=grobid_tei_xml.parse_citation_list_xml(text)\n",
    "            # Extract metadat and text and print the bibliographic details\n",
    "            ref_list=extract_bibliographic_details(grobid_biblios, print_choice=False)\n",
    "\n",
    "            doc = grobid_tei_xml.parse_document_xml(text)\n",
    "            title = doc.header.title\n",
    "            authors = ';'.join([a.full_name for a in doc.header.authors])\n",
    "            doi = str(doc.header.doi)\n",
    "            citations = str(len(doc.citations))\n",
    "            abstract = doc.abstract\n",
    "            body = doc.body\n",
    "\n",
    "            #----------------------------------------\n",
    "            print(f\"id_key: {id_key}\")\n",
    "            # Insert the information into the database\n",
    "            cursor.execute('''\n",
    "                INSERT INTO {table_name} (ID, Title, Authors, DOI, Citations, Abstract, Body, Refs)\n",
    "                VALUES (?, ?, ?, ?, ?, ?, ?, ?)\n",
    "            '''.format(table_name=table_name), (id_key, title, authors, doi, citations, abstract, body, ref_list))\n",
    "\n",
    "            # Commit the changes to the database\n",
    "            conn.commit()\n",
    "            # Increment the ID key for the next iteration\n",
    "            id_key += 1\n",
    "\n",
    "            # Retrieve the count of records from the table\n",
    "            cursor.execute(\"SELECT COUNT(*) FROM {table_name}\".format(table_name=table_name))\n",
    "            count = cursor.fetchone()[0]\n",
    "            print(\"record loop:\", count)\n",
    "elif choice.lower() == 'n':\n",
    "    print(\"\\n\\nActive database_pathname\", database_name)\n",
    "else:\n",
    "    print(\"Invalid choice! Please enter 'y' or 'n'.\")\n",
    "    \n",
    "# Retrieve the count of records from the table\n",
    "cursor.execute(\"SELECT COUNT(*) FROM {table_name}\".format(table_name=table_name))\n",
    "count = cursor.fetchone()[0]\n",
    "print(\"Total records:\", count)        \n",
    "# Close the connection\n",
    "conn.close()\n",
    "print(\"finished\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1b3b4f3-ca02-4fcd-b653-133fd323d9be",
   "metadata": {},
   "source": [
    "# Load and Check of SQL database: \n",
    "## Can also be used alone without Docker, once SQL database was created\n",
    " * database\n",
    "    * base_name='sample.db'\n",
    "    * database_name=os.path.join(cwd, \"docs\", pdf_collection_path, base_name) #database\n",
    " * Enter the record number or 'q' to quit: After: Enter 'n' for next choice or 'q' for quitEnter the record number or 'q' to quit: After: Enter 'n' for next choice or 'q' for quit# Check of SQL database: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1037e57-de4a-4bba-a8f2-b8f3d530882b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sqlite3\n",
    "from IPython.display import clear_output\n",
    "# Function to clear the console screen\n",
    "def clear_console():\n",
    "    if os.name == 'nt':  # for Windows\n",
    "        os.system('cls')\n",
    "    else:  # for Mac and Linux\n",
    "        os.system('clear')\n",
    "# Connect to the SQLite database\n",
    "conn = sqlite3.connect(database_name)\n",
    "# Retrieve the table name from the database\n",
    "cursor = conn.cursor()\n",
    "cursor.execute(\"SELECT name FROM sqlite_master WHERE type='table';\")\n",
    "table_name = cursor.fetchone()[0]\n",
    "# Create the loop\n",
    "while True:\n",
    "    # Clear the console screen\n",
    "    # Prompt the user to enter the record number or 'q' to quit\n",
    "    user_input = input(\"Enter the record number or 'q' to quit: After: Enter 'n' for next choice or 'q' for quit\")\n",
    "    \n",
    "    \n",
    "        \n",
    "    # Check if the user wants to quit\n",
    "    if user_input.lower() == 'q':\n",
    "        break\n",
    "    # Validate the user input as an integer\n",
    "    try:\n",
    "        chosen_record = int(user_input)\n",
    "    except ValueError:\n",
    "        print(\"Invalid input. Please enter a valid record number or 'q' to quit.\")\n",
    "        continue\n",
    "    # Retrieve the \"body\" column from the table for the chosen record\n",
    "    cursor.execute(\"SELECT body FROM {table_name} LIMIT 1 OFFSET {chosen_record}\".format(table_name=table_name, chosen_record=chosen_record - 1))\n",
    "    record_body = cursor.fetchone()\n",
    "    \n",
    "    # Check if the record exists\n",
    "    if record_body is None:\n",
    "        print(\"Record does not exist. Please enter a valid record number.\")\n",
    "        continue\n",
    "    # Clear the console screen again before printing the new text\n",
    "    \n",
    "   \n",
    "    # Print the body information of the chosen record\n",
    "    print(record_body[0])\n",
    "    user_input = input(\"Next: 'n' Quite: 'q'\")\n",
    "    \n",
    "    if user_input.lower() == 'n':\n",
    "        # Clear the screen (Windows-specific)\n",
    "        #print(\"clean\")\n",
    "        clear_output()\n",
    "    else:\n",
    "        break\n",
    "        \n",
    "        \n",
    "# Close the connection\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc948e52-7d3c-4990-b4d4-67f573a8c1a5",
   "metadata": {},
   "source": [
    "# adds metadata to chunk, seem to be not good for embedding\n",
    " * Metadata might be not useful here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e79aaed-9f9e-4524-a3f3-991e494ed0a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlite3\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain_text_splitters.base import Document\n",
    "\n",
    "# Connect to the SQLite database\n",
    "conn = sqlite3.connect(database_name)\n",
    "\n",
    "# Create the table if it doesn't exist\n",
    "cursor = conn.cursor()\n",
    "cursor.execute(\"SELECT ID, Title, Authors, DOI, Abstract, Body FROM {table_name}\".format(table_name=table_name))\n",
    "results = cursor.fetchall()\n",
    "\n",
    "# Define the desired chunk size\n",
    "desired_chunk_size = 2000\n",
    "# Define the chunk overlap\n",
    "chunk_overlap = 450\n",
    "\n",
    "# Initialize the text splitter\n",
    "r_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=desired_chunk_size,\n",
    "    chunk_overlap=chunk_overlap,\n",
    "    separators=[\"\\n\\n\", \"\\n\", \".\", \" \", \"\"]\n",
    ")\n",
    "\n",
    "# Loop through the results and process the data\n",
    "texts_db = []\n",
    "for result in results:\n",
    "    record_id, title, authors, doi, abstract, body = result\n",
    "    print(f\"Record_id {record_id} | title {title}\")\n",
    "    \n",
    "    # your code to retrieve abstract and body values from the database\n",
    "    if abstract is None:\n",
    "        abstract = \"\"\n",
    "    if body is None:\n",
    "        body = \"\"\n",
    "    \n",
    "    # Combine the abstract and body values\n",
    "    combined_text = abstract + \" \" + body\n",
    "    \n",
    "    # Create a Document object for the continuous large text\n",
    "    #document = Document(page_content=combined_text, metadata={'ID': record_id, 'Title': title, 'Authors': authors, 'DOI': doi})\n",
    "    #document = Document(page_content=combined_text, metadata={'ID': record_id})\n",
    "    document = Document(page_content=combined_text, metadata={'ID': record_id, 'Title': title})\n",
    "    # Chunk the document\n",
    "    r_splits = r_splitter.split_documents([document])\n",
    "    \n",
    "    # Add the chunked documents to the texts list, preserving the metadata\n",
    "   # Create Document objects for the chunks and add them to the texts list\n",
    "    for i, chunk in enumerate(r_splits):\n",
    "        chunk_document = Document(\n",
    "            page_content=chunk.page_content,\n",
    "            metadata={\n",
    "                'ID': chunk.metadata['ID'],\n",
    "                #'Page': i+1,\n",
    "                'Title': chunk.metadata['Title'],\n",
    "                #'Authors': chunk.metadata['Authors'],\n",
    "                #'DOI': chunk.metadata['DOI'],\n",
    "                #'Chunk size': len(chunk.page_content)\n",
    "            }\n",
    "        )\n",
    "        texts_db.append(chunk_document)\n",
    "\n",
    "# Close the connection\n",
    "conn.close()\n",
    "\n",
    "# Create the custom embedding object\n",
    "#embedding = CustomEmbedding2(texts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5cdfb95-a184-4ac0-afce-1ca82dd20ad0",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts_db"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42c285f4-fed9-45fd-8fa0-2085191f07e4",
   "metadata": {},
   "source": [
    "# Next Steps: Running LM Studio and Creating or Loading Vector Database\n",
    "1. Start LM Studio 0.2.27.\n",
    "2. Activate the Local Server Tab (\" <-> \").\n",
    "3. Select a model, e.g., \"Meta-Llama-3-8B-Instruct-Q8_0.gguf\".\n",
    "    - Choose the preset \"Llama3\".\n",
    "4. Choose a server port that does not conflict with the docker port.\n",
    "5. Load the \"Embedding Model\", e.g., \"nomic embed text v1 5 Q8_0.gguf\".\n",
    "    - Models were previously downloaded manually from Huggingface/Models/.\n",
    "6. Press \"Start Server\".\n",
    "7. Define the following paths in line 80-81:\n",
    "    ```python\n",
    "    \n",
    "    user_path = os.path.join(os.getcwd(), \"docs\", pdf_collection_path, \"sample\")\n",
    "    temp_dir = os.path.join(os.getcwd(), \"docs\", \"temp\")\n",
    "    ```\n",
    "8. Run the cell:\n",
    "    - If there is an error indicating the base is in use, ensure the path is correct and start from the cell \"Check SQL Database\" after restarting the kernel.\n",
    "9. If there is a connection error to LM, remove any other open proxy or browser proxy settings in \"Settings\", restart the kernel, and start from \"Check SQL Database\".\n",
    "10. After creating the vector database, set `new_vectordb` in line 80-81 to `False`.\n",
    "11. If \"new_vectordbnew_vectordb= False\" then can start directly from here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "760af7f8-f518-4654-91b3-52932fb4c606",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Default values\n",
    "#new_vectordb = True\n",
    "new_vectordb = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd54cdd8-1dad-46df-9dd5-81b0884a2184",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import Chroma\n",
    "from langchain.docstore.document import Document\n",
    "from openai import OpenAI\n",
    "from typing import List\n",
    "import os\n",
    "import shutil\n",
    "import tempfile\n",
    "import time\n",
    "\n",
    "class CustomEmbedding2:\n",
    "    def __init__(self):\n",
    "        self.embeddings = []  # List to store the embeddings\n",
    "\n",
    "    def embed_documents(self, texts: List[str]) -> List[List[float]]:\n",
    "        embeddings = [get_embedding(text) for text in texts]\n",
    "        self.embeddings = embeddings  # Store the embeddings in the `embeddings` attribute\n",
    "        return embeddings\n",
    "\n",
    "    def embed_query(self, text: str) -> List[float]:\n",
    "        embedding = get_embedding(text)\n",
    "        self.embeddings = [embedding]\n",
    "        return embedding\n",
    "\n",
    "    def get_embeddings(self) -> List[List[float]]:\n",
    "        return self.embeddings\n",
    "\n",
    "def get_embedding(text, model=\"TheBloke/nomic-embed-text\"):\n",
    "    text = text.replace(\"\\n\", \" \")\n",
    "    return client.embeddings.create(input=[text], model=model).data[0].embedding\n",
    "\n",
    "\n",
    "def create_vectordb(embedding, r_splits, new_vectordb=True, user_path=None, temp_dir=None):\n",
    "    cwd = os.getcwd()  # Get the current working directory\n",
    "\n",
    "    if new_vectordb:\n",
    "        if user_path is None:\n",
    "            if temp_dir is None:\n",
    "                # Create a temporary directory for persisting the vector database\n",
    "                temp_dir = tempfile.mkdtemp()\n",
    "                persist_directory = os.path.join(cwd, temp_dir, 'chroma_db')\n",
    "            else:\n",
    "                persist_directory = os.path.join(cwd, temp_dir, 'chroma_db')\n",
    "        else:\n",
    "            persist_directory = os.path.join(cwd, user_path, 'chroma_db')\n",
    "\n",
    "        # Remove the existing persist directory (if any)\n",
    "        if os.path.exists(persist_directory):\n",
    "            shutil.rmtree(persist_directory)\n",
    "\n",
    "        print(\"Creating a new vector database...\")\n",
    "        start_time = time.time()\n",
    "        vectordb = Chroma.from_documents(\n",
    "            documents=r_splits,\n",
    "            embedding=embedding,\n",
    "            persist_directory=persist_directory\n",
    "        )\n",
    "        end_time = time.time()\n",
    "        print(f\"New vector database created in {end_time - start_time:.2f} seconds. Directory: {persist_directory}\")\n",
    "\n",
    "    else:\n",
    "        if user_path is None:\n",
    "            print(\"Please provide a valid user path to load the existing vector database.\")\n",
    "            return None\n",
    "        else:\n",
    "            persist_directory = os.path.join(cwd, user_path, 'chroma_db')\n",
    "            if not os.path.exists(persist_directory):\n",
    "                print(f\"Vector database not found in the specified path: {persist_directory}\")\n",
    "                return None\n",
    "\n",
    "            print(\"Loading existing vector database...\")\n",
    "            start_time = time.time()\n",
    "            vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)\n",
    "            end_time = time.time()\n",
    "            print(f\"Existing vector database loaded in {end_time - start_time:.2f} seconds. Directory: {persist_directory}\")\n",
    "\n",
    "    return vectordb, persist_directory\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# Check if new_vectordb is True and set texts_db accordingly\n",
    "if new_vectordb:\n",
    "    texts_db = texts_db\n",
    "else:\n",
    "    texts_db = []\n",
    "\n",
    "user_path = os.path.join(os.getcwd(), \"docs\", pdf_collection_path, \"sample\")\n",
    "temp_dir = os.path.join(os.getcwd(), \"docs\", \"temp\")\n",
    "\n",
    "client = OpenAI(base_url=\"http://localhost:1238/v1\", api_key=\"lm-studio\")\n",
    "\n",
    "# Extract the text from the Document objects text_db \n",
    "#texts = [doc.page_content for doc in texts_db]\n",
    "# Create the custom embedding object\n",
    "embedding = CustomEmbedding2()\n",
    "\n",
    "\n",
    "vectordb, temp_dir = create_vectordb(embedding, texts_db, new_vectordb, user_path, temp_dir)\n",
    "print(\"\\nfinished\", temp_dir, \"\\n\", len(vectordb))\n",
    "print(vectordb._collection.count())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b7ef748-65ed-442a-8b75-c4c5eb0448e6",
   "metadata": {},
   "source": [
    "# Alternative Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9229f8b-46a5-4dd8-9012-2706b695731a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#from langchain.chat_models import ChatOpenAI\n",
    "from langchain.retrievers.multi_query import MultiQueryRetriever\n",
    "from langchain_openai import OpenAI\n",
    "\n",
    "llm  = OpenAI(base_url=\"http://localhost:1238/v1\", api_key=\"lm-studio\")\n",
    "#llm = ChatOpenAI(temperature=0)\n",
    "\n",
    "# we instantiated the retreiever above\n",
    "retriever_from_llm = MultiQueryRetriever.from_llm(\n",
    "    retriever=vectordb.as_retriever(), llm=llm\n",
    ")\n",
    "\n",
    "# Set logging for the queries\n",
    "import logging\n",
    "\n",
    "logging.basicConfig()\n",
    "logging.getLogger(\"langchain.retrievers.multi_query\").setLevel(logging.INFO)\n",
    "\n",
    "#unique_docs = retriever_from_llm.get_relevant_documents(query=question)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3f8b717-2eda-49a2-92e1-3373fcd4fdd8",
   "metadata": {},
   "source": [
    "## Based on  client.chat.completions.create  : \n",
    "    * Browser TAB\n",
    "    * Prompt Input\n",
    "    * Query\n",
    "    * Response\n",
    "    * Whole message sent to Model\n",
    "## Chose Retrieval: Standard or Multi (langChain): line 32\n",
    " * overall the responses are quite slow, because of retrieval of useful text snippets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae3811b2-7699-481c-a187-7260251acd9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import panel as pn\n",
    "from openai import OpenAI\n",
    "from langchain_community.vectorstores import Chroma\n",
    "#from langchain_community.embeddings import HuggingFaceEmbeddings\n",
    "from langchain.prompts import PromptTemplate\n",
    "import re\n",
    "import time\n",
    "\n",
    "\n",
    "# Define the template for the prompt\n",
    "template = \"\"\"\n",
    "Use the following context from the vector database and the previous conversation to answer the query. Incorporate your own knowledge and reasoning as an AI assistant:\n",
    "Previous Conversation: {history}\n",
    "Context: {context}\n",
    "Query: {query}\n",
    "Answer:\n",
    "{answer}\n",
    "\"\"\"\n",
    "\n",
    "# Create the prompt template\n",
    "prompt_template = PromptTemplate(\n",
    "    input_variables=[\"history\", \"context\", \"query\", \"answer\"],\n",
    "    template=template\n",
    ")\n",
    "\n",
    "\n",
    "history = []\n",
    "\n",
    "# Point to the local OpenAI server\n",
    "client = OpenAI(base_url=\"http://localhost:1238/v1\", api_key=\"lm-studio\")\n",
    "\n",
    "\n",
    "which_retriever='Multi'\n",
    "#which_retriever='Standard'\n",
    "\n",
    "if which_retriever =='Multi':\n",
    "    retriever=retriever_from_llm\n",
    "else:\n",
    "    # Set up the vector retriever\n",
    "    retriever = vectordb.as_retriever()\n",
    "\n",
    "    \n",
    "# Define a function to handle the user's input\n",
    "def handle_input2():\n",
    "    global is_first_run\n",
    "    wrap_width = 100\n",
    "\n",
    "    if is_first_run:\n",
    "        query = \"Hello, introduce yourself to someone opening this program for the first time. Be concise\"\n",
    "        is_first_run = False\n",
    "        some_context = \"\"\n",
    "    else:\n",
    "        output_box_retrievalTime.object = f\"Processing Retrieval {which_retriever}\"\n",
    "        output_box_gpt_responseTime = f\"Waiting for GPT\"\n",
    "        start_time = time.time()\n",
    "        query = input_box.value.strip()\n",
    "        search_results = retriever.get_relevant_documents(query)\n",
    "        #search_results = retriever_from_llm.get_relevant_documents(query)\n",
    "        some_context = \"\"\n",
    "        for result in search_results:\n",
    "            some_context += result.page_content + \"'Title-- '\"+result.metadata['Title']+\"--EndTitle | \\n\\n\"\n",
    "        \n",
    "        end_time = time.time()\n",
    "        # Calculate the elapsed time\n",
    "        elapsed_time = end_time - start_time\n",
    "        output_box_retrievalTime.object = f\"Time taken for Retrieval with: {elapsed_time:.2f} seconds\"\n",
    "    #history.append({\"role\": \"user\", \"content\": query)# + \" Database content: [\" + some_context + \"]\"})\n",
    "\n",
    "    if query:\n",
    "        # Generate the prompt using the template and input values\n",
    "        prompt = prompt_template.format(history=history, context=some_context, query=query, answer=\"\")\n",
    "        #prompt = prompt_template.format(history=history, answer=\"\")\n",
    "        output_box2.object=prompt\n",
    "        start_time = time.time()\n",
    "        completion = client.chat.completions.create(\n",
    "            model=\"LLMA/Meta-Llama-3-8B\",\n",
    "            messages=history + [{\"role\": \"user\", \"content\": prompt}],\n",
    "            temperature=0.7,\n",
    "            stream=True,\n",
    "        )\n",
    "\n",
    "        \n",
    "        new_message = {\"role\": \"assistant\", \"content\": \"\"}\n",
    "\n",
    "        full_response = \"\"\n",
    "        for chunk in completion:\n",
    "            if chunk.choices[0].delta.content:\n",
    "                new_message[\"content\"] += chunk.choices[0].delta.content\n",
    "                full_response += (chunk.choices[0].delta.content)\n",
    "\n",
    "        formatted_text = f\"<pre style='white-space: pre-wrap; width: {wrap_width}ch; font-family: Arial, sans-serif; font-size: 12px;'>**Query:**\\n{query}\\n\\n**Final Answer:**\\n{full_response}</pre>\"\n",
    "        output_box.object = formatted_text\n",
    "\n",
    "        history.append(new_message)\n",
    "        \n",
    "        input_box.value = \"\"\n",
    "        \n",
    "        end_time = time.time()\n",
    "        elapsed_time = end_time - start_time\n",
    "        output_box_responseTime.object = f\"Time taken for Model: {elapsed_time:.2f} seconds\"\n",
    "\n",
    "\n",
    "# Define a function to handle user input changes\n",
    "def handle_input_change(event):\n",
    "    handle_input2()\n",
    "\n",
    "# Load the Panel extension\n",
    "pn.extension()\n",
    "\n",
    "# Create an input box\n",
    "input_box = pn.widgets.TextAreaInput(width=400, height=100, styles={'overflow-y': 'scroll', 'text-align': 'center'})\n",
    "# Create the output box\n",
    "output_box = pn.pane.Markdown(\n",
    "    width=800,\n",
    "    height=400,\n",
    "    styles={\n",
    "        'overflow-y': 'scroll',\n",
    "        'color': 'black',                   # Text color\n",
    "        'background-color': '#f0f0f0',      # Background color (light gray)\n",
    "        'font-size': '14px',                # Font size\n",
    "        'font-family': 'Courier New',       # Font family\n",
    "        'padding': '5px',                   # Padding inside the pane\n",
    "        'border': '1px solid red',          # Border around the pane\n",
    "        'text-align': 'center'              # Center the text\n",
    "    }\n",
    ")\n",
    "\n",
    "output_box2 = pn.pane.Markdown(\n",
    "    height=400,\n",
    "    styles={\n",
    "        'overflow-y': 'scroll',\n",
    "        'overflow-x': 'scroll',\n",
    "        'color': 'black',        # Text color\n",
    "        'background-color': '#f0f0f4',  # Background color\n",
    "        'font-size': '14px',     # Font size\n",
    "        'font-family': 'Courier New',  # Font family\n",
    "        'padding': '5px',        # Padding inside the pane\n",
    "        'border': '1px solid blue'  # Border around the pane\n",
    "    }\n",
    ")\n",
    "\n",
    "\n",
    "output_box_retrievalTime = pn.pane.Markdown()\n",
    "output_box_responseTime = pn.pane.Markdown()\n",
    "# Watch for changes in the input box value\n",
    "input_box.param.watch(handle_input_change, 'value')\n",
    "\n",
    "# Create a variable to track if it's the first run\n",
    "is_first_run = True\n",
    "if is_first_run:\n",
    "        # Call handle_input2()\n",
    "        handle_input2()\n",
    "        is_first_run = False\n",
    "\n",
    "         \n",
    "# Create the dashboard layout using Panel\n",
    "dashboard_layout = pn.Column(\n",
    "    pn.pane.Markdown(\"# ChatGPT-like Conversation: RAG\"),\n",
    "    pn.Spacer(height=20),\n",
    "    pn.Row(input_box, output_box),\n",
    "    pn.Row(pn.Spacer(height=20)),\n",
    "    pn.Row(output_box_retrievalTime, output_box_responseTime),\n",
    "    pn.Row(output_box2),\n",
    "    name=\"ChatGPT Dashboard\"\n",
    ")\n",
    "\n",
    "# Create a dashboard object using the layout\n",
    "dashboard = pn.panel(dashboard_layout)\n",
    "\n",
    "# Display the dashboard in a new browser tab\n",
    "dashboard.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "073f50ae-1a12-4aae-aa90-7f5f947c971e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c87a3780-68d2-4fea-8fc5-e2984cdfd73f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "0fccb34c-790d-493f-a74f-5b2a12ba55c9",
   "metadata": {},
   "source": [
    "# Alternative Version: based on QA CHAIN with BROWSER TAB - \n",
    " * output still not consitent and short"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac895296-a285-4bbb-b26d-5ae9eebc19fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import panel as pn\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain.docstore.document import Document\n",
    "import textwrap\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "\n",
    "\n",
    "# Create a RAG model\n",
    "llm = OpenAI(base_url=\"http://localhost:1238/v1\", api_key=\"lm-studio\")\n",
    "retriever = vectordb.as_retriever()\n",
    "\n",
    "# Define a custom prompt template\n",
    "template = \"\"\"\n",
    "Use the following context from the vector database and the previous conversation to answer the query. Incorporate your own knowledge and reasoning as an AI assistant:\n",
    "\n",
    "Previous Conversation: {history}\n",
    "\n",
    "Context: {context}\n",
    "\n",
    "Query: {query}\n",
    "\n",
    "Answer:\n",
    "\"\"\"\n",
    "\n",
    "prompt_template = PromptTemplate(input_variables=[\"history\", \"context\", \"query\"], template=template)\n",
    "\n",
    "# Create a custom chain\n",
    "qa_chain = LLMChain(llm=llm, prompt=prompt_template)\n",
    "\n",
    "\n",
    "# Load the Panel extension\n",
    "pn.extension()\n",
    "\n",
    "import textwrap\n",
    "import panel as pn\n",
    "# Create an input box\n",
    "input_box = pn.widgets.TextAreaInput()\n",
    "# Create the output box\n",
    "output_box = pn.pane.Markdown()\n",
    "# Create a history list to store the conversation\n",
    "history = []\n",
    "# Define a function to handle the user's input\n",
    "def handle_input():\n",
    "    query = input_box.value.strip()\n",
    "    if query:\n",
    "        # Add user input to the history\n",
    "        #history.append({\"role\": \"user\", \"content\": query})\n",
    "        # Process the query and generate a response\n",
    "        context = retriever.get_relevant_documents(query)\n",
    "        result_str = qa_chain.run(\n",
    "            history=\"\\n\".join([f\"{msg['role']}: {msg['content']}\" for msg in history]),\n",
    "            query=query,\n",
    "            #max_tokens=8096,\n",
    "            max_tokens=-1,\n",
    "            context=\"\".join([doc.page_content for doc in context])\n",
    "        )\n",
    "        final_answer_start = result_str.find(\"Final Answer:\")\n",
    "        if final_answer_start != -1:\n",
    "            final_answer = result_str[final_answer_start + len(\"Final Answer:\"):].strip()\n",
    "            output_box.object = f\"**Query:**\\n{query}\\n\\n**Final Answer:**\\n{textwrap.fill(final_answer, width=80)}\"\n",
    "        else:\n",
    "            final_answer=result_str\n",
    "            output_box.object = f\"**Query:**\\n{query}\\n\\n**Answer:**\\n{textwrap.fill(result_str, width=80)}\"\n",
    "        \n",
    "        # Add assistant response to the history\n",
    "        history.append({\"role\": \"assistant\", \"content\": final_answer})\n",
    "        history.append({\"role\": \"user\", \"content\": query})\n",
    "        # Clear the input box\n",
    "        input_box.value = \"\"\n",
    "        \n",
    "        \n",
    "# Watch for changes in the input box value\n",
    "input_box.param.watch(lambda event: handle_input(), 'value')\n",
    "# Create the dashboard layout using Panel\n",
    "dashboard_layout = pn.Column(\n",
    "    pn.pane.Markdown(\"# ChatGPT-like Conversation\"),\n",
    "    pn.Spacer(height=20),\n",
    "    pn.Row(input_box, output_box),\n",
    "    name=\"ChatGPT Dashboard\"\n",
    ")\n",
    "# Create a dashboard object using the layout\n",
    "dashboard = pn.panel(dashboard_layout)\n",
    "# Display the dashboard in a new browser tab\n",
    "dashboard.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d9c0f4f-c064-4988-8e03-8d804d71717c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.19"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}