(self, in_chat: bool = True, stream: bool = True,
finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, return_img: bool = True)
| 1171 | self.chunk_list.append(chunk) |
| 1172 | |
| 1173 | def run_task(self, in_chat: bool = True, stream: bool = True, |
| 1174 | finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, return_img: bool = True): |
| 1175 | json_result: Dict[str, Any] = {'success': True} |
| 1176 | _session = None |
| 1177 | try: |
| 1178 | _session = session_maker() |
| 1179 | if self.ds: |
| 1180 | oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1 |
| 1181 | ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None |
| 1182 | |
| 1183 | self.filter_terminology_template(_session, oid, ds_id) |
| 1184 | |
| 1185 | self.filter_training_template(_session, oid, ds_id) |
| 1186 | |
| 1187 | self.filter_custom_prompts(_session, CustomPromptTypeEnum.GENERATE_SQL, oid, ds_id) |
| 1188 | |
| 1189 | self.init_messages(_session) |
| 1190 | |
| 1191 | # return id |
| 1192 | if in_chat: |
| 1193 | yield 'data:' + orjson.dumps({'type': 'id', 'id': self.get_record().id}).decode() + '\n\n' |
| 1194 | if self.get_record().regenerate_record_id: |
| 1195 | yield 'data:' + orjson.dumps({'type': 'regenerate_record_id', |
| 1196 | 'regenerate_record_id': self.get_record().regenerate_record_id}).decode() + '\n\n' |
| 1197 | yield 'data:' + orjson.dumps( |
| 1198 | {'type': 'question', 'question': self.get_record().question}).decode() + '\n\n' |
| 1199 | else: |
| 1200 | if stream: |
| 1201 | yield '> ' + self.trans('i18n_chat.record_id_in_mcp') + str(self.get_record().id) + '\n' |
| 1202 | yield '> ' + self.get_record().question + '\n\n' |
| 1203 | if not stream: |
| 1204 | json_result['record_id'] = self.get_record().id |
| 1205 | |
| 1206 | # select datasource if datasource is none |
| 1207 | if not self.ds: |
| 1208 | ds_res = self.select_datasource(_session) |
| 1209 | |
| 1210 | for chunk in ds_res: |
| 1211 | SQLBotLogUtil.info(chunk) |
| 1212 | if in_chat: |
| 1213 | yield 'data:' + orjson.dumps( |
| 1214 | {'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'), |
| 1215 | 'type': 'datasource-result'}).decode() + '\n\n' |
| 1216 | if in_chat: |
| 1217 | yield 'data:' + orjson.dumps({'id': self.ds.id, 'datasource_name': self.ds.name, |
| 1218 | 'engine_type': self.ds.type_name or self.ds.type, |
| 1219 | 'type': 'datasource'}).decode() + '\n\n' |
| 1220 | |
| 1221 | else: |
| 1222 | self.validate_history_ds(_session) |
| 1223 | |
| 1224 | # check connection |
| 1225 | connected = check_connection(ds=self.ds, trans=None) |
| 1226 | if not connected: |
| 1227 | raise SQLBotDBConnectionError('Connect DB failed') |
| 1228 | |
| 1229 | # generate sql |
| 1230 | sql_res = self.generate_sql(_session) |
no test coverage detected