fix: ruff formatting, allow-direct-references, and noqa for Kubeflow pipeline vars

2026-02-02 08:44:14 -05:00
parent 58465b77d8
commit 0462412353
5 changed files with 208 additions and 237 deletions
--- a/pipelines/voice_pipeline.py
+++ b/pipelines/voice_pipeline.py
@@ -14,13 +14,9 @@ from kfp import dsl
 from kfp import compiler


-@dsl.component(
-    base_image="python:3.13-slim",
-    packages_to_install=["httpx"]
-)
+@dsl.component(base_image="python:3.13-slim", packages_to_install=["httpx"])
 def transcribe_audio(
-    audio_b64: str,
-    whisper_url: str = "http://whisper-predictor.ai-ml.svc.cluster.local"
+    audio_b64: str, whisper_url: str = "http://whisper-predictor.ai-ml.svc.cluster.local"
 ) -> str:
    """Transcribe audio using Whisper STT service."""
    import base64
@@ -32,43 +28,35 @@ def transcribe_audio(
        response = client.post(
            f"{whisper_url}/v1/audio/transcriptions",
            files={"file": ("audio.wav", audio_bytes, "audio/wav")},
-            data={"model": "whisper-large-v3", "language": "en"}
+            data={"model": "whisper-large-v3", "language": "en"},
        )
        result = response.json()

    return result.get("text", "")


-@dsl.component(
-    base_image="python:3.13-slim",
-    packages_to_install=["httpx"]
-)
+@dsl.component(base_image="python:3.13-slim", packages_to_install=["httpx"])
 def generate_embeddings(
-    text: str,
-    embeddings_url: str = "http://embeddings-predictor.ai-ml.svc.cluster.local"
+    text: str, embeddings_url: str = "http://embeddings-predictor.ai-ml.svc.cluster.local"
 ) -> list:
    """Generate embeddings for RAG retrieval."""
    import httpx

    with httpx.Client(timeout=60.0) as client:
        response = client.post(
-            f"{embeddings_url}/embeddings",
-            json={"input": text, "model": "bge-small-en-v1.5"}
+            f"{embeddings_url}/embeddings", json={"input": text, "model": "bge-small-en-v1.5"}
        )
        result = response.json()

    return result["data"][0]["embedding"]


-@dsl.component(
-    base_image="python:3.13-slim",
-    packages_to_install=["pymilvus"]
-)
+@dsl.component(base_image="python:3.13-slim", packages_to_install=["pymilvus"])
 def retrieve_context(
    embedding: list,
    milvus_host: str = "milvus.ai-ml.svc.cluster.local",
    collection_name: str = "knowledge_base",
-    top_k: int = 5
+    top_k: int = 5,
 ) -> list:
    """Retrieve relevant documents from Milvus vector database."""
    from pymilvus import connections, Collection, utility
@@ -86,30 +74,29 @@ def retrieve_context(
        anns_field="embedding",
        param={"metric_type": "COSINE", "params": {"nprobe": 10}},
        limit=top_k,
-        output_fields=["text", "source"]
+        output_fields=["text", "source"],
    )

    documents = []
    for hits in results:
        for hit in hits:
-            documents.append({
-                "text": hit.entity.get("text"),
-                "source": hit.entity.get("source"),
-                "score": hit.distance
-            })
+            documents.append(
+                {
+                    "text": hit.entity.get("text"),
+                    "source": hit.entity.get("source"),
+                    "score": hit.distance,
+                }
+            )

    return documents


-@dsl.component(
-    base_image="python:3.13-slim",
-    packages_to_install=["httpx"]
-)
+@dsl.component(base_image="python:3.13-slim", packages_to_install=["httpx"])
 def rerank_documents(
    query: str,
    documents: list,
    reranker_url: str = "http://reranker-predictor.ai-ml.svc.cluster.local",
-    top_k: int = 3
+    top_k: int = 3,
 ) -> list:
    """Rerank documents using BGE reranker."""
    import httpx
@@ -123,30 +110,25 @@ def rerank_documents(
            json={
                "query": query,
                "documents": [doc["text"] for doc in documents],
-                "model": "bge-reranker-v2-m3"
-            }
+                "model": "bge-reranker-v2-m3",
+            },
        )
        result = response.json()

    # Sort by rerank score
    reranked = sorted(
-        zip(documents, result.get("scores", [0] * len(documents))),
-        key=lambda x: x[1],
-        reverse=True
+        zip(documents, result.get("scores", [0] * len(documents))), key=lambda x: x[1], reverse=True
    )[:top_k]

    return [doc for doc, score in reranked]


-@dsl.component(
-    base_image="python:3.13-slim",
-    packages_to_install=["httpx"]
-)
+@dsl.component(base_image="python:3.13-slim", packages_to_install=["httpx"])
 def generate_response(
    query: str,
    context: list,
    vllm_url: str = "http://llm-draft.ai-ml.svc.cluster.local:8000",
-    model: str = "mistralai/Mistral-7B-Instruct-v0.3"
+    model: str = "mistralai/Mistral-7B-Instruct-v0.3",
 ) -> str:
    """Generate response using vLLM."""
    import httpx
@@ -164,31 +146,22 @@ Keep responses concise and natural for speech synthesis."""

    messages = [
        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": user_content}
+        {"role": "user", "content": user_content},
    ]

    with httpx.Client(timeout=180.0) as client:
        response = client.post(
            f"{vllm_url}/v1/chat/completions",
-            json={
-                "model": model,
-                "messages": messages,
-                "max_tokens": 512,
-                "temperature": 0.7
-            }
+            json={"model": model, "messages": messages, "max_tokens": 512, "temperature": 0.7},
        )
        result = response.json()

    return result["choices"][0]["message"]["content"]


-@dsl.component(
-    base_image="python:3.13-slim",
-    packages_to_install=["httpx"]
-)
+@dsl.component(base_image="python:3.13-slim", packages_to_install=["httpx"])
 def synthesize_speech(
-    text: str,
-    tts_url: str = "http://tts-predictor.ai-ml.svc.cluster.local"
+    text: str, tts_url: str = "http://tts-predictor.ai-ml.svc.cluster.local"
 ) -> str:
    """Convert text to speech using TTS service."""
    import base64
@@ -197,11 +170,7 @@ def synthesize_speech(
    with httpx.Client(timeout=120.0) as client:
        response = client.post(
            f"{tts_url}/v1/audio/speech",
-            json={
-                "input": text,
-                "voice": "en_US-lessac-high",
-                "response_format": "wav"
-            }
+            json={"input": text, "voice": "en_US-lessac-high", "response_format": "wav"},
        )
        audio_b64 = base64.b64encode(response.content).decode("utf-8")

@@ -210,12 +179,9 @@ def synthesize_speech(

@dsl.pipeline(
    name="voice-assistant-rag-pipeline",
-    description="End-to-end voice assistant with RAG: STT -> Embeddings -> Milvus -> Rerank -> LLM -> TTS"
+    description="End-to-end voice assistant with RAG: STT -> Embeddings -> Milvus -> Rerank -> LLM -> TTS",
 )
-def voice_assistant_pipeline(
-    audio_b64: str,
-    collection_name: str = "knowledge_base"
-):
+def voice_assistant_pipeline(audio_b64: str, collection_name: str = "knowledge_base"):
    """
    Voice Assistant Pipeline with RAG

@@ -233,44 +199,28 @@ def voice_assistant_pipeline(
    embed_task.set_caching_options(enable_caching=True)

    # Step 3: Retrieve context from Milvus
-    retrieve_task = retrieve_context(
-        embedding=embed_task.output,
-        collection_name=collection_name
-    )
+    retrieve_task = retrieve_context(embedding=embed_task.output, collection_name=collection_name)

    # Step 4: Rerank documents
-    rerank_task = rerank_documents(
-        query=transcribe_task.output,
-        documents=retrieve_task.output
-    )
+    rerank_task = rerank_documents(query=transcribe_task.output, documents=retrieve_task.output)

    # Step 5: Generate response with context
-    llm_task = generate_response(
-        query=transcribe_task.output,
-        context=rerank_task.output
-    )
+    llm_task = generate_response(query=transcribe_task.output, context=rerank_task.output)

    # Step 6: Synthesize speech
-    tts_task = synthesize_speech(text=llm_task.output)
+    tts_task = synthesize_speech(text=llm_task.output)  # noqa: F841


-@dsl.pipeline(
-    name="text-to-speech-pipeline",
-    description="Simple text to speech pipeline"
-)
+@dsl.pipeline(name="text-to-speech-pipeline", description="Simple text to speech pipeline")
 def text_to_speech_pipeline(text: str):
    """Simple TTS pipeline for testing."""
-    tts_task = synthesize_speech(text=text)
+    tts_task = synthesize_speech(text=text)  # noqa: F841


@dsl.pipeline(
-    name="rag-query-pipeline", 
-    description="RAG query pipeline: Embed -> Retrieve -> Rerank -> LLM"
+    name="rag-query-pipeline", description="RAG query pipeline: Embed -> Retrieve -> Rerank -> LLM"
 )
-def rag_query_pipeline(
-    query: str,
-    collection_name: str = "knowledge_base"
-):
+def rag_query_pipeline(query: str, collection_name: str = "knowledge_base"):
    """
    RAG Query Pipeline (text input, no voice)

@@ -282,21 +232,14 @@ def rag_query_pipeline(
    embed_task = generate_embeddings(text=query)

    # Retrieve from Milvus
-    retrieve_task = retrieve_context(
-        embedding=embed_task.output,
-        collection_name=collection_name
-    )
+    retrieve_task = retrieve_context(embedding=embed_task.output, collection_name=collection_name)

    # Rerank
-    rerank_task = rerank_documents(
-        query=query,
-        documents=retrieve_task.output
-    )
+    rerank_task = rerank_documents(query=query, documents=retrieve_task.output)

    # Generate response
-    llm_task = generate_response(
-        query=query,
-        context=rerank_task.output
+    llm_task = generate_response(  # noqa: F841
+        query=query, context=rerank_task.output
    )


--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,6 +22,9 @@ dev = [
 requires = ["hatchling"]
 build-backend = "hatchling.build"

+[tool.hatch.metadata]
+allow-direct-references = true
+
 [tool.hatch.build.targets.wheel]
 packages = ["."]
 only-include = ["voice_assistant.py"]
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,11 +1,11 @@
 """
 Pytest configuration and fixtures for voice-assistant tests.
 """
+
 import asyncio
 import base64
 import os
-from typing import AsyncGenerator
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import MagicMock, patch

 import pytest

@@ -29,21 +29,54 @@ def event_loop():
 def sample_audio_b64():
    """Sample base64 encoded audio for testing."""
    # 16-bit PCM silence (44 bytes header + 1000 samples)
-    wav_header = bytes([
-        0x52, 0x49, 0x46, 0x46,  # "RIFF"
-        0x24, 0x08, 0x00, 0x00,  # File size
-        0x57, 0x41, 0x56, 0x45,  # "WAVE"
-        0x66, 0x6D, 0x74, 0x20,  # "fmt "
-        0x10, 0x00, 0x00, 0x00,  # Chunk size
-        0x01, 0x00,              # PCM format
-        0x01, 0x00,              # Mono
-        0x80, 0x3E, 0x00, 0x00,  # Sample rate (16000)
-        0x00, 0x7D, 0x00, 0x00,  # Byte rate
-        0x02, 0x00,              # Block align
-        0x10, 0x00,              # Bits per sample
-        0x64, 0x61, 0x74, 0x61,  # "data"
-        0x00, 0x08, 0x00, 0x00,  # Data size
-    ])
+    wav_header = bytes(
+        [
+            0x52,
+            0x49,
+            0x46,
+            0x46,  # "RIFF"
+            0x24,
+            0x08,
+            0x00,
+            0x00,  # File size
+            0x57,
+            0x41,
+            0x56,
+            0x45,  # "WAVE"
+            0x66,
+            0x6D,
+            0x74,
+            0x20,  # "fmt "
+            0x10,
+            0x00,
+            0x00,
+            0x00,  # Chunk size
+            0x01,
+            0x00,  # PCM format
+            0x01,
+            0x00,  # Mono
+            0x80,
+            0x3E,
+            0x00,
+            0x00,  # Sample rate (16000)
+            0x00,
+            0x7D,
+            0x00,
+            0x00,  # Byte rate
+            0x02,
+            0x00,  # Block align
+            0x10,
+            0x00,  # Bits per sample
+            0x64,
+            0x61,
+            0x74,
+            0x61,  # "data"
+            0x00,
+            0x08,
+            0x00,
+            0x00,  # Data size
+        ]
+    )
    silence = bytes([0x00] * 2048)
    return base64.b64encode(wav_header + silence).decode()

@@ -96,13 +129,14 @@ def mock_voice_request(sample_audio_b64):
@pytest.fixture
 def mock_clients():
    """Mock all service clients."""
-    with patch("voice_assistant.STTClient") as stt, \
-         patch("voice_assistant.EmbeddingsClient") as embeddings, \
-         patch("voice_assistant.RerankerClient") as reranker, \
-         patch("voice_assistant.LLMClient") as llm, \
-         patch("voice_assistant.TTSClient") as tts, \
-         patch("voice_assistant.MilvusClient") as milvus:
-        
+    with (
+        patch("voice_assistant.STTClient") as stt,
+        patch("voice_assistant.EmbeddingsClient") as embeddings,
+        patch("voice_assistant.RerankerClient") as reranker,
+        patch("voice_assistant.LLMClient") as llm,
+        patch("voice_assistant.TTSClient") as tts,
+        patch("voice_assistant.MilvusClient") as milvus,
+    ):
        yield {
            "stt": stt,
            "embeddings": embeddings,
--- a/tests/test_voice_assistant.py
+++ b/tests/test_voice_assistant.py
@@ -1,9 +1,9 @@
 """
 Unit tests for VoiceAssistant handler.
 """
-import base64
+
 import pytest
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, patch

 # Import after environment is set up in conftest
 from voice_assistant import VoiceAssistant, VoiceSettings
@@ -31,10 +31,7 @@ class TestVoiceSettings:
        monkeypatch.setenv("RAG_COLLECTION", "custom_collection")

        # Note: Would need to re-instantiate settings to pick up env vars
-        settings = VoiceSettings(
-            rag_top_k=20,
-            rag_collection="custom_collection"
-        )
+        settings = VoiceSettings(rag_top_k=20, rag_collection="custom_collection")

        assert settings.rag_top_k == 20
        assert settings.rag_collection == "custom_collection"
@@ -46,13 +43,14 @@ class TestVoiceAssistant:
    @pytest.fixture
    def handler(self):
        """Create handler with mocked clients."""
-        with patch("voice_assistant.STTClient"), \
-             patch("voice_assistant.EmbeddingsClient"), \
-             patch("voice_assistant.RerankerClient"), \
-             patch("voice_assistant.LLMClient"), \
-             patch("voice_assistant.TTSClient"), \
-             patch("voice_assistant.MilvusClient"):
-            
+        with (
+            patch("voice_assistant.STTClient"),
+            patch("voice_assistant.EmbeddingsClient"),
+            patch("voice_assistant.RerankerClient"),
+            patch("voice_assistant.LLMClient"),
+            patch("voice_assistant.TTSClient"),
+            patch("voice_assistant.MilvusClient"),
+        ):
            handler = VoiceAssistant()

            # Setup mock clients
@@ -167,13 +165,14 @@ class TestVoiceAssistant:
    @pytest.mark.asyncio
    async def test_setup_initializes_clients(self):
        """Test that setup initializes all clients."""
-        with patch("voice_assistant.STTClient") as stt_cls, \
-             patch("voice_assistant.EmbeddingsClient") as emb_cls, \
-             patch("voice_assistant.RerankerClient") as rer_cls, \
-             patch("voice_assistant.LLMClient") as llm_cls, \
-             patch("voice_assistant.TTSClient") as tts_cls, \
-             patch("voice_assistant.MilvusClient") as mil_cls:
-            
+        with (
+            patch("voice_assistant.STTClient") as stt_cls,
+            patch("voice_assistant.EmbeddingsClient") as emb_cls,
+            patch("voice_assistant.RerankerClient") as rer_cls,
+            patch("voice_assistant.LLMClient") as llm_cls,
+            patch("voice_assistant.TTSClient") as tts_cls,
+            patch("voice_assistant.MilvusClient") as mil_cls,
+        ):
            mil_cls.return_value.connect = AsyncMock()

            handler = VoiceAssistant()
--- a/voice_assistant.py
+++ b/voice_assistant.py
@@ -12,6 +12,7 @@ End-to-end voice assistant pipeline using handler-base:
 7. Synthesize speech with XTTS
 8. Publish result to NATS "voice.response.{request_id}"
 """
+
 import base64
 import logging
 from typing import Any, Optional
@@ -167,8 +168,7 @@ class VoiceAssistant(Handler):

            if self.voice_settings.include_sources:
                result["sources"] = [
-                    {"text": d["document"][:200], "score": d["score"]}
-                    for d in reranked[:3]
+                    {"text": d["document"][:200], "score": d["score"]} for d in reranked[:3]
                ]

            logger.info(f"Completed voice request {request_id}")
@@ -179,9 +179,7 @@ class VoiceAssistant(Handler):

            return result

-    async def _transcribe(
-        self, audio: bytes, language: Optional[str]
-    ) -> dict:
+    async def _transcribe(self, audio: bytes, language: Optional[str]) -> dict:
        """Transcribe audio to text."""
        with create_span("voice.stt"):
            return await self.stt.transcribe(audio, language=language)
@@ -191,9 +189,7 @@ class VoiceAssistant(Handler):
        with create_span("voice.embedding"):
            return await self.embeddings.embed_single(text)

-    async def _search_context(
-        self, embedding: list[float], collection: str
-    ) -> list[dict]:
+    async def _search_context(self, embedding: list[float], collection: str) -> list[dict]:
        """Search Milvus for relevant documents."""
        with create_span("voice.search"):
            return await self.milvus.search_with_texts(
@@ -202,9 +198,7 @@ class VoiceAssistant(Handler):
                text_field="text",
            )

-    async def _rerank_documents(
-        self, query: str, documents: list[dict]
-    ) -> list[dict]:
+    async def _rerank_documents(self, query: str, documents: list[dict]) -> list[dict]:
        """Rerank documents by relevance."""
        with create_span("voice.rerank"):
            texts = [d.get("text", "") for d in documents]
@@ -224,9 +218,7 @@ class VoiceAssistant(Handler):
    async def _synthesize_speech(self, text: str) -> str:
        """Synthesize speech and return base64."""
        with create_span("voice.tts"):
-            audio_bytes = await self.tts.synthesize(
-                text, language=self.voice_settings.tts_language
-            )
+            audio_bytes = await self.tts.synthesize(text, language=self.voice_settings.tts_language)
            return base64.b64encode(audio_bytes).decode()