groundx 2.0.15__py3-none-any.whl → 2.7.7__py3-none-any.whl

groundx/extract/classes/field.py

@@ -0,0 +1,88 @@
+import dateparser, typing
+
+from pydantic import BaseModel
+
+
+class ExtractedField(BaseModel):
+    confidence: typing.Optional[str] = None
+    conflicts: typing.List[typing.Any] = []
+    key: str
+
+    value: typing.Union[str, float, typing.List[typing.Any]] = ""
+
+    def __init__(
+        self,
+        value: typing.Union[str, float, typing.List[typing.Any]],
+        **data: typing.Any,
+    ):
+        super().__init__(**data)
+
+        self.set_value(value)
+
+    def contains(self, other: "ExtractedField") -> bool:
+        self_val = self.get_value()
+        other_val = other.get_value()
+        if not (isinstance(self_val, (str, float, int))):
+            raise Exception(f"unexpected self field value type [{type(self_val)}]")
+
+        if self.equal_to_value(other_val):
+            return True
+
+        if other_val in self.conflicts:
+            return True
+
+        return False
+
+    def equal_to_field(self, other: "ExtractedField") -> bool:
+        self_val = self.get_value()
+        other_val = other.get_value()
+        if not (isinstance(self_val, (str, float, int))):
+            raise Exception(f"unexpected self field value type [{type(self_val)}]")
+
+        return self.equal_to_value(other_val)
+
+    def equal_to_value(self, other: typing.Any) -> bool:
+        if not (isinstance(other, (str, float, int))):
+            raise Exception(f"unexpected value type [{type(other)}]")
+
+        exist = self.get_value()
+        if isinstance(exist, int):
+            exist = float(exist)
+        if isinstance(other, int):
+            other = float(other)
+        if isinstance(exist, str):
+            exist = exist.lower()
+        if isinstance(other, str):
+            other = other.lower()
+
+        return type(other) == type(exist) and other == exist
+
+    def get_value(self) -> typing.Union[str, float, typing.List[typing.Any]]:
+        return self.value
+
+    def remove_conflict(self, value: typing.Any) -> None:
+        if value in self.conflicts:
+            self.conflicts.remove(value)
+        if not self.equal_to_value(value):
+            self.conflicts.append(self.get_value())
+
+    def set_value(
+        self, value: typing.Union[str, float, typing.List[typing.Any]]
+    ) -> None:
+        if isinstance(value, int):
+            self.value = float(value)
+        elif isinstance(value, str) and "date" in self.key.lower():
+            try:
+                dt = dateparser.parse(value)
+                if dt is None:
+                    self.value = value
+                else:
+                    self.value = dt.strftime("%Y-%m-%d")
+            except Exception as e:
+                print(f"date error [{value}]: [{e}]")
+                self.value = value
+        else:
+            self.value = value
+
+
+ExtractedField.model_rebuild()

groundx/extract/classes/groundx.py

@@ -0,0 +1,147 @@
+import json, requests, typing
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from ..services.upload import Upload
+
+
+class GroundXDocument(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    base_url: str
+    document_id: str = Field(alias="documentID")
+    task_id: str = Field(alias="taskID")
+
+    def xray_path(self) -> str:
+        return f"layout/processed/{self.task_id}/{self.document_id}-xray.json"
+
+    def xray_url(self, base: typing.Optional[str] = None) -> str:
+        if not base:
+            base = self.base_url
+        if base.endswith("/"):
+            base = base[:-1]
+        return f"{base}/layout/processed/{self.task_id}/{self.document_id}-xray.json"
+
+    def xray(
+        self,
+        cache_dir: Path,
+        upload: typing.Optional[Upload] = None,
+        clear_cache: bool = False,
+        is_test: bool = False,
+        base: typing.Optional[str] = None,
+    ) -> "XRayDocument":
+        return XRayDocument.download(
+            self,
+            cache_dir=cache_dir,
+            upload=upload,
+            clear_cache=clear_cache,
+            is_test=is_test,
+            base=base,
+        )
+
+
+class GroundXResponse(BaseModel):
+    code: int
+    document_id: str = Field(alias="documentID")
+    model_id: int = Field(alias="modelID")
+    processor_id: int = Field(alias="processorID")
+    result_url: str = Field(alias="resultURL")
+    task_id: str = Field(alias="taskID")
+
+
+class BoundingBox(BaseModel):
+    bottomRightX: float
+    bottomRightY: float
+    topLeftX: float
+    topLeftY: float
+    corrected: typing.Optional[bool]
+    pageNumber: typing.Optional[int]
+
+
+class Chunk(BaseModel):
+    boundingBoxes: typing.Optional[typing.List[BoundingBox]] = []
+    chunk: typing.Optional[str] = None
+    contentType: typing.Optional[typing.List[str]] = []
+    json_: typing.Optional[typing.List[typing.Any]] = Field(None, alias="json")
+    multimodalUrl: typing.Optional[str] = None
+    narrative: typing.Optional[typing.List[str]] = None
+    pageNumbers: typing.Optional[typing.List[int]] = []
+    sectionSummary: typing.Optional[str] = None
+    suggestedText: typing.Optional[str] = None
+    text: typing.Optional[str] = None
+
+
+class DocumentPage(BaseModel):
+    chunks: typing.List[Chunk]
+    height: float
+    pageNumber: int
+    pageUrl: str
+    width: float
+
+
+class XRayDocument(BaseModel):
+    chunks: typing.List[Chunk]
+    documentPages: typing.List[DocumentPage] = []
+    sourceUrl: str
+    fileKeywords: typing.Optional[str] = None
+    fileName: typing.Optional[str] = None
+    fileType: typing.Optional[str] = None
+    fileSummary: typing.Optional[str] = None
+    language: typing.Optional[str] = None
+
+    @classmethod
+    def download(
+        cls,
+        gx_doc: GroundXDocument,
+        cache_dir: Path,
+        upload: typing.Optional[Upload] = None,
+        clear_cache: bool = False,
+        is_test: bool = False,
+        base: typing.Optional[str] = None,
+    ) -> "XRayDocument":
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        cache_file = cache_dir / f"{gx_doc.document_id}-xray.json"
+
+        if not clear_cache and cache_file.exists():
+            try:
+                with cache_file.open("r", encoding="utf-8") as f:
+                    payload = json.load(f)
+
+                return cls(**payload)
+            except Exception as e:
+                raise RuntimeError(
+                    f"Error loading cached X-ray JSON from {cache_file}: {e}"
+                )
+
+        if upload:
+            path = gx_doc.xray_path()
+            ru = upload.get_object(path)
+            if ru:
+                try:
+                    payload = json.loads(ru.decode("utf-8"))
+                    return cls(**payload)
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Error decoding X-ray JSON bytes from {path}: {e}"
+                    )
+
+        url = gx_doc.xray_url(base=base)
+        try:
+            resp = requests.get(url)
+            resp.raise_for_status()
+        except requests.RequestException as e:
+            raise RuntimeError(f"Error fetching X-ray JSON from {url}: {e}")
+
+        try:
+            payload = resp.json()
+        except ValueError as e:
+            raise RuntimeError(f"Invalid JSON returned from {url}: {e}")
+
+        if is_test is False:
+            try:
+                with cache_file.open("w", encoding="utf-8") as f:
+                    json.dump(payload, f)
+            except Exception as e:
+                print(f"Warning: failed to write X-ray JSON cache to {cache_file}: {e}")
+
+        return cls(**payload)

groundx/extract/classes/prompt.py

@@ -0,0 +1,36 @@
+import typing
+
+from pydantic import BaseModel
+
+from ..utility.classes import str_to_type_sequence
+
+
+class Prompt(BaseModel):
+    attr_name: str
+    prompt: str
+    type: typing.Union[str, typing.List[str]]
+
+    class Config:
+        validate_by_name = True
+
+    def valid_value(self, value: typing.Any) -> bool:
+        ty = self.type
+
+        types: typing.List[typing.Type[typing.Any]] = []
+        if isinstance(ty, list):
+            for t in ty:
+                if t == "int" or t == "float":
+                    types.extend([int, float])
+                elif t == "str":
+                    types.append(str)
+
+            return isinstance(value, tuple(types))
+
+        exp = str_to_type_sequence(ty)
+        for et in exp:
+            if et in (int, float):
+                types.extend([int, float])
+            else:
+                types.append(et)
+        types = list(dict.fromkeys(types))
+        return isinstance(value, tuple(types))

groundx/extract/classes/test_document.py

@@ -0,0 +1,109 @@
+import pytest, typing, unittest
+
+pytest.importorskip("PIL")
+
+from io import BytesIO
+from pathlib import Path
+from PIL import Image
+from unittest.mock import patch
+
+from .document import Document, DocumentRequest
+from .test_groundx import TestXRay
+
+
+def DR(**data: typing.Any) -> DocumentRequest:
+    return DocumentRequest.model_validate(data)
+
+
+def test_doc() -> Document:
+    return Document.from_request(
+        cache_dir=Path("./cache"),
+        base_url="",
+        req=test_request(),
+    )
+
+
+def test_request() -> DocumentRequest:
+    return DR(documentID="D", fileName="F", modelID=1, processorID=1, taskID="T")
+
+
+class TestDocument(unittest.TestCase):
+    def setUp(self) -> None:
+        patcher = patch(
+            "groundx.extract.classes.document.GroundXDocument.xray", autospec=True
+        )
+        self.mock_xray = patcher.start()
+        self.addCleanup(patcher.stop)
+        self.mock_xray.return_value = TestXRay("http://test.co", [])
+
+    def test_init_name(self) -> None:
+        st1: Document = test_doc()
+        self.assertEqual(st1.file_name, "F")
+        st2: Document = Document.from_request(
+            cache_dir=Path("./cache"),
+            base_url="",
+            req=DR(
+                documentID="D", fileName="F.pdf", modelID=1, processorID=1, taskID="T"
+            ),
+        )
+        self.assertEqual(st2.file_name, "F.pdf")
+        st3: Document = Document.from_request(
+            cache_dir=Path("./cache"),
+            base_url="",
+            req=DR(documentID="D", fileName="F.", modelID=1, processorID=1, taskID="T"),
+        )
+        self.assertEqual(st3.file_name, "F.")
+
+
+class TestDocumentRequest(unittest.TestCase):
+    def test_load_images_cached(self) -> None:
+        urls: typing.List[str] = [
+            "http://example.com/page1.png",
+            "http://example.com/page2.png",
+        ]
+
+        red_img = Image.new("RGB", (10, 10), color="red")
+        buf = BytesIO()
+        red_img.save(buf, format="PNG")
+
+        st = test_request()
+        st.page_images = [red_img, red_img]
+        st.page_image_dict = {
+            urls[0]: 0,
+            urls[1]: 1,
+        }
+        st.load_images(urls)
+        self.assertEqual(len(st.page_images), 2)
+        self.assertEqual(len(st.page_image_dict), 2)
+
+    def test_load_images_download(self) -> None:
+        urls = ["http://example.com/page1.png", "http://example.com/page2.png"]
+
+        red_img = Image.new("RGB", (10, 10), color="red")
+        buf = BytesIO()
+        red_img.save(buf, format="PNG")
+        img_bytes = buf.getvalue()
+
+        class TestResp:
+            content = img_bytes
+
+            def raise_for_status(self) -> None:
+                pass
+
+        with patch("requests.get", return_value=TestResp()):
+            st = test_request()
+            st.load_images(urls)
+
+            self.assertEqual(len(st.page_images), 2)
+            self.assertEqual(len(st.page_image_dict), 2)
+            for img in st.page_images:
+                self.assertIsInstance(img, Image.Image)
+                self.assertEqual(img.size, (10, 10))
+
+    def test_load_images_error(self) -> None:
+        urls = ["http://example.com/page1.png", "http://example.com/page2.png"]
+
+        st = test_request()
+        st.load_images(urls, should_sleep=False)
+        self.assertEqual(len(st.page_images), 0)
+        self.assertEqual(len(st.page_image_dict), 0)

groundx/extract/classes/test_field.py

@@ -0,0 +1,43 @@
+import pytest, typing, unittest
+
+pytest.importorskip("dateparser")
+
+from .field import ExtractedField
+
+
+def TestField(
+    name: str,
+    value: typing.Union[str, float, typing.List[typing.Any]],
+    conflicts: typing.List[typing.Any] = [],
+) -> ExtractedField:
+    return ExtractedField(
+        key=name.replace("_", " "),
+        value=value,
+        conflicts=conflicts,
+    )
+
+
+class TestExtractedField(unittest.TestCase):
+    def test_equalToValue_string(self):
+        ef = TestField("test", "hello")
+        self.assertTrue(ef.equal_to_value("hello"))
+        self.assertFalse(ef.equal_to_value("world"))
+
+    def test_equalToValue_int_float_equivalence(self):
+        ef = TestField("test", int(10))
+        self.assertTrue(ef.equal_to_value(10.0))
+        self.assertTrue(ef.equal_to_value(10))
+
+    def test_equalToValue_mismatch(self):
+        ef = TestField("test", 3.14)
+        self.assertFalse(ef.equal_to_value(2.71))
+
+    def test_set_value_dates(self):
+        ef1 = TestField("test date", "3/29/25")
+        self.assertEqual(ef1.get_value(), "2025-03-29")
+        ef2 = TestField("test date", "2025-03-29")
+        self.assertEqual(ef2.get_value(), "2025-03-29")
+
+
+if __name__ == "__main__":
+    unittest.main()

groundx/extract/classes/test_groundx.py

@@ -0,0 +1,223 @@
+import requests, typing, unittest
+from pathlib import Path
+from unittest.mock import patch
+
+from pydantic import ValidationError
+
+from .groundx import (
+    GroundXDocument,
+    XRayDocument,
+    Chunk,
+    BoundingBox,
+    DocumentPage,
+)
+
+
+class TestChunk:
+    def __init__(self, json_str: str):
+        self.sectionSummary = None
+        self.suggestedText = json_str
+
+
+class TestDocumentPage:
+    def __init__(self, page_url: str):
+        self.pageUrl = page_url
+
+
+class TestXRay:
+    def __init__(
+        self,
+        source_url: str,
+        chunks: typing.Optional[typing.List[TestChunk]] = [],
+        document_pages: typing.Optional[typing.List[str]] = [],
+    ):
+        self.chunks = chunks
+        self.documentPages: typing.List[TestDocumentPage] = []
+        if document_pages is not None:
+            for p in document_pages:
+                self.documentPages.append(TestDocumentPage(p))
+        self.sourceUrl = source_url
+
+
+def GD(**data: typing.Any) -> GroundXDocument:
+    return GroundXDocument.model_validate(data)
+
+
+def test_xray(gx: GroundXDocument) -> XRayDocument:
+    return XRayDocument.download(
+        gx, cache_dir=Path("./cache"), base="https://upload.test", is_test=True
+    )
+
+
+class TestGroundX(unittest.TestCase):
+    def make_dummy_response(
+        self,
+        payload: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        status_ok: bool = True,
+        json_error: bool = False,
+    ) -> typing.Any:
+        class DummyResponse:
+            def raise_for_status(self):
+                if not status_ok:
+                    raise requests.HTTPError("HTTP error!")
+
+            def json(self):
+                if json_error:
+                    raise ValueError("Bad JSON!")
+                return payload
+
+        return DummyResponse()
+
+    def test_xray_url(self):
+        gx = GD(base_url="", documentID="doc123", taskID="taskABC")
+        expected = "https://upload.test/layout/processed/taskABC/doc123-xray.json"
+        self.assertEqual(gx.xray_url(base="https://upload.test"), expected)
+
+    def test_download_success(self):
+        payload: typing.Dict[str, typing.Any] = {
+            "chunks": [],
+            "documentPages": [],
+            "sourceUrl": "https://example.com/foo.pdf",
+        }
+        dummy = self.make_dummy_response(payload=payload, status_ok=True)
+        with patch("requests.get", return_value=dummy):
+            gx = GD(base_url="", documentID="D", taskID="T")
+            xdoc = test_xray(gx)
+            self.assertIsInstance(xdoc, XRayDocument)
+            self.assertEqual(xdoc.chunks, [])
+            self.assertEqual(xdoc.documentPages, [])
+            self.assertEqual(xdoc.sourceUrl, payload["sourceUrl"])
+
+    def test_download_request_exception(self):
+        with patch("requests.get", side_effect=requests.RequestException("no network")):
+            gx = GD(base_url="", documentID="D", taskID="T")
+            with self.assertRaises(RuntimeError) as cm:
+                test_xray(gx)
+            self.assertIn("Error fetching X-ray JSON", str(cm.exception))
+
+    def test_download_http_error(self):
+        dummy = self.make_dummy_response(payload={}, status_ok=False)
+        with patch("requests.get", return_value=dummy):
+            gx = GD(base_url="", documentID="D", taskID="T")
+            with self.assertRaises(RuntimeError) as cm:
+                test_xray(gx)
+            self.assertIn("HTTP error!", str(cm.exception))
+
+    def test_download_json_error(self):
+        dummy = self.make_dummy_response(payload=None, status_ok=True, json_error=True)
+        with patch("requests.get", return_value=dummy):
+            gx = GD(base_url="", documentID="D", taskID="T")
+            with self.assertRaises(RuntimeError) as cm:
+                test_xray(gx)
+            self.assertIn("Invalid JSON returned", str(cm.exception))
+
+    def test_validation_error_on_missing_required_fields(self) -> None:
+        payload: typing.Dict[str, typing.Any] = {
+            "documentPages": [],
+            "sourceUrl": "https://example.com/foo.pdf",
+        }
+        dummy = self.make_dummy_response(payload=payload, status_ok=True)
+        with patch("requests.get", return_value=dummy):
+            gx = GD(base_url="", documentID="D", taskID="T")
+            with self.assertRaises(ValidationError) as cm:
+                test_xray(gx)
+            self.assertIn("Field required", str(cm.exception))
+
+    def test_xray_method_delegates_to_download(self) -> None:
+        gx = GD(base_url="", documentID="X", taskID="Y")
+
+        sentinel = object()
+        with patch.object(XRayDocument, "download", return_value=sentinel):
+            result = gx.xray(
+                cache_dir=Path("./cache"), base="https://upload.test", is_test=True
+            )
+            self.assertIs(result, sentinel)
+
+    def test_chunk_json_alias(self) -> None:
+        raw: typing.Dict[str, typing.Any] = {
+            "boundingBoxes": [],
+            "chunk": "id123",
+            "contentType": [],
+            "json": [{"foo": "bar"}],
+            "multimodalUrl": None,
+            "narrative": None,
+            "pageNumbers": [],
+            "sectionSummary": None,
+            "suggestedText": None,
+            "text": None,
+        }
+        chunk = Chunk.model_validate(raw)
+        self.assertEqual(chunk.json_, [{"foo": "bar"}])
+
+        self.assertNotIn("json':", chunk.model_dump_json().replace('"json"', ""))
+
+    def test_roundtrip_xray_to_models(self):
+        payload: dict[str, typing.Any] = {
+            "chunks": [
+                {
+                    "boundingBoxes": [
+                        {
+                            "bottomRightX": 10.0,
+                            "bottomRightY": 20.0,
+                            "topLeftX": 1.0,
+                            "topLeftY": 2.0,
+                            "corrected": True,
+                            "pageNumber": 1,
+                        }
+                    ],
+                    "chunk": "foo",
+                    "contentType": ["paragraph"],
+                    "json": [{"a": 1}],
+                    "multimodalUrl": None,
+                    "narrative": ["narr1"],
+                    "pageNumbers": [1],
+                    "sectionSummary": None,
+                    "suggestedText": None,
+                    "text": "hello",
+                }
+            ],
+            "documentPages": [
+                {
+                    "chunks": [],
+                    "height": 500,
+                    "pageNumber": 1,
+                    "pageUrl": "https://page.jpg",
+                    "width": 400,
+                }
+            ],
+            "sourceUrl": "https://doc.pdf",
+            "fileKeywords": "kw",
+            "fileName": "file.pdf",
+            "fileType": "pdf",
+            "fileSummary": "sum",
+            "language": "en",
+        }
+        dummy = self.make_dummy_response(payload=payload, status_ok=True)
+        with patch("requests.get", return_value=dummy):
+            gx = GD(base_url="", documentID="D", taskID="T")
+            xdoc = test_xray(gx)
+
+            self.assertEqual(xdoc.fileType, "pdf")
+            self.assertEqual(xdoc.fileName, "file.pdf")
+            self.assertEqual(xdoc.fileKeywords, "kw")
+            self.assertEqual(xdoc.language, "en")
+
+            self.assertEqual(len(xdoc.chunks), 1)
+            chunk = xdoc.chunks[0]
+            self.assertIsInstance(chunk, Chunk)
+            self.assertEqual(chunk.chunk, "foo")
+            bb: typing.Optional[BoundingBox] = None
+            if chunk.boundingBoxes is not None and len(chunk.boundingBoxes) > 0:
+                bb = chunk.boundingBoxes[0]
+            self.assertIsInstance(bb, BoundingBox)
+            assert bb is not None, "BoundingBox should not be None"
+            self.assertTrue(bb.corrected)
+
+            self.assertEqual(len(xdoc.documentPages), 1)
+            page = xdoc.documentPages[0]
+            self.assertIsInstance(page, DocumentPage)
+            self.assertEqual(page.pageUrl, "https://page.jpg")
+
+
+if __name__ == "__main__":
+    unittest.main()