Source code for langchain_community.embeddings.oracleai

# Authors:
#   Harichandan Roy (hroy)
#   David Jiang (ddjiang)
#
# -----------------------------------------------------------------------------
# oracleai.py
# -----------------------------------------------------------------------------

from __future__ import annotations

import json
import logging
import traceback
from typing import TYPE_CHECKING, Any, Dict, List, Optional

from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseModel, Extra

if TYPE_CHECKING:
    from oracledb import Connection

logger = logging.getLogger(__name__)

"""OracleEmbeddings class"""


class OracleEmbeddings(BaseModel, Embeddings):
    """获取嵌入向量"""

    """Oracle数据库连接"""
    conn: Any
    """嵌入参数"""
    params: Dict[str, Any]
    """代理"""
    proxy: Optional[str] = None

    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)

    class Config:
        """此pydantic对象的配置。"""

        extra = Extra.forbid

    """    1 - 用户需要具有创建存储过程、创建挖掘模型、创建任何目录权限。
    2 - 将创建存储过程、创建挖掘模型、创建任何目录的权限授予<用户>;"""

    @staticmethod
    def load_onnx_model(
        conn: Connection, dir: str, onnx_file: str, model_name: str
    ) -> None:
        """将一个ONNX模型加载到Oracle数据库中。
参数:
    conn: Oracle连接,
    dir: Oracle目录,
    onnx_file: ONNX文件名,
    model_name: 模型的名称。
"""

        try:
            if conn is None or dir is None or onnx_file is None or model_name is None:
                raise Exception("Invalid input")

            cursor = conn.cursor()
            cursor.execute(
                """
                begin
                    dbms_data_mining.drop_model(model_name => :model, force => true);
                    SYS.DBMS_VECTOR.load_onnx_model(:path, :filename, :model, 
                        json('{"function" : "embedding", 
                            "embeddingOutput" : "embedding", 
                            "input": {"input": ["DATA"]}}'));
                end;""",
                path=dir,
                filename=onnx_file,
                model=model_name,
            )

            cursor.close()

        except Exception as ex:
            logger.info(f"An exception occurred :: {ex}")
            traceback.print_exc()
            cursor.close()
            raise

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """使用OracleEmbeddings计算文档嵌入。
参数：
    texts：要嵌入的文本列表。
返回：
    每个输入文本的嵌入列表。
"""

        try:
            import oracledb
        except ImportError as e:
            raise ImportError(
                "Unable to import oracledb, please install with "
                "`pip install -U oracledb`."
            ) from e

        if texts is None:
            return None

        embeddings: List[List[float]] = []
        try:
            # returns strings or bytes instead of a locator
            oracledb.defaults.fetch_lobs = False
            cursor = self.conn.cursor()

            if self.proxy:
                cursor.execute(
                    "begin utl_http.set_proxy(:proxy); end;", proxy=self.proxy
                )

            for text in texts:
                cursor.execute(
                    "select t.* "
                    + "from dbms_vector_chain.utl_to_embeddings(:content, "
                    + "json(:params)) t",
                    content=text,
                    params=json.dumps(self.params),
                )

                for row in cursor:
                    if row is None:
                        embeddings.append([])
                    else:
                        rdata = json.loads(row[0])
                        # dereference string as array
                        vec = json.loads(rdata["embed_vector"])
                        embeddings.append(vec)

            cursor.close()
            return embeddings
        except Exception as ex:
            logger.info(f"An exception occurred :: {ex}")
            traceback.print_exc()
            cursor.close()
            raise

    def embed_query(self, text: str) -> List[float]:
        """使用OracleEmbeddings计算查询嵌入。
参数：
    text：要嵌入的文本。
返回：
    文本的嵌入。
"""
        return self.embed_documents([text])[0]


# uncomment the following code block to run the test

"""
# A sample unit test.

''' get the Oracle connection '''
conn = oracledb.connect(
    user="",
    password="",
    dsn="")
print("Oracle connection is established...")

''' params '''
embedder_params = {"provider":"database", "model":"demo_model"}
proxy = ""

''' instance '''
embedder = OracleEmbeddings(conn=conn, params=embedder_params, proxy=proxy)

embed = embedder.embed_query("Hello World!")
print(f"Embedding generated by OracleEmbeddings: {embed}")

conn.close()
print("Connection is closed.")

"""