pyxecm 2.0.2__py3-none-any.whl → 2.0.4__py3-none-any.whl

pyxecm/customizer/api/terminal/router.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import contextlib
+import json
 import os
 import pty
 import signal
@@ -15,13 +16,16 @@ router = APIRouter(tags=["terminal"])
 
 
 @router.websocket("/ws/terminal")
-async def ws_terminal(websocket: WebSocket, pod: str = Query(...), command: str = Query(...)) -> None:
+async def ws_terminal(
+    websocket: WebSocket, pod: str = Query(...), command: str = Query(...), container: str = Query(None)
+) -> None:
     """Websocket to connect to a shell session in a pod.
 
     Args:
         websocket (WebSocket): WebSocket to connect to the shell session.
         pod (str): pod name to connect to.
         command (str): command to be executed.
+        container (str, optional): container name to connect to.
 
     """
     await websocket.accept()
@@ -30,7 +34,7 @@ async def ws_terminal(websocket: WebSocket, pod: str = Query(...), command: str
         # Wait for the first message to be the token
         token = await websocket.receive_text()
 
-        user = await get_current_user(token)
+        user = await get_current_user(token=token, api_key=token)
         authrorized = await get_authorized_user(user)
 
         if not authrorized:
@@ -45,11 +49,17 @@ async def ws_terminal(websocket: WebSocket, pod: str = Query(...), command: str
     except WebSocketDisconnect:
         return
 
-    process = ["bash"] if pod == "customizer" else ["kubectl", "exec", "-it", pod, "--", command]
+    container = ["-c", container] if container else []
+
+    process = [command] if pod == "customizer" else ["kubectl", "exec", "-it", pod, *container, "--", command]
 
     pid, fd = pty.fork()
+    import fcntl
+    import struct
+    import termios
+
     if pid == 0:
-        subprocess.run(process, check=False)  # noqa: ASYNC221
+        process = subprocess.run(process, check=False)  # noqa: ASYNC221
 
     async def read_from_pty() -> None:
         loop = asyncio.get_event_loop()
@@ -64,24 +74,39 @@ async def ws_terminal(websocket: WebSocket, pod: str = Query(...), command: str
         try:
             while True:
                 data = await websocket.receive_text()
-                os.write(fd, data.encode())
+                if data.startswith("{") and data.endswith("}"):
+                    message = json.loads(data)
+                    if message.get("type") == "resize":
+                        rows = message.get("rows")
+                        cols = message.get("cols")
+
+                        packed_size = struct.pack("HHHH", rows, cols, 0, 0)
+                        fcntl.ioctl(fd, termios.TIOCSWINSZ, packed_size)
+
+                else:
+                    os.write(fd, data.encode())
         except Exception:  # noqa: S110
             pass
 
-    # Launch read/write tasks
-    read_task = asyncio.create_task(read_from_pty())
-    write_task = asyncio.create_task(write_to_pty())
+    try:
+        # Launch read/write tasks
+        read_task = asyncio.create_task(read_from_pty())
+        write_task = asyncio.create_task(write_to_pty())
+
+        done, pending = await asyncio.wait([read_task, write_task], return_when=asyncio.FIRST_COMPLETED)
 
-    done, pending = await asyncio.wait([read_task, write_task], return_when=asyncio.FIRST_COMPLETED)
+        # Cancel other task
+        for task in pending:
+            task.cancel()
+    finally:
+        with contextlib.suppress(Exception):
+            await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
 
-    # Cancel other task
-    for task in pending:
-        task.cancel()
+        with contextlib.suppress(Exception):
+            process.terminate()
 
-    try:  # noqa: SIM105
-        os.kill(pid, signal.SIGKILL)
-    except ProcessLookupError:
-        pass  # Already exited
+        with contextlib.suppress(Exception):
+            os.kill(pid, signal.SIGKILL)
 
-    with contextlib.suppress(Exception):
-        os.close(fd)
+        with contextlib.suppress(Exception):
+            os.close(fd)

pyxecm/customizer/api/v1_csai/models.py

@@ -0,0 +1,18 @@
+"""Define Models for Payload."""
+
+from pydantic import BaseModel
+
+
+class CSAIEmbedMetadata(BaseModel):
+    """Defines Data Model for embeding metadata for documents and workspaces."""
+
+    node_id: int = None
+    crawl: bool = False
+    wait_for_completion: bool = False
+    message_override: dict = None
+    timeout: float = 30.0
+    document_metadata: bool = False
+    images: bool = False
+    image_prompt: str = ""
+    workspace_metadata: bool = True
+    remove_existing: bool = False

pyxecm/customizer/api/v1_csai/router.py

@@ -9,15 +9,40 @@ from fastapi.responses import JSONResponse
 
 from pyxecm.customizer.api.auth.functions import get_authorized_user
 from pyxecm.customizer.api.auth.models import User
-from pyxecm.customizer.api.common.functions import get_k8s_object, get_settings
+from pyxecm.customizer.api.common.functions import get_k8s_object, get_otcs_object, get_settings
 from pyxecm.customizer.api.settings import CustomizerAPISettings
+from pyxecm.customizer.api.v1_csai.models import CSAIEmbedMetadata
 from pyxecm.customizer.k8s import K8s
+from pyxecm.otcs import OTCS
 
 router = APIRouter(prefix="/api/v1/csai", tags=["csai"])
 
 logger = logging.getLogger("pyxecm.customizer.api.v1_csai")
 
 
+@router.post("/metadata")
+def embed_metadata(
+    user: Annotated[User, Depends(get_authorized_user)],  # noqa: ARG001
+    otcs_object: Annotated[OTCS, Depends(get_otcs_object)],
+    body: Annotated[CSAIEmbedMetadata, Body()],
+) -> JSONResponse:
+    """Embed the Metadata of the given objects.
+
+    Args:
+        user (Annotated[User, Depends): User required for authentication
+        otcs_object (Annotated[OTCS, Depends(get_otcs_object)]): OTCS object to interact with OTCS
+        body (Annotated[CSAIEmbedMetadata, Body): Request body
+
+    Returns:
+        JSONResponse: JSONResponse with success=true/false
+
+    """
+
+    success = otcs_object.aviator_embed_metadata(**body.model_dump())
+
+    return JSONResponse({"success": success})
+
+
 @router.get("")
 def get_csai_config_data(
     user: Annotated[User, Depends(get_authorized_user)],

pyxecm/customizer/api/v1_otcs/router.py

@@ -35,10 +35,15 @@ async def put_otcs_logs(
 
     if "all" in hosts:
         hosts = []
-        for sts in ["otcs-admin", "otcs-frontend", "otcs-backend-search"]:
+        for sts in ["otcs-admin", "otcs-frontend", "otcs-backend-search", "otcs-da"]:
             try:
-                sts_replicas = k8s_object.get_stateful_set_scale(sts).status.replicas
-                hosts.extend([f"{sts}-{i}" for i in range(sts_replicas)])
+                sts_replicas = k8s_object.get_stateful_set_scale(sts)
+
+                if sts_replicas is None:
+                    logger.debug("Cannot get statefulset {sts}")
+                    continue
+
+                hosts.extend([f"{sts}-{i}" for i in range(sts_replicas.status.replicas)])
             except Exception as e:
                 raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR) from e
 
@@ -118,10 +123,15 @@ async def get_otcs_log_files(
     response = {"status": {host: bool(otcs_logs_lock[host].locked()) for host in otcs_logs_lock}, "files": files}
 
     # Extend response with all hosts
-    for sts in ["otcs-admin", "otcs-frontend", "otcs-backend-search"]:
+    for sts in ["otcs-admin", "otcs-frontend", "otcs-backend-search", "otcs-da"]:
         try:
-            sts_replicas = k8s_object.get_stateful_set_scale(sts).status.replicas
-            for i in range(sts_replicas):
+            sts_replicas = k8s_object.get_stateful_set_scale(sts)
+
+            if sts_replicas is None:
+                logger.debug("Cannot get statefulset {sts}")
+                continue
+
+            for i in range(sts_replicas.status.replicas):
                 host = f"{sts}-{i}"
 
                 if host in otcs_logs_lock:

pyxecm/customizer/api/v1_payload/functions.py

@@ -129,16 +129,22 @@ def import_payload(
         payload_dir = payload
 
     if payload_dir is None:
-        import_payload_file(payload, enabled, dependencies)
+        try:
+            import_payload_file(payload, enabled, dependencies)
+        except PayloadImportError as exc:
+            logger.error(exc)
+            logger.debug(exc, exc_info=True)
         return
+
     elif not os.path.isdir(payload_dir):
         return
 
     for filename in sorted(os.listdir(payload_dir)):
         try:
             import_payload_file(os.path.join(payload_dir, filename), enabled, dependencies)
-        except PayloadImportError:
-            logger.error("Payload import failed")
+        except PayloadImportError as exc:
+            logger.error(exc)
+            logger.debug(exc, exc_info=True)
 
 
 def prepare_dependencies(dependencies: list) -> list | None: