chore: fix runtime lint

cristipufu · cristipufu · commit 6c822a445090 · 2025-06-13T21:12:44.000+03:00
diff --git a/src/uipath_llamaindex/_cli/_runtime/_hitl.py b/src/uipath_llamaindex/_cli/_runtime/_hitl.py
@@ -90,7 +90,7 @@ async def read(cls, resume_trigger: UiPathResumeTrigger) -> Optional[str]:
                     return job.output_arguments
 
             case UiPathResumeTriggerType.API:
-                if resume_trigger.api_resume.inbox_id:
+                if resume_trigger.api_resume and resume_trigger.api_resume.inbox_id:
                     return await _get_api_payload(resume_trigger.api_resume.inbox_id)
 
             case _:
diff --git a/src/uipath_llamaindex/_cli/_runtime/_runtime.py b/src/uipath_llamaindex/_cli/_runtime/_runtime.py
@@ -12,7 +12,10 @@
     JsonPickleSerializer,
 )
 from llama_index.core.workflow.handler import WorkflowHandler
-from openinference.instrumentation.llama_index import LlamaIndexInstrumentor, get_current_span
+from openinference.instrumentation.llama_index import (
+    LlamaIndexInstrumentor,
+    get_current_span,
+)
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
@@ -23,14 +26,13 @@
     UiPathRuntimeResult,
     UiPathRuntimeStatus,
 )
+from uipath.tracing import TracingManager
 
 from .._tracing._oteladapter import LlamaIndexExporter
 from ._context import UiPathLlamaIndexRuntimeContext
 from ._exception import UiPathLlamaIndexRuntimeError
 from ._hitl import HitlProcessor, HitlReader
 
-from uipath.tracing import TracingManager
-
 logger = logging.getLogger(__name__)
 
 
@@ -74,14 +76,20 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                 if os.path.exists(self.state_file_path):
                     os.remove(self.state_file_path)
 
+            if self.context.workflow is None:
+                return None
+
             start_event_class = self.context.workflow._start_event_class
-            ev = start_event_class(**self.context.input_json)
+            ev = start_event_class(**(self.context.input_json or {}))
             await self.load_workflow_context()
 
+            if self.context.workflow_context is None:
+                return None
+
             handler: WorkflowHandler = self.context.workflow.run(
                 start_event=ev if self.context.resume else None,
                 ctx=self.context.workflow_context,
-                **self.context.input_json,
+                **(self.context.input_json or {}),
             )
 
             resume_trigger: Optional[UiPathResumeTrigger] = None
@@ -94,9 +102,10 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                     if self.context.resume and not response_applied:
                         # If we are resuming, we need to apply the response to the event stream.
                         response_applied = True
-                        self.context.workflow_context.send_event(
-                            await self.get_response_event()
-                        )
+                        response_event = await self.get_response_event()
+                        if response_event:
+                            # If we have a response event, send it to the workflow context.
+                            self.context.workflow_context.send_event(response_event)
                     else:
                         resume_trigger = await hitl_processor.create_resume_trigger()
                         break
@@ -244,6 +253,9 @@ async def load_workflow_context(self):
         """
         logger.debug(f"Resumed: {self.context.resume} Input: {self.context.input_json}")
 
+        if self.context.workflow is None:
+            return
+
         if not self.context.resume:
             self.context.workflow_context = Context(self.context.workflow)
             return
@@ -277,7 +289,7 @@ async def get_response_event(self) -> Optional[HumanResponseEvent]:
         """
         if self.context.input_json:
             # If input_json is provided, use it to create a HumanResponseEvent
-            return HumanResponseEvent(**self.context.input_json)
+            return HumanResponseEvent(**(self.context.input_json or {}))
         # If resumed_trigger is set, fetch the feedback
         if self.context.resumed_trigger:
             feedback = await HitlReader.read(self.context.resumed_trigger)
diff --git a/src/uipath_llamaindex/_cli/cli_init.py b/src/uipath_llamaindex/_cli/cli_init.py
@@ -5,7 +5,7 @@
 from typing import Any, Dict
 
 from llama_index.core.workflow import StopEvent, Workflow
-from llama_index.core.workflow.drawing import StepConfig
+from llama_index.core.workflow.drawing import StepConfig  # type: ignore
 from llama_index.core.workflow.events import (
     HumanResponseEvent,
     InputRequiredEvent,
@@ -129,11 +129,13 @@ def draw_all_possible_flows_mermaid(
 
     # Track event types to avoid duplicates
     event_types = {}
+    current_stop_event = (
+        None  # Only one kind of `StopEvent` is allowed in a `Workflow`.
+    )
+    step_config: StepConfig | None = None
 
-    # Only one kind of `StopEvent` is allowed in a `Workflow`.
-    current_stop_event = None
     for _, step_func in steps.items():
-        step_config: StepConfig = getattr(step_func, "__step_config", None)
+        step_config = getattr(step_func, "__step_config", None)
         if step_config is None:
             continue
 
@@ -227,12 +229,13 @@ def draw_all_possible_flows_mermaid(
             event_id = f"event_{clean_id(event_name)}"
 
             if step_name == "_done" and issubclass(event_type, StopEvent):
-                stop_event_name = current_stop_event.__name__
-                stop_event_id = f"event_{clean_id(stop_event_name)}"
-                edge = f"{stop_event_id} --> {step_id}"
-                if edge not in edges:
-                    edges.add(edge)
-                    mermaid_diagram.append(f"    {edge}")
+                if current_stop_event:
+                    stop_event_name = current_stop_event.__name__
+                    stop_event_id = f"event_{clean_id(stop_event_name)}"
+                    edge = f"{stop_event_id} --> {step_id}"
+                    if edge not in edges:
+                        edges.add(edge)
+                        mermaid_diagram.append(f"    {edge}")
             else:
                 edge = f"{event_id} --> {step_id}"
                 if edge not in edges:
diff --git a/src/uipath_llamaindex/_cli/cli_run.py b/src/uipath_llamaindex/_cli/cli_run.py
@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from os import environ as env
-from typing import Any, Dict, Optional
+from typing import Optional
 
 from dotenv import load_dotenv
 from uipath._cli._runtime._contracts import UiPathTraceContext
diff --git a/src/uipath_llamaindex/embeddings/_openai_embeddings.py b/src/uipath_llamaindex/embeddings/_openai_embeddings.py
@@ -1,48 +1,46 @@
-import os
-from enum import Enum
-from typing import Any, Union
-
-from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
-
-
-class OpenAIEmbeddingModel(Enum):
-    TEXT_EMBEDDING_3_LARGE = "text-embedding-3-large"
-    TEXT_EMBEDDING_ADA_002 = "text-embedding-ada-002"
-
-
-class UiPathOpenAIEmbedding(AzureOpenAIEmbedding):
-    def __init__(
-        self,
-        model: Union[
-            str, OpenAIEmbeddingModel
-        ] = OpenAIEmbeddingModel.TEXT_EMBEDDING_ADA_002,
-        api_version: str = "2024-10-21",
-        **kwargs: Any,
-    ):
-        default_headers_dict = {
-            "X-UIPATH-STREAMING-ENABLED": "false",
-            "X-UiPath-LlmGateway-RequestingProduct": "uipath-python-sdk",
-            "X-UiPath-LlmGateway-RequestingFeature": "llama-index-agent",
-        }
-
-        model_value = model.value if isinstance(model, OpenAIEmbeddingModel) else model
-
-        base_url = os.environ.get(
-            "UIPATH_URL", "EMPTY"
-        ).rstrip("/")
-        
-        if base_url == "EMPTY":
-            raise ValueError(
-                "UIPATH_URL environment variable is not set. Please run uipath auth."
-            )
-
-        defaults = {
-            "model": model_value,
-            "deployment_name": model_value,
-            "azure_endpoint": f"{base_url}/llmgateway_/",
-            "api_key": os.environ.get("UIPATH_ACCESS_TOKEN"),
-            "api_version": api_version,
-            "default_headers": default_headers_dict,
-        }
-        final_kwargs = {**defaults, **kwargs}
-        super().__init__(**final_kwargs)
+import os
+from enum import Enum
+from typing import Any, Union
+
+from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding  # type: ignore
+
+
+class OpenAIEmbeddingModel(Enum):
+    TEXT_EMBEDDING_3_LARGE = "text-embedding-3-large"
+    TEXT_EMBEDDING_ADA_002 = "text-embedding-ada-002"
+
+
+class UiPathOpenAIEmbedding(AzureOpenAIEmbedding):
+    def __init__(
+        self,
+        model: Union[
+            str, OpenAIEmbeddingModel
+        ] = OpenAIEmbeddingModel.TEXT_EMBEDDING_ADA_002,
+        api_version: str = "2024-10-21",
+        **kwargs: Any,
+    ):
+        default_headers_dict = {
+            "X-UIPATH-STREAMING-ENABLED": "false",
+            "X-UiPath-LlmGateway-RequestingProduct": "uipath-python-sdk",
+            "X-UiPath-LlmGateway-RequestingFeature": "llama-index-agent",
+        }
+
+        model_value = model.value if isinstance(model, OpenAIEmbeddingModel) else model
+
+        base_url = os.environ.get("UIPATH_URL", "EMPTY").rstrip("/")
+
+        if base_url == "EMPTY":
+            raise ValueError(
+                "UIPATH_URL environment variable is not set. Please run uipath auth."
+            )
+
+        defaults = {
+            "model": model_value,
+            "deployment_name": model_value,
+            "azure_endpoint": f"{base_url}/llmgateway_/",
+            "api_key": os.environ.get("UIPATH_ACCESS_TOKEN"),
+            "api_version": api_version,
+            "default_headers": default_headers_dict,
+        }
+        final_kwargs = {**defaults, **kwargs}
+        super().__init__(**final_kwargs)
diff --git a/src/uipath_llamaindex/llms/_openai_llms.py b/src/uipath_llamaindex/llms/_openai_llms.py
@@ -1,54 +1,54 @@
-import os
-from enum import Enum
-from typing import Any, Union
-
-from llama_index.llms.azure_openai import AzureOpenAI
-
-
-class OpenAIModel(Enum):
-    GPT_4_1_2025_04_14 = "gpt-4.1-2025-04-14"
-    GPT_4_1_MINI_2025_04_14 = "gpt-4.1-mini-2025-04-14"
-    GPT_4_1_NANO_2025_04_14 = "gpt-4.1-nano-2025-04-14"
-    GPT_4O_2024_05_13 = "gpt-4o-2024-05-13"
-    GPT_4O_2024_08_06 = "gpt-4o-2024-08-06"
-    GPT_4O_2024_11_20 = "gpt-4o-2024-11-20"
-    GPT_4O_MINI_2024_07_18 = "gpt-4o-mini-2024-07-18"
-    O3_MINI_2025_01_31 = "o3-mini-2025-01-31"
-    TEXT_DAVINCI_003 = "text-davinci-003"
-
-
-# Define your custom AzureOpenAI class with default settings
-class UiPathOpenAI(AzureOpenAI):
-    def __init__(
-        self,
-        model: Union[str, OpenAIModel] = OpenAIModel.GPT_4O_MINI_2024_07_18,
-        api_version: str = "2024-10-21",
-        **kwargs: Any,
-    ):
-        default_headers_dict = {
-            "X-UIPATH-STREAMING-ENABLED": "false",
-            "X-UiPath-LlmGateway-RequestingProduct": "uipath-python-sdk",
-            "X-UiPath-LlmGateway-RequestingFeature": "llama-index-agent",
-        }
-        model_value = model.value if isinstance(model, OpenAIModel) else model
-
-        base_url = os.environ.get(
-            "UIPATH_URL", "EMPTY"
-        ).rstrip("/")
-        
-        if base_url == "EMPTY":
-            raise ValueError(
-                "UIPATH_URL environment variable is not set. Please run uipath auth."
-            )
-
-        defaults = {
-            "model": model_value,
-            "deployment_name": model_value,
-            "azure_endpoint": f"{base_url}/llmgateway_/",
-            "api_key": os.environ.get("UIPATH_ACCESS_TOKEN"),
-            "api_version": api_version,
-            "is_chat_model": True,
-            "default_headers": default_headers_dict,
-        }
-        final_kwargs = {**defaults, **kwargs}
-        super().__init__(**final_kwargs)
+import os
+from enum import Enum
+from typing import Any, Union
+
+from llama_index.llms.azure_openai import AzureOpenAI  # type: ignore
+
+
+class OpenAIModel(Enum):
+    GPT_4_1_2025_04_14 = "gpt-4.1-2025-04-14"
+    GPT_4_1_MINI_2025_04_14 = "gpt-4.1-mini-2025-04-14"
+    GPT_4_1_NANO_2025_04_14 = "gpt-4.1-nano-2025-04-14"
+    GPT_4O_2024_05_13 = "gpt-4o-2024-05-13"
+    GPT_4O_2024_08_06 = "gpt-4o-2024-08-06"
+    GPT_4O_2024_11_20 = "gpt-4o-2024-11-20"
+    GPT_4O_MINI_2024_07_18 = "gpt-4o-mini-2024-07-18"
+    O3_MINI_2025_01_31 = "o3-mini-2025-01-31"
+    TEXT_DAVINCI_003 = "text-davinci-003"
+
+
+# Define your custom AzureOpenAI class with default settings
+class UiPathOpenAI(AzureOpenAI):
+    def __init__(
+        self,
+        model: Union[str, OpenAIModel] = OpenAIModel.GPT_4O_MINI_2024_07_18,
+        api_version: str = "2024-10-21",
+        **kwargs: Any,
+    ):
+        default_headers_dict = {
+            "X-UIPATH-STREAMING-ENABLED": "false",
+            "X-UiPath-LlmGateway-RequestingProduct": "uipath-python-sdk",
+            "X-UiPath-LlmGateway-RequestingFeature": "llama-index-agent",
+        }
+        model_value = model.value if isinstance(model, OpenAIModel) else model
+
+        base_url = os.environ.get(
+            "UIPATH_URL", "EMPTY"
+        ).rstrip("/")
+
+        if base_url == "EMPTY":
+            raise ValueError(
+                "UIPATH_URL environment variable is not set. Please run uipath auth."
+            )
+
+        defaults = {
+            "model": model_value,
+            "deployment_name": model_value,
+            "azure_endpoint": f"{base_url}/llmgateway_/",
+            "api_key": os.environ.get("UIPATH_ACCESS_TOKEN"),
+            "api_version": api_version,
+            "is_chat_model": True,
+            "default_headers": default_headers_dict,
+        }
+        final_kwargs = {**defaults, **kwargs}
+        super().__init__(**final_kwargs)
