DB-GPT/dbgpt/vis/client.py
明天 f866580703
feat: dbgpts modules bug fix (#1921)
Co-authored-by: 途杨 <tuyang.yhj@antgroup.com>
Co-authored-by: lhwan <1017484907@qq.com>
Co-authored-by: hustcc <i@hust.cc>
2024-08-30 10:03:58 +08:00

59 lines
1.6 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""Client for vis protocol."""
from typing import Dict, Type
from .base import Vis
from .tags.vis_agent_message import VisAgentMessages
from .tags.vis_agent_plans import VisAgentPlans
from .tags.vis_api_response import VisApiResponse
from .tags.vis_app_link import VisAppLink
from .tags.vis_chart import VisChart
from .tags.vis_code import VisCode
from .tags.vis_dashboard import VisDashboard
from .tags.vis_plugin import VisPlugin
class VisClient:
"""Client for vis protocol."""
def __init__(self):
"""Client for vis protocol."""
self._vis_tag: Dict[str, Vis] = {}
def register(self, vis_cls: Type[Vis]):
"""Register the vis protocol."""
self._vis_tag[vis_cls.vis_tag()] = vis_cls()
def get(self, tag_name):
"""Get the vis protocol by tag name."""
if tag_name not in self._vis_tag:
raise ValueError(f"Vis protocol tags not yet supported[{tag_name}]")
return self._vis_tag[tag_name]
def tag_names(self):
"""Return the tag names of the vis protocol."""
self._vis_tag.keys()
vis_client = VisClient()
vis_client.register(VisCode)
vis_client.register(VisChart)
vis_client.register(VisDashboard)
vis_client.register(VisAgentPlans)
vis_client.register(VisAgentMessages)
vis_client.register(VisPlugin)
vis_client.register(VisAppLink)
vis_client.register(VisApiResponse)
def vis_name_change(vis_message: str) -> str:
"""Change vis tag name use new name."""
replacements = {
"```vis-chart": "```vis-db-chart",
}
for old_tag, new_tag in replacements.items():
vis_message = vis_message.replace(old_tag, new_tag)
return vis_message