mirror of https://github.com/hwchase17/langchain
Add quip loader (#12259)
- **Description:** implement [quip](https://quip.com) loader - **Issue:** https://github.com/langchain-ai/langchain/issues/10352 - **Dependencies:** No - pass make format, make lint, make test --------- Co-authored-by: Hao Fan <h_fan@apple.com> Co-authored-by: Bagatur <baskaryan@gmail.com>pull/12619/head^2
Frank
11 months ago
committed by
GitHub
parent
c9a6940d58
commit
bf5805bb32
@ -0,0 +1,66 @@
|
||||
{
|
||||
"metadata": {
|
||||
"vscode": {
|
||||
"interpreter": {
|
||||
"hash": "cc99336516f23363341912c6723b01ace86f02e26b4290be1efc0677e2e2ec24"
|
||||
}
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python",
|
||||
"display_name": "Python (Pyodide)",
|
||||
"language": "python"
|
||||
},
|
||||
"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.9.1"
|
||||
}
|
||||
},
|
||||
"nbformat_minor": 4,
|
||||
"nbformat": 4,
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": "# Quip\n\n>[Quip](https://quip.com) is a collaborative productivity software suite for mobile and Web. It allows groups of people to create and edit documents and spreadsheets as a group, typically for business purposes.\n\nA loader for `Quip` docs.\n\nPlease refer [here](https://quip.com/dev/automation/documentation/current#section/Authentication/Get-Access-to-Quip's-APIs) to know how to get personal access token. \n\nSpecify a list `folder_ids` and/or `thread_ids` to load in the corresponding docs into Document objects, if both are specified, loader will get all `thread_ids` belong to this folder based on `folder_ids`, combine with passed `thread_ids`, the union of both sets will be returned.\n\n* How to know folder_id ? \n go to quip folder, right click folder and copy link, extract suffix from link as folder_id. Hint: `https://example.quip.com/<folder_id>`\n* How to know thread_id ? \n thread_id is the document id. Go to quip doc, right click doc and copy link, extract suffix from link as thread_id. Hint: `https://exmaple.quip.com/<thread_id>`\n \nYou can also set `include_all_folders` as `True` will fetch group_folder_ids and \nYou can also specify a boolean `include_attachments` to include attachments, this is set to False by default, if set to True all attachments will be downloaded and QuipLoader will extract the text from the attachments and add it to the Document object. Currently supported attachment types are: `PDF`, `PNG`, `JPEG/JPG`, `SVG`, `Word` and `Excel`. Also you can sepcify a boolean `include_comments` to include comments in document, this is set to False by default, if set to True all comments in document will be fetched and QuipLoader will add them to Document objec.\n",
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": "Before using QuipLoader make sure you have the latest version of the quip-api package installed:",
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": "#!pip install quip-api",
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"execution_count": 1,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": "## Examples",
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": "### Personal Access Token",
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": "from langchain.document_loaders import QuipLoader\n\nloader = QuipLoader(api_url=\"https://platform.quip.com\", access_token=\"change_me\", request_timeout=60)\ndocuments = loader.load(\n folder_ids={'123','456'}, thread_ids={'abc', 'efg'} include_attachments=False, include_comments=False\n)",
|
||||
"metadata": {},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
}
|
||||
]
|
||||
}
|
||||
@ -0,0 +1,232 @@
|
||||
import logging
|
||||
import re
|
||||
import xml.etree.cElementTree
|
||||
import xml.sax.saxutils
|
||||
from io import BytesIO
|
||||
from typing import List, Optional, Sequence
|
||||
from xml.etree.ElementTree import ElementTree
|
||||
|
||||
from langchain.docstore.document import Document
|
||||
from langchain.document_loaders.base import BaseLoader
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_MAXIMUM_TITLE_LENGTH = 64
|
||||
|
||||
|
||||
class QuipLoader(BaseLoader):
|
||||
"""Load `Quip` pages.
|
||||
|
||||
Port of https://github.com/quip/quip-api/tree/master/samples/baqup
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, api_url: str, access_token: str, request_timeout: Optional[int] = 60
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
api_url: https://platform.quip.com
|
||||
access_token: token of access quip API. Please refer:
|
||||
https://quip.com/dev/automation/documentation/current#section/Authentication/Get-Access-to-Quip's-APIs
|
||||
request_timeout: timeout of request, default 60s.
|
||||
"""
|
||||
try:
|
||||
from quip_api.quip import QuipClient
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"`quip_api` package not found, please run " "`pip install quip_api`"
|
||||
)
|
||||
|
||||
self.quip_client = QuipClient(
|
||||
access_token=access_token, base_url=api_url, request_timeout=request_timeout
|
||||
)
|
||||
|
||||
def load(
|
||||
self,
|
||||
folder_ids: Optional[List[str]] = None,
|
||||
thread_ids: Optional[List[str]] = None,
|
||||
max_docs: Optional[int] = 1000,
|
||||
include_all_folders: bool = False,
|
||||
include_comments: bool = False,
|
||||
include_images: bool = False,
|
||||
) -> List[Document]:
|
||||
"""
|
||||
Args:
|
||||
:param folder_ids: List of specific folder IDs to load, defaults to None
|
||||
:param thread_ids: List of specific thread IDs to load, defaults to None
|
||||
:param max_docs: Maximum number of docs to retrieve in total, defaults 1000
|
||||
:param include_all_folders: Include all folders that your access_token
|
||||
can access, but doesn't include your private folder
|
||||
:param include_comments: Include comments, defaults to False
|
||||
:param include_images: Include images, defaults to False
|
||||
"""
|
||||
if not folder_ids and not thread_ids and not include_all_folders:
|
||||
raise ValueError(
|
||||
"Must specify at least one among `folder_ids`, `thread_ids` "
|
||||
"or set `include_all`_folders as True"
|
||||
)
|
||||
|
||||
thread_ids = thread_ids or []
|
||||
|
||||
if folder_ids:
|
||||
for folder_id in folder_ids:
|
||||
self.get_thread_ids_by_folder_id(folder_id, 0, thread_ids)
|
||||
|
||||
if include_all_folders:
|
||||
user = self.quip_client.get_authenticated_user()
|
||||
if "group_folder_ids" in user:
|
||||
self.get_thread_ids_by_folder_id(
|
||||
user["group_folder_ids"], 0, thread_ids
|
||||
)
|
||||
if "shared_folder_ids" in user:
|
||||
self.get_thread_ids_by_folder_id(
|
||||
user["shared_folder_ids"], 0, thread_ids
|
||||
)
|
||||
|
||||
thread_ids = list(set(thread_ids[:max_docs]))
|
||||
return self.process_threads(thread_ids, include_images, include_comments)
|
||||
|
||||
def get_thread_ids_by_folder_id(
|
||||
self, folder_id: str, depth: int, thread_ids: List[str]
|
||||
) -> None:
|
||||
"""Get thread ids by folder id and update in thread_ids"""
|
||||
from quip_api.quip import HTTPError, QuipError
|
||||
|
||||
try:
|
||||
folder = self.quip_client.get_folder(folder_id)
|
||||
except QuipError as e:
|
||||
if e.code == 403:
|
||||
logging.warning(
|
||||
f"depth {depth}, Skipped over restricted folder {folder_id}, {e}"
|
||||
)
|
||||
else:
|
||||
logging.warning(
|
||||
f"depth {depth}, Skipped over folder {folder_id} "
|
||||
f"due to unknown error {e.code}"
|
||||
)
|
||||
return
|
||||
except HTTPError as e:
|
||||
logging.warning(
|
||||
f"depth {depth}, Skipped over folder {folder_id} "
|
||||
f"due to HTTP error {e.code}"
|
||||
)
|
||||
return
|
||||
|
||||
title = folder["folder"].get("title", "Folder %s" % folder_id)
|
||||
|
||||
logging.info(f"depth {depth}, Processing folder {title}")
|
||||
for child in folder["children"]:
|
||||
if "folder_id" in child:
|
||||
self.get_thread_ids_by_folder_id(
|
||||
child["folder_id"], depth + 1, thread_ids
|
||||
)
|
||||
elif "thread_id" in child:
|
||||
thread_ids.append(child["thread_id"])
|
||||
|
||||
def process_threads(
|
||||
self, thread_ids: Sequence[str], include_images: bool, include_messages: bool
|
||||
) -> List[Document]:
|
||||
"""Process a list of thread into a list of documents."""
|
||||
docs = []
|
||||
for thread_id in thread_ids:
|
||||
doc = self.process_thread(thread_id, include_images, include_messages)
|
||||
if doc is not None:
|
||||
docs.append(doc)
|
||||
return docs
|
||||
|
||||
def process_thread(
|
||||
self, thread_id: str, include_images: bool, include_messages: bool
|
||||
) -> Optional[Document]:
|
||||
thread = self.quip_client.get_thread(thread_id)
|
||||
thread_id = thread["thread"]["id"]
|
||||
title = thread["thread"]["title"]
|
||||
link = thread["thread"]["link"]
|
||||
update_ts = thread["thread"]["updated_usec"]
|
||||
sanitized_title = QuipLoader._sanitize_title(title)
|
||||
|
||||
logger.info(
|
||||
f"processing thread {thread_id} title {sanitized_title} "
|
||||
f"link {link} update_ts {update_ts}"
|
||||
)
|
||||
|
||||
if "html" in thread:
|
||||
# Parse the document
|
||||
try:
|
||||
tree = self.quip_client.parse_document_html(thread["html"])
|
||||
except xml.etree.cElementTree.ParseError as e:
|
||||
logger.error(f"Error parsing thread {title} {thread_id}, skipping, {e}")
|
||||
return None
|
||||
|
||||
metadata = {
|
||||
"title": sanitized_title,
|
||||
"update_ts": update_ts,
|
||||
"id": thread_id,
|
||||
"source": link,
|
||||
}
|
||||
|
||||
# Download each image and replace with the new URL
|
||||
text = ""
|
||||
if include_images:
|
||||
text = self.process_thread_images(tree)
|
||||
|
||||
if include_messages:
|
||||
text = text + "/n" + self.process_thread_messages(thread_id)
|
||||
|
||||
return Document(
|
||||
page_content=thread["html"] + text,
|
||||
metadata=metadata,
|
||||
)
|
||||
return None
|
||||
|
||||
def process_thread_images(self, tree: ElementTree) -> str:
|
||||
text = ""
|
||||
|
||||
try:
|
||||
from PIL import Image
|
||||
from pytesseract import pytesseract
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"`Pillow or pytesseract` package not found, "
|
||||
"please run "
|
||||
"`pip install Pillow` or `pip install pytesseract`"
|
||||
)
|
||||
|
||||
for img in tree.iter("img"):
|
||||
src = img.get("src")
|
||||
if not src or not src.startswith("/blob"):
|
||||
continue
|
||||
_, _, thread_id, blob_id = src.split("/")
|
||||
blob_response = self.quip_client.get_blob(thread_id, blob_id)
|
||||
try:
|
||||
image = Image.open(BytesIO(blob_response.read()))
|
||||
text = text + "\n" + pytesseract.image_to_string(image)
|
||||
except OSError as e:
|
||||
logger.error(f"failed to convert image to text, {e}")
|
||||
raise e
|
||||
return text
|
||||
|
||||
def process_thread_messages(self, thread_id: str) -> str:
|
||||
max_created_usec = None
|
||||
messages = []
|
||||
while True:
|
||||
chunk = self.quip_client.get_messages(
|
||||
thread_id, max_created_usec=max_created_usec, count=100
|
||||
)
|
||||
messages.extend(chunk)
|
||||
if chunk:
|
||||
max_created_usec = chunk[-1]["created_usec"] - 1
|
||||
else:
|
||||
break
|
||||
messages.reverse()
|
||||
|
||||
texts = [message["text"] for message in messages]
|
||||
|
||||
return "\n".join(texts)
|
||||
|
||||
@staticmethod
|
||||
def _sanitize_title(title: str) -> str:
|
||||
sanitized_title = re.sub(r"\s", " ", title)
|
||||
sanitized_title = re.sub(r"(?u)[^- \w.]", "", sanitized_title)
|
||||
if len(sanitized_title) > _MAXIMUM_TITLE_LENGTH:
|
||||
sanitized_title = sanitized_title[:_MAXIMUM_TITLE_LENGTH]
|
||||
return sanitized_title
|
||||
@ -0,0 +1,179 @@
|
||||
from typing import Dict
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from langchain.document_loaders.quip import QuipLoader
|
||||
from langchain.schema import Document
|
||||
|
||||
try:
|
||||
from quip_api.quip import QuipClient # noqa: F401
|
||||
|
||||
quip_installed = True
|
||||
except ImportError:
|
||||
quip_installed = False
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_quip(): # type: ignore
|
||||
# mock quip_client
|
||||
with patch("quip_api.quip.QuipClient") as mock_quip:
|
||||
yield mock_quip
|
||||
|
||||
|
||||
@pytest.mark.requires("quip_api")
|
||||
class TestQuipLoader:
|
||||
API_URL = "https://example-api.quip.com"
|
||||
DOC_URL_PREFIX = ("https://example.quip.com",)
|
||||
ACCESS_TOKEN = "api_token"
|
||||
|
||||
MOCK_FOLDER_IDS = ["ABC"]
|
||||
MOCK_THREAD_IDS = ["ABC", "DEF"]
|
||||
|
||||
def test_quip_loader_initialization(self, mock_quip: MagicMock) -> None:
|
||||
QuipLoader(self.API_URL, access_token=self.ACCESS_TOKEN, request_timeout=60)
|
||||
mock_quip.assert_called_once_with(
|
||||
access_token=self.ACCESS_TOKEN, base_url=self.API_URL, request_timeout=60
|
||||
)
|
||||
|
||||
def test_quip_loader_load_date_invalid_args(self) -> None:
|
||||
quip_loader = QuipLoader(
|
||||
self.API_URL, access_token=self.ACCESS_TOKEN, request_timeout=60
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Must specify at least one among `folder_ids`, `thread_ids` or "
|
||||
"set `include_all`_folders as True",
|
||||
):
|
||||
quip_loader.load()
|
||||
|
||||
def test_quip_loader_load_data_by_folder_id(self, mock_quip: MagicMock) -> None:
|
||||
mock_quip.get_folder.side_effect = [
|
||||
self._get_mock_folder(self.MOCK_FOLDER_IDS[0])
|
||||
]
|
||||
mock_quip.get_thread.side_effect = [
|
||||
self._get_mock_thread(self.MOCK_THREAD_IDS[0]),
|
||||
self._get_mock_thread(self.MOCK_THREAD_IDS[1]),
|
||||
]
|
||||
|
||||
quip_loader = self._get_mock_quip_loader(mock_quip)
|
||||
documents = quip_loader.load(folder_ids=[self.MOCK_FOLDER_IDS[0]])
|
||||
assert mock_quip.get_folder.call_count == 1
|
||||
assert mock_quip.get_thread.call_count == 2
|
||||
assert len(documents) == 2
|
||||
assert all(isinstance(doc, Document) for doc in documents)
|
||||
assert (
|
||||
documents[0].metadata.get("source")
|
||||
== f"https://example.quip.com/{self.MOCK_THREAD_IDS[0]}"
|
||||
)
|
||||
assert (
|
||||
documents[1].metadata.get("source")
|
||||
== f"https://example.quip.com/{self.MOCK_THREAD_IDS[1]}"
|
||||
)
|
||||
|
||||
def test_quip_loader_load_data_all_folder(self, mock_quip: MagicMock) -> None:
|
||||
mock_quip.get_authenticated_user.side_effect = [
|
||||
self._get_mock_authenticated_user()
|
||||
]
|
||||
|
||||
mock_quip.get_folder.side_effect = [
|
||||
self._get_mock_folder(self.MOCK_FOLDER_IDS[0]),
|
||||
]
|
||||
|
||||
mock_quip.get_thread.side_effect = [
|
||||
self._get_mock_thread(self.MOCK_THREAD_IDS[0]),
|
||||
self._get_mock_thread(self.MOCK_THREAD_IDS[1]),
|
||||
]
|
||||
|
||||
quip_loader = self._get_mock_quip_loader(mock_quip)
|
||||
documents = quip_loader.load(include_all_folders=True)
|
||||
assert mock_quip.get_folder.call_count == 1
|
||||
assert mock_quip.get_thread.call_count == 2
|
||||
assert len(documents) == 2
|
||||
assert all(isinstance(doc, Document) for doc in documents)
|
||||
assert (
|
||||
documents[0].metadata.get("source")
|
||||
== f"https://example.quip.com/{self.MOCK_THREAD_IDS[0]}"
|
||||
)
|
||||
assert (
|
||||
documents[1].metadata.get("source")
|
||||
== f"https://example.quip.com/{self.MOCK_THREAD_IDS[1]}"
|
||||
)
|
||||
|
||||
def test_quip_loader_load_data_by_thread_id(self, mock_quip: MagicMock) -> None:
|
||||
mock_quip.get_thread.side_effect = [
|
||||
self._get_mock_thread(self.MOCK_THREAD_IDS[0]),
|
||||
self._get_mock_thread(self.MOCK_THREAD_IDS[1]),
|
||||
]
|
||||
|
||||
quip_loader = self._get_mock_quip_loader(mock_quip)
|
||||
documents = quip_loader.load(thread_ids=self.MOCK_THREAD_IDS)
|
||||
|
||||
assert mock_quip.get_folder.call_count == 0
|
||||
assert mock_quip.get_thread.call_count == 2
|
||||
assert len(documents) == 2
|
||||
assert all(isinstance(doc, Document) for doc in documents)
|
||||
assert (
|
||||
documents[0].metadata.get("source")
|
||||
== f"https://example.quip.com/{self.MOCK_THREAD_IDS[0]}"
|
||||
)
|
||||
assert (
|
||||
documents[1].metadata.get("source")
|
||||
== f"https://example.quip.com/{self.MOCK_THREAD_IDS[1]}"
|
||||
)
|
||||
|
||||
def _get_mock_quip_loader(self, mock_quip: MagicMock) -> QuipLoader:
|
||||
quip_loader = QuipLoader(
|
||||
self.API_URL, access_token=self.ACCESS_TOKEN, request_timeout=60
|
||||
)
|
||||
quip_loader.quip_client = mock_quip
|
||||
return quip_loader
|
||||
|
||||
def _get_mock_folder(self, folder_id: str) -> Dict:
|
||||
return {
|
||||
"folder": {
|
||||
"title": "runbook",
|
||||
"creator_id": "testing",
|
||||
"folder_type": "shared",
|
||||
"parent_id": "ABCD",
|
||||
"inherit_mode": "inherit",
|
||||
"color": "manila",
|
||||
"id": f"{folder_id}",
|
||||
"created_usec": 1668405728528904,
|
||||
"updated_usec": 1697356632672453,
|
||||
"link": "https://example.quip.com/YPH9OAR2Eu5",
|
||||
},
|
||||
"member_ids": [],
|
||||
"children": [
|
||||
{"thread_id": "ABC"},
|
||||
{"thread_id": "DEF"},
|
||||
],
|
||||
}
|
||||
|
||||
def _get_mock_thread(self, thread_id: str) -> Dict:
|
||||
return {
|
||||
"thread": {
|
||||
"author_id": "testing",
|
||||
"thread_class": "document",
|
||||
"owning_company_id": "ABC",
|
||||
"id": f"{thread_id}",
|
||||
"created_usec": 1690873126670055,
|
||||
"updated_usec": 1690874891638991,
|
||||
"title": f"Unit Test Doc {thread_id}",
|
||||
"link": f"https://example.quip.com/{thread_id}",
|
||||
"document_id": "ABC",
|
||||
"type": "document",
|
||||
"is_template": False,
|
||||
"is_deleted": False,
|
||||
},
|
||||
"user_ids": [],
|
||||
"shared_folder_ids": ["ABC"],
|
||||
"expanded_user_ids": ["ABCDEFG"],
|
||||
"invited_user_emails": [],
|
||||
"access_levels": {"ABCD": {"access_level": "OWN"}},
|
||||
"html": "<h1 id='temp:C:ABCD'>How to write Python Test </h1>",
|
||||
}
|
||||
|
||||
def _get_mock_authenticated_user(self) -> Dict:
|
||||
return {"shared_folder_ids": self.MOCK_FOLDER_IDS, "id": "Test"}
|
||||
Loading…
Reference in New Issue