Merge remote-tracking branch 'origin/main'

This commit is contained in:
wangjiahao
2025-05-16 16:06:54 +08:00
8 changed files with 375 additions and 387 deletions
-276
View File
@@ -1,276 +0,0 @@
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
+6 -1
View File
@@ -1,6 +1,6 @@
from fastapi import APIRouter
from ..crud.datasource import get_datasource_list, check_status, create_ds, update_ds, delete_ds, getTables, getFields, \
execSql, update_table_and_fields, getTablesByDs, chooseTables
execSql, update_table_and_fields, getTablesByDs, chooseTables, preview
from common.core.deps import SessionDep
from ..models.datasource import CoreDatasource, CreateDatasource, EditObj, CoreTable
from ..crud.table import get_tables_by_ds_id
@@ -73,3 +73,8 @@ async def field_list(session: SessionDep, id: int):
@router.post("/editLocalComment")
async def edit_local(session: SessionDep, data: EditObj):
update_table_and_fields(session, data)
@router.post("/previewData/{id}")
async def edit_local(session: SessionDep, id: int, data: EditObj):
return preview(session, id, data)
@@ -167,3 +167,11 @@ def update_table_and_fields(session: SessionDep, data: EditObj):
update_table(session, data.table)
for field in data.fields:
update_field(session, field)
def preview(session: SessionDep, id: int, data: EditObj):
ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
sql: str = ""
if ds.type == "mysql":
sql = f"""SELECT {", ".join([f.field_name for f in data.fields if f.checked])} FROM {data.table.table_name} LIMIT 100"""
return exec_sql(ds, sql)
@@ -0,0 +1,250 @@
import logging
from typing import Dict, List
from langchain.schema import Document
from langchain_openai import OpenAIEmbeddings
from langchain.vectorstores.pgvector import PGVector
from sqlalchemy import Tuple
from apps.chat.schemas.schema_engine import SchemaEngine
from apps.datasource.models.datasource import CoreDatasource
from apps.db.db import get_uri
from langchain_community.utilities import SQLDatabase
from sqlmodel import Session, select
from apps.system.models.system_model import AiModelDetail
from common.core.db import engine
class DsEmbeddingManage:
def __init__(self, ds: CoreDatasource):
self.ds = ds
db_uri = get_uri(ds)
self.db = SQLDatabase.from_uri(db_uri)
# Get AI model configuration from database
with Session(engine) as session:
aimodel = session.exec(select(AiModelDetail).where(
AiModelDetail.status == True,
AiModelDetail.api_key.is_not(None)
)).first()
if not aimodel:
raise RuntimeError("No available AI model configuration found")
# Initialize embedding model with AI model config
self.embedding_model = OpenAIEmbeddings(
model=aimodel.name,
openai_api_base=aimodel.endpoint if aimodel.endpoint else None,
api_key=aimodel.api_key,
)
def getSchema(self):
schema_engine = SchemaEngine(engine=self.db._engine)
mschema = schema_engine.mschema
mschema_str = mschema.to_mschema()
return mschema_str
def mschema_to_documents(self, mschema_str: str) -> List[Document]:
"""Convert M-Schema string to list of Documents
Args:
mschema_str: M-Schema format string
Returns:
List of Documents for vector storage
"""
docs = []
current_table = ""
table_content = []
table_line = ""
current_table_col_docs = []
# Split schema into lines
lines = mschema_str.split('\n')
for line in lines:
# Skip empty lines and header sections
if not line.strip() or line.startswith('【'):
continue
# Process table header
if line.startswith('# Table:'):
# Store previous table if exists
if current_table and table_content:
doc_content = f"Table: {current_table}\n" + '\n'.join(table_content)
table_doc = Document(
page_content=doc_content,
metadata={
"type": "table",
"name": current_table,
"origin": table_line
}
)
docs.append(table_doc)
docs.extend(current_table_col_docs)
# Start new table
current_table = line.split('# Table: ')[1].split(',')[0].strip()
table_content = []
current_table_col_docs = []
table_line = line
# If there's a table comment, add it
if ',' in line:
table_content.append(f"Description: {line.split(',')[1].strip()}. Contains columns: ")
continue
# Add column information to current table
if line.strip() and not line.startswith('[') and not line.startswith(']') and line.startswith('('):
# Clean up the line
column_line = line.strip('(),')
column_info_list = column_line.split(',')
column_name = column_info_list[0].split(':')[0]
table_content.append(column_name + ",\n")
column_content = column_line
column_doc = Document(
page_content=column_content,
metadata = {
"type": "column",
"name": column_name,
"table": current_table,
"origin": line
}
)
current_table_col_docs.append(column_doc)
# Don't forget to add the last table
if current_table and table_content:
doc_content = f"Table: {current_table}\n" + '\n'.join(table_content)
docs.append(Document(
page_content=doc_content,
metadata={
"type": "table",
"name": current_table,
"origin": table_line
}
))
docs.extend(current_table_col_docs)
return docs
def store_schema_embeddings(self):
"""Store schema embeddings in PGVector"""
try:
# Get schema string
mschema_str = self.getSchema()
# Convert to documents
docs = self.mschema_to_documents(mschema_str)
# print(docs)
# Store in PGVector
PGVector.from_documents(
embedding=self.embedding_model,
documents=docs,
collection_name="schema_embeddings",
connection_string=str(self.db._engine.url),
pre_delete_collection=True
)
return f"Successfully stored {len(docs)} schema documents in vector store"
except Exception as e:
raise RuntimeError(f"Failed to store schema embeddings: {str(e)}")
def get_relevant_schema(
self,
query: str,
collection_name: str = "schema_embeddings",
top_k: int = 5
) -> 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:
List of dictionaries containing relevant schema information
"""
try:
store = PGVector(
collection_name=collection_name,
connection_string=str(self.db._engine.url),
embedding_function=self.embedding_model,
)
# Execute similarity search
docs = store.similarity_search_with_score(query, k=top_k)
relevant_tables = {}
for doc, score in docs:
metadata = doc.metadata
item = {
"name": metadata["name"],
"type": metadata["type"],
"origin": metadata["origin"],
"score": float(score)
}
if metadata["type"] == "table":
table_name = metadata["name"]
item["columns"] = []
if table_name not in relevant_tables:
relevant_tables[table_name] = item
else:
relevant_tables[table_name]["score"] = max(relevant_tables[table_name]["score"], item["score"])
relevant_tables[table_name]["origin"] = item["origin"]
elif metadata["type"] == "column":
table_name = f"{metadata['table']}"
if table_name not in relevant_tables:
relevant_tables[table_name] = {
"name": table_name,
"type": "table",
"score": float(score) * 0.9,
"columns": [item]
}
else:
relevant_tables[table_name]["columns"].append(item)
# Convert dictionary to sorted list
relevant_tables_list = list(relevant_tables.values())
# Sort by score in descending order
relevant_tables_list.sort(key=lambda x: x["score"], reverse=True)
return relevant_tables_list
except Exception as e:
logging.error(f"Error during schema search: {str(e)}")
raise
def generate_schema_context(self, relevant_tables_list: List[Dict]) -> str:
context = "## Relevant Database Schema Information:\n\n【Schema】\n"
if relevant_tables_list:
for table in sorted(relevant_tables_list, key=lambda x: -x["score"]):
if table['origin'] is not None:
context += f"\n{table['origin']}"
else:
context += f"# Table: {table['name']}"
context += f" (relevance score: {table['score']:.2f})\n"
columns: List[Dict] = table["columns"]
# context += f"[\n{',\n'.join(f'({column["origin"].strip('(),')} , relevance score: {column["score"]:.2f})' for column in columns)}]\n"
context += "\n[\n" + ",\n".join(column["origin"] for column in columns) + "\n]\n"
return context
def get_schema_with_query(self,
query: str,
collection_name: str = "schema_embeddings",
top_k: int = 5):
list = self.get_relevant_schema(query, collection_name, top_k)
context = self.generate_schema_context(list)
print(context)
return context
def execute(self):
"""Main execution method"""
return self.store_schema_embeddings()
+5 -3
View File
@@ -9,10 +9,12 @@ def delete_field_by_ds_id(session: SessionDep, id: int):
def get_fields_by_table_id(session: SessionDep, id: int):
return session.query(CoreField).filter(CoreField.table_id == id).order_by(
CoreField.field_name.asc()).all()
return session.query(CoreField).filter(CoreField.table_id == id).all()
def update_field(session: SessionDep, item: CoreField):
session.add(item)
record = session.query(CoreField).filter(CoreField.id == item.id).first()
record.checked = item.checked
record.custom_comment = item.custom_comment
session.add(record)
session.commit()
+4 -1
View File
@@ -14,5 +14,8 @@ def get_tables_by_ds_id(session: SessionDep, id: int):
def update_table(session: SessionDep, item: CoreTable):
session.add(item)
record = session.query(CoreTable).filter(CoreTable.id == item.id).first()
record.checked = item.checked
record.custom_comment = item.custom_comment
session.add(record)
session.commit()
+4 -1
View File
@@ -11,5 +11,8 @@ export const datasourceApi = {
getFields: (id: Number, table_name: string) => request.post(`/datasource/getFields/${id}/${table_name}`),
execSql: (id: Number, sql: string) => request.post(`/datasource/execSql/${id}/${sql}`),
chooseTables: (id: Number, data: any) => request.post(`/datasource/chooseTables/${id}`, data),
tableList: (id: Number) => request.post(`/datasource/tableList/${id}`)
tableList: (id: Number) => request.post(`/datasource/tableList/${id}`),
fieldList: (id: Number) => request.post(`/datasource/fieldList/${id}`),
edit: (data: any) => request.post('/datasource/editLocalComment', data),
previewData: (id: Number, data: any) => request.post(`/datasource/previewData/${id}`, data),
}
+98 -105
View File
@@ -5,145 +5,132 @@
<el-button text :icon="ArrowLeft" @click="back()" />
{{ props.dsName }}
</div>
<el-button type="primary" @click="save()">
Save
</el-button>
</el-row>
<el-row class="data">
<el-table-v2
:columns="tableColumns"
:data="tableData"
:width="1000"
:height="800"
fixed
/>
<el-row class="container" :gutter="20">
<el-col :span="4">
Tables
<el-input v-model="searchValue" placeholder="Search" />
<div>
<div v-for="item in tableList" class="table-item" @click="clickTable(item)">{{ item.table_name }}</div>
</div>
</el-col>
<el-col :span="20">
<div v-if="fieldList.length === 0">No data, select a table from left</div>
<div v-else>
<div style="display: flex;justify-content: space-between;align-items: center;">
<div style="display: flex;justify-content: start;align-items: center;">
<span>{{ currentTable.table_name }}</span>
<span>&nbsp;|&nbsp;Comment:&nbsp;</span>
<span>
<el-input v-model="currentTable.custom_comment" />
</span>
</div>
<el-button type="primary" @click="save()">
Save
</el-button>
</div>
<el-tabs v-model="activeName" class="demo-tabs" @tab-click="handleClick">
<el-tab-pane label="Table Schema" name="schema">
<el-table :data="fieldList" style="width: 100%">
<el-table-column prop="field_name" label="Name" width="180" />
<el-table-column prop="field_type" label="Type" width="180" />
<el-table-column prop="field_comment" label="Comment" />
<el-table-column label="Custom Comment">
<template #default="scope">
<div style="display: flex; align-items: center">
<el-input v-model="scope.row.custom_comment" />
</div>
</template>
</el-table-column>
<el-table-column label="Status" width="180">
<template #default="scope">
<div style="display: flex; align-items: center">
<el-switch
v-model="scope.row.checked"
size="small"
/>
</div>
</template>
</el-table-column>
</el-table>
</el-tab-pane>
<el-tab-pane label="Preview" name="preview">
<div>Preview 100 items</div>
<el-table :data="previewData.data" style="width: 100%;height: 600px;">
<el-table-column v-for="c in previewData.fields" :prop="c" :label="c"/>
</el-table>
</el-tab-pane>
</el-tabs>
</div>
</el-col>
</el-row>
<el-drawer
v-model="drawer"
direction="btt"
:destroy-on-close="true"
:close-on-click-modal="false"
size="80%"
>
<template #header>
<div>Detail</div>
</template>
<template #default>
<el-row :gutter="20">
<el-table-v2
:columns="fieldColumns"
:data="fieldData"
:width="1000"
:height="800"
fixed
/>
</el-row>
</template>
<template #footer>
<div style="flex: auto">
<el-button @click="cancelClick">Cancel</el-button>
<el-button type="primary" @click="confirmClick">Confirm</el-button>
</div>
</template>
</el-drawer>
</div>
</template>
<script lang="tsx" setup>
import { ref } from 'vue'
import { datasourceApi } from '@/api/datasource'
import { ElButton } from 'element-plus'
import { h } from 'vue'
import { onMounted } from 'vue'
import { ArrowLeft } from '@element-plus/icons-vue'
import type { TabsPaneContext } from 'element-plus'
const props = defineProps({
dsId: { type: [Number], required: true },
dsName: { type: [String], required: true }
})
const drawer = ref<boolean>(false)
const tableColumns = ref<any>([
{
key:"tableName",
dataKey:"tableName",
title: 'Table Name',
width: 150,
},
{
key:"tableComment",
dataKey:"tableComment",
title: 'Table Comment',
width: 300,
},
{
key: 'operations',
title: 'Operations',
cellRenderer: ({ rowData }: { rowData: any }) => {
return h(ElButton,{
onClick: () => showFields(rowData.tableName)
},"Show Fields")
},
width: 150,
align: 'center',
}
])
const tableData = ref<any>([])
const fieldColumns = ref<any>([
{
key:"fieldName",
dataKey:"fieldName",
title: 'Field Name',
width: 150,
},
{
key:"fieldType",
dataKey:"fieldType",
title: 'Field Type',
width: 150,
},
{
key:"fieldComment",
dataKey:"fieldComment",
title: 'Field Comment',
width: 300,
}
])
const fieldData = ref<any>([])
const dsId = ref<Number>(0)
const searchValue = ref('')
const tableList = ref<any>([])
const currentTable = ref<any>({})
const fieldList = ref<any>([])
const previewData = ref<any>({})
const activeName = ref('schema')
const buildData = () => {
return {"table": currentTable.value, "fields": fieldList.value}
}
const back = () => {
history.back()
}
const save = () => {
}
const showFields = (tableName: string) =>{
drawer.value = true
datasourceApi.getFields(dsId.value, tableName).then((res) => {
fieldData.value = res
datasourceApi.edit(buildData()).then(() => {
ElMessage({
message: 'Save success',
type: 'success',
showClose: true
})
})
}
const cancelClick = () => {
drawer.value = false
const clickTable = (table: any) => {
currentTable.value = table
datasourceApi.fieldList(table.id).then((res) => {
fieldList.value = res
datasourceApi.previewData(dsId.value, buildData()).then((res) => {
previewData.value = res
})
})
}
const confirmClick = () => {
// save something
cancelClick()
const handleClick = (tab: TabsPaneContext) => {
if(tab.paneName === 'preview') {
datasourceApi.previewData(dsId.value, buildData()).then((res) => {
previewData.value = res
})
}
}
onMounted(() => {
dsId.value = props.dsId
datasourceApi.getTables(props.dsId).then((res) => {
tableData.value = res
fieldList.value = []
datasourceApi.tableList(props.dsId).then((res) => {
tableList.value = res
})
})
defineExpose({ open })
</script>
<style lang="less" scoped>
@@ -158,5 +145,11 @@ defineExpose({ open })
display: flex;
align-items: center;
}
.container{
height: 100%;
.table-item{
margin: 10px 0;
}
}
}
</style>