feat(core): Support simple DB query for sdk (#917)

Co-authored-by: chengfangyin2 <chengfangyin3@jd.com>

dbgpt/core/__init__.py

@@ -11,7 +11,7 @@ from dbgpt.core.interface.message import (
     OnceConversation,
 )
 from dbgpt.core.interface.prompt import PromptTemplate, PromptTemplateOperator
-from dbgpt.core.interface.output_parser import BaseOutputParser
+from dbgpt.core.interface.output_parser import BaseOutputParser, SQLOutputParser
 from dbgpt.core.interface.serialization import Serializable, Serializer
 from dbgpt.core.interface.cache import (
     CacheKey,
@@ -33,6 +33,7 @@ __ALL__ = [
     "PromptTemplate",
     "PromptTemplateOperator",
     "BaseOutputParser",
+    "SQLOutputParser",
     "Serializable",
     "Serializer",
     "CacheKey",

dbgpt/core/awel/task/task_impl.py

@@ -53,7 +53,7 @@ class SimpleTaskOutput(TaskOutput[T], Generic[T]):
 
     @property
     def is_empty(self) -> bool:
-        return not self._data
+        return self._data is None
 
     async def _apply_func(self, func) -> Any:
         if asyncio.iscoroutinefunction(func):

dbgpt/core/interface/output_parser.py

@@ -251,3 +251,13 @@ def _parse_model_response(response: ResponseTye):
     else:
         raise ValueError(f"Unsupported response type {type(response)}")
     return resp_obj_ex
+
+
+class SQLOutputParser(BaseOutputParser):
+    def __init__(self, is_stream_out: bool = False, **kwargs):
+        super().__init__(is_stream_out=is_stream_out, **kwargs)
+
+    def parse_model_nostream_resp(self, response: ResponseTye, sep: str):
+        model_out_text = super().parse_model_nostream_resp(response, sep)
+        clean_str = super().parse_prompt_response(model_out_text)
+        return json.loads(clean_str, strict=True)

dbgpt/core/interface/retriever.py

@@ -0,0 +1,23 @@
+from abc import abstractmethod
+from dbgpt.core.awel import MapOperator
+from dbgpt.core.awel.task.base import IN, OUT
+
+
+class RetrieverOperator(MapOperator[IN, OUT]):
+    """The Abstract Retriever Operator."""
+
+    async def map(self, input_value: IN) -> OUT:
+        """Map input value to output value.
+
+        Args:
+            input_value (IN): The input value.
+
+        Returns:
+            OUT: The output value.
+        """
+        # The retrieve function is blocking, so we need to wrap it in a blocking_func_to_async.
+        return await self.blocking_func_to_async(self.retrieve, input_value)
+
+    @abstractmethod
+    def retrieve(self, input_value: IN) -> OUT:
+        """Retrieve data for input value."""