import json import re from abc import ABC, abstractmethod from typing import Dict, NamedTuple import pandas as pd from pilot.utils import build_logger from pilot.out_parser.base import BaseOutputParser, T from pilot.configs.model_config import LOGDIR from pilot.configs.config import Config from pilot.scene.chat_db.data_loader import DbDataLoader CFG = Config() class SqlAction(NamedTuple): sql: str thoughts: Dict logger = build_logger("webserver", LOGDIR + "DbChatOutputParser.log") class DbChatOutputParser(BaseOutputParser): def __init__(self, sep: str, is_stream_out: bool): super().__init__(sep=sep, is_stream_out=is_stream_out) def parse_prompt_response(self, model_out_text): clean_str = super().parse_prompt_response(model_out_text) print("clean prompt response:", clean_str) response = json.loads(clean_str) for key in sorted(response): if key.strip() == "sql": sql = response[key] if key.strip() == "thoughts": thoughts = response[key] return SqlAction(sql, thoughts) def parse_view_response(self, speak, data) -> str: ### tool out data to table view data_loader = DbDataLoader() if len(data) <= 1: data.insert(0, ["result"]) df = pd.DataFrame(data[1:], columns=data[0]) if not CFG.NEW_SERVER_MODE and not CFG.SERVER_LIGHT_MODE: table_style = """""" html_table = df.to_html(index=False, escape=False) html = f"{table_style}{html_table}" view_text = f"##### {str(speak)}" + "\n" + html.replace("\n", " ") return view_text else: return data_loader.get_table_view_by_conn(data, speak) @property def _type(self) -> str: return "sql_chat"