Merge pull request #26 from weni-ai/feature/links

load and index urls

Dockerfile

@@ -9,6 +9,10 @@ COPY pyproject.toml poetry.lock ./
 RUN poetry config virtualenvs.create false && \
  poetry install --no-dev
 
+RUN apt update && apt install libmagic1 -y
+RUN poetry add python-magic
+RUN python -m nltk.downloader punkt averaged_perceptron_tagger -d /usr/share/nltk_data
+
 COPY . .
 
 EXPOSE 8000

app/handlers/nexus.py

@@ -16,10 +16,18 @@ def __init__(self) -> None:
 
  def index_succedded(self, task_succeded: bool, nexus_task_uuid: str, file_type: str) -> None:
  endpoint = f'{self.base_url}/api/v1/content-base-file'
+
+ if file_type == "txt":
+ ftype = "text"
+ elif file_type == "urls":
+ ftype = "link"
+ else:
+ ftype = "file"
+
  data = { 
  "status": int(task_succeded),
  "task_uuid": nexus_task_uuid,
- "file_type": "text" if file_type == "txt" else "file",
+ "file_type": ftype,
  }
  response = requests.patch(url=endpoint, data=json.dumps(data), headers=self.headers)
  response.raise_for_status()

app/loaders/__init__.py

@@ -1,9 +1,24 @@
 
 import os
-from app.loaders.loaders import (DataLoader, txt_loader, pdf_loader, docx_loader, xlsx_loader, DataLoaderCls, PDFLoader, DocxLoader, TxtLoader, XlsxLoader)
+from app.loaders.loaders import (
+ DataLoader,
+ txt_loader,
+ pdf_loader,
+ docx_loader,
+ xlsx_loader,
+)
+from app.loaders.loaders import (
+ DataLoaderCls,
+ PDFLoader,
+ DocxLoader,
+ TxtLoader,
+ XlsxLoader,
+ URLsLoader,
+)
 from langchain.schema.document import Document
 from typing import List
 from app.text_splitters import ITextSplitter
+
 supported_loaders = {
  'txt': txt_loader,
  'pdf': pdf_loader,
@@ -20,9 +35,9 @@
  'txt': TxtLoader,
  'xlsx': XlsxLoader,
  'xls': XlsxLoader,
+ 'urls': URLsLoader,
 }
 
-
 def load_file_and_get_raw_text(file_name: str, file_type: str) -> str:
  file_path = f'{os.environ.get("FILE_PATH")}/{file_name}'
  loader = supported_loaders.get(file_type)
@@ -47,6 +62,5 @@ def load_file_url_and_split_text(file_url: str, file_type: str, text_splitter: I
  load_type = kwargs.get("load_type", None)
 
  loader = supported_loaders_cls.get(file_type)
-
  data_loader = DataLoaderCls(loader=loader, file=file_url, load_type=load_type)
  return data_loader.load_and_split_text(text_splitter)

app/loaders/loaders.py

@@ -5,11 +5,10 @@
 
 from langchain.document_loaders import (
  TextLoader, PyPDFLoader, UnstructuredExcelLoader,
- UnstructuredWordDocumentLoader, Docx2txtLoader,
- PDFMinerLoader
+ UnstructuredWordDocumentLoader, Docx2txtLoader, UnstructuredURLLoader, PDFMinerLoader
 )
 from langchain.schema.document import Document
-from typing import Callable, List
+from typing import Callable, List, Union
 from app.text_splitters import ITextSplitter
 
 
@@ -65,7 +64,7 @@ def txt_loader(file: str) -> Callable:
 
 class TxtLoader(DocumentLoader):
  def _get_file(self, file: str):
- if os.environ.get("AWS_STORAGE_BUCKET_NAME") in file:
+ if os.environ.get("AWS_STORAGE_BUCKET_NAME") in file: # pragma: no cover
  response = requests.get(file)
  if response.status_code == 200:
  file_path = f"/tmp/{uuid.uuid4()}.txt"
@@ -189,7 +188,7 @@ def __init__(self, file:str) -> None:
  tmp_file, _ = self._get_temp_file(file)
  self.loader = UnstructuredExcelLoader(tmp_file, mode="single")
 
- def _get_temp_file(self, file_url: str):
+ def _get_temp_file(self, file_url: str): # pragma: no cover
  result = urlparse(file_url)
  filename = result.path.strip("/")
  file_path, message = urlretrieve(file_url, f"/tmp/{filename}")
@@ -211,3 +210,31 @@ def load_and_split_text(self, text_splitter: ITextSplitter) -> List[Document]:
  for chunk in text_chunks:
  split_pages.append(Document(page_content=chunk, metadata=metadatas))
  return split_pages
+
+
+class URLsLoader(DocumentLoader):
+ def _urls(self, urls: Union[List[str], str]):
+ if isinstance(urls, str):
+ return [urls]
+ return urls
+
+ def __init__(self, urls: Union[List[str], str]) -> None:
+ self.urls = self._urls(urls)
+ self.loader = UnstructuredURLLoader(urls=self.urls)
+
+ def load(self) -> List[Document]:
+ return self.loader.load()
+
+ def load_and_split_text(self, text_splitter: ITextSplitter) -> List[Document]:
+ split_pages = []
+
+ pages = self.loader.load_and_split()
+ for page in pages:
+ page_content = page.page_content.lower()
+ metadatas = page.metadata
+ metadatas.update({"full_page": page_content})
+
+ text_chunks = text_splitter.split_text(page_content)
+ for chunk in text_chunks:
+ split_pages.append(Document(page_content=chunk, metadata=metadatas))
+ return split_pages

app/tests/test_document_loader.py

@@ -5,6 +5,8 @@
  PDFLoader,
  DocxLoader,
  TxtLoader,
+ URLsLoader,
+ XlsxLoader,
  pdf_loader,
  txt_loader,
  docx_loader,
@@ -122,18 +124,43 @@ def test_load_xlsx(self):
  raw_text = data_loader.raw_text()
  self.assertEqual(type(raw_text), str)
 
+ @mock.patch("app.loaders.loaders.XlsxLoader._get_temp_file")
+ def test_load_xlsx_cls(self, mock_file_url):
+ file_path = f'{self.path}/{self.file_name}.xlsx'
+ mock_file_url.return_value = (file_path, "")
+ xlsx_loader = XlsxLoader(file_path)
+ split_pages: List[Document] = xlsx_loader.load_and_split_text(self.text_splitter)
+ self.assertEqual(list, type(split_pages))
+
  def test_pdf_loader_cls(self):
  file_path = f'{self.path}/{self.file_name}.pdf'
  pdf_loader = PDFLoader(file_path)
  split_pages: List[Document] = pdf_loader.load_and_split_text(self.text_splitter)
  self.assertEqual(list, type(split_pages))
+
+ def test_urls_loader_cls(self):
+ urls_loader = URLsLoader("https://en.wikipedia.org/wiki/Unit_testing")
+ split_pages: List[Document] = urls_loader.load()
+ self.assertEqual(list, type(split_pages))
+
+ def test_urls_loader_and_split_cls(self):
+ urls_loader = URLsLoader("https://en.wikipedia.org/wiki/Unit_testing")
+ split_pages: List[Document] = urls_loader.load_and_split_text(self.text_splitter)
+ self.assertEqual(list, type(split_pages))
+
+ def test_urls_list_loader_and_split_cls(self):
+ urls = ["https://en.wikipedia.org/wiki/Unit_testing"]
+ urls_loader = URLsLoader(urls)
+ split_pages: List[Document] = urls_loader.load_and_split_text(self.text_splitter)
+ self.assertEqual(list, type(split_pages))
 
  def test_docx_loader_cls(self):
  file_path = f'{self.path}/{self.file_name}.docx'
  docx_loader = DocxLoader(file_path)
  split_pages: List[Document] = docx_loader.load_and_split_text(self.text_splitter)
  self.assertEqual(list, type(split_pages))
 
+ @mock.patch.dict(os.environ, {"AWS_STORAGE_BUCKET_NAME": "file-path"})
  def test_txt_loader_cls(self):
  file_path = f'{self.path}/{self.file_name}.txt'
  docx_loader = TxtLoader(file_path)

app/util.py

@@ -7,11 +7,11 @@ class ContentHandler(EmbeddingsContentHandler):
  content_type = "application/json"
  accepts = "application/json"
 
- def transform_input(self, inputs: list[str], model_kwargs: dict) -> bytes:
+ def transform_input(self, inputs: list[str], model_kwargs: dict) -> bytes: # pragma: no cover
  input_str = json.dumps({"inputs": inputs, **model_kwargs})
  return input_str.encode("utf-8")
 
- def transform_output(self, output: bytes) -> list[list[float]]:
+ def transform_output(self, output: bytes) -> list[list[float]]: # pragma: no cover
  response_json = json.loads(output.read().decode("utf-8"))
  return response_json["vectors"]
 

pyproject.toml

@@ -27,6 +27,7 @@ pydantic = "2.3.0"
 celery = "^5.3.6"
 redis = "^5.0.1"
 sentry-sdk = {extras = ["fastapi"], version = "^1.35.0"}
+xlrd = "^2.0.1"
 pdfminer-six = "^20231228"