Embedchain json url support (#878)

Co-authored-by: Deven Patel <deven298@yahoo.com>
This commit is contained in:
Deven Patel
2023-10-30 16:19:11 -07:00
committed by GitHub
parent 68dc274f72
commit 5255a37c93
2 changed files with 84 additions and 8 deletions

View File

@@ -1,9 +1,14 @@
import hashlib
import json
import os
import re
import requests
from embedchain.loaders.base_loader import BaseLoader
VALID_URL_PATTERN = "^https:\/\/[0-9A-z.]+.[0-9A-z.]+.[a-z]+\/.*\.json$"
class JSONLoader(BaseLoader):
@staticmethod
@@ -20,18 +25,32 @@ class JSONLoader(BaseLoader):
loader = LLHBUBJSONLoader()
if not isinstance(content, str) and not os.path.isfile(content):
if not isinstance(content, str):
print(f"Invaid content input. Provide the correct path to the json file saved locally in {content}")
data = []
data_content = []
with open(content, "r") as json_file:
json_data = json.load(json_file)
docs = loader.load_data(json_data)
for doc in docs:
doc_content = doc.text
data.append({"content": doc_content, "meta_data": {"url": content}})
data_content.append(doc_content)
# Load json data from various sources. TODO: add support for dictionary
if os.path.isfile(content):
with open(content, "r") as json_file:
json_data = json.load(json_file)
elif re.match(VALID_URL_PATTERN, content):
response = requests.get(content)
if response.status_code == 200:
json_data = response.json()
else:
raise ValueError(
f"Loading data from the given url: {content} failed. \
Make sure the url is working."
)
else:
raise ValueError(f"Invalid content to load json data from: {content}")
docs = loader.load_data(json_data)
for doc in docs:
doc_content = doc.text
data.append({"content": doc_content, "meta_data": {"url": content}})
data_content.append(doc_content)
doc_id = hashlib.sha256((content + ", ".join(data_content)).encode()).hexdigest()
return {"doc_id": doc_id, "data": data}

View File

@@ -1,5 +1,8 @@
import hashlib
import pytest
from llama_index.readers.schema.base import Document
from embedchain.loaders.json import JSONLoader
@@ -32,3 +35,57 @@ def test_load_data(mocker):
expected_doc_id = hashlib.sha256((content + ", ".join(["content1", "content2"])).encode()).hexdigest()
assert result["doc_id"] == expected_doc_id
def test_load_data_url(mocker):
content = "https://example.com/posts.json"
mocker.patch("os.path.isfile", return_value=False) # Mocking os.path.isfile to simulate a URL case
mocker.patch(
"llama_hub.jsondata.base.JSONDataReader.load_data",
return_value=[Document(text="content1"), Document(text="content2")],
)
mock_response = mocker.Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"document1": "content1", "document2": "content2"}
mocker.patch("requests.get", return_value=mock_response)
result = JSONLoader.load_data(content)
assert "doc_id" in result
assert "data" in result
expected_data = [
{"content": "content1", "meta_data": {"url": content}},
{"content": "content2", "meta_data": {"url": content}},
]
assert result["data"] == expected_data
expected_doc_id = hashlib.sha256((content + ", ".join(["content1", "content2"])).encode()).hexdigest()
assert result["doc_id"] == expected_doc_id
def test_load_data_invalid_content(mocker):
mocker.patch("os.path.isfile", return_value=False)
mocker.patch("requests.get")
content = "123"
with pytest.raises(ValueError, match="Invalid content to load json data from"):
JSONLoader.load_data(content)
def test_load_data_invalid_url(mocker):
mocker.patch("os.path.isfile", return_value=False)
mock_response = mocker.Mock()
mock_response.status_code = 404
mocker.patch("requests.get", return_value=mock_response)
content = "http://invalid-url.com/"
with pytest.raises(ValueError, match=f"Invalid content to load json data from: {content}"):
JSONLoader.load_data(content)