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

groundx/extract/classes/test_prompt.py

@@ -0,0 +1,68 @@
+import typing, unittest
+
+from .prompt import Prompt
+
+
+def TestPrompt(
+    name: str,
+    ty: typing.Union[str, typing.List[str]],
+) -> Prompt:
+    return Prompt(
+        attr_name=name,
+        prompt=name.replace("_", "-"),
+        type=ty,
+    )
+
+
+class TestPromptValidValue(unittest.TestCase):
+    def test_single_type_str(self):
+        p = TestPrompt("field1", "str")
+        self.assertTrue(p.valid_value("hello"))
+        self.assertFalse(p.valid_value(123))
+        self.assertFalse(p.valid_value((1, 2, 3)))
+        self.assertFalse(p.valid_value([1, 2, 3]))
+
+    def test_single_type_int(self):
+        p = TestPrompt("field1", "int")
+        self.assertFalse(p.valid_value("hello"))
+        self.assertTrue(p.valid_value(123))
+        self.assertTrue(p.valid_value(12.3))
+        self.assertFalse(p.valid_value((1, 2, 3)))
+        self.assertFalse(p.valid_value([1, 2, 3]))
+
+    def test_single_type_float(self):
+        p = TestPrompt("field1", "float")
+        self.assertFalse(p.valid_value("hello"))
+        self.assertTrue(p.valid_value(123))
+        self.assertTrue(p.valid_value(12.3))
+        self.assertTrue(p.valid_value(123.0))
+        self.assertFalse(p.valid_value((1, 2, 3)))
+        self.assertFalse(p.valid_value([1, 2, 3]))
+
+    def test_single_type_list(self):
+        p = TestPrompt("field1", "list")
+        self.assertFalse(p.valid_value("hello"))
+        self.assertFalse(p.valid_value(123))
+        self.assertFalse(p.valid_value(12.3))
+        self.assertFalse(p.valid_value(123.0))
+        self.assertFalse(p.valid_value((1, 2, 3)))
+        self.assertTrue(p.valid_value([1, 2, 3]))
+
+    def test_list_of_types_success_and_failure(self):
+        p = TestPrompt("field2", ["str", "float"])
+        self.assertTrue(p.valid_value("hello"))
+        self.assertTrue(p.valid_value(123))
+        self.assertTrue(p.valid_value(12.3))
+        self.assertTrue(p.valid_value(123.0))
+        self.assertFalse(p.valid_value((1, 2, 3)))
+        self.assertFalse(p.valid_value([1, 2, 3]))
+
+    def test_repr_contains_fields(self):
+        p = TestPrompt("field_5", "int")
+        rep = repr(p)
+        self.assertIn("field_5", rep)
+        self.assertIn("field-5", rep)
+
+
+if __name__ == "__main__":
+    unittest.main()

groundx/extract/post_process/__init__.py

@@ -0,0 +1,7 @@
+from .post_process import check_map, check_valid
+
+
+__all__ = [
+    "check_map",
+    "check_valid",
+]

groundx/extract/post_process/post_process.py

@@ -0,0 +1,33 @@
+import typing
+
+
+def check_map(
+    fty: str,
+    sty: str,
+    val: str,
+    mp: typing.Dict[str, typing.Dict[str, str]],
+    should_warn: bool = True,
+) -> typing.Optional[str]:
+    if sty not in mp:
+        sty = ""
+    if sty not in mp:
+        return None
+
+    vl = val.lower().strip()
+
+    nmp = mp[sty]
+    if vl not in nmp:
+        if should_warn:
+            print(f"[arcadia-v1] {fty} not found [{sty}] [{vl}]")
+        return None
+
+    return nmp[vl]
+
+
+def check_valid(sty: str, val: str, valid: typing.Dict[str, typing.List[str]]) -> bool:
+    vl = val.lower().strip()
+
+    if sty not in valid:
+        sty = ""
+
+    return sty in valid and vl in valid[sty]

groundx/extract/services/__init__.py

@@ -0,0 +1,14 @@
+from .logger import Logger
+from .sheets_client import SheetsClient
+from .ratelimit import RateLimit
+from .status import Status
+from .upload import Upload
+
+
+__all__ = [
+    "Logger",
+    "RateLimit",
+    "SheetsClient",
+    "Status",
+    "Upload",
+]

groundx/extract/services/csv.py

@@ -0,0 +1,76 @@
+import csv, typing
+from pathlib import Path
+
+
+def append_row(
+    csv_path: Path,
+    headers: typing.List[str],
+    row: typing.Dict[str, str],
+) -> None:
+    with csv_path.open("a", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=headers)
+        writer.writerow(row)
+
+
+def extraction_row(
+    record: typing.Mapping[str, typing.Any], keys_in_order: typing.Sequence[str]
+) -> typing.List[typing.Any]:
+    return [record.get(k, "") for k in keys_in_order]
+
+
+def find_rows(
+    query: typing.Dict[str, str],
+    csv_path: str,
+) -> typing.List[typing.Dict[str, str]]:
+    with open(csv_path, newline="", encoding="utf-8") as f:
+        reader = csv.DictReader(f)
+
+        rows: typing.List[typing.Dict[str, str]] = []
+        for row in reader:
+            matches: typing.List[str] = []
+            for k, v in query.items():
+                if str(row.get(k)) == str(v):
+                    matches.append(k)
+
+            if len(matches) == len(query):
+                rows.append(row)
+
+        return rows
+
+
+def load_row(
+    key: str,
+    match: typing.List[str],
+    csv_path: typing.Optional[Path] = None,
+    rows: typing.Optional[typing.List[typing.Dict[str, str]]] = None,
+) -> typing.Optional[typing.Dict[str, str]]:
+    if csv_path is None and rows is None:
+        raise Exception("csv_path and rows are None")
+
+    if rows is None and csv_path:
+        rows = load_rows(csv_path)
+
+    if not rows:
+        raise Exception("rows are None")
+
+    return next((r for r in rows if r.get(key) in match), None)
+
+
+def load_rows(csv_path: Path) -> typing.List[typing.Dict[str, str]]:
+    rows: typing.List[typing.Dict[str, str]] = []
+    with csv_path.open("r", newline="") as csvfile:
+        reader = csv.DictReader(csvfile)
+        for row in reader:
+            rows.append(row)
+
+    return rows
+
+
+def save_rows(
+    csv_path: Path, headers: typing.List[str], rows: typing.List[typing.Dict[str, str]]
+) -> None:
+    with csv_path.open("w", newline="") as csvfile:
+        writer = csv.DictWriter(csvfile, fieldnames=headers)
+        writer.writeheader()
+        for r in rows:
+            writer.writerow(r)

groundx/extract/services/logger.py

@@ -0,0 +1,126 @@
+import logging, logging.config, typing
+
+from .logging_cfg import logging_config
+
+
+class Logger:
+    def __init__(
+        self,
+        name: str,
+        level: str,
+    ) -> None:
+        logging.config.dictConfig(logging_config(name, level))
+
+        self.logger = logging.getLogger(name)
+
+    def debug_msg(
+        self,
+        msg: str,
+        name: typing.Optional[str] = None,
+        document_id: typing.Optional[str] = None,
+        task_id: typing.Optional[str] = None,
+    ) -> None:
+        self.print_msg("DEBUG", msg, name, document_id, task_id)
+
+    def error_msg(
+        self,
+        msg: str,
+        name: typing.Optional[str] = None,
+        document_id: typing.Optional[str] = None,
+        task_id: typing.Optional[str] = None,
+    ) -> None:
+        self.print_msg("ERROR", msg, name, document_id, task_id)
+
+    def info_msg(
+        self,
+        msg: str,
+        name: typing.Optional[str] = None,
+        document_id: typing.Optional[str] = None,
+        task_id: typing.Optional[str] = None,
+    ) -> None:
+        self.print_msg("INFO", msg, name, document_id, task_id)
+
+    def report_error(
+        self,
+        api_key: str,
+        callback_url: str,
+        req: typing.Optional[typing.Dict[str, typing.Any]],
+        msg: str,
+    ) -> None:
+        import requests
+
+        self.error_msg(msg)
+
+        if req is None or callback_url == "":
+            return
+
+        requests.post(
+            callback_url,
+            json=req,
+            headers={"X-API-Key": api_key},
+        )
+
+    def report_result(
+        self,
+        api_key: str,
+        callback_url: str,
+        req: typing.Dict[str, typing.Any],
+    ):
+        import requests
+
+        if callback_url == "":
+            return
+
+        self.info_msg("calling back to [%s]" % (callback_url))
+
+        requests.post(
+            callback_url,
+            json=req,
+            headers={"X-API-Key": api_key},
+        )
+
+    def warning_msg(
+        self,
+        msg: str,
+        name: typing.Optional[str] = None,
+        document_id: typing.Optional[str] = None,
+        task_id: typing.Optional[str] = None,
+    ) -> None:
+        self.print_msg("WARNING", msg, name, document_id, task_id)
+
+    def print_msg(
+        self,
+        level: str,
+        msg: str,
+        name: typing.Optional[str] = None,
+        document_id: typing.Optional[str] = None,
+        task_id: typing.Optional[str] = None,
+    ) -> None:
+        prefix = ""
+        if name:
+            if prefix != "":
+                prefix += " "
+            prefix += f"[{name}]"
+        if document_id:
+            if prefix != "":
+                prefix += " "
+            prefix += f"d [{document_id}]"
+        if task_id:
+            if prefix != "":
+                prefix += " "
+            prefix += f"t [{task_id}]"
+
+        text = ""
+        if prefix != "":
+            text += f"{prefix} "
+        text += f"\n\n\t>> {msg}\n"
+
+        lvl = level.upper()
+        if lvl == "ERROR":
+            self.logger.error(text)
+        elif lvl in ("WARN", "WARNING"):
+            self.logger.warning(text)
+        elif lvl == "INFO":
+            self.logger.info(text)
+        else:
+            self.logger.debug(text)

groundx/extract/services/logging_cfg.py

@@ -0,0 +1,53 @@
+import typing
+
+
+def logging_config(name: str, level: str) -> typing.Dict[str, typing.Any]:
+    return {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "default": {
+                "format": "%(asctime)s - [%(process)d] - %(levelname)s - %(message)s",
+            },
+        },
+        "handlers": {
+            "default": {
+                "level": level.upper(),
+                "formatter": "default",
+                "class": "logging.StreamHandler",
+                "stream": "ext://sys.stdout",
+            },
+        },
+        "loggers": {
+            "": {
+                "handlers": ["default"],
+                "level": "WARNING",
+            },
+            name: {
+                "handlers": ["default"],
+                "level": level.upper(),
+                "propagate": False,
+            },
+            "gunicorn.error": {
+                "level": "INFO",
+                "handlers": ["default"],
+                "propagate": False,
+            },
+            "gunicorn.access": {
+                "level": "WARNING",
+                "handlers": ["default"],
+                "propagate": False,
+            },
+            "uvicorn": {"level": "INFO", "handlers": ["default"], "propagate": False},
+            "uvicorn.error": {
+                "level": "INFO",
+                "handlers": ["default"],
+                "propagate": False,
+            },
+            "uvicorn.access": {
+                "level": "WARNING",
+                "handlers": ["default"],
+                "propagate": False,
+            },
+        },
+    }

groundx/extract/services/ratelimit.py

@@ -0,0 +1,104 @@
+import typing
+
+from dataclasses import asdict
+from fastapi import Request, HTTPException
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import JSONResponse, Response
+from starlette.types import ASGIApp
+
+from ..classes.api import ProcessResponse
+from .logger import Logger
+from .status import Status
+from ..settings.settings import ContainerSettings
+from .utility import get_gunicorn_threads, get_thread_id, get_worker_id
+
+
+class RateLimit(BaseHTTPMiddleware):
+    def __init__(
+        self,
+        app: ASGIApp,
+        settings: ContainerSettings,
+        logger: Logger,
+    ) -> None:
+        super().__init__(app)
+
+        self.worker_id = get_worker_id()
+        num_threads = get_gunicorn_threads()
+        if num_threads > 1:
+            num_threads = num_threads - 1
+
+        self.status = Status(
+            settings,
+            logger,
+        )
+
+        self.settings = settings
+        self.logger = logger
+
+        self.thread_ids: typing.Dict[str, typing.Any] = {}
+
+        self.status.set_worker_available(self.worker_id)
+
+        self.logger.info_msg(
+            f"[{self.settings.service}] ratelimit init [{num_threads}]"
+        )
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: typing.Callable[[Request], typing.Awaitable[Response]],
+    ) -> Response:
+        thread_id, self.thread_ids = get_thread_id(self.thread_ids)
+        wasSet = False
+
+        try:
+            if request.url.path == "/health":
+                response = await call_next(request)
+
+                self.status.refresh_worker(self.worker_id)
+
+                available, total = self.status.get_worker_state(self.worker_id)
+
+                response = self.status.set_headers(
+                    response, self.worker_id, available, total
+                )
+
+                return response
+
+            api_key = request.headers.get("X-API-Key") or request.headers.get(
+                "Authorization"
+            )
+            if api_key and api_key.startswith("Bearer "):
+                api_key = api_key.split("Bearer ")[1]
+            if not api_key or api_key not in self.settings.get_valid_api_keys():
+                raise HTTPException(status_code=403, detail="Invalid API key")
+
+            request.state.api_key = api_key
+
+            wasSet = True
+            self.status.set_worker_unavailable(self.worker_id)
+
+            response = await call_next(request)
+
+            wasSet = False
+            self.status.set_worker_available(self.worker_id)
+
+            available, total = self.status.get_service_state()
+
+            response.headers.update(
+                {
+                    "X-RateLimit-Limit-Requests": str(total),
+                    "X-RateLimit-Remaining-Requests": str(max(0, available)),
+                    "X-Worker-ID": f"{self.worker_id}:{thread_id}",
+                }
+            )
+
+            return response
+        except HTTPException as exc:
+            if wasSet:
+                self.status.set_worker_available(self.worker_id)
+
+            return JSONResponse(
+                status_code=exc.status_code,
+                content=asdict(ProcessResponse(message=exc.detail)),
+            )

groundx/extract/services/sheets_client.py

@@ -0,0 +1,160 @@
+import json, os, typing
+from pathlib import Path
+
+from google.oauth2 import service_account
+from googleapiclient.discovery import (
+    build,  # pyright: ignore[reportUnknownVariableType]
+)
+import gspread
+
+from ..settings.settings import ContainerSettings, GCP_CREDENTIALS
+
+SPREADSHEET_MIME = "application/vnd.google-apps.spreadsheet"
+
+
+class SheetsClient:
+    client: gspread.Client
+    # drive: DriveResource
+    drive: typing.Any
+    settings: ContainerSettings
+
+    def __init__(
+        self,
+        settings: ContainerSettings,
+        scopes: typing.Optional[typing.List[str]] = None,
+    ):
+        self.scopes: typing.List[str] = scopes or [
+            "https://www.googleapis.com/auth/spreadsheets",
+            "https://www.googleapis.com/auth/drive",
+        ]
+        self.settings = settings
+
+        creds_dict = _load_credentials_from_env()
+        if not creds_dict:
+            raise ValueError(f"{GCP_CREDENTIALS} does not load valid credentials")
+
+        creds = service_account.Credentials.from_service_account_info(
+            creds_dict, scopes=scopes
+        )
+        self.drive = build("drive", "v3", credentials=creds)
+
+        auth_scopes = self.scopes
+        if scopes:
+            auth_scopes = scopes
+
+        self.client = gspread.service_account_from_dict(creds_dict, scopes=auth_scopes)
+
+    def create_headers_if_missing(
+        self, ws: gspread.Worksheet, headers: typing.List[str]
+    ) -> None:
+        existing = ws.row_values(1)
+        if not existing:
+            ws.insert_row(headers, 1)
+
+    def find_sheet_by_name(
+        self,
+        spreadsheet_name: str,
+        drive_id: str,
+    ) -> typing.Optional[str]:
+        cln = spreadsheet_name.replace("'", "\\'")
+
+        q = f"name = '{cln}' and mimeType = '{SPREADSHEET_MIME}' and trashed = false"
+
+        resp = (
+            self.drive.files()
+            .list(
+                q=q,
+                corpora="drive",
+                driveId=drive_id,
+                includeItemsFromAllDrives=True,
+                supportsAllDrives=True,
+                fields="files(id, name)",
+            )
+            .execute()
+        )
+        files = resp.get("files", [])
+        return files[0].get("id") if files else None
+
+    def open_or_create_spreadsheet(
+        self,
+        spreadsheet_name: str,
+        drive_id: str,
+        sheet_1_title: typing.Optional[str] = None,
+    ) -> gspread.Spreadsheet:
+        file_id = self.find_sheet_by_name(spreadsheet_name, drive_id)
+        if file_id:
+            return self.client.open_by_key(file_id)
+
+        if self.settings.google_sheets_template_id:
+            created = (
+                self.drive.files()
+                .copy(
+                    fileId=self.settings.google_sheets_template_id,
+                    body={"name": spreadsheet_name, "parents": [drive_id]},
+                    supportsAllDrives=True,
+                    fields="id,name,parents,driveId",
+                )
+                .execute()
+            )
+        else:
+            created = (
+                self.drive.files()
+                .create(
+                    body={
+                        "name": spreadsheet_name,
+                        "mimeType": SPREADSHEET_MIME,
+                        "parents": [drive_id],
+                    },
+                    supportsAllDrives=True,
+                    fields="id,name,parents,driveId",
+                )
+                .execute()
+            )
+
+        cid = created.get("id")
+        if not cid:
+            raise Exception(f"create spreadsheet failed\n{created}")
+
+        sh = self.client.open_by_key(cid)
+
+        if sheet_1_title:
+            sh.sheet1.update_title(sheet_1_title)
+
+        return sh
+
+    def open_or_create_worksheet(
+        self,
+        sh: gspread.Spreadsheet,
+        title: str,
+        headers: typing.List[str],
+        rows: int = 1000,
+    ) -> gspread.Worksheet:
+        cols = len(headers)
+        try:
+            ws = sh.worksheet(title)
+            self.create_headers_if_missing(ws, headers)
+        except gspread.WorksheetNotFound:
+            ws = sh.add_worksheet(title=title, rows=rows, cols=cols)
+            ws.append_row(headers)
+
+        return ws
+
+
+def _load_credentials_from_env() -> typing.Optional[typing.Dict[str, typing.Any]]:
+    raw = os.environ.get(GCP_CREDENTIALS)
+    if not raw:
+        if Path("./gcv.json").exists():
+            with open("./gcv.json") as f:
+                data = f.read()
+            return json.loads(data)
+
+        return None
+
+    try:
+        creds = json.loads(raw)
+        if not isinstance(creds, dict):
+            raise ValueError(f"{GCP_CREDENTIALS} is not type dict [{type(creds)}]")
+
+        return typing.cast(typing.Dict[str, typing.Any], creds)
+    except Exception as e:
+        raise ValueError(f"{GCP_CREDENTIALS} is set but not valid JSON: {e}") from e