| 59 | |
| 60 | |
| 61 | def process(receiver, max_round): |
| 62 | container = Container() |
| 63 | while True: |
| 64 | data_item, session, sender = receiver.recv() |
| 65 | |
| 66 | if data_item is None and session is None and sender is None: |
| 67 | break |
| 68 | |
| 69 | entry = data_item |
| 70 | # container = self.container |
| 71 | init = build_sql(entry, container.conn) |
| 72 | container.execute(init) |
| 73 | db = entry['table']['table_name'] |
| 74 | session.inject({"role": "user", "content": big_prompt}) |
| 75 | session.inject({"role": "agent", "content": "Ok."}) |
| 76 | prompt = entry["description"] + "\n" + entry["add_description"] |
| 77 | session.inject({"role": "user", "content": prompt}) |
| 78 | res = session.action() |
| 79 | try: |
| 80 | action = re.search(r"Action: (.*?)\n", res) |
| 81 | rounds = 0 |
| 82 | while action and action.group(1) == "Operation" and rounds < max_round: |
| 83 | res = re.search(r"```sql\n([\s\S]*?)\n```", res) |
| 84 | if not res: |
| 85 | answer = "" |
| 86 | break |
| 87 | sql = res.group(1).strip() |
| 88 | sql = sql.replace("\n", " ") |
| 89 | response = container.execute(sql, db) |
| 90 | if response: |
| 91 | session.inject({"role": "user", "content": response}) |
| 92 | else: |
| 93 | session.inject({"role": "user", "content": ""}) |
| 94 | res = session.action() |
| 95 | action = re.search(r"Action: (.*?)\n", res) |
| 96 | rounds += 1 |
| 97 | else: |
| 98 | answer = re.search(r"\nFinal Answer:(.*)", res) |
| 99 | if answer: |
| 100 | answer = answer.group(1) |
| 101 | else: |
| 102 | answer = "" |
| 103 | except Exception as e: |
| 104 | error = str(e) |
| 105 | answer = "" |
| 106 | else: |
| 107 | error = "" |
| 108 | if data_item["type"][0] in ("INSERT", "DELETE", "UPDATE"): |
| 109 | columns = ",".join([f"`{escape(column['name'], container.conn)}`" |
| 110 | for column in entry["table"]["table_info"]["columns"]]) |
| 111 | md5_query = f"select md5(group_concat(rowhash order by rowhash)) as hash " \ |
| 112 | f"from( SELECT substring(MD5(CONCAT_WS(',', {columns})), 1, 5) AS rowhash FROM `{db}`) as sub;" |
| 113 | answer = container.execute(md5_query, db) |
| 114 | container.execute(f"drop database `{db}`") |
| 115 | sender.send({ |
| 116 | "answer": str(answer), |
| 117 | "type": entry["type"][0], |
| 118 | "history": session.history, |