From faa79c099427296bb6e9a9467daf5d4664fd314c Mon Sep 17 00:00:00 2001 From: fit2cloud-chenyw Date: Fri, 16 May 2025 15:28:19 +0800 Subject: [PATCH] feat: Schema embedding --- backend/apps/chat/schemas/embedding.py | 276 ------------------- backend/apps/datasource/crud/ds_embedding.py | 250 +++++++++++++++++ 2 files changed, 250 insertions(+), 276 deletions(-) delete mode 100644 backend/apps/chat/schemas/embedding.py create mode 100644 backend/apps/datasource/crud/ds_embedding.py diff --git a/backend/apps/chat/schemas/embedding.py b/backend/apps/chat/schemas/embedding.py deleted file mode 100644 index a25ce71a..00000000 --- a/backend/apps/chat/schemas/embedding.py +++ /dev/null @@ -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 diff --git a/backend/apps/datasource/crud/ds_embedding.py b/backend/apps/datasource/crud/ds_embedding.py new file mode 100644 index 00000000..7fe50ee8 --- /dev/null +++ b/backend/apps/datasource/crud/ds_embedding.py @@ -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() + +