onshape integration and metadata extraction

app.py

@@ -1,4 +1,6 @@
 import os
+import json
+import asyncio
 from enum import Enum
 import platform
 import re
@@ -24,6 +26,8 @@ from torch import Tensor
 
 from llm_service import LLMService
 from mv_utils_zs import Realistic_Projection
+from onshape.onshape_translation import OnshapeTranslation
+from onshape.onshape_download import OnshapeDownload
 
 os.environ.get("GRADIO_TEMP_DIR", "gradio_cache")  # You must set it in `.env` file also
 os_name = platform.system()
@@ -38,11 +42,15 @@ else:
 # The Gradio 3D Model component default accept
 GRADIO_3D_MODEL_DEFAULT_FORMAT = [".obj", ".glb", ".gltf", ".stl", ".splat", ".ply"]
 USER_REQUIRE_FORMAT = [".3dxml", ".step"]
-FREECAD_LOW_LEVEL_FORMAT = [".step", ".igs", ".iges"]
+FREECAD_LOW_LEVEL_FORMAT = [".step", ".igs", ".iges", ".stp"]
+ONSHAPE_SUPPORTED_FORMAT = [".prt", ".asm", ".jt"]
 FREECAD_NATIVE_FORMAT = [".fcstd"]
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
 SIMILARITY_SCORE_THRESHOLD = 0.7
 
+
+llm_service = LLMService.from_partner()
+
 ####################################################################################################################
 # Transform high-level to low-level
 ####################################################################################################################
@@ -112,7 +120,7 @@ Mesh.export(to_export, "{obj_path}")
 
 
 # Dummy converter from STEP/3DXML to OBJ (replace with real converter)
-def convert_to_obj(file: str) -> str:
+async def convert_to_obj(file: str) -> str:
     if file is None:
         return None
     logger.info(f"Converting {file} to .obj")
@@ -131,6 +139,55 @@ def convert_to_obj(file: str) -> str:
         raise Exception(f"Do nothing at convert_to_obj with file {file}")
 
 
+async def onshape_converter(
+    input_file_path: str,
+    output_file: str | None = None,
+    did: str = "ef42d7639096f3e61a4d4f07",
+    wid: str = "5fcd0f25ce3dee08bbb823bf",
+    format_name: str = "STEP",
+) -> Dict:
+    """
+    Convert proprietary 3D file to open-source format using the Onshape API.
+    """
+    file_path = input_file_path
+
+    # Upload file and translate it to the desired format
+    translator = OnshapeTranslation(did, wid, file_path, format_name)
+    response = translator.upload_and_translate()
+
+    # Check the translation status via polling
+    response = translator.get_translation_status(response.id)
+    while response.request_state not in ["DONE", "FAILED"]:
+        logger.info(
+            f"Waiting for translation to complete. Current state: {response.request_state}"
+        )
+        response = translator.get_translation_status(response.id)
+        await asyncio.sleep(6)
+    logger.success(f"Translation completed with state: {response.request_state}")
+
+    # If translation failed, raise an error
+    if response.request_state == "FAILED":
+        logger.error(f"Translation failed: {response.failure_reason}")
+        raise gr.Error(f"Translation failed: {response.failure_reason}")
+
+    # Download the translated file
+    ## you can find it in `resultElementIds` when `requestState` of `TranslationStatusResponse` is `DONE`
+    assert (
+        response.result_element_ids is not None and len(response.result_element_ids) > 0
+    ), "No result element IDs found in translation response"
+    eid = response.result_element_ids[0]
+    prefix_path, ext = os.path.splitext(file_path)
+    if output_file is None:
+        output_file = f"{prefix_path}_{eid}.{format_name.lower()}"
+    downloader = OnshapeDownload(did, wid, eid, output_file)
+    downloader.download()
+
+    return {
+        "eid": eid,
+        "output_file": output_file,
+    }
+
+
 ####################################################################################################################
 # Feature Extraction
 ####################################################################################################################
@@ -534,7 +591,7 @@ If metadata filtering is required, delegate to the **Keyword Search Agent** by c
             get_description_of_model_to_analysis,
         ],
         handoffs=[keyword_search_agent],
-    )
+    )  # type:ignore
 
     # Prepare the prompt for the Datum Agent
     prompt_input = f"""An user is watching a 3D object and wants to query it. 
@@ -671,19 +728,67 @@ def extract_step_metadata(file_path):
     return metadata
 
 
+async def extract_step_metadata_using_llm(file_path: str) -> Dict:
+    logger.info("Extracting STEP metadata using LLM")
+    metadata = {}
+
+    try:
+        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+            content = f.read()
+
+            # Trim to the HEADER section ending with ENDSEC;
+            endsec_index = content.find("ENDSEC;")
+            if endsec_index != -1:
+                content = content[:endsec_index].strip() + "\nENDSEC;"
+                logger.info("Using trimmed content up to ENDSEC;")  # \n```{content}\n")
+            else:
+                logger.warning("No ENDSEC; found in the STEP file, using full content.")
+
+            # Prepare prompt for LLM
+            system_prompt = """You are a STEP file expert. Given the HEADER section of a STEP file, extract the following fields in JSON format:\n
+- Description (from FILE_DESCRIPTION)\n
+- Description_Level\n
+- FileName\n
+- Created\n
+- Authors (as a comma-separated string)\n
+- Organizations (as a comma-separated string)\n
+- Preprocessor\n
+- OriginatingSystem\n
+- Authorization\n
+- Schema\n\n
+Only return a valid JSON object with these fields.
+Here is the content of the STEP file:\n
+content = ```step\n{content}\n```
+"""
+
+            # Ask the LLM
+            raw_response = await llm_service.chat_with_text(
+                prompt=system_prompt.format(content=content),
+                return_as_json=True,
+            )
+
+            dict_response = json.loads(raw_response)
+            return dict_response  # Or: return dict_to_markdown(dict_response)
+
+    except Exception as e:
+        logger.error(f"Failed to extract STEP metadata with LLM: {e}")
+
+    return metadata
+
+
 def dict_to_markdown(metadata: dict) -> str:
     return "\n".join(f"{key}: {value}" for key, value in metadata.items())
 
 
 # Dummy parser - Replace with real parser
-def parse_3d_file(original_filepath: str) -> Dict[str, Any]:
+async def parse_3d_file(original_filepath: str) -> Dict[str, Any]:
     if original_filepath is None:
         return "No file"
     if original_filepath.endswith((".3dxml", ".3DXML")):
         meta = extract_header_from_3dxml(original_filepath)
         return meta
     elif original_filepath.endswith((".step", ".STEP")):
-        meta = extract_step_metadata(original_filepath)
+        meta = await extract_step_metadata_using_llm(original_filepath)
         return meta
     logger.warning(f"No metadata found in the file {original_filepath}")
     return {}
@@ -692,6 +797,7 @@ def parse_3d_file(original_filepath: str) -> Dict[str, Any]:
 def render_3D_metadata(
     original_filepath: str, obj_path: str, embedding_dict: dict
 ) -> Tuple[str, str]:
+    logger.info(f"Rendering 3D metadata for {original_filepath} and {obj_path}")
     return (
         embedding_dict.get(obj_path, {}).get("metadata", "No metadata found!"),
         embedding_dict.get(obj_path, {}).get("description", "No description found!"),
@@ -741,8 +847,6 @@ def aggregate_images(
     return aggregate_img
 
 
-llm_service = LLMService.from_partner()
-
 DESCRIPTION_AGGREGATED_DEPTH_MAP_PROMPT = """You are a manufacturing expert analyzing 3D objects for production purposes. Given a set of multi-view depth maps of a single object, extract all possible special features relevant to manufacturing.
 
 Your output must follow the structured format provided below and be as complete and specific as possible, even if some features are inferred or uncertain.
@@ -865,7 +969,12 @@ def normalize_metadata(metadata: Dict[str, Any]) -> Dict[str, object]:
     return normalized
 
 
-async def accumulate_and_embedding(input_files, file_list, embedding_dict):
+async def accumulate_and_embedding(
+    input_files: List[str],
+    file_list: List[str],
+    embedding_dict: Dict[str, Any],
+    converting_store_map: Dict[str, str],
+):
     # accumulate
     if not isinstance(input_files, list):
         input_files = [input_files]
@@ -883,10 +992,36 @@ async def accumulate_and_embedding(input_files, file_list, embedding_dict):
     # embedding
     for file_path in new_files:
         logger.info("Processing new upload file:", file_path)
-        obj_path = convert_to_obj(file_path)
+
+        # If proprietary file, translate first
+        prefix_path, ext = os.path.splitext(file_path)
+        if ext.lower() in ONSHAPE_SUPPORTED_FORMAT:
+            response = await onshape_converter(input_file_path=file_path)
+            step_path = response.get("output_file", "")  # type:str
+            logger.info(
+                f"Converted {file_path} to {step_path} using Onshape converter."
+            )
+        else:
+            step_path = None
+
+        # Convert to obj
+        if step_path is not None:
+            obj_path = await convert_to_obj(step_path)
+            logger.info(f"Converted {step_path} to {obj_path} using FreeCAD converter.")
+        else:
+            obj_path = await convert_to_obj(file_path)
+            logger.info(f"Converted {file_path} to {obj_path}.")
+
+        # Generate embeddings for the 3D object
         embeddings = await embedding_3d_object(obj_path)
+
         # Extract metadata from the 3D file
-        metadata_extraction = parse_3d_file(original_filepath=file_path)
+        if step_path is not None:
+            metadata_extraction = await parse_3d_file(original_filepath=step_path)
+            logger.info(f"Extracted metadata from STEP file: {metadata_extraction}")
+        else:
+            metadata_extraction = await parse_3d_file(original_filepath=file_path)
+
         # Extract geometric features from the 3D object such as volume, dimention, surface
         metadata_aggregation = extract_geometric_features(obj_path)
         metadata = (
@@ -915,12 +1050,15 @@ async def accumulate_and_embedding(input_files, file_list, embedding_dict):
         embedding_dict[obj_path]["image_embedding"] = embeddings["image_embedding"]
         embedding_dict[obj_path]["text_embedding"] = text_embedding
 
+        # Store mapping of original file path to converted obj path
+        converting_store_map[file_path] = obj_path
+
         # if os.environ.get("ENVIRONMENT") == "local":
         #     # save to local file
         #     torch.save(embedding_dict, "embedding_dict.pt")
         #     logger.info("Saved embedding_dict to local file.")
 
-    return all_files, gr.update(choices=all_files), embedding_dict
+    return all_files, gr.update(choices=all_files), embedding_dict, converting_store_map
 
 
 def select_file(filename, file_list):
@@ -932,15 +1070,20 @@ def select_file(filename, file_list):
     return "File not found."
 
 
-def render_3D_object(filepath) -> Tuple[str, str]:
+async def render_3D_object(filepath, converting_store_map) -> Tuple[str, str]:
     _, ext = os.path.splitext(filepath)
     ext = ext.lower()
     if ext in tuple(GRADIO_3D_MODEL_DEFAULT_FORMAT):
         return filepath, filepath
     if ext in tuple(
-        USER_REQUIRE_FORMAT + FREECAD_LOW_LEVEL_FORMAT + FREECAD_NATIVE_FORMAT
+        USER_REQUIRE_FORMAT
+        + FREECAD_LOW_LEVEL_FORMAT
+        + FREECAD_NATIVE_FORMAT
+        + ONSHAPE_SUPPORTED_FORMAT
     ):
-        return convert_to_obj(filepath), filepath
+        if filepath in converting_store_map:
+            return converting_store_map[filepath], filepath
+        return await convert_to_obj(filepath), filepath
     return filepath, filepath
 
 
@@ -953,6 +1096,7 @@ valid_file_types = list(
         + USER_REQUIRE_FORMAT
         + FREECAD_NATIVE_FORMAT
         + FREECAD_LOW_LEVEL_FORMAT
+        + ONSHAPE_SUPPORTED_FORMAT
     )
 )
 valid_file_types = valid_file_types + [t.upper() for t in valid_file_types]
@@ -961,6 +1105,7 @@ with gr.Blocks() as demo:
         file_state = gr.State(sample_files)
         ###################################### !IMPORTANT #############################################################
         embedding_store = gr.State({})  ####### !IMPORTANT. This is in memory vector database ##########################
+        converting_store_map = gr.State({})  ####### !IMPORTANT. This is in memory vector database ##########################
         file_input = gr.File(
             file_count="multiple",
             label="Upload files (You can append more)",
@@ -972,8 +1117,8 @@ with gr.Blocks() as demo:
             file_dropdown = gr.Dropdown(
                 label="Select a file to process", choices=sample_files, interactive=True
             )
-            description_render = gr.Textbox(label="Description", lines=6)
             metadata_render = gr.Textbox(label="Metadata", lines=6)
+            description_render = gr.Textbox(label="Description", lines=6)
         with gr.Column(scale=1):
             model_render = gr.Model3D(label="3D", height=500, interactive=False)
             model_hidden_filepath = gr.Textbox(visible=False)
@@ -1044,8 +1189,8 @@ with gr.Blocks() as demo:
 
     file_input.change(
         fn=accumulate_and_embedding,
-        inputs=[file_input, file_state, embedding_store],
-        outputs=[file_state, file_dropdown, embedding_store],
+        inputs=[file_input, file_state, embedding_store, converting_store_map],
+        outputs=[file_state, file_dropdown, embedding_store, converting_store_map],
     )
     # query button
     query_button.click(
@@ -1069,16 +1214,24 @@ with gr.Blocks() as demo:
     )
     # model query
     model_q_1_btn.click(
-        render_3D_object, model_q_1, [model_render, model_hidden_filepath]
+        render_3D_object,
+        [model_q_1, converting_store_map],
+        [model_render, model_hidden_filepath],
     )
     model_q_2_btn.click(
-        render_3D_object, model_q_2, [model_render, model_hidden_filepath]
+        render_3D_object,
+        [model_q_2, converting_store_map],
+        [model_render, model_hidden_filepath],
     )
     model_q_3_btn.click(
-        render_3D_object, model_q_3, [model_render, model_hidden_filepath]
+        render_3D_object,
+        [model_q_3, converting_store_map],
+        [model_render, model_hidden_filepath],
     )
     model_q_4_btn.click(
-        render_3D_object, model_q_4, [model_render, model_hidden_filepath]
+        render_3D_object,
+        [model_q_4, converting_store_map],
+        [model_render, model_hidden_filepath],
     )
     # sim button
     sim_button.click(
@@ -1101,20 +1254,30 @@ with gr.Blocks() as demo:
     )
     # model similarity
     model_s_1_btn.click(
-        render_3D_object, model_s_1, [model_render, model_hidden_filepath]
+        render_3D_object,
+        [model_s_1, converting_store_map],
+        [model_render, model_hidden_filepath],
     )
     model_s_2_btn.click(
-        render_3D_object, model_s_2, [model_render, model_hidden_filepath]
+        render_3D_object,
+        [model_s_2, converting_store_map],
+        [model_render, model_hidden_filepath],
     )
     model_s_3_btn.click(
-        render_3D_object, model_s_3, [model_render, model_hidden_filepath]
+        render_3D_object,
+        [model_s_3, converting_store_map],
+        [model_render, model_hidden_filepath],
     )
     model_s_4_btn.click(
-        render_3D_object, model_s_4, [model_render, model_hidden_filepath]
+        render_3D_object,
+        [model_s_4, converting_store_map],
+        [model_render, model_hidden_filepath],
     )
     # drop down
     file_dropdown.change(
-        render_3D_object, file_dropdown, [model_render, model_hidden_filepath]
+        render_3D_object,
+        [file_dropdown, converting_store_map],
+        [model_render, model_hidden_filepath],
     )
     # parse metadata
     model_hidden_filepath.change(

onshape/onshape_download.py

@@ -1,6 +1,6 @@
 import requests
 from loguru import logger
-from onshape_base import OnshapeBase
+from onshape.onshape_base import OnshapeBase
 
 
 class OnshapeDownload(OnshapeBase):

onshape/onshape_translation.py

@@ -2,8 +2,8 @@ import time
 
 import requests
 from loguru import logger
-from onshape_base import OnshapeBase
-from onshape_schema import TranslationResponse, TranslationStatusResponse
+from onshape.onshape_base import OnshapeBase
+from onshape.onshape_schema import TranslationResponse, TranslationStatusResponse
 
 
 class OnshapeTranslation(OnshapeBase):