add RAG part

.gitignore

@@ -11,3 +11,8 @@ wheels/
 
 .DS_Store
 gradio_cached_examples
+
+.env
+local_qdrant
+
+.vscode

pyproject.toml

@@ -7,9 +7,14 @@ authors = [
 ]
 dependencies = [
     "gradio>=4.19.2",
-    "langchain>=0.1.9",
     "gradio-pdf>=0.0.5",
     "loguru>=0.7.2",
+    "pypdf>=4.0.2",
+    "langchain>=0.1.9",
+    "openai>=1.12.0",
+    "tiktoken>=0.6.0",
+    "qdrant-client>=1.7.3",
+    "requests>=2.31.0",
 ]
 readme = "README.md"
 requires-python = ">= 3.8"

requirements-dev.lock

@@ -22,6 +22,7 @@ annotated-types==0.6.0
 anyio==4.3.0
     # via httpx
     # via langchain-core
+    # via openai
     # via starlette
 asttokens==2.4.1
     # via icecream
@@ -51,6 +52,8 @@ dataclasses-json==0.6.4
     # via langchain-community
 decorator==5.1.1
     # via ipython
+distro==1.9.0
+    # via openai
 executing==2.0.1
     # via icecream
     # via stack-data
@@ -75,17 +78,30 @@ gradio-client==0.10.1
     # via gradio
 gradio-pdf==0.0.5
     # via pdfchat
+grpcio==1.62.0
+    # via grpcio-tools
+    # via qdrant-client
+grpcio-tools==1.62.0
+    # via qdrant-client
 h11==0.14.0
     # via httpcore
     # via uvicorn
+h2==4.1.0
+    # via httpx
+hpack==4.0.0
+    # via h2
 httpcore==1.0.4
     # via httpx
 httpx==0.27.0
     # via gradio
     # via gradio-client
+    # via openai
+    # via qdrant-client
 huggingface-hub==0.20.3
     # via gradio
     # via gradio-client
+hyperframe==6.0.1
+    # via h2
 icecream==2.1.3
 idna==3.6
     # via anyio
@@ -150,6 +166,9 @@ numpy==1.26.4
     # via langchain-community
     # via matplotlib
     # via pandas
+    # via qdrant-client
+openai==1.12.0
+    # via pdfchat
 orjson==3.9.15
     # via gradio
     # via langsmith
@@ -171,8 +190,12 @@ pexpect==4.9.0
 pillow==10.2.0
     # via gradio
     # via matplotlib
+portalocker==2.8.2
+    # via qdrant-client
 prompt-toolkit==3.0.43
     # via ipython
+protobuf==4.25.3
+    # via grpcio-tools
 ptyprocess==0.7.0
     # via pexpect
 pure-eval==0.2.2
@@ -183,6 +206,8 @@ pydantic==2.6.2
     # via langchain
     # via langchain-core
     # via langsmith
+    # via openai
+    # via qdrant-client
 pydantic-core==2.16.3
     # via pydantic
 pydub==0.25.1
@@ -193,6 +218,8 @@ pygments==2.17.2
     # via rich
 pyparsing==3.1.1
     # via matplotlib
+pypdf==4.0.2
+    # via pdfchat
 python-dateutil==2.8.2
     # via matplotlib
     # via pandas
@@ -206,15 +233,21 @@ pyyaml==6.0.1
     # via langchain
     # via langchain-community
     # via langchain-core
+qdrant-client==1.7.3
+    # via pdfchat
 referencing==0.33.0
     # via jsonschema
     # via jsonschema-specifications
+regex==2023.12.25
+    # via tiktoken
 requests==2.31.0
     # via huggingface-hub
     # via langchain
     # via langchain-community
     # via langchain-core
     # via langsmith
+    # via pdfchat
+    # via tiktoken
 rich==13.7.0
     # via typer
 rpds-py==0.18.0
@@ -224,6 +257,8 @@ ruff==0.2.2
     # via gradio
 semantic-version==2.10.0
     # via gradio
+setuptools==69.1.1
+    # via grpcio-tools
 shellingham==1.5.4
     # via typer
 six==1.16.0
@@ -232,6 +267,7 @@ six==1.16.0
 sniffio==1.3.0
     # via anyio
     # via httpx
+    # via openai
 sqlalchemy==2.0.27
     # via langchain
     # via langchain-community
@@ -243,12 +279,15 @@ tenacity==8.2.3
     # via langchain
     # via langchain-community
     # via langchain-core
+tiktoken==0.6.0
+    # via pdfchat
 tomlkit==0.12.0
     # via gradio
 toolz==0.12.1
     # via altair
 tqdm==4.66.2
     # via huggingface-hub
+    # via openai
 traitlets==5.14.1
     # via ipython
     # via matplotlib-inline
@@ -259,6 +298,7 @@ typing-extensions==4.9.0
     # via gradio
     # via gradio-client
     # via huggingface-hub
+    # via openai
     # via pydantic
     # via pydantic-core
     # via sqlalchemy
@@ -269,6 +309,7 @@ typing-inspect==0.9.0
 tzdata==2024.1
     # via pandas
 urllib3==2.2.1
+    # via qdrant-client
     # via requests
 uvicorn==0.27.1
     # via gradio

requirements.lock

@@ -22,6 +22,7 @@ annotated-types==0.6.0
 anyio==4.3.0
     # via httpx
     # via langchain-core
+    # via openai
     # via starlette
 attrs==23.2.0
     # via aiohttp
@@ -45,6 +46,8 @@ cycler==0.12.1
 dataclasses-json==0.6.4
     # via langchain
     # via langchain-community
+distro==1.9.0
+    # via openai
 fastapi==0.110.0
     # via gradio
 ffmpy==0.3.2
@@ -66,17 +69,30 @@ gradio-client==0.10.1
     # via gradio
 gradio-pdf==0.0.5
     # via pdfchat
+grpcio==1.62.0
+    # via grpcio-tools
+    # via qdrant-client
+grpcio-tools==1.62.0
+    # via qdrant-client
 h11==0.14.0
     # via httpcore
     # via uvicorn
+h2==4.1.0
+    # via httpx
+hpack==4.0.0
+    # via h2
 httpcore==1.0.4
     # via httpx
 httpx==0.27.0
     # via gradio
     # via gradio-client
+    # via openai
+    # via qdrant-client
 huggingface-hub==0.20.3
     # via gradio
     # via gradio-client
+hyperframe==6.0.1
+    # via h2
 idna==3.6
     # via anyio
     # via httpx
@@ -135,6 +151,9 @@ numpy==1.26.4
     # via langchain-community
     # via matplotlib
     # via pandas
+    # via qdrant-client
+openai==1.12.0
+    # via pdfchat
 orjson==3.9.15
     # via gradio
     # via langsmith
@@ -152,12 +171,18 @@ pandas==2.2.1
 pillow==10.2.0
     # via gradio
     # via matplotlib
+portalocker==2.8.2
+    # via qdrant-client
+protobuf==4.25.3
+    # via grpcio-tools
 pydantic==2.6.2
     # via fastapi
     # via gradio
     # via langchain
     # via langchain-core
     # via langsmith
+    # via openai
+    # via qdrant-client
 pydantic-core==2.16.3
     # via pydantic
 pydub==0.25.1
@@ -166,6 +191,8 @@ pygments==2.17.2
     # via rich
 pyparsing==3.1.1
     # via matplotlib
+pypdf==4.0.2
+    # via pdfchat
 python-dateutil==2.8.2
     # via matplotlib
     # via pandas
@@ -179,15 +206,21 @@ pyyaml==6.0.1
     # via langchain
     # via langchain-community
     # via langchain-core
+qdrant-client==1.7.3
+    # via pdfchat
 referencing==0.33.0
     # via jsonschema
     # via jsonschema-specifications
+regex==2023.12.25
+    # via tiktoken
 requests==2.31.0
     # via huggingface-hub
     # via langchain
     # via langchain-community
     # via langchain-core
     # via langsmith
+    # via pdfchat
+    # via tiktoken
 rich==13.7.0
     # via typer
 rpds-py==0.18.0
@@ -197,6 +230,8 @@ ruff==0.2.2
     # via gradio
 semantic-version==2.10.0
     # via gradio
+setuptools==69.1.1
+    # via grpcio-tools
 shellingham==1.5.4
     # via typer
 six==1.16.0
@@ -204,6 +239,7 @@ six==1.16.0
 sniffio==1.3.0
     # via anyio
     # via httpx
+    # via openai
 sqlalchemy==2.0.27
     # via langchain
     # via langchain-community
@@ -213,12 +249,15 @@ tenacity==8.2.3
     # via langchain
     # via langchain-community
     # via langchain-core
+tiktoken==0.6.0
+    # via pdfchat
 tomlkit==0.12.0
     # via gradio
 toolz==0.12.1
     # via altair
 tqdm==4.66.2
     # via huggingface-hub
+    # via openai
 typer==0.9.0
     # via gradio
 typing-extensions==4.9.0
@@ -226,6 +265,7 @@ typing-extensions==4.9.0
     # via gradio
     # via gradio-client
     # via huggingface-hub
+    # via openai
     # via pydantic
     # via pydantic-core
     # via sqlalchemy
@@ -236,6 +276,7 @@ typing-inspect==0.9.0
 tzdata==2024.1
     # via pandas
 urllib3==2.2.1
+    # via qdrant-client
     # via requests
 uvicorn==0.27.1
     # via gradio

requirements.txt

@@ -21,6 +21,7 @@ annotated-types==0.6.0
 anyio==4.3.0
     # via httpx
     # via langchain-core
+    # via openai
     # via starlette
 attrs==23.2.0
     # via aiohttp
@@ -44,6 +45,8 @@ cycler==0.12.1
 dataclasses-json==0.6.4
     # via langchain
     # via langchain-community
+distro==1.9.0
+    # via openai
 fastapi==0.110.0
     # via gradio
 ffmpy==0.3.2
@@ -65,17 +68,30 @@ gradio-client==0.10.1
     # via gradio
 gradio-pdf==0.0.5
     # via pdfchat
+grpcio==1.62.0
+    # via grpcio-tools
+    # via qdrant-client
+grpcio-tools==1.62.0
+    # via qdrant-client
 h11==0.14.0
     # via httpcore
     # via uvicorn
+h2==4.1.0
+    # via httpx
+hpack==4.0.0
+    # via h2
 httpcore==1.0.4
     # via httpx
 httpx==0.27.0
     # via gradio
     # via gradio-client
+    # via openai
+    # via qdrant-client
 huggingface-hub==0.20.3
     # via gradio
     # via gradio-client
+hyperframe==6.0.1
+    # via h2
 idna==3.6
     # via anyio
     # via httpx
@@ -134,6 +150,9 @@ numpy==1.26.4
     # via langchain-community
     # via matplotlib
     # via pandas
+    # via qdrant-client
+openai==1.12.0
+    # via pdfchat
 orjson==3.9.15
     # via gradio
     # via langsmith
@@ -151,12 +170,18 @@ pandas==2.2.1
 pillow==10.2.0
     # via gradio
     # via matplotlib
+portalocker==2.8.2
+    # via qdrant-client
+protobuf==4.25.3
+    # via grpcio-tools
 pydantic==2.6.2
     # via fastapi
     # via gradio
     # via langchain
     # via langchain-core
     # via langsmith
+    # via openai
+    # via qdrant-client
 pydantic-core==2.16.3
     # via pydantic
 pydub==0.25.1
@@ -165,6 +190,8 @@ pygments==2.17.2
     # via rich
 pyparsing==3.1.1
     # via matplotlib
+pypdf==4.0.2
+    # via pdfchat
 python-dateutil==2.8.2
     # via matplotlib
     # via pandas
@@ -178,15 +205,21 @@ pyyaml==6.0.1
     # via langchain
     # via langchain-community
     # via langchain-core
+qdrant-client==1.7.3
+    # via pdfchat
 referencing==0.33.0
     # via jsonschema
     # via jsonschema-specifications
+regex==2023.12.25
+    # via tiktoken
 requests==2.31.0
     # via huggingface-hub
     # via langchain
     # via langchain-community
     # via langchain-core
     # via langsmith
+    # via pdfchat
+    # via tiktoken
 rich==13.7.0
     # via typer
 rpds-py==0.18.0
@@ -196,6 +229,8 @@ ruff==0.2.2
     # via gradio
 semantic-version==2.10.0
     # via gradio
+setuptools==69.1.1
+    # via grpcio-tools
 shellingham==1.5.4
     # via typer
 six==1.16.0
@@ -203,6 +238,7 @@ six==1.16.0
 sniffio==1.3.0
     # via anyio
     # via httpx
+    # via openai
 sqlalchemy==2.0.27
     # via langchain
     # via langchain-community
@@ -212,12 +248,15 @@ tenacity==8.2.3
     # via langchain
     # via langchain-community
     # via langchain-core
+tiktoken==0.6.0
+    # via pdfchat
 tomlkit==0.12.0
     # via gradio
 toolz==0.12.1
     # via altair
 tqdm==4.66.2
     # via huggingface-hub
+    # via openai
 typer==0.9.0
     # via gradio
 typing-extensions==4.9.0
@@ -225,6 +264,7 @@ typing-extensions==4.9.0
     # via gradio
     # via gradio-client
     # via huggingface-hub
+    # via openai
     # via pydantic
     # via pydantic-core
     # via sqlalchemy
@@ -235,6 +275,7 @@ typing-inspect==0.9.0
 tzdata==2024.1
     # via pandas
 urllib3==2.2.1
+    # via qdrant-client
     # via requests
 uvicorn==0.27.1
     # via gradio

src/pdfchat/app.py

@@ -1,12 +1,52 @@
+import os
+import re
 import time
 from dataclasses import dataclass
 from pathlib import Path
 
 import gradio as gr
+import requests
 from gradio_pdf import PDF
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.prompts import PromptTemplate
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.vectorstores.qdrant import Qdrant
 from loguru import logger
-
-MODEL_CALM2 = "cyberagent/calm2"
+from pypdf import PdfReader
+
+MODEL_CALM2 = "cyberagent/calm2-7b-chat"
+text_splitter = CharacterTextSplitter(
+    separator="\n\n",
+    chunk_size=1000,
+    chunk_overlap=0,
+)
+QDRANT_MODE = "local"
+if QDRANT_MODE == "local":
+    QDRANT_CLIENT_CONFIG = {
+        "path": "./local_qdrant",
+    }
+elif QDRANT_MODE == "cloud":
+    QDRANT_CLIENT_CONFIG = {
+        "url": os.environ.get("QDRANT_URL"),
+        "api_key": os.environ.get("QDRANT_API_KEY"),
+    }
+    if not QDRANT_CLIENT_CONFIG["url"] or not QDRANT_CLIENT_CONFIG["api_key"]:
+        raise ValueError(
+            "Please set the QDRANT_URL and QDRANT_API_KEY environment variables."
+        )
+COLLECTION_NAME = "pdfchat"
+PROMPT_TEMPLATE = """以下の文脈を利用して、最後の質問に答えなさい。
+答えがわからない場合は、わからないと答えてください。
+
+【文脈】
+{context}
+
+【質問】
+{question}
+
+【答え】
+"""
+LLM_URL = os.environ.get("PDFCHAT_LLM_URL")
 
 
 @dataclass
@@ -46,34 +86,85 @@ def open_file(file_path: str) -> str:
     if file_path.suffix == ".txt":
         text = file_path.read_text()
     elif file_path.suffix == ".pdf":
-        text = "WARNING: PDF file is not supported yet."
+        text = parse_pdf(file_path)
     else:
         text = "WARNING: Unsupported file format."
 
     return text
 
 
-def get_response(query: str, document: str | None) -> str:
-    response = ""
-    if not document:
-        response = "No document is uploaded. Please upload a document."
-    else:
-        response = f"Your document: {document}"
+def parse_pdf(file_path: Path, backend="pypdf") -> str:
+    reader = PdfReader(file_path)
+    contents = "".join([page.extract_text() for page in reader.pages])
+
+    contents = re.sub(r"[ 　]+\n[ 　]+", "\n", contents)
+    contents = re.sub(r"[ 　]+\n", "\n", contents)
+    contents = re.sub(r"\n[ 　]+", "\n", contents)
+    contents = re.sub(r"[^。\n]\n", "", contents)
+    contents = re.sub(r"[\w、（）]\n[\w、（）]", "", contents)
+    contents = re.sub(r"\n{3,}", "\n\n", contents)
+
+    return contents
+
+
+def get_response(prompt: str) -> str:
+    response = requests.post(
+        LLM_URL,
+        json={
+            "prompt": prompt,
+            "max_new_tokens": 2048,
+        },
+    ).json()
 
     return response
 
 
-def bot(history: ChatHistory, query: str, file_path: str) -> ChatHistory:
+def retrieve_relevant_documents(query: str, document: str | None) -> list[str]:
+    if not document:
+        return "No document is uploaded. Please upload a document."
+
+    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+    if not OPENAI_API_KEY:
+        raise ValueError("Please set the OPENAI_API_KEY environment variable.")
+
+    documents = text_splitter.split_text(document)
+    embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
+
+    db = Qdrant.from_texts(
+        texts=documents,
+        embedding=embeddings,
+        **QDRANT_CLIENT_CONFIG,
+    )
+    retriever = db.as_retriever()
+    relevant_documents = [
+        doc.page_content for doc in retriever.get_relevant_documents(query)
+    ]
+
+    return relevant_documents
+
+
+def build_prompt(query: str, context: str) -> str:
+    prompt = PromptTemplate(
+        template=PROMPT_TEMPLATE,
+        input_variables=["context", "question"],
+    ).format(context=context, question=query)
+
+    return prompt
+
+
+def main(history: ChatHistory, query: str, file_path: str | None) -> ChatHistory:
     history = ChatHistory(history)
     document = open_file(file_path) if file_path else None
-    response = get_response(query, document)
-    history.add_chat(Chat(query=query, response=response))
+    relevant_documents = retrieve_relevant_documents(query=query, document=document)
+    prompt = build_prompt(query=query, context="\n\n".join(relevant_documents))
+    response_message = get_response(prompt)["message"]
+    history.add_chat(Chat(query=query, response=response_message))
     logger.info(history)
 
     history.clear_last_response()
-    for char in response:
+    for char in response_message:
         history[-1].response += char
-        time.sleep(0.02)
+        time.sleep(0.01)
         yield history
 
 
@@ -121,7 +212,7 @@ with gr.Blocks() as app:
                 )
                 submit_button = gr.Button("Submit", variant="primary", size="sm")
                 submit = submit_button.click(
-                    fn=bot,
+                    fn=main,
                     inputs=[chatbot, text_box, file_box],
                     outputs=chatbot,
                 )