From b203584cbe4d018982d1d4c57403ee8499ab843d Mon Sep 17 00:00:00 2001 From: fit2cloud-chenyw Date: Wed, 23 Jul 2025 11:51:19 +0800 Subject: [PATCH] perf: Delete embedding --- backend/apps/datasource/crud/ds_embedding.py | 250 ------------------- 1 file changed, 250 deletions(-) delete mode 100644 backend/apps/datasource/crud/ds_embedding.py diff --git a/backend/apps/datasource/crud/ds_embedding.py b/backend/apps/datasource/crud/ds_embedding.py deleted file mode 100644 index 7fe50ee8..00000000 --- a/backend/apps/datasource/crud/ds_embedding.py +++ /dev/null @@ -1,250 +0,0 @@ -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() - -