diff --git a/.gitignore b/.gitignore index 6c37af14..752ee168 100644 --- a/.gitignore +++ b/.gitignore @@ -182,4 +182,5 @@ cython_debug/ # PyPI configuration file .pypirc -.DS_Store \ No newline at end of file +.DS_Store +test.py \ No newline at end of file diff --git a/backend/apps/chat/schemas/embedding.py b/backend/apps/chat/schemas/embedding.py new file mode 100644 index 00000000..a25ce71a --- /dev/null +++ b/backend/apps/chat/schemas/embedding.py @@ -0,0 +1,276 @@ +import os +from typing import List, Dict, Tuple, Optional +from sqlalchemy import inspect, text +from sqlmodel import SQLModel, create_engine, Session +from langchain_openai import OpenAIEmbeddings +from langchain.vectorstores import PGVector +from langchain.schema import Document +import logging +from apps.chat.schemas.chat_base_schema import LLMConfig +from common.core.config import settings +from common.core.db import engine + +class SchemaEmbeddingManager: + def __init__(self, config: LLMConfig): + """Initialize SchemaEmbeddingManager + + Args: + db_uri: Database connection URI in format postgresql://user:password@host:port/database + embedding_model: Embedding model to use + """ + # self.db_uri = db_uri + self.embedding_model = OpenAIEmbeddings( + model=config.model_name, + openai_api_base=config.api_base_url, + api_key=config.api_key + ) + # self.engine = self._create_engine() + self.engine = engine + self._setup_vector_extension() + + + def _setup_vector_extension(self): + """Set up PgVector extension using SQLModel""" + try: + with Session(self.engine) as session: + # Enable vector extension + session.exec(text("CREATE EXTENSION IF NOT EXISTS vector;")) + session.commit() + + # Create vector type if needed + session.exec(text(""" + DO $$ + BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'vector') THEN + CREATE TYPE vector AS (x float8[]); + END IF; + END + $$; + """)) + session.commit() + logging.info("Vector extension setup completed successfully") + except Exception as e: + logging.error(f"Failed to setup vector extension: {str(e)}") + raise + + def get_session(self) -> Session: + """Get a new database session""" + return Session(self.engine) + + def extract_schema_info(self, schema_name: str = "public") -> List[Dict]: + """ + Extract database schema information and structure it + + Args: + schema_name: Schema name to extract, defaults to public + + Returns: + List of dictionaries containing schema information + """ + metadata = SQLModel.metadata + metadata.reflect(bind=self.engine, schema=schema_name) + + schema_info = [] + inspector = inspect(self.engine) + + for table_name in inspector.get_table_names(schema=schema_name): + table = metadata.tables[f"{schema_name}.{table_name}"] + table_comment = inspector.get_table_comment(table_name, schema=schema_name) + + # Table-level information + table_info = { + "type": "table", + "schema": schema_name, + "name": table_name, + "description": table_comment.get('text') if table_comment else f"Table {table_name} in schema {schema_name}", + "columns": [] + } + + # Column-level information + for column in inspector.get_columns(table_name, schema=schema_name): + col_name = column['name'] + # Get column comment directly from column info + col_comment = column.get('comment', f"Column {col_name} in table {table_name}") + + column_info = { + "type": "column", + "schema": schema_name, + "table": table_name, + "name": col_name, + "data_type": str(column['type']), + "description": col_comment, + "is_primary_key": col_name in [pk['name'] for pk in inspector.get_pk_constraint(table_name, schema=schema_name).get('constrained_columns', [])], + "foreign_key": None # Will be updated below + } + + # Get foreign key information + for fk in inspector.get_foreign_keys(table_name, schema=schema_name): + if col_name in fk['constrained_columns']: + column_info["foreign_key"] = f"{fk['referred_table']}.{fk['referred_columns'][0]}" + break + + table_info["columns"].append(column_info) + + schema_info.append(table_info) + + return schema_info + + def generate_schema_documents(self, schema_info: List[Dict]) -> List[Document]: + """ + Convert schema information into LangChain Document objects + + Args: + schema_info: Schema information from extract_schema_info + + Returns: + List of LangChain Document objects + """ + documents = [] + + for table in schema_info: + # Table-level document + table_content = f"Table {table['schema']}.{table['name']}: {table['description']}. " + table_content += f"Contains columns: {', '.join(col['name'] for col in table['columns'])}" + + table_doc = Document( + page_content=table_content, + metadata={ + "type": "table", + "schema": table["schema"], + "name": table["name"], + "description": table["description"], + "full_name": f"{table['schema']}.{table['name']}" + } + ) + documents.append(table_doc) + + # Column-level documents + for column in table["columns"]: + column_content = f"Column {table['schema']}.{table['name']}.{column['name']}: {column['description']}. " + column_content += f"Data type: {column['data_type']}. " + if column["is_primary_key"]: + column_content += "This is a primary key. " + if column["foreign_key"]: + column_content += f"Foreign key to {column['foreign_key']}." + + column_doc = Document( + page_content=column_content, + metadata={ + "type": "column", + "schema": table["schema"], + "table": table["name"], + "name": column["name"], + "description": column["description"], + "data_type": column["data_type"], + "full_name": f"{table['schema']}.{table['name']}.{column['name']}" + } + ) + documents.append(column_doc) + + return documents + + def store_embeddings(self, documents: List[Document], collection_name: str = "schema_embeddings"): + """Store schema document embeddings in PgVector + + Args: + documents: Documents from generate_schema_documents + collection_name: Vector collection name + """ + try: + PGVector.from_documents( + embedding=self.embedding_model, + documents=documents, + collection_name=collection_name, + connection_string=self.db_uri, + pre_delete_collection=True # Optionally recreate collection + ) + logging.info(f"Successfully stored {len(documents)} embeddings in collection {collection_name}") + except Exception as e: + logging.error(f"Failed to store embeddings: {str(e)}") + raise + + def get_relevant_schema( + self, + query: str, + collection_name: str = "schema_embeddings", + top_k: int = 5 + ) -> Tuple[List[Dict], List[Dict]]: + """Get relevant schema information based on user query + + Args: + query: User query text + collection_name: Vector collection name + top_k: Number of most relevant items to return + + Returns: + Tuple containing (relevant_tables_list, relevant_columns_list) + """ + try: + store = PGVector( + collection_name=collection_name, + connection_string=self.db_uri, + embedding_function=self.embedding_model, + ) + + # Execute similarity search + docs = store.similarity_search_with_score(query, k=top_k) + + relevant_tables = [] + relevant_columns = [] + seen_tables = set() + + for doc, score in docs: + metadata = doc.metadata + item = { + "name": metadata["full_name"], + "description": metadata["description"], + "type": metadata["type"], + "score": float(score) + } + + if metadata["type"] == "table" and metadata["full_name"] not in seen_tables: + relevant_tables.append(item) + seen_tables.add(metadata["full_name"]) + elif metadata["type"] == "column": + relevant_columns.append(item) + # Ensure associated table is also included + table_name = f"{metadata['schema']}.{metadata['table']}" + if table_name not in seen_tables: + relevant_tables.append({ + "name": table_name, + "description": f"Table containing column {metadata['name']}", + "type": "table", + "score": float(score) * 0.9 # Slightly lower table score + }) + seen_tables.add(table_name) + + return relevant_tables, relevant_columns + + except Exception as e: + logging.error(f"Error during schema search: {str(e)}") + raise + + def generate_schema_context(self, relevant_tables: List[Dict], relevant_columns: List[Dict]) -> str: + """ + Generate concise schema context for LLM usage + + Args: + relevant_tables: List of relevant tables + relevant_columns: List of relevant columns + + Returns: + Formatted schema context string + """ + context = "## Relevant Database Schema Information:\n\n" + + if relevant_tables: + context += "### Tables:\n" + for table in sorted(relevant_tables, key=lambda x: -x["score"]): + context += f"- {table['name']}: {table['description']} (relevance score: {table['score']:.2f})\n" + + if relevant_columns: + context += "\n### Columns:\n" + for column in sorted(relevant_columns, key=lambda x: -x["score"]): + context += f"- {column['name']}: {column['description']} (relevance score: {column['score']:.2f})\n" + + return context diff --git a/backend/apps/chat/schemas/llm.py b/backend/apps/chat/schemas/llm.py index 3771f258..f9815704 100644 --- a/backend/apps/chat/schemas/llm.py +++ b/backend/apps/chat/schemas/llm.py @@ -2,6 +2,7 @@ from langchain_community.utilities import SQLDatabase from langgraph.prebuilt import create_react_agent from langchain_core.prompts import ChatPromptTemplate from apps.chat.schemas.chat_base_schema import LLMConfig, LLMFactory +from apps.chat.schemas.schema_engine import SchemaEngine from apps.datasource.models.datasource import CoreDatasource from apps.db.db import exec_sql, get_uri from common.core.config import settings @@ -103,9 +104,13 @@ class AgentService: async def async_generate(self, question: str) -> AsyncGenerator[str, None]: chain = self.prompt | self.agent_executor - schema = self.db.get_table_info() + # schema = self.db.get_table_info() - async for chunk in chain.astream({"schema": schema, "question": question}): + schema_engine = SchemaEngine(engine=self.db._engine) + mschema = schema_engine.mschema + mschema_str = mschema.to_mschema() + + async for chunk in chain.astream({"schema": mschema_str, "question": question}): if not isinstance(chunk, dict): continue diff --git a/backend/apps/chat/schemas/m_schema.py b/backend/apps/chat/schemas/m_schema.py new file mode 100644 index 00000000..058e2f94 --- /dev/null +++ b/backend/apps/chat/schemas/m_schema.py @@ -0,0 +1,195 @@ +from apps.chat.schemas.utils import examples_to_str, read_json, write_json +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union + + +class MSchema: + def __init__(self, db_id: str = 'Anonymous', schema: Optional[str] = None): + self.db_id = db_id + self.schema = schema + self.tables = {} + self.foreign_keys = [] + + def add_table(self, name, fields={}, comment=None): + self.tables[name] = {"fields": fields.copy(), 'examples': [], 'comment': comment} + + def add_field(self, table_name: str, field_name: str, field_type: str = "", + primary_key: bool = False, nullable: bool = True, default: Any = None, + autoincrement: bool = False, comment: str = "", examples: list = [], **kwargs): + self.tables[table_name]["fields"][field_name] = { + "type": field_type, + "primary_key": primary_key, + "nullable": nullable, + "default": default if default is None else f'{default}', + "autoincrement": autoincrement, + "comment": comment, + "examples": examples.copy(), + **kwargs} + + def add_foreign_key(self, table_name, field_name, ref_schema, ref_table_name, ref_field_name): + self.foreign_keys.append([table_name, field_name, ref_schema, ref_table_name, ref_field_name]) + + def get_field_type(self, field_type, simple_mode=True)->str: + if not simple_mode: + return field_type + else: + return field_type.split("(")[0] + + def has_table(self, table_name: str) -> bool: + if table_name in self.tables.keys(): + return True + else: + return False + + def has_column(self, table_name: str, field_name: str) -> bool: + if self.has_table(table_name): + if field_name in self.tables[table_name]["fields"].keys(): + return True + else: + return False + else: + return False + + def get_field_info(self, table_name: str, field_name: str) -> Dict: + try: + return self.tables[table_name]['fields'][field_name] + except: + return {} + + def single_table_mschema(self, table_name: str, selected_columns: List = None, + example_num=3, show_type_detail=False) -> str: + table_info = self.tables.get(table_name, {}) + output = [] + table_comment = table_info.get('comment', '') + if table_comment is not None and table_comment != 'None' and len(table_comment) > 0: + if self.schema is not None and len(self.schema) > 0: + output.append(f"# Table: {self.schema}.{table_name}, {table_comment}") + else: + output.append(f"# Table: {table_name}, {table_comment}") + else: + if self.schema is not None and len(self.schema) > 0: + output.append(f"# Table: {self.schema}.{table_name}") + else: + output.append(f"# Table: {table_name}") + + field_lines = [] + for field_name, field_info in table_info['fields'].items(): + if selected_columns is not None and field_name.lower() not in selected_columns: + continue + + raw_type = self.get_field_type(field_info['type'], not show_type_detail) + field_line = f"({field_name}:{raw_type.upper()}" + if field_info['comment'] != '': + field_line += f", {field_info['comment'].strip()}" + else: + pass + + is_primary_key = field_info.get('primary_key', False) + if is_primary_key: + field_line += f", Primary Key" + + if len(field_info.get('examples', [])) > 0 and example_num > 0: + examples = field_info['examples'] + examples = [s for s in examples if s is not None] + examples = examples_to_str(examples) + if len(examples) > example_num: + examples = examples[:example_num] + + if raw_type in ['DATE', 'TIME', 'DATETIME', 'TIMESTAMP']: + examples = [examples[0]] + elif len(examples) > 0 and max([len(s) for s in examples]) > 20: + if max([len(s) for s in examples]) > 50: + examples = [] + else: + examples = [examples[0]] + else: + pass + if len(examples) > 0: + example_str = ', '.join([str(example) for example in examples]) + field_line += f", Examples: [{example_str}]" + else: + pass + else: + field_line += "" + field_line += ")" + + field_lines.append(field_line) + output.append('[') + output.append(',\n'.join(field_lines)) + output.append(']') + + return '\n'.join(output) + + def to_mschema(self, selected_tables: List = None, selected_columns: List = None, + example_num=3, show_type_detail=False) -> str: + """ + Convert to a MSchema string. + + Args: + selected_tables: None means select all tables + selected_columns: None means select all columns, format: ['table_name.column_name'] + example_num: Number of examples to show for each field + show_type_detail: Whether to show detailed type information + + Returns: + Formatted schema string + """ + output = [] + + output.append(f"【DB_ID】 {self.db_id}") + output.append(f"【Schema】") + + # Process selected tables and columns + if selected_tables is not None: + selected_tables = [s.lower() for s in selected_tables] + if selected_columns is not None: + selected_columns = [s.lower() for s in selected_columns] + selected_tables = [s.split('.')[0].lower() for s in selected_columns] + + # Process each table + for table_name, table_info in self.tables.items(): + if selected_tables is None or table_name.lower() in selected_tables: + cur_table_type = table_info.get('type', 'table') + column_names = list(table_info['fields'].keys()) + if selected_columns is not None: + cur_selected_columns = [c.lower() for c in column_names if f"{table_name}.{c}".lower() in selected_columns] + else: + cur_selected_columns = selected_columns + output.append(self.single_table_mschema(table_name, cur_selected_columns, example_num, show_type_detail)) + + # Add foreign key information, skip when table_type is view + if self.foreign_keys: + output.append("【Foreign keys】") + for fk in self.foreign_keys: + ref_schema = fk[2] + table1, column1, _, table2, column2 = fk + if selected_tables is None or \ + (table1.lower() in selected_tables and table2.lower() in selected_tables): + output.append(f"{fk[0]}.{fk[1]}={fk[3]}.{fk[4]}") + + return '\n'.join(output) + + def dump(self): + """ + Dump schema information to dictionary + + Returns: + Dictionary containing schema information + """ + schema_dict = { + "db_id": self.db_id, + "schema": self.schema, + "tables": self.tables, + "foreign_keys": self.foreign_keys + } + return schema_dict + + def save(self, file_path: str): + schema_dict = self.dump() + write_json(file_path, schema_dict) + + def load(self, file_path: str): + data = read_json(file_path) + self.db_id = data.get("db_id", "Anonymous") + self.schema = data.get("schema", None) + self.tables = data.get("tables", {}) + self.foreign_keys = data.get("foreign_keys", []) diff --git a/backend/apps/chat/schemas/schema_engine.py b/backend/apps/chat/schemas/schema_engine.py new file mode 100644 index 00000000..63910503 --- /dev/null +++ b/backend/apps/chat/schemas/schema_engine.py @@ -0,0 +1,126 @@ +import json, os +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union +from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, select, text +from sqlalchemy.engine import Engine +from llama_index.core import SQLDatabase +from apps.chat.schemas.utils import read_json, write_json, save_raw_text, examples_to_str +from apps.chat.schemas.m_schema import MSchema + + +class SchemaEngine(SQLDatabase): + def __init__(self, engine: Engine, schema: Optional[str] = None, metadata: Optional[MetaData] = None, + ignore_tables: Optional[List[str]] = None, include_tables: Optional[List[str]] = None, + sample_rows_in_table_info: int = 3, indexes_in_table_info: bool = False, + custom_table_info: Optional[dict] = None, view_support: bool = False, max_string_length: int = 300, + mschema: Optional[MSchema] = None, db_name: Optional[str] = ''): + super().__init__(engine, schema, metadata, ignore_tables, include_tables, sample_rows_in_table_info, + indexes_in_table_info, custom_table_info, view_support, max_string_length) + + self._db_name = db_name + # Dictionary to store table names and their corresponding schema + self._tables_schemas: Dict[str, str] = {} + + # If a schema is specified, filter by that schema and store that value for every table. + """ if schema: + self._usable_tables = [ + table_name for table_name in self._usable_tables + if self._inspector.has_table(table_name, schema) + ] + for table_name in self._usable_tables: + self._tables_schemas[table_name] = schema + else: + all_tables = [] + # Iterate through all available schemas + for s in self.get_schema_names(): + tables = self._inspector.get_table_names(schema=s) + all_tables.extend(tables) + for table in tables: + self._tables_schemas[table] = s + self._usable_tables = all_tables """ + d_schema = self.default_schema_name() + for table_name in self._usable_tables: + self._tables_schemas[table_name] = d_schema + + self._dialect = engine.dialect.name + if mschema is not None: + self._mschema = mschema + else: + self._mschema = MSchema(db_id=db_name) + self.init_mschema() + + @property + def mschema(self) -> MSchema: + """Return M-Schema""" + return self._mschema + + def get_pk_constraint(self, table_name: str) -> Dict: + return self._inspector.get_pk_constraint(table_name, self._tables_schemas[table_name] )['constrained_columns'] + + def get_table_comment(self, table_name: str): + try: + return self._inspector.get_table_comment(table_name, self._tables_schemas[table_name])['text'] + except: # sqlite does not support comments + return '' + + def default_schema_name(self) -> Optional[str]: + return self._inspector.default_schema_name + + def get_schema_names(self) -> List[str]: + return self._inspector.get_schema_names() + + def get_foreign_keys(self, table_name: str): + return self._inspector.get_foreign_keys(table_name, self._tables_schemas[table_name]) + + def get_unique_constraints(self, table_name: str): + return self._inspector.get_unique_constraints(table_name, self._tables_schemas[table_name]) + + def fectch_distinct_values(self, table_name: str, column_name: str, max_num: int = 5): + table = Table(table_name, self.metadata_obj, autoload_with=self._engine, schema=self._tables_schemas[table_name]) + # Construct SELECT DISTINCT query + query = select(table.c[column_name]).distinct().limit(max_num) + values = [] + with self._engine.connect() as connection: + result = connection.execute(query) + distinct_values = result.fetchall() + for value in distinct_values: + if value[0] is not None and value[0] != '': + values.append(value[0]) + return values + + def init_mschema(self): + for table_name in self._usable_tables: + table_comment = self.get_table_comment(table_name) + table_comment = '' if table_comment is None else table_comment.strip() + table_with_schema = self._tables_schemas[table_name] + '.' + table_name + self._mschema.add_table(table_with_schema, fields={}, comment=table_comment) + pks = self.get_pk_constraint(table_name) + + fks = self.get_foreign_keys(table_name) + for fk in fks: + referred_schema = fk['referred_schema'] + for c, r in zip(fk['constrained_columns'], fk['referred_columns']): + self._mschema.add_foreign_key(table_with_schema, c, referred_schema, fk['referred_table'], r) + + fields = self._inspector.get_columns(table_name, schema=self._tables_schemas[table_name]) + for field in fields: + field_type = f"{field['type']!s}" + field_name = field['name'] + primary_key = field_name in pks + field_comment = field.get("comment", None) + field_comment = "" if field_comment is None else field_comment.strip() + autoincrement = field.get('autoincrement', False) + default = field.get('default', None) + if default is not None: + default = f'{default}' + + try: + examples = self.fectch_distinct_values(table_name, field_name, 5) + except: + examples = [] + examples = examples_to_str(examples) + + self._mschema.add_field( + table_with_schema, field_name, field_type=field_type, primary_key=primary_key, + nullable=field['nullable'], default=default, autoincrement=autoincrement, + comment=field_comment, examples=examples + ) \ No newline at end of file diff --git a/backend/apps/chat/schemas/utils.py b/backend/apps/chat/schemas/utils.py new file mode 100644 index 00000000..efa6ecf3 --- /dev/null +++ b/backend/apps/chat/schemas/utils.py @@ -0,0 +1,80 @@ +import datetime +import decimal +import re +import json + + +def write_json(path, data): + with open(path, 'w', encoding='utf-8') as f: + json.dump(data, f, ensure_ascii=False, indent=2) + + +def read_json(path): + with open(path, 'r', encoding='utf-8') as f: + data = json.load(f) + return data + + +def read_text(filename)->str: + data = [] + with open(filename, 'r', encoding='utf-8') as file: + for line in file.readlines(): + line = line.strip() + data.append(line) + return data + + +def save_raw_text(filename, content): + with open(filename, 'w', encoding='utf-8') as file: + file.write(content) + + +def read_map_file(path): + data = {} + with open(path, 'r', encoding='utf-8') as f: + for line in f.readlines(): + line = line.strip().split('\t') + data[line[0]] = line[1].split('、') + data[line[0]].append(line[0]) + return data + + +def save_json(target_file,js,indent=4): + with open(target_file, 'w', encoding='utf-8') as f: + json.dump(js, f, ensure_ascii=False, indent=indent) + +def is_email(string): + pattern = r'^[\w\.-]+@[\w\.-]+\.\w+$' + match = re.match(pattern, string) + if match: + return True + else: + return False + + +def examples_to_str(examples: list) -> list[str]: + """ + from examples to a list of str + """ + values = examples + for i in range(len(values)): + if isinstance(values[i], datetime.date): + values = [values[i]] + break + elif isinstance(values[i], datetime.datetime): + values = [values[i]] + break + elif isinstance(values[i], decimal.Decimal): + values[i] = str(float(values[i])) + elif is_email(str(values[i])): + values = [] + break + elif 'http://' in str(values[i]) or 'https://' in str(values[i]): + values = [] + break + elif values[i] is not None and not isinstance(values[i], str): + pass + elif values[i] is not None and '.com' in values[i]: + pass + + return [str(v) for v in values if v is not None and len(str(v)) > 0] \ No newline at end of file diff --git a/frontend/components.d.ts b/frontend/components.d.ts index 505c5408..bb9f32f9 100644 --- a/frontend/components.d.ts +++ b/frontend/components.d.ts @@ -11,7 +11,9 @@ declare module 'vue' { ElAvatar: typeof import('element-plus/es')['ElAvatar'] ElButton: typeof import('element-plus/es')['ElButton'] ElCheckbox: typeof import('element-plus/es')['ElCheckbox'] + ElCol: typeof import('element-plus/es')['ElCol'] ElDialog: typeof import('element-plus/es')['ElDialog'] + ElDrawer: typeof import('element-plus/es')['ElDrawer'] ElDropdown: typeof import('element-plus/es')['ElDropdown'] ElDropdownItem: typeof import('element-plus/es')['ElDropdownItem'] ElDropdownMenu: typeof import('element-plus/es')['ElDropdownMenu'] @@ -23,10 +25,12 @@ declare module 'vue' { ElMenuItem: typeof import('element-plus/es')['ElMenuItem'] ElOption: typeof import('element-plus/es')['ElOption'] ElPagination: typeof import('element-plus/es')['ElPagination'] + ElRow: typeof import('element-plus/es')['ElRow'] ElSelect: typeof import('element-plus/es')['ElSelect'] ElSlider: typeof import('element-plus/es')['ElSlider'] ElTable: typeof import('element-plus/es')['ElTable'] ElTableColumn: typeof import('element-plus/es')['ElTableColumn'] + ElTableV2: typeof import('element-plus/es')['ElTableV2'] ElTabPane: typeof import('element-plus/es')['ElTabPane'] ElTabs: typeof import('element-plus/es')['ElTabs'] ElTooltip: typeof import('element-plus/es')['ElTooltip']