text2sql_backend / src /scheme.py
LightRT's picture
Project Completion Commit
52adb86
raw
history blame
873 Bytes
from sqlalchemy import create_engine , inspect
from langchain_core.documents import Document
def create_scheme(database_url : str) -> list[dict] :
scheme_docs = []
engine = create_engine(database_url)
inspector = inspect(engine)
tables = inspector.get_table_names()
for table in tables :
clean_table = table.replace("_" , " ")
columns = inspector.get_columns(table)
clean_columns = []
for col in columns :
clean_col = col['name'].replace("_" , " ")
clean_columns.append(clean_col)
doc = f"Table: {clean_table}.\nColumns: {', '.join(clean_columns)}"
scheme_docs.append(
Document(
page_content=doc,
metadata={
"table_name" : table
}
)
)
return scheme_docs