I3eka commented on code in PR #43132: URL: https://github.com/apache/superset/pull/43132#discussion_r4385019069
########## superset/ai/orchestrator.py: ########## @@ -0,0 +1,647 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Runs one assistant turn end to end. + +Sits between the HTTP layer and the runtime: loads the conversation, assembles +the prompt, resolves the tools the chosen profile allows, drives the runtime, +publishes every event to the bus, and records the outcome on the assistant +message. + +Deliberately independent of *where* it runs. The same function body serves the +inline path and the Celery path, which is what makes the execution mode a +configuration choice rather than two implementations that drift apart. +""" + +from __future__ import annotations + +import asyncio +import logging +import uuid as uuid_module +from collections.abc import AsyncIterator, Iterator +from dataclasses import dataclass +from typing import Any + +from superset.ai.events import ( + cancelled_event, + done_event, + error_event, + GENERIC_ERROR_MESSAGE, + session_event, + StreamEvent, +) +from superset.ai.llm.base import Message +from superset.ai.telemetry import bind_run, current_run, start_run +from superset.ai.types import MessageRole, MessageStatus, RunOutcome, StreamEventType +from superset.utils.decorators import transaction + +logger = logging.getLogger(__name__) + +#: Cache key prefix for a run's cancellation flag. A flag rather than a signal +#: because a worker cannot be interrupted mid-call reliably; the runtime checks +#: this between steps. +_CANCEL_PREFIX = "ai-cancel-" + +#: How long a cancellation request stays meaningful. +_CANCEL_TTL_SECONDS = 900 + +#: Stored when a run is stopped before it produced any answer, so the +#: transcript still records that the turn happened. +_STOPPED_WITHOUT_ANSWER = "_Stopped before an answer was produced._" + +#: Stored when a run exhausted its time budget without saying anything. Phrased +#: as something the user can act on, because retrying is usually the right move. +_TIMED_OUT_WITHOUT_ANSWER = ( + "The assistant ran out of time before it could answer. Please try again." +) + +#: Ceiling on the page context recorded on a message. Well below the prompt's own +#: limit: this is stored per turn and read back with the whole transcript. +_RECORDED_CONTEXT_LIMIT = 4_000 + + +@dataclass +class TurnRequest: + """One unit of work: answer the latest message on a thread.""" + + thread_uuid: str + user_id: int + run_id: str + #: Assistant message row to fill in. Created before the run starts so a + #: client that reconnects has something to attach to. + assistant_message_uuid: str + profile_key: str | None = None + #: Concrete model to pin, overriding the profile's tier. + model: str | None = None + #: What the user had on screen when they asked. Supplied by the client, + #: which is the only party that knows which tab is open, what is typed in + #: the editor and which filters are applied. + page_context: dict[str, Any] | None = None + + def to_payload(self) -> dict[str, Any]: + """Serialise for the task broker.""" + return { + "thread_uuid": self.thread_uuid, + "user_id": self.user_id, + "run_id": self.run_id, + "assistant_message_uuid": self.assistant_message_uuid, + "profile_key": self.profile_key, + "model": self.model, + "page_context": self.page_context, + } + + @classmethod + def from_payload(cls, payload: dict[str, Any]) -> TurnRequest: + """Rebuild from a broker payload.""" + return cls(**payload) + + +def new_run_id() -> str: + """Identifier for one run, used as the event-stream key.""" + return str(uuid_module.uuid4()) + + +#: Runs cancelled in this process. +#: +#: Held alongside the cache rather than instead of it. Superset's default cache +#: is a null cache, which accepts a write and discards it — so a cache-only +#: implementation would leave cancellation silently broken on a default install, +#: with the button appearing to work and nothing stopping. This set makes inline +#: execution correct with no cache at all; the cache is what carries a +#: cancellation across processes for worker execution. +_CANCELLED_LOCALLY: set[str] = set() + + +def request_cancel(run_id: str) -> None: + """ + Ask a run to stop. + + Cooperative by design: the flag is recorded here and observed by the runtime + between steps. A run blocked inside a single long model call or query will + not notice until that call returns, which is a real limit worth documenting + rather than hiding. + """ + from superset.extensions import cache_manager + + _CANCELLED_LOCALLY.add(run_id) + try: + cache_manager.cache.set( Review Comment: Follow-up correction in a9f7d3f94e: the shared backend was correct, but the raw Redis SET value was not. redis-py rejects Python `True`; `request_cancel` caught that error, so the earlier mock-only test missed that no flag was stored. The value is now `b"1"`. Added a regression through the real `RedisCacheBackend` methods and redis-py's actual encoder (only network transport is replaced). It fails on the previous head and passes after the fix, including get/clear and TTL checks. All 69 cancellation/orchestrator/event-bus tests and pre-commit checks for 122 PR files pass. No live Redis server or deployment was used. This corrects my earlier validation claim; the PR remains Draft / ON HOLD behind #42805 and #42806. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
