refactor: Refactor for core SDK (#1092)

Author: fangyinc

Makefile

@@ -81,6 +81,19 @@ clean: ## Clean up the environment
 	find . -type d -name '.pytest_cache' -delete
 	find . -type d -name '.coverage' -delete
 
+.PHONY: clean-dist
+clean-dist: ## Clean up the distribution
+	rm -rf dist/ *.egg-info build/
+
+.PHONY: package
+package: clean-dist ## Package the project for distribution
+	IS_DEV_MODE=false python setup.py sdist bdist_wheel
+
+.PHONY: upload
+upload: package ## Upload the package to PyPI
+	# upload to testpypi: twine upload --repository testpypi dist/*
+	twine upload dist/*
+
 .PHONY: help
 help:  ## Display this help screen
 	@echo "Available commands:"

dbgpt/__init__.py

@@ -1,12 +1,18 @@
-from dbgpt.component import BaseComponent, SystemApp
-
-__ALL__ = ["SystemApp", "BaseComponent"]
+"""DB-GPT: Next Generation Data Interaction Solution with LLMs.
+"""
+from dbgpt import _version  # noqa: E402
+from dbgpt.component import BaseComponent, SystemApp  # noqa: F401
 
 _CORE_LIBS = ["core", "rag", "model", "agent", "datasource", "vis", "storage", "train"]
 _SERVE_LIBS = ["serve"]
 _LIBS = _CORE_LIBS + _SERVE_LIBS
 
 
+__version__ = _version.version
+
+__ALL__ = ["__version__", "SystemApp", "BaseComponent"]
+
+
 def __getattr__(name: str):
     # Lazy load
     import importlib

dbgpt/_version.py

@@ -0,0 +1 @@
+version = "0.4.7"

dbgpt/agent/agents/expand/retrieve_summary_assistant_agent.py

@@ -579,7 +579,7 @@ def _num_token_from_text(self, text: str, model: str = "gpt-3.5-turbo-0613"):
     from dbgpt.agent.agents.agent import AgentContext
     from dbgpt.agent.agents.user_proxy_agent import UserProxyAgent
     from dbgpt.core.interface.llm import ModelMetadata
-    from dbgpt.model import OpenAILLMClient
+    from dbgpt.model.proxy import OpenAILLMClient
 
     llm_client = OpenAILLMClient()
     context: AgentContext = AgentContext(

dbgpt/app/chat_adapter.py

@@ -9,7 +9,7 @@
 from typing import Dict, List, Tuple
 
 from dbgpt.core.interface.message import ModelMessage, ModelMessageRoleType
-from dbgpt.model.conversation import Conversation, get_conv_template
+from dbgpt.model.llm.conversation import Conversation, get_conv_template
 
 
 class BaseChatAdpter:
@@ -21,7 +21,7 @@ def match(self, model_path: str):
 
     def get_generate_stream_func(self, model_path: str):
         """Return the generate stream handler func"""
-        from dbgpt.model.inference import generate_stream
+        from dbgpt.model.llm.inference import generate_stream
 
         return generate_stream
 

dbgpt/app/scene/base_chat.py

@@ -171,13 +171,13 @@ def llm_client(self) -> LLMClient:
 
     async def call_llm_operator(self, request: ModelRequest) -> ModelOutput:
         llm_task = build_cached_chat_operator(self.llm_client, False, CFG.SYSTEM_APP)
-        return await llm_task.call(call_data={"data": request})
+        return await llm_task.call(call_data=request)
 
     async def call_streaming_operator(
         self, request: ModelRequest
     ) -> AsyncIterator[ModelOutput]:
         llm_task = build_cached_chat_operator(self.llm_client, True, CFG.SYSTEM_APP)
-        async for out in await llm_task.call_stream(call_data={"data": request}):
+        async for out in await llm_task.call_stream(call_data=request):
             yield out
 
     def do_action(self, prompt_response):
@@ -251,11 +251,9 @@ async def _build_model_request(self) -> ModelRequest:
             str_history=self.prompt_template.str_history,
             request_context=req_ctx,
         )
-        node_input = {
-            "data": ChatComposerInput(
-                messages=self.history_messages, prompt_dict=input_values
-            )
-        }
+        node_input = ChatComposerInput(
+            messages=self.history_messages, prompt_dict=input_values
+        )
         # llm_messages = self.generate_llm_messages()
         model_request: ModelRequest = await node.call(call_data=node_input)
         model_request.context.cache_enable = self.model_cache_enable

dbgpt/app/scene/operator/app_operator.py

@@ -87,7 +87,7 @@ async def map(self, input_value: ChatComposerInput) -> ModelRequest:
         end_node: BaseOperator = self._sub_compose_dag.leaf_nodes[0]
         # Sub dag, use the same dag context in the parent dag
         messages = await end_node.call(
-            call_data={"data": input_value}, dag_ctx=self.current_dag_context
+            call_data=input_value, dag_ctx=self.current_dag_context
         )
         span_id = self._request_context.span_id
         model_request = ModelRequest.build_request(

dbgpt/component.py

@@ -1,3 +1,7 @@
+"""Component module for dbgpt.
+
+Manages the lifecycle and registration of components.
+"""
 from __future__ import annotations
 
 import asyncio

dbgpt/core/awel/__init__.py

@@ -22,6 +22,7 @@
     JoinOperator,
     MapOperator,
     ReduceStreamOperator,
+    TriggerOperator,
 )
 from .operator.stream_operator import (
     StreamifyAbsOperator,
@@ -50,6 +51,7 @@
     "BaseOperator",
     "JoinOperator",
     "ReduceStreamOperator",
+    "TriggerOperator",
     "MapOperator",
     "BranchOperator",
     "InputOperator",
@@ -150,4 +152,6 @@ def setup_dev_environment(
         for trigger in dag.trigger_nodes:
             trigger_manager.register_trigger(trigger)
     trigger_manager.after_register()
-    uvicorn.run(app, host=host, port=port)
+    if trigger_manager.keep_running():
+        # Should keep running
+        uvicorn.run(app, host=host, port=port)

dbgpt/core/awel/operator/base.py

@@ -28,7 +28,7 @@
 
 F = TypeVar("F", bound=FunctionType)
 
-CALL_DATA = Union[Dict, Dict[str, Dict]]
+CALL_DATA = Union[Dict[str, Any], Any]
 
 
 class WorkflowRunner(ABC, Generic[T]):
@@ -197,6 +197,8 @@ async def call(
         Returns:
             OUT: The output of the node after execution.
         """
+        if call_data:
+            call_data = {"data": call_data}
         out_ctx = await self._runner.execute_workflow(
             self, call_data, exist_dag_ctx=dag_ctx
         )
@@ -242,6 +244,8 @@ async def call_stream(
         Returns:
             AsyncIterator[OUT]: An asynchronous iterator over the output stream.
         """
+        if call_data:
+            call_data = {"data": call_data}
         out_ctx = await self._runner.execute_workflow(
             self, call_data, streaming_call=True, exist_dag_ctx=dag_ctx
         )

dbgpt/core/awel/task/base.py

@@ -28,6 +28,14 @@ def __bool__(self):
 SKIP_DATA = _EMPTY_DATA_TYPE()
 PLACEHOLDER_DATA = _EMPTY_DATA_TYPE()
 
+
+def is_empty_data(data: Any):
+    """Check if the data is empty."""
+    if isinstance(data, _EMPTY_DATA_TYPE):
+        return data in (EMPTY_DATA, SKIP_DATA)
+    return False
+
+
 MapFunc = Union[Callable[[IN], OUT], Callable[[IN], Awaitable[OUT]]]
 ReduceFunc = Union[Callable[[IN], OUT], Callable[[IN], Awaitable[OUT]]]
 StreamFunc = Callable[[IN], Awaitable[AsyncIterator[OUT]]]

dbgpt/core/awel/task/task_impl.py

@@ -24,7 +24,6 @@
     EMPTY_DATA,
     OUT,
     PLACEHOLDER_DATA,
-    SKIP_DATA,
     InputContext,
     InputSource,
     MapFunc,
@@ -37,6 +36,7 @@
     TaskState,
     TransformFunc,
     UnStreamFunc,
+    is_empty_data,
 )
 
 logger = logging.getLogger(__name__)
@@ -99,7 +99,7 @@ def new_output(self) -> TaskOutput[T]:
     @property
     def is_empty(self) -> bool:
         """Return True if the output data is empty."""
-        return self._data == EMPTY_DATA or self._data == SKIP_DATA
+        return is_empty_data(self._data)
 
     @property
     def is_none(self) -> bool:
@@ -171,7 +171,7 @@ def is_stream(self) -> bool:
     @property
     def is_empty(self) -> bool:
         """Return True if the output data is empty."""
-        return self._data == EMPTY_DATA or self._data == SKIP_DATA
+        return is_empty_data(self._data)
 
     @property
     def is_none(self) -> bool:
@@ -330,7 +330,7 @@ def _read_data(self, task_ctx: TaskContext) -> Any:
         """
         call_data = task_ctx.call_data
         data = call_data.get("data", EMPTY_DATA) if call_data else EMPTY_DATA
-        if data == EMPTY_DATA:
+        if is_empty_data(data):
             raise ValueError("No call data for current SimpleCallDataInputSource")
         return data
 

dbgpt/core/awel/trigger/http_trigger.py

@@ -1,12 +1,8 @@
 """Http trigger for AWEL."""
-from __future__ import annotations
-
 import logging
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union, cast
 
-from starlette.requests import Request
-
 from dbgpt._private.pydantic import BaseModel
 
 from ..dag.base import DAG
@@ -15,9 +11,10 @@
 
 if TYPE_CHECKING:
     from fastapi import APIRouter
+    from starlette.requests import Request
 
-RequestBody = Union[Type[Request], Type[BaseModel], Type[str]]
-StreamingPredictFunc = Callable[[Union[Request, BaseModel]], bool]
+    RequestBody = Union[Type[Request], Type[BaseModel], Type[str]]
+    StreamingPredictFunc = Callable[[Union[Request, BaseModel]], bool]
 
 logger = logging.getLogger(__name__)
 
@@ -32,9 +29,9 @@ def __init__(
         self,
         endpoint: str,
         methods: Optional[Union[str, List[str]]] = "GET",
-        request_body: Optional[RequestBody] = None,
+        request_body: Optional["RequestBody"] = None,
         streaming_response: bool = False,
-        streaming_predict_func: Optional[StreamingPredictFunc] = None,
+        streaming_predict_func: Optional["StreamingPredictFunc"] = None,
         response_model: Optional[Type] = None,
         response_headers: Optional[Dict[str, str]] = None,
         response_media_type: Optional[str] = None,
@@ -69,6 +66,7 @@ def mount_to_router(self, router: "APIRouter") -> None:
             router (APIRouter): The router to mount the trigger.
         """
         from fastapi import Depends
+        from starlette.requests import Request
 
         methods = [self._methods] if isinstance(self._methods, str) else self._methods
 
@@ -114,8 +112,10 @@ async def route_function(body=Depends(_request_body_dependency)):
 
 
 async def _parse_request_body(
-    request: Request, request_body_cls: Optional[RequestBody]
+    request: "Request", request_body_cls: Optional["RequestBody"]
 ):
+    from starlette.requests import Request
+
     if not request_body_cls:
         return None
     if request_body_cls == Request:
@@ -152,7 +152,7 @@ async def _trigger_dag(
         raise ValueError("HttpTrigger just support one leaf node in dag")
     end_node = cast(BaseOperator, leaf_nodes[0])
     if not streaming_response:
-        return await end_node.call(call_data={"data": body})
+        return await end_node.call(call_data=body)
     else:
         headers = response_headers
         media_type = response_media_type if response_media_type else "text/event-stream"
@@ -163,7 +163,7 @@ async def _trigger_dag(
                 "Connection": "keep-alive",
                 "Transfer-Encoding": "chunked",
             }
-        generator = await end_node.call_stream(call_data={"data": body})
+        generator = await end_node.call_stream(call_data=body)
         background_tasks = BackgroundTasks()
         background_tasks.add_task(dag._after_dag_end)
         return StreamingResponse(

dbgpt/core/awel/trigger/trigger_manager.py

@@ -24,6 +24,14 @@ class TriggerManager(ABC):
     def register_trigger(self, trigger: Any) -> None:
         """Register a trigger to current manager."""
 
+    def keep_running(self) -> bool:
+        """Whether keep running.
+
+        Returns:
+            bool: Whether keep running, True means keep running, False means stop.
+        """
+        return False
+
 
 class HttpTriggerManager(TriggerManager):
     """Http trigger manager.
@@ -64,6 +72,8 @@ def register_trigger(self, trigger: Any) -> None:
             self._trigger_map[trigger_id] = trigger
 
     def _init_app(self, system_app: SystemApp):
+        if not self.keep_running():
+            return
         logger.info(
             f"Include router {self._router} to prefix path {self._router_prefix}"
         )
@@ -72,6 +82,14 @@ def _init_app(self, system_app: SystemApp):
             raise RuntimeError("System app not initialized")
         app.include_router(self._router, prefix=self._router_prefix, tags=["AWEL"])
 
+    def keep_running(self) -> bool:
+        """Whether keep running.
+
+        Returns:
+            bool: Whether keep running, True means keep running, False means stop.
+        """
+        return len(self._trigger_map) > 0
+
 
 class DefaultTriggerManager(TriggerManager, BaseComponent):
     """Default trigger manager for AWEL.
@@ -105,3 +123,11 @@ def after_register(self) -> None:
         """After register, init the trigger manager."""
         if self.system_app:
             self.http_trigger._init_app(self.system_app)
+
+    def keep_running(self) -> bool:
+        """Whether keep running.
+
+        Returns:
+            bool: Whether keep running, True means keep running, False means stop.
+        """
+        return self.http_trigger.keep_running()

dbgpt/core/interface/operator/composer_operator.py

@@ -70,7 +70,7 @@ async def map(self, input_value: ChatComposerInput) -> ModelRequest:
         end_node: BaseOperator = cast(BaseOperator, self._sub_compose_dag.leaf_nodes[0])
         # Sub dag, use the same dag context in the parent dag
         return await end_node.call(
-            call_data={"data": input_value}, dag_ctx=self.current_dag_context
+            call_data=input_value, dag_ctx=self.current_dag_context
         )
 
     def _build_composer_dag(self) -> DAG:

dbgpt/core/interface/operator/prompt_operator.py

@@ -150,7 +150,7 @@ class PromptBuilderOperator(
                     )
                 )
 
-            single_input = {"data": {"dialect": "mysql"}}
+            single_input = {"dialect": "mysql"}
             single_expected_messages = [
                 ModelMessage(
                     content="Please write a mysql SQL count the length of a field",

dbgpt/model/__init__.py

@@ -1,9 +1,12 @@
-from dbgpt.model.cluster.client import DefaultLLMClient
+try:
+    from dbgpt.model.cluster.client import DefaultLLMClient
+except ImportError as exc:
+    # logging.warning("Can't import dbgpt.model.DefaultLLMClient")
+    DefaultLLMClient = None
 
-# from dbgpt.model.utils.chatgpt_utils import OpenAILLMClient
-from dbgpt.model.proxy.llms.chatgpt import OpenAILLMClient
 
-__ALL__ = [
-    "DefaultLLMClient",
-    "OpenAILLMClient",
-]
+_exports = []
+if DefaultLLMClient:
+    _exports.append("DefaultLLMClient")
+
+__ALL__ = _exports