fix: ruff formatting and allow-direct-references for handler-base dep

This commit is contained in:
2026-02-02 08:44:51 -05:00
parent 7c7a147db6
commit 4fbde95eb7
4 changed files with 82 additions and 90 deletions

View File

@@ -8,6 +8,7 @@ Bridges NATS events to workflow engines using handler-base:
3. Monitor execution and publish status updates 3. Monitor execution and publish status updates
4. Publish completion to "ai.pipeline.status.{request_id}" 4. Publish completion to "ai.pipeline.status.{request_id}"
""" """
import logging import logging
from typing import Any, Optional from typing import Any, Optional
from datetime import datetime from datetime import datetime
@@ -137,9 +138,7 @@ class PipelineBridge(Handler):
try: try:
if engine == "argo": if engine == "argo":
run_id = await self._submit_argo( run_id = await self._submit_argo(pipeline["template"], parameters, request_id)
pipeline["template"], parameters, request_id
)
else: else:
run_id = await self._submit_kubeflow( run_id = await self._submit_kubeflow(
pipeline["pipeline_id"], parameters, request_id pipeline["pipeline_id"], parameters, request_id
@@ -155,9 +154,7 @@ class PipelineBridge(Handler):
} }
# Publish status update # Publish status update
await self.nats.publish( await self.nats.publish(f"ai.pipeline.status.{request_id}", result)
f"ai.pipeline.status.{request_id}", result
)
logger.info(f"Pipeline {pipeline_name} submitted: {run_id}") logger.info(f"Pipeline {pipeline_name} submitted: {run_id}")
return result return result
@@ -170,9 +167,7 @@ class PipelineBridge(Handler):
"error": str(e), "error": str(e),
} }
async def _submit_argo( async def _submit_argo(self, template: str, parameters: dict, request_id: str) -> str:
self, template: str, parameters: dict, request_id: str
) -> str:
"""Submit workflow to Argo Workflows.""" """Submit workflow to Argo Workflows."""
with create_span("pipeline.submit.argo") as span: with create_span("pipeline.submit.argo") as span:
if span: if span:
@@ -191,10 +186,7 @@ class PipelineBridge(Handler):
"spec": { "spec": {
"workflowTemplateRef": {"name": template}, "workflowTemplateRef": {"name": template},
"arguments": { "arguments": {
"parameters": [ "parameters": [{"name": k, "value": str(v)} for k, v in parameters.items()]
{"name": k, "value": str(v)}
for k, v in parameters.items()
]
}, },
}, },
} }
@@ -208,9 +200,7 @@ class PipelineBridge(Handler):
result = response.json() result = response.json()
return result["metadata"]["name"] return result["metadata"]["name"]
async def _submit_kubeflow( async def _submit_kubeflow(self, pipeline_id: str, parameters: dict, request_id: str) -> str:
self, pipeline_id: str, parameters: dict, request_id: str
) -> str:
"""Submit run to Kubeflow Pipelines.""" """Submit run to Kubeflow Pipelines."""
with create_span("pipeline.submit.kubeflow") as span: with create_span("pipeline.submit.kubeflow") as span:
if span: if span:
@@ -220,10 +210,7 @@ class PipelineBridge(Handler):
"name": f"{pipeline_id}-{request_id[:8]}", "name": f"{pipeline_id}-{request_id[:8]}",
"pipeline_spec": { "pipeline_spec": {
"pipeline_id": pipeline_id, "pipeline_id": pipeline_id,
"parameters": [ "parameters": [{"name": k, "value": str(v)} for k, v in parameters.items()],
{"name": k, "value": str(v)}
for k, v in parameters.items()
],
}, },
} }

View File

@@ -24,6 +24,9 @@ dev = [
requires = ["hatchling"] requires = ["hatchling"]
build-backend = "hatchling.build" build-backend = "hatchling.build"
[tool.hatch.metadata]
allow-direct-references = true
[tool.hatch.build.targets.wheel] [tool.hatch.build.targets.wheel]
packages = ["."] packages = ["."]
only-include = ["pipeline_bridge.py"] only-include = ["pipeline_bridge.py"]

View File

@@ -1,9 +1,10 @@
""" """
Pytest configuration and fixtures for pipeline-bridge tests. Pytest configuration and fixtures for pipeline-bridge tests.
""" """
import asyncio import asyncio
import os import os
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import MagicMock
import pytest import pytest

View File

@@ -1,6 +1,7 @@
""" """
Unit tests for PipelineBridge handler. Unit tests for PipelineBridge handler.
""" """
import pytest import pytest
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch