From e9208bf5b0df8a05ba571251229cac5687c98504 Mon Sep 17 00:00:00 2001
From: Benjamin Bolte <ben@bolte.cc>
Date: Fri, 30 Aug 2024 14:35:36 -0700
Subject: [PATCH] better config options

---
 frontend/src/gen/api.ts      | 24 +++++++++++++++++++++++-
 store/app/routers/onshape.py | 30 +++++++++++++++++++++++++-----
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/frontend/src/gen/api.ts b/frontend/src/gen/api.ts
index 163c7459..38621651 100644
--- a/frontend/src/gen/api.ts
+++ b/frontend/src/gen/api.ts
@@ -700,6 +700,17 @@ export interface components {
             /** Api Key */
             api_key: string;
         };
+        /** Body_pull_onshape_document_onshape_pull__listing_id__get */
+        Body_pull_onshape_document_onshape_pull__listing_id__get: {
+            /** Suffix To Joint Effort */
+            suffix_to_joint_effort?: {
+                [key: string]: number;
+            } | null;
+            /** Suffix To Joint Velocity */
+            suffix_to_joint_velocity?: {
+                [key: string]: number;
+            } | null;
+        };
         /** Body_upload_artifacts_upload__listing_id__post */
         Body_upload_artifacts_upload__listing_id__post: {
             /** Files */
@@ -1892,6 +1903,13 @@ export interface operations {
         parameters: {
             query?: {
                 token?: string | null;
+                default_prismatic_joint_effort?: number;
+                default_prismatic_joint_velocity?: number;
+                default_revolute_joint_effort?: number;
+                default_revolute_joint_velocity?: number;
+                voxel_size?: number;
+                convex_collision_meshes?: boolean;
+                add_mjcf?: boolean;
             };
             header?: never;
             path: {
@@ -1899,7 +1917,11 @@ export interface operations {
             };
             cookie?: never;
         };
-        requestBody?: never;
+        requestBody?: {
+            content: {
+                "application/json": components["schemas"]["Body_pull_onshape_document_onshape_pull__listing_id__get"];
+            };
+        };
         responses: {
             /** @description Successful Response */
             200: {
diff --git a/store/app/routers/onshape.py b/store/app/routers/onshape.py
index 8ce94ad0..ebaae8fb 100644
--- a/store/app/routers/onshape.py
+++ b/store/app/routers/onshape.py
@@ -6,6 +6,7 @@
 
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 from fastapi.responses import StreamingResponse
+from kol.onshape.config import ConverterConfig
 from pydantic.main import BaseModel
 
 from store.app.db import Crud
@@ -40,9 +41,7 @@ async def set_onshape_document(
 
 
 async def pull_onshape_document_generator(
-    listing_id: str,
-    user: User,
-    crud: Crud,
+    listing_id: str, user: User, crud: Crud, *, config: ConverterConfig | None = None
 ) -> AsyncIterable[str]:
     # Gets the listing and makes sure the user has permission to write to it.
     listing = await crud.get_listing(listing_id)
@@ -53,7 +52,7 @@ async def pull_onshape_document_generator(
     if not await can_write_listing(user, listing):
         raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="User cannot write to this listing")
     yield f"event: message\ndata: {json.dumps({'message': 'Starting download', 'level': 'success'})}\n\n"
-    async for event in crud.download_onshape_document(listing, onshape_url):
+    async for event in crud.download_onshape_document(listing, onshape_url, config=config):
         yield event
     yield "event: finish\ndata: finish\n\n"
 
@@ -64,6 +63,15 @@ async def pull_onshape_document(
     request: Request,
     crud: Annotated[Crud, Depends(Crud.get)],
     token: str | None = None,
+    default_prismatic_joint_effort: float = 80.0,
+    default_prismatic_joint_velocity: float = 5.0,
+    default_revolute_joint_effort: float = 80.0,
+    default_revolute_joint_velocity: float = 5.0,
+    suffix_to_joint_effort: dict[str, float] | None = None,
+    suffix_to_joint_velocity: dict[str, float] | None = None,
+    voxel_size: float = 0.002,
+    convex_collision_meshes: bool = False,
+    add_mjcf: bool = True,
 ) -> StreamingResponse:
     # Because the default EventStream implementation doesn't provide an easy
     # way to pass the token in the header, we have to pass it as a query
@@ -77,7 +85,19 @@ async def pull_onshape_document(
         raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Permission denied")
     user = await crud.get_user(api_key.user_id, throw_if_missing=True)
 
+    config = ConverterConfig(
+        default_prismatic_joint_effort=default_prismatic_joint_effort,
+        default_prismatic_joint_velocity=default_prismatic_joint_velocity,
+        default_revolute_joint_effort=default_revolute_joint_effort,
+        default_revolute_joint_velocity=default_revolute_joint_velocity,
+        suffix_to_joint_effort=suffix_to_joint_effort or {},
+        suffix_to_joint_velocity=suffix_to_joint_velocity or {},
+        voxel_size=voxel_size,
+        convex_collision_meshes=convex_collision_meshes,
+        add_mjcf=add_mjcf,
+    )
+
     return StreamingResponse(
-        content=pull_onshape_document_generator(listing_id, user, crud),
+        content=pull_onshape_document_generator(listing_id, user, crud, config=config),
         media_type="text/event-stream",
     )