mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-13 21:47:12 +00:00
Document loader for Cube Semantic Layer (#6882)
### Description This pull request introduces the "Cube Semantic Layer" document loader, which demonstrates the retrieval of Cube's data model metadata in a format suitable for passing to LLMs as embeddings. This enhancement aims to provide contextual information and improve the understanding of data. Twitter handle: @the_cube_dev --------- Co-authored-by: rlm <pexpresss31@gmail.com>
This commit is contained in:
86
tests/unit_tests/document_loaders/test_cube_semantic.py
Normal file
86
tests/unit_tests/document_loaders/test_cube_semantic.py
Normal file
@@ -0,0 +1,86 @@
|
||||
from typing import List
|
||||
from unittest import TestCase
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import requests
|
||||
|
||||
from langchain.docstore.document import Document
|
||||
from langchain.document_loaders import CubeSemanticLoader
|
||||
|
||||
|
||||
class TestCubeSemanticLoader(TestCase):
|
||||
@patch.object(requests, "get")
|
||||
def test_load_success(self, mock_get: MagicMock) -> None:
|
||||
# Arrange
|
||||
cube_api_url: str = "https://example.com/cube_api"
|
||||
cube_api_token: str = "abc123"
|
||||
mock_response: MagicMock = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response_json: dict = {
|
||||
"cubes": [
|
||||
{
|
||||
"type": "view",
|
||||
"name": "cube1",
|
||||
"measures": [{"type": "sum", "name": "sales", "title": "Sales"}],
|
||||
"dimensions": [
|
||||
{
|
||||
"type": "string",
|
||||
"name": "product_name",
|
||||
"title": "Product Name",
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_response.json.return_value = mock_response_json
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
expected_docs: List[Document] = [
|
||||
Document(
|
||||
page_content=(
|
||||
"table name: cube1, "
|
||||
"column name: sales, "
|
||||
"column data type: sum, "
|
||||
"column title: Sales, "
|
||||
"column description: None"
|
||||
),
|
||||
metadata={
|
||||
"table_name": "cube1",
|
||||
"column_name": "sales",
|
||||
"column_data_type": "sum",
|
||||
"column_title": "Sales",
|
||||
"column_description": "None",
|
||||
},
|
||||
),
|
||||
Document(
|
||||
page_content=(
|
||||
"table name: cube1, "
|
||||
"column name: product_name, "
|
||||
"column data type: string, "
|
||||
"column title: Product Name, "
|
||||
"column description: None"
|
||||
),
|
||||
metadata={
|
||||
"table_name": "cube1",
|
||||
"column_name": "product_name",
|
||||
"column_data_type": "string",
|
||||
"column_title": "Product Name",
|
||||
"column_description": "None",
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
loader: CubeSemanticLoader = CubeSemanticLoader(cube_api_url, cube_api_token)
|
||||
|
||||
# Act
|
||||
result: List[Document] = loader.load()
|
||||
|
||||
# Assert
|
||||
self.assertEqual(result, expected_docs)
|
||||
mock_get.assert_called_once_with(
|
||||
cube_api_url,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": cube_api_token,
|
||||
},
|
||||
)
|
Reference in New Issue
Block a user