mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-08-29 13:20:26 +00:00
122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
"""Dashboard Action Module."""
|
||
|
||
import json
|
||
import logging
|
||
from typing import List, Optional
|
||
|
||
from dbgpt._private.pydantic import BaseModel, Field, model_to_dict
|
||
from dbgpt.vis.tags.vis_dashboard import Vis, VisDashboard
|
||
|
||
from ...core.action.base import Action, ActionOutput
|
||
from ...resource.base import AgentResource, ResourceType
|
||
from ...resource.database import DBResource
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ChartItem(BaseModel):
|
||
"""Chart item model."""
|
||
|
||
title: str = Field(
|
||
...,
|
||
description="The title of the current analysis chart.",
|
||
)
|
||
display_type: str = Field(
|
||
...,
|
||
description="The chart rendering method selected for SQL. If you don’t know "
|
||
"what to output, just output 'response_table' uniformly.",
|
||
)
|
||
sql: str = Field(
|
||
..., description="Executable sql generated for the current target/problem"
|
||
)
|
||
thought: str = Field(..., description="Summary of thoughts to the user")
|
||
|
||
def to_dict(self):
|
||
"""Convert to dict."""
|
||
return model_to_dict(self)
|
||
|
||
|
||
class DashboardAction(Action[List[ChartItem]]):
|
||
"""Dashboard action class."""
|
||
|
||
def __init__(self, **kwargs):
|
||
"""Dashboard action init."""
|
||
super().__init__(**kwargs)
|
||
self._render_protocol = VisDashboard()
|
||
|
||
@property
|
||
def resource_need(self) -> Optional[ResourceType]:
|
||
"""Return the resource type needed for the action."""
|
||
return ResourceType.DB
|
||
|
||
@property
|
||
def render_protocol(self) -> Optional[Vis]:
|
||
"""Return the render protocol."""
|
||
return self._render_protocol
|
||
|
||
@property
|
||
def out_model_type(self):
|
||
"""Return the output model type."""
|
||
return List[ChartItem]
|
||
|
||
async def run(
|
||
self,
|
||
ai_message: str,
|
||
resource: Optional[AgentResource] = None,
|
||
rely_action_out: Optional[ActionOutput] = None,
|
||
need_vis_render: bool = True,
|
||
**kwargs,
|
||
) -> ActionOutput:
|
||
"""Perform the action."""
|
||
try:
|
||
input_param = self._input_convert(ai_message, List[ChartItem])
|
||
except Exception as e:
|
||
logger.exception(str(e))
|
||
return ActionOutput(
|
||
is_exe_success=False,
|
||
content="The requested correctly structured answer could not be found.",
|
||
)
|
||
if not isinstance(input_param, list):
|
||
return ActionOutput(
|
||
is_exe_success=False,
|
||
content="The requested correctly structured answer could not be found.",
|
||
)
|
||
chart_items: List[ChartItem] = input_param
|
||
try:
|
||
db_resources: List[DBResource] = DBResource.from_resource(self.resource)
|
||
if not db_resources:
|
||
raise ValueError("The database resource is not found！")
|
||
|
||
db = db_resources[0]
|
||
|
||
if not db:
|
||
raise ValueError("The database resource is not found！")
|
||
|
||
chart_params = []
|
||
for chart_item in chart_items:
|
||
chart_dict = {}
|
||
try:
|
||
sql_df = await db.query_to_df(chart_item.sql)
|
||
chart_dict = chart_item.to_dict()
|
||
|
||
chart_dict["data"] = sql_df
|
||
except Exception as e:
|
||
logger.warning(f"Sql execute failed！{str(e)}")
|
||
chart_dict["err_msg"] = str(e)
|
||
chart_params.append(chart_dict)
|
||
if not self.render_protocol:
|
||
raise ValueError("The render protocol is not initialized!")
|
||
view = await self.render_protocol.display(charts=chart_params)
|
||
return ActionOutput(
|
||
is_exe_success=True,
|
||
content=json.dumps(
|
||
[chart_item.to_dict() for chart_item in chart_items]
|
||
),
|
||
view=view,
|
||
)
|
||
except Exception as e:
|
||
logger.exception("Dashboard generate Failed！")
|
||
return ActionOutput(
|
||
is_exe_success=False, content=f"Dashboard action run failed!{str(e)}"
|
||
)
|