langchain/libs/community/langchain_community/document_loaders/nuclia.py
Erick Friis 3a2eb6e12b
infra: add print rule to ruff (#16221)
Added noqa for existing prints. Can slowly remove / will prevent more
being intro'd
2024-02-09 16:13:30 -08:00

34 lines
1.1 KiB
Python

import json
import uuid
from typing import List
from langchain_core.documents import Document
from langchain_community.document_loaders.base import BaseLoader
from langchain_community.tools.nuclia.tool import NucliaUnderstandingAPI
class NucliaLoader(BaseLoader):
"""Load from any file type using `Nuclia Understanding API`."""
def __init__(self, path: str, nuclia_tool: NucliaUnderstandingAPI):
self.nua = nuclia_tool
self.id = str(uuid.uuid4())
self.nua.run({"action": "push", "id": self.id, "path": path, "text": None})
def load(self) -> List[Document]:
"""Load documents."""
data = self.nua.run(
{"action": "pull", "id": self.id, "path": None, "text": None}
)
if not data:
return []
obj = json.loads(data)
text = obj["extracted_text"][0]["body"]["text"]
print(text) # noqa: T201
metadata = {
"file": obj["file_extracted_data"][0],
"metadata": obj["field_metadata"][0],
}
return [Document(page_content=text, metadata=metadata)]