feat(awel): AWEL supports http trigger and add some AWEL examples (#815)

- AWEL supports http trigger.
- Disassemble the KBQA into some atomic operator.
- Add some AWEL examples.

examples/awel/simple_chat_dag_example.py

@@ -0,0 +1,54 @@
+"""AWEL: Simple chat dag example
+
+    Example:
+
+    .. code-block:: shell
+
+        curl -X POST http://127.0.0.1:5000/api/v1/awel/trigger/examples/simple_chat \
+        -H "Content-Type: application/json" -d '{
+            "model": "proxyllm",
+            "user_input": "hello"
+        }'
+"""
+from typing import Dict
+from pydantic import BaseModel, Field
+
+from pilot.awel import DAG, HttpTrigger, MapOperator
+from pilot.scene.base_message import ModelMessage
+from pilot.model.base import ModelOutput
+from pilot.model.operator.model_operator import ModelOperator
+
+
+class TriggerReqBody(BaseModel):
+    model: str = Field(..., description="Model name")
+    user_input: str = Field(..., description="User input")
+
+
+class RequestHandleOperator(MapOperator[TriggerReqBody, Dict]):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    async def map(self, input_value: TriggerReqBody) -> Dict:
+        hist = []
+        hist.append(ModelMessage.build_human_message(input_value.user_input))
+        hist = list(h.dict() for h in hist)
+        params = {
+            "prompt": input_value.user_input,
+            "messages": hist,
+            "model": input_value.model,
+            "echo": False,
+        }
+        print(f"Receive input value: {input_value}")
+        return params
+
+
+with DAG("dbgpt_awel_simple_dag_example") as dag:
+    # Receive http request and trigger dag to run.
+    trigger = HttpTrigger(
+        "/examples/simple_chat", methods="POST", request_body=TriggerReqBody
+    )
+    request_handle_task = RequestHandleOperator()
+    model_task = ModelOperator()
+    # type(out) == ModelOutput
+    model_parse_task = MapOperator(lambda out: out.to_dict())
+    trigger >> request_handle_task >> model_task >> model_parse_task

examples/awel/simple_dag_example.py

@@ -0,0 +1,32 @@
+"""AWEL: Simple dag example
+
+    Example:
+
+    .. code-block:: shell
+
+        curl -X GET http://127.0.0.1:5000/api/v1/awel/trigger/examples/hello\?name\=zhangsan
+
+"""
+from pydantic import BaseModel, Field
+
+from pilot.awel import DAG, HttpTrigger, MapOperator
+
+
+class TriggerReqBody(BaseModel):
+    name: str = Field(..., description="User name")
+    age: int = Field(18, description="User age")
+
+
+class RequestHandleOperator(MapOperator[TriggerReqBody, str]):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    async def map(self, input_value: TriggerReqBody) -> str:
+        print(f"Receive input value: {input_value}")
+        return f"Hello, {input_value.name}, your age is {input_value.age}"
+
+
+with DAG("simple_dag_example") as dag:
+    trigger = HttpTrigger("/examples/hello", request_body=TriggerReqBody)
+    map_node = RequestHandleOperator()
+    trigger >> map_node

examples/awel/simple_rag_example.py

@@ -0,0 +1,70 @@
+"""AWEL: Simple rag example
+
+    Example:
+
+    .. code-block:: shell
+
+        curl -X POST http://127.0.0.1:5000/api/v1/awel/trigger/examples/simple_rag \
+        -H "Content-Type: application/json" -d '{
+            "conv_uid": "36f0e992-8825-11ee-8638-0242ac150003",
+            "model_name": "proxyllm",
+            "chat_mode": "chat_knowledge",
+            "user_input": "What is DB-GPT?",
+            "select_param": "default"
+        }'
+
+"""
+
+from pilot.awel import HttpTrigger, DAG, MapOperator
+from pilot.scene.operator._experimental import (
+    ChatContext,
+    PromptManagerOperator,
+    ChatHistoryStorageOperator,
+    ChatHistoryOperator,
+    EmbeddingEngingOperator,
+    BaseChatOperator,
+)
+from pilot.scene.base import ChatScene
+from pilot.openapi.api_view_model import ConversationVo
+from pilot.model.base import ModelOutput
+from pilot.model.operator.model_operator import ModelOperator
+
+
+class RequestParseOperator(MapOperator[ConversationVo, ChatContext]):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    async def map(self, input_value: ConversationVo) -> ChatContext:
+        return ChatContext(
+            current_user_input=input_value.user_input,
+            model_name=input_value.model_name,
+            chat_session_id=input_value.conv_uid,
+            select_param=input_value.select_param,
+            chat_scene=ChatScene.ChatKnowledge,
+        )
+
+
+with DAG("simple_rag_example") as dag:
+    trigger_task = HttpTrigger(
+        "/examples/simple_rag", methods="POST", request_body=ConversationVo
+    )
+    req_parse_task = RequestParseOperator()
+    prompt_task = PromptManagerOperator()
+    history_storage_task = ChatHistoryStorageOperator()
+    history_task = ChatHistoryOperator()
+    embedding_task = EmbeddingEngingOperator()
+    chat_task = BaseChatOperator()
+    model_task = ModelOperator()
+    output_parser_task = MapOperator(lambda out: out.to_dict()["text"])
+
+    (
+        trigger_task
+        >> req_parse_task
+        >> prompt_task
+        >> history_storage_task
+        >> history_task
+        >> embedding_task
+        >> chat_task
+        >> model_task
+        >> output_parser_task
+    )

pilot/awel/__init__.py

@@ -1,8 +1,17 @@
-"""Agentic Workflow Expression Language (AWEL)"""
+"""Agentic Workflow Expression Language (AWEL)
+
+Note:
+
+AWEL is still an experimental feature and only opens the lowest level API. 
+The stability of this API cannot be guaranteed at present.
+
+"""
+
+from pilot.component import SystemApp
 
 from .dag.base import DAGContext, DAG
 
-from .operator.base import BaseOperator, WorkflowRunner, initialize_awel
+from .operator.base import BaseOperator, WorkflowRunner
 from .operator.common_operator import (
     JoinOperator,
     ReduceStreamOperator,
@@ -28,6 +37,7 @@
     SimpleStreamTaskOutput,
     _is_async_iterator,
 )
+from .trigger.http_trigger import HttpTrigger
 from .runner.local_runner import DefaultWorkflowRunner
 
 __all__ = [
@@ -57,4 +67,21 @@
     "StreamifyAbsOperator",
     "UnstreamifyAbsOperator",
     "TransformStreamAbsOperator",
+    "HttpTrigger",
 ]
+
+
+def initialize_awel(system_app: SystemApp, dag_filepath: str):
+    from .dag.dag_manager import DAGManager
+    from .dag.base import DAGVar
+    from .trigger.trigger_manager import DefaultTriggerManager
+    from .operator.base import initialize_runner
+
+    DAGVar.set_current_system_app(system_app)
+
+    system_app.register(DefaultTriggerManager)
+    dag_manager = DAGManager(system_app, dag_filepath)
+    system_app.register_instance(dag_manager)
+    initialize_runner(DefaultWorkflowRunner())
+    # Load all dags
+    dag_manager.load_dags()

pilot/awel/base.py

@@ -0,0 +1,7 @@
+from abc import ABC, abstractmethod
+
+
+class Trigger(ABC):
+    @abstractmethod
+    async def trigger(self) -> None:
+        """Trigger the workflow or a specific operation in the workflow."""

pilot/awel/dag/base.py

@@ -1,14 +1,20 @@
 from abc import ABC, abstractmethod
-from typing import Optional, Dict, List, Sequence, Union, Any
+from typing import Optional, Dict, List, Sequence, Union, Any, Set
 import uuid
 import contextvars
 import threading
 import asyncio
+import logging
 from collections import deque
+from functools import cache
+from concurrent.futures import Executor
 
+from pilot.component import SystemApp
 from ..resource.base import ResourceGroup
 from ..task.base import TaskContext
 
+logger = logging.getLogger(__name__)
+
 DependencyType = Union["DependencyMixin", Sequence["DependencyMixin"]]
 
 
@@ -96,6 +102,8 @@ def __rlshift__(self, nodes: DependencyType) -> "DependencyMixin":
 class DAGVar:
     _thread_local = threading.local()
     _async_local = contextvars.ContextVar("current_dag_stack", default=deque())
+    _system_app: SystemApp = None
+    _executor: Executor = None
 
     @classmethod
     def enter_dag(cls, dag) -> None:
@@ -138,18 +146,48 @@ def get_current_dag(cls) -> Optional["DAG"]:
                 return cls._thread_local.current_dag_stack[-1]
             return None
 
+    @classmethod
+    def get_current_system_app(cls) -> SystemApp:
+        if not cls._system_app:
+            raise RuntimeError("System APP not set for DAGVar")
+        return cls._system_app
+
+    @classmethod
+    def set_current_system_app(cls, system_app: SystemApp) -> None:
+        if cls._system_app:
+            logger.warn("System APP has already set, nothing to do")
+        else:
+            cls._system_app = system_app
+
+    @classmethod
+    def get_executor(cls) -> Executor:
+        return cls._executor
+
+    @classmethod
+    def set_executor(cls, executor: Executor) -> None:
+        cls._executor = executor
+
 
 class DAGNode(DependencyMixin, ABC):
     resource_group: Optional[ResourceGroup] = None
     """The resource group of current DAGNode"""
 
     def __init__(
-        self, dag: Optional["DAG"] = None, node_id: str = None, node_name: str = None
+        self,
+        dag: Optional["DAG"] = None,
+        node_id: Optional[str] = None,
+        node_name: Optional[str] = None,
+        system_app: Optional[SystemApp] = None,
+        executor: Optional[Executor] = None,
     ) -> None:
         super().__init__()
         self._upstream: List["DAGNode"] = []
         self._downstream: List["DAGNode"] = []
         self._dag: Optional["DAG"] = dag or DAGVar.get_current_dag()
+        self._system_app: Optional[SystemApp] = (
+            system_app or DAGVar.get_current_system_app()
+        )
+        self._executor: Optional[Executor] = executor or DAGVar.get_executor()
         if not node_id and self._dag:
             node_id = self._dag._new_node_id()
         self._node_id: str = node_id
@@ -159,6 +197,10 @@ def __init__(
     def node_id(self) -> str:
         return self._node_id
 
+    @property
+    def system_app(self) -> SystemApp:
+        return self._system_app
+
     def set_node_id(self, node_id: str) -> None:
         self._node_id = node_id
 
@@ -178,7 +220,7 @@ def node_name(self) -> str:
         return self._node_name
 
     @property
-    def dag(self) -> "DAGNode":
+    def dag(self) -> "DAG":
         return self._dag
 
     def set_upstream(self, nodes: DependencyType) -> "DAGNode":
@@ -254,17 +296,69 @@ class DAG:
     def __init__(
         self, dag_id: str, resource_group: Optional[ResourceGroup] = None
     ) -> None:
+        self._dag_id = dag_id
         self.node_map: Dict[str, DAGNode] = {}
+        self._root_nodes: Set[DAGNode] = None
+        self._leaf_nodes: Set[DAGNode] = None
+        self._trigger_nodes: Set[DAGNode] = None
 
     def _append_node(self, node: DAGNode) -> None:
         self.node_map[node.node_id] = node
+        # clear cached nodes
+        self._root_nodes = None
+        self._leaf_nodes = None
 
     def _new_node_id(self) -> str:
         return str(uuid.uuid4())
 
+    @property
+    def dag_id(self) -> str:
+        return self._dag_id
+
+    def _build(self) -> None:
+        from ..operator.common_operator import TriggerOperator
+
+        nodes = set()
+        for _, node in self.node_map.items():
+            nodes = nodes.union(_get_nodes(node))
+        self._root_nodes = list(set(filter(lambda x: not x.upstream, nodes)))
+        self._leaf_nodes = list(set(filter(lambda x: not x.downstream, nodes)))
+        self._trigger_nodes = list(
+            set(filter(lambda x: isinstance(x, TriggerOperator), nodes))
+        )
+
+    @property
+    def root_nodes(self) -> List[DAGNode]:
+        if not self._root_nodes:
+            self._build()
+        return self._root_nodes
+
+    @property
+    def leaf_nodes(self) -> List[DAGNode]:
+        if not self._leaf_nodes:
+            self._build()
+        return self._leaf_nodes
+
+    @property
+    def trigger_nodes(self):
+        if not self._trigger_nodes:
+            self._build()
+        return self._trigger_nodes
+
     def __enter__(self):
         DAGVar.enter_dag(self)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         DAGVar.exit_dag()
+
+
+def _get_nodes(node: DAGNode, is_upstream: Optional[bool] = True) -> set[DAGNode]:
+    nodes = set()
+    if not node:
+        return nodes
+    nodes.add(node)
+    stream_nodes = node.upstream if is_upstream else node.downstream
+    for node in stream_nodes:
+        nodes = nodes.union(_get_nodes(node, is_upstream))
+    return nodes