import datetime import json from sqlalchemy import ( Boolean, Column, String, Integer, ) from sqlalchemy.dialects.mysql import LONGTEXT from utils.context import fmt_tasks_tree from utils.logging import logging from models.base import Base, BaseMixin from consts.consts import ContextTypes, SessionStatus class Context(Base, BaseMixin): __tablename__ = "context" _type = Column(String(16), nullable=False) name = Column(String(32), nullable=False) org_id = Column(String(64), nullable=False, index=True) project_id = Column(String(64), nullable=False) cluster_id = Column(String(64), nullable=False, default=0, index=True) db_name = Column(String(64), nullable=False, default="") session_id = Column(String(64), nullable=False, default="", index=True) creator = Column(String(32), nullable=False) data = Column(LONGTEXT, nullable=False) data_context_id = Column(Integer) # session context will have a ref to data context domain_context_id = Column(Integer) # both session context and data context will have a ref to domain context database_uri = Column(String(4096), index=True, default="") identifier = Column(String(64), nullable=False, default="", index=True) identifier_type = Column(String(16), nullable=False, default="") thread_id = Column(String(64), nullable=False, default="") description = Column(String(4096), nullable=False, default="") is_default = Column(Boolean, nullable=False, default=False) @classmethod def create( cls, session, _type, name, org_id, project_id, cluster_id, db_name, session_id, creator, data=None, data_context_id=None, domain_context_id=None, database_uri="", identifier="", identifier_type="", thread_id="", description="", ): data = data or {} org_id = str(org_id) project_id = str(project_id) cluster_id = str(cluster_id) item = cls( _type=_type, name=name, org_id=org_id, project_id=project_id, cluster_id=cluster_id, db_name=db_name, session_id=session_id, creator=creator, data=json.dumps(data), data_context_id=data_context_id, domain_context_id=domain_context_id, database_uri=database_uri, identifier=identifier, identifier_type=identifier_type, thread_id=thread_id, description=description, ) session.add(item) session.flush() return item @classmethod def get_by_database_uri(cls, session, database_uri): return session.query(cls).filter(cls.database_uri == database_uri).order_by(cls.created_at.desc()).first() @classmethod def get_domain_context_by_org_id(cls, session, org_id): org_id = str(org_id) return session.query(cls).filter(cls.org_id == org_id, cls._type == ContextTypes.DOMAIN).first() @classmethod def get_session_context_by_id(cls, session, context_id): return session.query(cls).filter( cls.id == context_id, cls._type.in_((ContextTypes.SESSION, ContextTypes.SESSION_CHAT, ContextTypes.SESSION_FIX_REFINE, ContextTypes.SESSION_ONETIME)), ).first() @classmethod def get_session_context_by_session_id(cls, session, session_id): return session.query(cls).filter( cls.session_id == session_id, cls._type.in_((ContextTypes.SESSION, ContextTypes.SESSION_CHAT, ContextTypes.SESSION_FIX_REFINE, ContextTypes.SESSION_ONETIME)), ).first() @classmethod def get_data_context_by_cluster_and_db(cls, session, cluster_id, db_name): cluster_id = str(cluster_id) return session.query(cls).filter( cls.cluster_id == cluster_id, cls.db_name == db_name, cls._type == ContextTypes.DATA, ).first() @classmethod def get_all_by_org_id_and_type(cls, session, org_id, _type): org_id = str(org_id) return session.query(cls).filter(cls.org_id == org_id, cls._type == _type).order_by(cls.created_at.desc()).all() @classmethod def get_all_by_org_id(cls, session, org_id): org_id = str(org_id) return session.query(cls).filter(cls.org_id == org_id).order_by(cls.created_at.desc()).all() def load_data(self, escape=False): data = json.loads(self.data) if escape: data.pop("database_uri", None) data.pop("public_key", None) return data def update_data(self, data): self.data = json.dumps(data) def reset_session_context(self, s, thread_id: str): logging.info(f"Reset session context {self.id} from {self.thread_id} to thread {thread_id}") self.thread_id = thread_id self.created_at = datetime.datetime.now() self.data = json.dumps({}) s.flush() def is_session_expired(self): if not self.thread_id: return False from svc.openai import check_thread_status # avoid circyle import return check_thread_status(self.identifier or 0, self.thread_id) def trim_tasks_tree(self): _data = self.load_data() tasks_tree = _data.get("tasks_tree", {}) data = [] for query_id, query_structure in tasks_tree.items(): for task_id, task_structure in query_structure.items(): data.append({ "type": "chat2sql", "input": { "raw_question": task_structure["task"], }, "output": { "sql": task_structure["sql"], "clarified_task": task_structure.get("clarified_task", ""), "sql_error": task_structure.get("sql_error", ""), "description": task_structure.get("description", ""), "created_at": task_structure.get("created_at", 0), "chart_options": task_structure.get("chartOptions", {}), "refine_list": task_structure.get("refine_list", []), "question_id": query_id, "task_id": task_id }, }) data.sort(key=lambda x: x["output"]["created_at"], reverse=True) return data def to_session_context(self, simple=True): """format session context, for chat session only""" if self._type != ContextTypes.SESSION_CHAT: raise ValueError("Not a session context") payload = { "meta": { "session_context_id": self.id, "cluster_id": self.cluster_id, "name": self.name, "creator": self.creator, "database": self.db_name, "org_id": self.org_id, "created_at": int(self.created_at.timestamp()), "updated_at": int(self.updated_at.timestamp()), "status": SessionStatus.EXPIRED if self.is_session_expired() else SessionStatus.ACTIVE, } } if not simple: payload["data"] = self.trim_tasks_tree() return payload def fmt_session_context_v2(self, simple=False): if self._type not in (ContextTypes.SESSION_CHAT, ContextTypes.SESSION): return {} result = { "session_id": self.id, "meta": { "org_id": self.org_id, "creator": self.creator, "name": self.name, "created_at": int(self.created_at.timestamp()), "updated_at": int(self.updated_at.timestamp()), }, } if not simple: messages = [] for answer_id, tasks_tree in self.load_data().get("tasks_tree", {}).items(): root_task = tasks_tree["0"] message = { "created_at": root_task.get("created_at"), "question": { "question": root_task["task"], "feedback_answer_id": root_task.get("feedback_answer_id", ""), "feedback_task_id": root_task.get("feedback_task_id", ""), }, "answers": [ fmt_tasks_tree(answer_id, tasks_tree), ], } messages.append(message) result["messages"] = messages return result def fmt_data_summary_v2(self, simple=False): if self._type != ContextTypes.DATA: return {} data = self.load_data() result = { "data_summary_id": self.id, "status": data.get("status", "init"), "default": self.is_default, "cluster_id": self.cluster_id, "database": self.db_name, "description": { "user": self.description, "system": data.get("summary", ""), }, "summary": data.get("short_summary", ""), } if not simple: tables = {} for table_name, table_data in data.get("db_schema", {}).items(): columns = {} for column_name, column_data in table_data.get("columns", {}).items(): columns[column_name] = { "name": column_name, "description": column_data.get("description", ""), } tables[table_name] = { "name": table_name, "description": table_data.get("description", ""), "columns": columns, } relationships = {} for table_name, relations in data.get("table_relationship", {}).items(): refined_relations = [] for relation in relations: refined_relations.append({ "referencing_table": relation.get("referencing_table", ""), "referencing_table_column": relation.get("foreign_key_column", ""), "referenced_table": relation.get("referenced_table", ""), "referenced_table_column": relation.get("primary_key_column", ""), }) relationships[table_name] = refined_relations result["tables"] = tables result["keywords"] = data.get("summary_keywords", []) result["relationships"] = relationships return result def to_json(self, escape=False): return { "id": self.id, "created_at": int(self.created_at.timestamp()), "updated_at": int(self.updated_at.timestamp()), "type": self._type, "name": self.name, "org_id": self.org_id, "cluster_id": self.cluster_id, "db_name": self.db_name, "session_id": self.session_id, "creator": self.creator, "domain_context_id": self.domain_context_id, "data_context_id": self.data_context_id, "data": self.load_data(escape), } @classmethod def get_all_data_context_by_cluster_id_and_database(cls, session, cluster_id, database): cluster_id = str(cluster_id) return session.query(cls).filter( cls.cluster_id == cluster_id, cls.db_name == database, cls._type == ContextTypes.DATA, ).order_by(cls.id.desc()).all() @classmethod def get_all_data_context_by_database_uri(cls, session, database_uri, org_id): return session.query(cls).filter( cls.database_uri == database_uri, cls._type == ContextTypes.DATA, cls.org_id == org_id, ).order_by(cls.id.desc()).all() @classmethod def get_all_session_contexts_by_identifier_and_type(cls, session, identifier_type, identifier, limit, offset, _type=ContextTypes.SESSION_CHAT): q = session.query(cls).filter( cls.identifier == str(identifier), cls.identifier_type == str(identifier_type), cls._type == _type, ) count = q.count() items = q.order_by(cls.id.desc()).limit(limit).offset(offset).all() return items, count def update_description_and_default(self, session, description, is_default): # reset all other data context to non-default if is_default: session.query(Context).filter( Context.org_id == self.org_id, Context._type == ContextTypes.DATA, Context.cluster_id == self.cluster_id, Context.db_name == self.db_name, ).update({"is_default": False}) self.description = description self.is_default = is_default session.flush() def update_table_description(self, session, table_name, description): data = self.load_data() if "db_schema" not in data: return if table_name not in data["db_schema"]: return data["db_schema"][table_name]["description"] = description self.update_data(data) def update_columns_description(self, session, table_name, columns, update_type): data = self.load_data() if "db_schema" not in data: logging.info("No db_schema in data") return if table_name not in data["db_schema"]: logging.info(f"No table %s in db_schema: %s", table_name, data["db_schema"].keys()) return for i in columns: column_name = i.get("name") description = i.get("description") if column_name not in data["db_schema"][table_name]["columns"]: logging.info(f"No column {column_name} in table {table_name}") continue if update_type == "override": data["db_schema"][table_name]["columns"][column_name]["description"] = description # TODO: support update_type == "merge" logging.info("Update columns description: %s", columns) self.update_data(data)