# Copyright (c) 2025 Oracle and/or its affiliates. # # This program is free software; you can redistribute it and/or modify # it under the terms of the GNU General Public License, version 2.0, as # published by the Free Software Foundation. # # This program is designed to work with certain software (including # but not limited to OpenSSL) that is licensed under separate terms, # as designated in a particular file or component or in included license # documentation. The authors of MySQL hereby grant you an # additional permission to link the program and your derivative works # with the separately licensed software that they have either included with # the program or referenced in the documentation. # # Without limiting anything contained in the foregoing, this file, # which is part of MySQL Connector/Python, is also subject to the # Universal FOSS Exception, version 1.0, a copy of which can be found at # http://oss.oracle.com/licenses/universal-foss-exception. # # This program is distributed in the hope that it will be useful, but # WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. # See the GNU General Public License, version 2.0, for more details. # # You should have received a copy of the GNU General Public License # along with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA """Embeddings integration utilities for MySQL Connector/Python. Provides MyEmbeddings class to generate embeddings via MySQL HeatWave using ML_EMBED_TABLE and ML_EMBED_ROW. """ from typing import Dict, List, Optional import pandas as pd from langchain_core.embeddings import Embeddings from pydantic import PrivateAttr from mysql.ai.utils import ( atomic_transaction, execute_sql, format_value_sql, source_schema, sql_table_from_df, sql_table_to_df, temporary_sql_tables, ) from mysql.connector.abstracts import MySQLConnectionAbstract class MyEmbeddings(Embeddings): """ Embedding generator class that uses a MySQL database to compute embeddings for input text. This class batches input text into temporary SQL tables, invokes MySQL's ML_EMBED_TABLE to generate embeddings, and retrieves the results as lists of floats. Attributes: _db_connection (MySQLConnectionAbstract): MySQL connection used for all database operations. schema_name (str): Name of the database schema to use. options_placeholder (str): SQL-ready placeholder string for ML_EMBED_TABLE options. options_params (dict): Dictionary of concrete option values to be passed as SQL parameters. """ _db_connection: MySQLConnectionAbstract = PrivateAttr() def __init__( self, db_connection: MySQLConnectionAbstract, options: Optional[Dict] = None ): """ Initialize MyEmbeddings with a database connection and optional embedding parameters. References: https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-embed-row.html A full list of supported options can be found under "options" NOTE: The supported "options" are the intersection of the options provided in https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-embed-row.html https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-embed-table.html Args: db_connection: Active MySQL connector database connection. options: Optional dictionary of options for embedding operations. Raises: ValueError: If the schema name is not valid DatabaseError: If a database connection issue occurs. If an operational error occurs during execution. """ super().__init__() self._db_connection = db_connection self.schema_name = source_schema(db_connection) options = options or {} self.options_placeholder, self.options_params = format_value_sql(options) def embed_documents(self, texts: List[str]) -> List[List[float]]: """ Generate embeddings for a list of input texts using the MySQL ML embedding procedure. References: https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-embed-table.html Args: texts: List of input strings to embed. Returns: List of lists of floats, with each inner list containing the embedding for a text. Raises: DatabaseError: If provided options are invalid or unsupported. If a database connection issue occurs. If an operational error occurs during execution. ValueError: If one or more text entries were unable to be embedded. Implementation notes: - Creates a temporary table to pass input text to the MySQL embedding service. - Adds a primary key to ensure results preserve input order. - Calls ML_EMBED_TABLE and fetches the resulting embeddings. - Deletes the temporary table after use to avoid polluting the database. - Embedding vectors are extracted from the "embeddings" column of the result table. """ if not texts: return [] df = pd.DataFrame({"id": range(len(texts)), "text": texts}) with ( atomic_transaction(self._db_connection) as cursor, temporary_sql_tables(self._db_connection) as temporary_tables, ): qualified_table_name, table_name = sql_table_from_df( cursor, self.schema_name, df ) temporary_tables.append((self.schema_name, table_name)) # ML_EMBED_TABLE expects input/output columns and options as parameters embed_query = ( "CALL sys.ML_EMBED_TABLE(" f"'{qualified_table_name}.text', " f"'{qualified_table_name}.embeddings', " f"{self.options_placeholder}" ")" ) execute_sql(cursor, embed_query, params=self.options_params) # Read back all columns, including "embeddings" df_embeddings = sql_table_to_df(cursor, self.schema_name, table_name) if df_embeddings["embeddings"].isnull().any() or any( e is None for e in df_embeddings["embeddings"] ): raise ValueError( "Failure to generate embeddings for one or more text entry." ) # Convert fetched embeddings to lists of floats embeddings = df_embeddings["embeddings"].tolist() embeddings = [list(e) for e in embeddings] return embeddings def embed_query(self, text: str) -> List[float]: """ Generate an embedding for a single text string. References: https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-embed-row.html Args: text: The input string to embed. Returns: List of floats representing the embedding vector. Raises: DatabaseError: If provided options are invalid or unsupported. If a database connection issue occurs. If an operational error occurs during execution. Example: >>> MyEmbeddings(db_conn).embed_query("Hello world") [0.1, 0.2, ...] """ with atomic_transaction(self._db_connection) as cursor: execute_sql( cursor, f'SELECT sys.ML_EMBED_ROW("%s", {self.options_placeholder})', params=(text, *self.options_params), ) return list(cursor.fetchone()[0])