cosmotech-acceleration-library 1.0.0__py3-none-any.whl

cosmotech/coal/cosmotech_api/dataset/download/file.py

@@ -0,0 +1,216 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+import csv
+import io
+import json
+import os
+import tempfile
+import time
+from pathlib import Path
+from typing import Dict, List, Any, Optional, Union, Tuple
+
+from cosmotech_api import WorkspaceApi
+from openpyxl import load_workbook
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech.coal.cosmotech_api.connection import get_api_client
+
+
+def download_file_dataset(
+    organization_id: str,
+    workspace_id: str,
+    file_name: str,
+    target_folder: Optional[Union[str, Path]] = None,
+    read_files: bool = True,
+) -> Tuple[Dict[str, Any], Path]:
+    """
+    Download file dataset.
+
+    Args:
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        file_name: File name to download
+        target_folder: Optional folder to save files (if None, uses temp dir)
+        read_files: Whether to read file contents
+
+    Returns:
+        Tuple of (content dict, folder path)
+    """
+    start_time = time.time()
+    LOGGER.info(T("coal.services.dataset.download_started").format(dataset_type="File"))
+    LOGGER.debug(
+        T("coal.services.dataset.file_downloading").format(
+            organization_id=organization_id,
+            workspace_id=workspace_id,
+            file_name=file_name,
+        )
+    )
+
+    # Create temp directory for downloaded files
+    if target_folder is None:
+        tmp_dataset_dir = tempfile.mkdtemp()
+    else:
+        tmp_dataset_dir = Path(target_folder)
+        tmp_dataset_dir.mkdir(parents=True, exist_ok=True)
+        tmp_dataset_dir = str(tmp_dataset_dir)
+
+    LOGGER.debug(T("coal.services.dataset.using_folder").format(folder=tmp_dataset_dir))
+
+    content = dict()
+
+    with get_api_client()[0] as api_client:
+        api_ws = WorkspaceApi(api_client)
+
+        # Find all files matching the pattern
+        list_start = time.time()
+        LOGGER.debug(T("coal.services.dataset.listing_workspace_files"))
+        all_api_files = api_ws.find_all_workspace_files(organization_id, workspace_id)
+
+        existing_files = list(_f.file_name for _f in all_api_files if _f.file_name.startswith(file_name))
+        list_time = time.time() - list_start
+
+        LOGGER.debug(T("coal.services.dataset.workspace_files_found").format(count=len(existing_files)))
+        LOGGER.debug(T("coal.common.timing.operation_completed").format(operation="file listing", time=list_time))
+
+        if not existing_files:
+            LOGGER.warning(T("coal.services.dataset.no_files_found").format(file_name=file_name))
+            return content, Path(tmp_dataset_dir)
+
+        # Download and process each file
+        for _file_name in existing_files:
+            download_start = time.time()
+            LOGGER.debug(T("coal.services.dataset.downloading_file").format(file_name=_file_name))
+
+            dl_file = api_ws.download_workspace_file(
+                organization_id=organization_id,
+                workspace_id=workspace_id,
+                file_name=_file_name,
+            )
+
+            target_file = os.path.join(tmp_dataset_dir, _file_name.split("/")[-1])
+            with open(target_file, "wb") as tmp_file:
+                tmp_file.write(dl_file)
+
+            download_time = time.time() - download_start
+            LOGGER.debug(T("coal.services.dataset.file_downloaded").format(file_name=_file_name, path=target_file))
+            LOGGER.debug(
+                T("coal.common.timing.operation_completed").format(
+                    operation=f"download {_file_name}", time=download_time
+                )
+            )
+
+            if not read_files:
+                continue
+
+            # Process file based on type
+            process_start = time.time()
+
+            if ".xls" in _file_name:
+                LOGGER.debug(T("coal.services.dataset.processing_excel").format(file_name=target_file))
+                wb = load_workbook(target_file, data_only=True)
+
+                for sheet_name in wb.sheetnames:
+                    sheet = wb[sheet_name]
+                    content[sheet_name] = list()
+                    headers = next(sheet.iter_rows(max_row=1, values_only=True))
+
+                    def item(_row: tuple) -> dict:
+                        return {k: v for k, v in zip(headers, _row)}
+
+                    row_count = 0
+                    for r in sheet.iter_rows(min_row=2, values_only=True):
+                        row = item(r)
+                        new_row = dict()
+
+                        for key, value in row.items():
+                            try:
+                                converted_value = json.load(io.StringIO(value))
+                            except (json.decoder.JSONDecodeError, TypeError):
+                                converted_value = value
+
+                            if converted_value is not None:
+                                new_row[key] = converted_value
+
+                        if new_row:
+                            content[sheet_name].append(new_row)
+                            row_count += 1
+
+                    LOGGER.debug(
+                        T("coal.services.dataset.sheet_processed").format(sheet_name=sheet_name, rows=row_count)
+                    )
+
+            elif ".csv" in _file_name:
+                LOGGER.debug(T("coal.services.dataset.processing_csv").format(file_name=target_file))
+                with open(target_file, "r") as file:
+                    current_filename = os.path.basename(target_file)[: -len(".csv")]
+                    content[current_filename] = list()
+
+                    row_count = 0
+                    for csv_row in csv.DictReader(file):
+                        csv_row: dict
+                        new_row = dict()
+
+                        for key, value in csv_row.items():
+                            try:
+                                # Try to convert any json row to dict object
+                                converted_value = json.load(io.StringIO(value))
+                            except json.decoder.JSONDecodeError:
+                                converted_value = value
+
+                            if converted_value == "":
+                                converted_value = None
+
+                            if converted_value is not None:
+                                new_row[key] = converted_value
+
+                        content[current_filename].append(new_row)
+                        row_count += 1
+
+                    LOGGER.debug(
+                        T("coal.services.dataset.csv_processed").format(file_name=current_filename, rows=row_count)
+                    )
+
+            elif ".json" in _file_name:
+                LOGGER.debug(T("coal.services.dataset.processing_json").format(file_name=target_file))
+                with open(target_file, "r") as _file:
+                    current_filename = os.path.basename(target_file)
+                    content[current_filename] = json.load(_file)
+
+                    if isinstance(content[current_filename], dict):
+                        item_count = len(content[current_filename])
+                    elif isinstance(content[current_filename], list):
+                        item_count = len(content[current_filename])
+                    else:
+                        item_count = 1
+
+                    LOGGER.debug(
+                        T("coal.services.dataset.json_processed").format(file_name=current_filename, items=item_count)
+                    )
+
+            else:
+                LOGGER.debug(T("coal.services.dataset.processing_text").format(file_name=target_file))
+                with open(target_file, "r") as _file:
+                    current_filename = os.path.basename(target_file)
+                    content[current_filename] = "\n".join(line for line in _file)
+
+                    line_count = content[current_filename].count("\n") + 1
+                    LOGGER.debug(
+                        T("coal.services.dataset.text_processed").format(file_name=current_filename, lines=line_count)
+                    )
+
+            process_time = time.time() - process_start
+            LOGGER.debug(
+                T("coal.common.timing.operation_completed").format(operation=f"process {_file_name}", time=process_time)
+            )
+
+    elapsed_time = time.time() - start_time
+    LOGGER.info(T("coal.common.timing.operation_completed").format(operation="File download", time=elapsed_time))
+    LOGGER.info(T("coal.services.dataset.download_completed").format(dataset_type="File"))
+
+    return content, Path(tmp_dataset_dir)

cosmotech/coal/cosmotech_api/dataset/download/twingraph.py

@@ -0,0 +1,188 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+import time
+import tempfile
+from pathlib import Path
+from typing import Dict, List, Any, Optional, Union, Tuple
+
+from cosmotech_api import (
+    DatasetApi,
+    DatasetTwinGraphQuery,
+    TwinGraphQuery,
+    TwingraphApi,
+)
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech.coal.cosmotech_api.connection import get_api_client
+from cosmotech.coal.cosmotech_api.dataset.utils import get_content_from_twin_graph_data
+from cosmotech.coal.cosmotech_api.dataset.converters import convert_dataset_to_files
+
+
+def download_twingraph_dataset(
+    organization_id: str,
+    dataset_id: str,
+    target_folder: Optional[Union[str, Path]] = None,
+) -> Tuple[Dict[str, Any], Path]:
+    """
+    Download dataset from TwinGraph.
+
+    Args:
+        organization_id: Organization ID
+        dataset_id: Dataset ID
+        target_folder: Optional folder to save files (if None, uses temp dir)
+
+    Returns:
+        Tuple of (content dict, folder path)
+    """
+    start_time = time.time()
+    LOGGER.info(T("coal.services.dataset.download_started").format(dataset_type="TwinGraph"))
+    LOGGER.debug(
+        T("coal.services.dataset.twingraph_downloading").format(organization_id=organization_id, dataset_id=dataset_id)
+    )
+
+    with get_api_client()[0] as api_client:
+        dataset_api = DatasetApi(api_client)
+
+        # Query nodes
+        nodes_start = time.time()
+        LOGGER.debug(T("coal.services.dataset.twingraph_querying_nodes").format(dataset_id=dataset_id))
+        nodes_query = DatasetTwinGraphQuery(query="MATCH(n) RETURN n")
+
+        nodes = dataset_api.twingraph_query(
+            organization_id=organization_id,
+            dataset_id=dataset_id,
+            dataset_twin_graph_query=nodes_query,
+        )
+
+        nodes_time = time.time() - nodes_start
+        LOGGER.debug(T("coal.services.dataset.twingraph_nodes_found").format(count=len(nodes)))
+        LOGGER.debug(T("coal.common.timing.operation_completed").format(operation="nodes query", time=nodes_time))
+
+        # Query edges
+        edges_start = time.time()
+        LOGGER.debug(T("coal.services.dataset.twingraph_querying_edges").format(dataset_id=dataset_id))
+        edges_query = DatasetTwinGraphQuery(query="MATCH(n)-[r]->(m) RETURN n as src, r as rel, m as dest")
+
+        edges = dataset_api.twingraph_query(
+            organization_id=organization_id,
+            dataset_id=dataset_id,
+            dataset_twin_graph_query=edges_query,
+        )
+
+        edges_time = time.time() - edges_start
+        LOGGER.debug(T("coal.services.dataset.twingraph_edges_found").format(count=len(edges)))
+        LOGGER.debug(T("coal.common.timing.operation_completed").format(operation="edges query", time=edges_time))
+
+        # Process results
+        process_start = time.time()
+        content = get_content_from_twin_graph_data(nodes, edges, True)
+        process_time = time.time() - process_start
+
+        LOGGER.debug(T("coal.common.timing.operation_completed").format(operation="data processing", time=process_time))
+
+    # Convert to files if target_folder is provided
+    if target_folder:
+        dataset_info = {
+            "type": "twincache",
+            "content": content,
+            "name": f"TwinGraph Dataset {dataset_id}",
+        }
+        target_folder = convert_dataset_to_files(dataset_info, target_folder)
+    else:
+        target_folder = tempfile.mkdtemp()
+
+    elapsed_time = time.time() - start_time
+    LOGGER.info(T("coal.common.timing.operation_completed").format(operation="TwinGraph download", time=elapsed_time))
+    LOGGER.info(T("coal.services.dataset.download_completed").format(dataset_type="TwinGraph"))
+
+    return content, Path(target_folder)
+
+
+def download_legacy_twingraph_dataset(
+    organization_id: str,
+    cache_name: str,
+    target_folder: Optional[Union[str, Path]] = None,
+) -> Tuple[Dict[str, Any], Path]:
+    """
+    Download dataset from legacy TwinGraph.
+
+    Args:
+        organization_id: Organization ID
+        cache_name: Twin cache name
+        target_folder: Optional folder to save files (if None, uses temp dir)
+
+    Returns:
+        Tuple of (content dict, folder path)
+    """
+    start_time = time.time()
+    LOGGER.info(T("coal.services.dataset.download_started").format(dataset_type="Legacy TwinGraph"))
+    LOGGER.debug(
+        T("coal.services.dataset.legacy_twingraph_downloading").format(
+            organization_id=organization_id, cache_name=cache_name
+        )
+    )
+
+    with get_api_client()[0] as api_client:
+        api_instance = TwingraphApi(api_client)
+
+        # Query nodes
+        nodes_start = time.time()
+        LOGGER.debug(T("coal.services.dataset.legacy_twingraph_querying_nodes").format(cache_name=cache_name))
+        _query_nodes = TwinGraphQuery(query="MATCH(n) RETURN n")
+
+        nodes = api_instance.query(
+            organization_id=organization_id,
+            graph_id=cache_name,
+            twin_graph_query=_query_nodes,
+        )
+
+        nodes_time = time.time() - nodes_start
+        LOGGER.debug(T("coal.services.dataset.legacy_twingraph_nodes_found").format(count=len(nodes)))
+        LOGGER.debug(T("coal.common.timing.operation_completed").format(operation="nodes query", time=nodes_time))
+
+        # Query relationships
+        rel_start = time.time()
+        LOGGER.debug(T("coal.services.dataset.legacy_twingraph_querying_relations").format(cache_name=cache_name))
+        _query_rel = TwinGraphQuery(query="MATCH(n)-[r]->(m) RETURN n as src, r as rel, m as dest")
+
+        rel = api_instance.query(
+            organization_id=organization_id,
+            graph_id=cache_name,
+            twin_graph_query=_query_rel,
+        )
+
+        rel_time = time.time() - rel_start
+        LOGGER.debug(T("coal.services.dataset.legacy_twingraph_relations_found").format(count=len(rel)))
+        LOGGER.debug(T("coal.common.timing.operation_completed").format(operation="relations query", time=rel_time))
+
+        # Process results
+        process_start = time.time()
+        content = get_content_from_twin_graph_data(nodes, rel, False)
+        process_time = time.time() - process_start
+
+        LOGGER.debug(T("coal.common.timing.operation_completed").format(operation="data processing", time=process_time))
+
+    # Convert to files if target_folder is provided
+    if target_folder:
+        dataset_info = {
+            "type": "twincache",
+            "content": content,
+            "name": f"Legacy TwinGraph Dataset {cache_name}",
+        }
+        target_folder = convert_dataset_to_files(dataset_info, target_folder)
+    else:
+        target_folder = tempfile.mkdtemp()
+
+    elapsed_time = time.time() - start_time
+    LOGGER.info(
+        T("coal.common.timing.operation_completed").format(operation="Legacy TwinGraph download", time=elapsed_time)
+    )
+    LOGGER.info(T("coal.services.dataset.download_completed").format(dataset_type="Legacy TwinGraph"))
+
+    return content, Path(target_folder)

cosmotech/coal/cosmotech_api/dataset/utils.py

@@ -0,0 +1,132 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+from typing import Dict, List, Any
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def get_content_from_twin_graph_data(
+    nodes: List[Dict], relationships: List[Dict], restore_names: bool = False
+) -> Dict[str, List[Dict]]:
+    """
+    Extract content from twin graph data.
+
+    When restore_names is True, the "id" value inside the "properties" field in the cypher query response is used
+    instead of the numerical id found in the "id" field. When restore_names is set to False, this function
+    keeps the previous behavior implemented when adding support for twingraph in v2 (default: False)
+
+    Example with a sample of cypher response:
+    [{
+      n: {
+        id: "50"  <-- this id is used if restore_names is False
+        label: "Customer"
+        properties: {
+          Satisfaction: 0
+          SurroundingSatisfaction: 0
+          Thirsty: false
+          id: "Lars_Coret"  <-- this id is used if restore_names is True
+        }
+        type: "NODE"
+      }
+    }]
+
+    Args:
+        nodes: List of node data from cypher query
+        relationships: List of relationship data from cypher query
+        restore_names: Whether to use property ID instead of node ID
+
+    Returns:
+        Dict mapping entity types to lists of entities
+    """
+    LOGGER.debug(
+        T("coal.services.dataset.processing_graph_data").format(
+            nodes_count=len(nodes),
+            relationships_count=len(relationships),
+            restore_names=restore_names,
+        )
+    )
+
+    content = dict()
+    # build keys
+    for item in relationships:
+        content[item["src"]["label"]] = list()
+        content[item["dest"]["label"]] = list()
+        content[item["rel"]["label"]] = list()
+
+    # Process nodes
+    for item in nodes:
+        label = item["n"]["label"]
+        props = item["n"]["properties"].copy()  # Create a copy to avoid modifying the original
+        if not restore_names:
+            props.update({"id": item["n"]["id"]})
+        content.setdefault(label, list())
+        content[label].append(props)
+
+    # Process relationships
+    for item in relationships:
+        src = item["src"]
+        dest = item["dest"]
+        rel = item["rel"]
+        props = rel["properties"].copy()  # Create a copy to avoid modifying the original
+        content[rel["label"]].append(
+            {
+                "id": rel["id"],
+                "source": src["properties"]["id"] if restore_names else src["id"],
+                "target": dest["properties"]["id"] if restore_names else dest["id"],
+                **props,
+            }
+        )
+
+    # Log the number of entities by type
+    for entity_type, entities in content.items():
+        LOGGER.debug(T("coal.services.dataset.entity_count").format(entity_type=entity_type, count=len(entities)))
+
+    return content
+
+
+def sheet_to_header(sheet_content: List[Dict]) -> List[str]:
+    """
+    Extract header fields from sheet content.
+
+    Args:
+        sheet_content: List of dictionaries representing sheet rows
+
+    Returns:
+        List of field names with id, source, and target fields first if present
+    """
+    LOGGER.debug(T("coal.services.dataset.extracting_headers").format(rows=len(sheet_content)))
+
+    fieldnames = []
+    has_src = False
+    has_id = False
+
+    for r in sheet_content:
+        for k in r.keys():
+            if k not in fieldnames:
+                if k in ["source", "target"]:
+                    has_src = True
+                elif k == "id":
+                    has_id = True
+                else:
+                    fieldnames.append(k)
+
+    # Ensure source/target and id fields come first
+    if has_src:
+        fieldnames = ["source", "target"] + fieldnames
+    if has_id:
+        fieldnames = ["id"] + fieldnames
+
+    LOGGER.debug(
+        T("coal.services.dataset.headers_extracted").format(
+            count=len(fieldnames),
+            fields=", ".join(fieldnames[:5]) + ("..." if len(fieldnames) > 5 else ""),
+        )
+    )
+
+    return fieldnames

cosmotech/coal/cosmotech_api/parameters.py

@@ -0,0 +1,48 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+"""
+Parameter handling functions.
+
+This module provides functions for handling parameters in solution templates.
+"""
+
+import json
+import os
+import pathlib
+from csv import DictWriter
+from typing import List, Dict, Any
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def write_parameters(
+    parameter_folder: str, parameters: List[Dict[str, Any]], write_csv: bool, write_json: bool
+) -> None:
+    """
+    Write parameters to CSV and/or JSON files.
+
+    Args:
+        parameter_folder: The folder to write the parameters to
+        parameters: The parameters to write
+        write_csv: Whether to write the parameters to a CSV file
+        write_json: Whether to write the parameters to a JSON file
+    """
+    if write_csv:
+        tmp_parameter_file = os.path.join(parameter_folder, "parameters.csv")
+        LOGGER.info(T("coal.cosmotech_api.runner.generating_file").format(file=tmp_parameter_file))
+        with open(tmp_parameter_file, "w") as _file:
+            _w = DictWriter(_file, fieldnames=["parameterId", "value", "varType", "isInherited"])
+            _w.writeheader()
+            _w.writerows(parameters)
+
+    if write_json:
+        tmp_parameter_file = os.path.join(parameter_folder, "parameters.json")
+        LOGGER.info(T("coal.cosmotech_api.runner.generating_file").format(file=tmp_parameter_file))
+        with open(tmp_parameter_file, "w") as _file:
+            json.dump(parameters, _file, indent=2)

cosmotech/coal/cosmotech_api/run.py

@@ -0,0 +1,25 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+from typing import Any
+from typing import Optional
+
+import cosmotech_api
+
+
+def get_run_metadata(
+    api_client: cosmotech_api.api_client.ApiClient,
+    organization_id: str,
+    workspace_id: str,
+    runner_id: str,
+    run_id: str,
+    include: Optional[list[str]] = None,
+    exclude: Optional[list[str]] = None,
+) -> dict[str, Any]:
+    run_api = cosmotech_api.RunApi(api_client)
+
+    run: cosmotech_api.Run = run_api.get_run(organization_id, workspace_id, runner_id, run_id)
+    return run.model_dump(by_alias=True, exclude_none=True, include=include, exclude=exclude, mode="json")