Merge remote-tracking branch 'origin/main'

.gitignore

@@ -182,4 +182,5 @@ cython_debug/
 # PyPI configuration file
 .pypirc
 
-.DS_Store
+.DS_Store
+test.py

backend/alembic/versions/0a6f11be9be4_005_table_and_field.py

@@ -0,0 +1,49 @@
+"""005_table_and_field
+
+Revision ID: 0a6f11be9be4
+Revises: 8fe654655905
+Create Date: 2025-05-15 10:20:25.686576
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+
+
+# revision identifiers, used by Alembic.
+revision = '0a6f11be9be4'
+down_revision = '8fe654655905'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('core_field',
+    sa.Column('id', sa.Integer(), sa.Identity(always=True), nullable=False),
+    sa.Column('ds_id', sa.BigInteger(), nullable=True),
+    sa.Column('table_id', sa.BigInteger(), nullable=True),
+    sa.Column('checked', sa.Boolean(), nullable=False),
+    sa.Column('field_name', sa.Text(), nullable=True),
+    sa.Column('field_type', sqlmodel.sql.sqltypes.AutoString(length=128), nullable=True),
+    sa.Column('field_comment', sa.Text(), nullable=True),
+    sa.Column('custom_comment', sa.Text(), nullable=True),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_table('core_table',
+    sa.Column('id', sa.Integer(), sa.Identity(always=True), nullable=False),
+    sa.Column('ds_id', sa.BigInteger(), nullable=True),
+    sa.Column('checked', sa.Boolean(), nullable=False),
+    sa.Column('table_name', sa.Text(), nullable=True),
+    sa.Column('table_comment', sa.Text(), nullable=True),
+    sa.Column('custom_comment', sa.Text(), nullable=True),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('core_table')
+    op.drop_table('core_field')
+    # ### end Alembic commands ###

backend/apps/chat/api/chat.py

@@ -3,11 +3,10 @@ from fastapi.responses import StreamingResponse
 from sqlmodel import select
 from apps.chat.schemas.chat_base_schema import LLMConfig
 from apps.chat.schemas.chat_schema import ChatQuestion
-from apps.chat.schemas.llm import AgentService, LLMService
+from apps.chat.schemas.llm import AgentService
 from apps.datasource.models.datasource import CoreDatasource
 from apps.system.models.system_model import AiModelDetail
 from common.core.deps import SessionDep
-from sse_starlette.sse import EventSourceResponse
 import json
 import asyncio
 

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

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
                 

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", [])

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
+                )

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]

backend/apps/datasource/api/datasource.py

@@ -1,7 +1,7 @@
 from fastapi import APIRouter
 from ..crud.datasource import get_datasource_list, check_status, create_ds, update_ds, delete_ds, getTables, getFields, execSql
 from common.core.deps import SessionDep
-from ..models.datasource import CoreDatasource
+from ..models.datasource import CoreDatasource, CreateDatasource
 
 router = APIRouter(tags=["datasource"], prefix="/datasource")
 
@@ -17,7 +17,7 @@ async def check(session: SessionDep, ds: CoreDatasource):
 
 
 @router.post("/add", response_model=CoreDatasource)
-async def add(session: SessionDep, ds: CoreDatasource):
+async def add(session: SessionDep, ds: CreateDatasource):
     return create_ds(session, ds)
 
 

backend/apps/datasource/crud/datasource.py

@@ -1,9 +1,13 @@
 from sqlmodel import select
-from ..models.datasource import CoreDatasource, DatasourceConf
+from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema
 import datetime
 from common.core.deps import SessionDep
 from apps.db.db import get_session, get_tables, get_fields, exec_sql
-from sqlalchemy import text
+from sqlalchemy import text, and_
+from common.utils.utils import deepcopy_ignore_extra
+from typing import List
+from ..crud.table import delete_table_by_ds_id
+from ..crud.field import delete_field_by_ds_id
 
 
 def get_datasource_list(session: SessionDep) -> CoreDatasource:
@@ -26,21 +30,30 @@ def check_status(session: SessionDep, ds: CoreDatasource):
     return False
 
 
-def create_ds(session: SessionDep, ds: CoreDatasource):
+def create_ds(session: SessionDep, create_ds: CreateDatasource):
+    ds = CoreDatasource()
+    deepcopy_ignore_extra(create_ds, ds)
     ds.create_time = datetime.datetime.now()
     status = check_status(session, ds)
     ds.status = "Success" if status is True else "Fail"
     record = CoreDatasource(**ds.model_dump())
-    # get tables and fields
-    if status:
-        pass
     session.add(record)
+    session.flush()
+    session.refresh(record)
+    ds.id = record.id
     session.commit()
+
+    # save tables and fields
+    if status:
+        sync_table(session, ds, create_ds.tables)
+
     return ds
 
 
 def update_ds(session: SessionDep, ds: CoreDatasource):
     ds.id = int(ds.id)
+    status = check_status(session, ds)
+    ds.status = "Success" if status is True else "Fail"
     record = session.exec(select(CoreDatasource).where(CoreDatasource.id == ds.id)).first()
     update_data = ds.model_dump(exclude_unset=True)
     for field, value in update_data.items():
@@ -54,6 +67,8 @@ def delete_ds(session: SessionDep, id: int):
     term = session.exec(select(CoreDatasource).where(CoreDatasource.id == id)).first()
     session.delete(term)
     session.commit()
+    delete_table_by_ds_id(session, id)
+    delete_field_by_ds_id(session, id)
     return {
         "message": f"Datasource with ID {id} deleted successfully."
     }
@@ -74,3 +89,65 @@ def getFields(session: SessionDep, id: int, table_name: str):
 def execSql(session: SessionDep, id: int, sql: str):
     ds = session.exec(select(CoreDatasource).where(CoreDatasource.id == id)).first()
     return exec_sql(ds, sql)
+
+
+def sync_table(session: SessionDep, ds: CoreDatasource, tables: List[CoreTable]):
+    id_list = []
+    for item in tables:
+        statement = select(CoreTable).where(and_(CoreTable.ds_id == ds.id, CoreTable.table_name == item.table_name))
+        record = session.exec(statement).first()
+        # update exist table, only update table_comment
+        if record is not None:
+            item.id = record.id
+            id_list.append(record.id)
+
+            record.table_comment = item.table_comment
+            session.add(record)
+            session.commit()
+        else:
+            # save new table
+            table = CoreTable(ds_id=ds.id, checked=True, table_name=item.table_name, table_comment=item.table_comment,
+                              custom_comment=item.table_comment)
+            session.add(table)
+            session.flush()
+            session.refresh(table)
+            item.id = table.id
+            id_list.append(table.id)
+            session.commit()
+
+        # sync field
+        fields = getFields(session, ds.id, item.table_name)
+        sync_fields(session, ds, item, fields)
+
+    if len(id_list) > 0:
+        session.query(CoreTable).filter(and_(CoreTable.ds_id == ds.id, CoreTable.id.not_in(id_list))).delete(
+            synchronize_session=False)
+
+
+def sync_fields(session: SessionDep, ds: CoreDatasource, table: CoreTable, fields: List[ColumnSchema]):
+    id_list = []
+    for item in fields:
+        statement = select(CoreField).where(
+            and_(CoreField.table_id == table.id, CoreField.field_name == item.fieldName))
+        record = session.exec(statement).first()
+        if record is not None:
+            item.id = record.id
+            id_list.append(record.id)
+
+            record.field_comment = item.fieldComment
+            session.add(record)
+            session.commit()
+        else:
+            field = CoreField(ds_id=ds.id, table_id=table.id, checked=True, field_name=item.fieldName,
+                              field_type=item.fieldType, field_comment=item.fieldComment,
+                              custom_comment=item.fieldComment)
+            session.add(field)
+            session.flush()
+            session.refresh(field)
+            item.id = field.id
+            id_list.append(field.id)
+            session.commit()
+
+    if len(id_list) > 0:
+        session.query(CoreField).filter(and_(CoreField.table_id == table.id, CoreField.id.not_in(id_list))).delete(
+            synchronize_session=False)

backend/apps/datasource/crud/field.py

@@ -0,0 +1,7 @@
+from common.core.deps import SessionDep
+from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema
+
+
+def delete_field_by_ds_id(session: SessionDep, id: int):
+    session.query(CoreField).filter(CoreField.ds_id == id).delete(synchronize_session=False)
+    session.commit()

backend/apps/datasource/crud/table.py

@@ -0,0 +1,7 @@
+from common.core.deps import SessionDep
+from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema
+
+
+def delete_table_by_ds_id(session: SessionDep, id: int):
+    session.query(CoreTable).filter(CoreTable.ds_id == id).delete(synchronize_session=False)
+    session.commit()

backend/apps/datasource/models/datasource.py

@@ -2,7 +2,7 @@ from sqlmodel import SQLModel, Field
 from sqlalchemy import Column, Text, BigInteger, DateTime, Integer, Identity
 from datetime import datetime
 from pydantic import BaseModel
-from typing import List
+from typing import List, Optional
 
 
 class CoreDatasource(SQLModel, table=True):
@@ -17,20 +17,39 @@ class CoreDatasource(SQLModel, table=True):
     status: str = Field(max_length=64, nullable=True)
 
 
-class CoreTable(BaseModel):
-    checked: str = True
-    table_name: str = ''
-    table_comment: str = ''
-    custom_comment: str = ''
-    table_fields: List = []
+class CoreTable(SQLModel, table=True):
+    __tablename__ = "core_table"
+    id: int = Field(sa_column=Column(Integer, Identity(always=True), nullable=False, primary_key=True))
+    ds_id:int = Field(sa_column=Column(BigInteger()))
+    checked: bool = Field(default=True)
+    table_name: str = Field(sa_column=Column(Text))
+    table_comment: str = Field(sa_column=Column(Text))
+    custom_comment: str = Field(sa_column=Column(Text))
 
 
-class CoreField(BaseModel):
-    checked: str = True
-    field_name: str = ''
-    field_type: str = ''
-    field_comment: str = ''
-    custom_comment: str = ''
+class CoreField(SQLModel, table=True):
+    __tablename__ = "core_field"
+    id: int = Field(sa_column=Column(Integer, Identity(always=True), nullable=False, primary_key=True))
+    ds_id: int = Field(sa_column=Column(BigInteger()))
+    table_id: int = Field(sa_column=Column(BigInteger()))
+    checked: bool = Field(default=True)
+    field_name: str = Field(sa_column=Column(Text))
+    field_type: str = Field(max_length=128, nullable=True)
+    field_comment: str = Field(sa_column=Column(Text))
+    custom_comment: str = Field(sa_column=Column(Text))
+
+
+# datasource create obj
+class CreateDatasource(BaseModel):
+    id: int = None
+    name: str = ''
+    description: str = ''
+    type: str = ''
+    configuration: str = ''
+    create_time: Optional[datetime] = None
+    create_by: int = 0
+    status: str = ''
+    tables: List[CoreTable] = []
 
 
 class DatasourceConf(BaseModel):

backend/common/utils/utils.py

@@ -37,3 +37,13 @@ def verify_password_reset_token(token: str) -> str | None:
         return str(decoded_token["sub"])
     except InvalidTokenError:
         return None
+
+
+def deepcopy_ignore_extra(src, dest):
+    import copy
+    for attr in vars(src):
+        if hasattr(dest, attr):
+            src_value = getattr(src, attr)
+            dest_value = copy.deepcopy(src_value)  # deep copy
+            setattr(dest, attr, dest_value)
+    return dest

backend/pyproject.toml

@@ -26,9 +26,9 @@ dependencies = [
     "langgraph>=0.3,<0.4",
     "vllm>=0.8.5",
     "dashscope>=1.14.0,<2.0.0",
-    "sse-starlette>=1.8.0,<2.0.0",
     "pymysql (>=1.1.1,<2.0.0)",
     "cryptography (>=44.0.3,<45.0.0)",
+    "llama_index>=0.12.35"
 ]
 [[tool.uv.index]]
 url = "https://pypi.tuna.tsinghua.edu.cn/simple"

frontend/components.d.ts

@@ -11,6 +11,7 @@ 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']

frontend/src/views/ds/form.vue

@@ -137,6 +137,7 @@ const save = async(formEl: FormInstance | undefined) => {
           emit('refresh')
         })
       } else {
+        form.value.tables = [{"table_name":"core_dataset_table"},{"table_name":"pre_user"}]
         datasourceApi.add(form.value).then((res: any) => {
           console.log(res)
           close()