codekingpro/portable-devtools

1# Authors:2#   Harichandan Roy (hroy)3#   David Jiang (ddjiang)4#5# -----------------------------------------------------------------------------6# oracleai.py7# -----------------------------------------------------------------------------8 9from __future__ import annotations10 11import hashlib12import json13import logging14import os15import random16import struct17import time18import traceback19from html.parser import HTMLParser20from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union21 22from langchain_core.document_loaders import BaseLoader23from langchain_core.documents import Document24from langchain_text_splitters import TextSplitter25 26if TYPE_CHECKING:27    from oracledb import Connection28 29logger = logging.getLogger(__name__)30 31"""ParseOracleDocMetadata class"""32 33 34class ParseOracleDocMetadata(HTMLParser):35    """Parse Oracle doc metadata..."""36 37    def __init__(self) -> None:38        super().__init__()39        self.reset()40        self.match = False41        self.metadata: Dict[str, Any] = {}42 43    def handle_starttag(self, tag: str, attrs: List[Tuple[str, Optional[str]]]) -> None:44        if tag == "meta":45            entry: Optional[str] = ""46            for name, value in attrs:47                if name == "name":48                    entry = value49                if name == "content":50                    if entry:51                        self.metadata[entry] = value52        elif tag == "title":53            self.match = True54 55    def handle_data(self, data: str) -> None:56        if self.match:57            self.metadata["title"] = data58            self.match = False59 60    def get_metadata(self) -> Dict[str, Any]:61        return self.metadata62 63 64"""OracleDocReader class"""65 66 67class OracleDocReader:68    """Read a file"""69 70    @staticmethod71    def generate_object_id(input_string: Union[str, None] = None) -> str:72        out_length = 32  # output length73        hash_len = 8  # hash value length74 75        if input_string is None:76            input_string = "".join(77                random.choices(78                    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789",79                    k=16,80                )81            )82 83        # timestamp84        timestamp = int(time.time())85        timestamp_bin = struct.pack(">I", timestamp)  # 4 bytes86 87        # hash_value88        hashval_bin = hashlib.sha256(input_string.encode()).digest()89        hashval_bin = hashval_bin[:hash_len]  # 8 bytes90 91        # counter92        counter_bin = struct.pack(">I", random.getrandbits(32))  # 4 bytes93 94        # binary object id95        object_id = timestamp_bin + hashval_bin + counter_bin  # 16 bytes96        object_id_hex = object_id.hex()  # 32 bytes97        object_id_hex = object_id_hex.zfill(98            out_length99        )  # fill with zeros if less than 32 bytes100 101        object_id_hex = object_id_hex[:out_length]102 103        return object_id_hex104 105    @staticmethod106    def read_file(107        conn: Connection, file_path: str, params: dict108    ) -> Union[Document, None]:109        """Read a file using OracleReader110        Args:111            conn: Oracle Connection,112            file_path: Oracle Directory,113            params: ONNX file name.114        Returns:115            Plain text and metadata as Langchain Document.116        """117 118        metadata: Dict[str, Any] = {}119        try:120            import oracledb121        except ImportError as e:122            raise ImportError(123                "Unable to import oracledb, please install with "124                "`pip install -U oracledb`."125            ) from e126        try:127            oracledb.defaults.fetch_lobs = False128            cursor = conn.cursor()129 130            with open(file_path, "rb") as f:131                data = f.read()132 133            if data is None:134                return Document(page_content="", metadata=metadata)135 136            mdata = cursor.var(oracledb.DB_TYPE_CLOB)137            text = cursor.var(oracledb.DB_TYPE_CLOB)138            cursor.execute(139                """140                declare141                    input blob;142                begin143                    input := :blob;144                    :mdata := dbms_vector_chain.utl_to_text(input, json(:pref));145                    :text := dbms_vector_chain.utl_to_text(input);146                end;""",147                blob=data,148                pref=json.dumps(params),149                mdata=mdata,150                text=text,151            )152            cursor.close()153 154            if mdata is None:155                metadata = {}156            else:157                doc_data = str(mdata.getvalue())158                if doc_data.startswith("<!DOCTYPE html") or doc_data.startswith(159                    "<HTML>"160                ):161                    p = ParseOracleDocMetadata()162                    p.feed(doc_data)163                    metadata = p.get_metadata()164 165            doc_id = OracleDocReader.generate_object_id(conn.username + "$" + file_path)166            metadata["_oid"] = doc_id167            metadata["_file"] = file_path168 169            if text is None:170                return Document(page_content="", metadata=metadata)171            else:172                return Document(page_content=str(text.getvalue()), metadata=metadata)173 174        except Exception as ex:175            logger.info(f"An exception occurred :: {ex}")176            logger.info(f"Skip processing {file_path}")177            cursor.close()178            return None179 180 181"""OracleDocLoader class"""182 183 184class OracleDocLoader(BaseLoader):185    """Read documents using OracleDocLoader186    Args:187        conn: Oracle Connection,188        params: Loader parameters.189    """190 191    def __init__(self, conn: Connection, params: Dict[str, Any], **kwargs: Any):192        self.conn = conn193        self.params = json.loads(json.dumps(params))194        super().__init__(**kwargs)195 196    def load(self) -> List[Document]:197        """Load data into LangChain Document objects..."""198        try:199            import oracledb200        except ImportError as e:201            raise ImportError(202                "Unable to import oracledb, please install with "203                "`pip install -U oracledb`."204            ) from e205 206        ncols = 0207        results: List[Document] = []208        metadata: Dict[str, Any] = {}209        m_params = {"plaintext": "false"}210        try:211            # extract the parameters212            if self.params is not None:213                self.file = self.params.get("file")214                self.dir = self.params.get("dir")215                self.owner = self.params.get("owner")216                self.tablename = self.params.get("tablename")217                self.colname = self.params.get("colname")218            else:219                raise Exception("Missing loader parameters")220 221            oracledb.defaults.fetch_lobs = False222 223            if self.file:224                doc = OracleDocReader.read_file(self.conn, self.file, m_params)225 226                if doc is None:227                    return results228 229                results.append(doc)230 231            if self.dir:232                skip_count = 0233                for file_name in os.listdir(self.dir):234                    file_path = os.path.join(self.dir, file_name)235                    if os.path.isfile(file_path):236                        doc = OracleDocReader.read_file(self.conn, file_path, m_params)237 238                        if doc is None:239                            skip_count = skip_count + 1240                            logger.info(f"Total skipped: {skip_count}\n")241                        else:242                            results.append(doc)243 244            if self.tablename:245                try:246                    if self.owner is None or self.colname is None:247                        raise Exception("Missing owner or column name or both.")248 249                    cursor = self.conn.cursor()250                    self.mdata_cols = self.params.get("mdata_cols")251                    if self.mdata_cols is not None:252                        if len(self.mdata_cols) > 3:253                            raise Exception(254                                "Exceeds the max number of columns "255                                + "you can request for metadata."256                            )257 258                        # execute a query to get column data types259                        sql = (260                            "select column_name, data_type from all_tab_columns "261                            + "where owner = :ownername and "262                            + "table_name = :tablename"263                        )264                        cursor.execute(265                            sql,266                            ownername=self.owner.upper(),267                            tablename=self.tablename.upper(),268                        )269 270                        # cursor.execute(sql)271                        rows = cursor.fetchall()272                        for row in rows:273                            if row[0] in self.mdata_cols:274                                if row[1] not in [275                                    "NUMBER",276                                    "BINARY_DOUBLE",277                                    "BINARY_FLOAT",278                                    "LONG",279                                    "DATE",280                                    "TIMESTAMP",281                                    "VARCHAR2",282                                ]:283                                    raise Exception(284                                        "The datatype for the column requested "285                                        + "for metadata is not supported."286                                    )287 288                    self.mdata_cols_sql = ", rowid"289                    if self.mdata_cols is not None:290                        for col in self.mdata_cols:291                            self.mdata_cols_sql = self.mdata_cols_sql + ", " + col292 293                    # [TODO] use bind variables294                    sql = (295                        "select dbms_vector_chain.utl_to_text(t."296                        + self.colname297                        + ", json('"298                        + json.dumps(m_params)299                        + "')) mdata, dbms_vector_chain.utl_to_text(t."300                        + self.colname301                        + ") text"302                        + self.mdata_cols_sql303                        + " from "304                        + self.owner305                        + "."306                        + self.tablename307                        + " t"308                    )309 310                    cursor.execute(sql)311                    for row in cursor:312                        metadata = {}313 314                        if row is None:315                            doc_id = OracleDocReader.generate_object_id(316                                self.conn.username317                                + "$"318                                + self.owner319                                + "$"320                                + self.tablename321                                + "$"322                                + self.colname323                            )324                            metadata["_oid"] = doc_id325                            results.append(Document(page_content="", metadata=metadata))326                        else:327                            if row[0] is not None:328                                data = str(row[0])329                                if data.startswith("<!DOCTYPE html") or data.startswith(330                                    "<HTML>"331                                ):332                                    p = ParseOracleDocMetadata()333                                    p.feed(data)334                                    metadata = p.get_metadata()335 336                            doc_id = OracleDocReader.generate_object_id(337                                self.conn.username338                                + "$"339                                + self.owner340                                + "$"341                                + self.tablename342                                + "$"343                                + self.colname344                                + "$"345                                + str(row[2])346                            )347                            metadata["_oid"] = doc_id348                            metadata["_rowid"] = row[2]349 350                            # process projected metadata cols351                            if self.mdata_cols is not None:352                                ncols = len(self.mdata_cols)353 354                            for i in range(0, ncols):355                                metadata[self.mdata_cols[i]] = row[i + 2]356 357                            if row[1] is None:358                                results.append(359                                    Document(page_content="", metadata=metadata)360                                )361                            else:362                                results.append(363                                    Document(364                                        page_content=str(row[1]), metadata=metadata365                                    )366                                )367                except Exception as ex:368                    logger.info(f"An exception occurred :: {ex}")369                    traceback.print_exc()370                    cursor.close()371                    raise372 373            return results374        except Exception as ex:375            logger.info(f"An exception occurred :: {ex}")376            traceback.print_exc()377            raise378 379 380class OracleTextSplitter(TextSplitter):381    """Splitting text using Oracle chunker."""382 383    def __init__(self, conn: Connection, params: Dict[str, Any], **kwargs: Any) -> None:384        """Initialize."""385        self.conn = conn386        self.params = params387        super().__init__(**kwargs)388        try:389            import json390 391            try:392                import oracledb393            except ImportError as e:394                raise ImportError(395                    "Unable to import oracledb, please install with "396                    "`pip install -U oracledb`."397                ) from e398 399            self._oracledb = oracledb400            self._json = json401        except ImportError:402            raise ImportError(403                "oracledb or json or both are not installed. "404                + "Please install them. "405                + "Recommendations: `pip install oracledb`. "406            )407 408    def split_text(self, text: str) -> List[str]:409        """Split incoming text and return chunks."""410 411        try:412            import oracledb413        except ImportError as e:414            raise ImportError(415                "Unable to import oracledb, please install with "416                "`pip install -U oracledb`."417            ) from e418 419        splits = []420 421        try:422            # returns strings or bytes instead of a locator423            self._oracledb.defaults.fetch_lobs = False424 425            cursor = self.conn.cursor()426 427            cursor.setinputsizes(content=oracledb.CLOB)428            cursor.execute(429                "select t.column_value from "430                + "dbms_vector_chain.utl_to_chunks(:content, json(:params)) t",431                content=text,432                params=self._json.dumps(self.params),433            )434 435            while True:436                row = cursor.fetchone()437                if row is None:438                    break439                d = self._json.loads(row[0])440                splits.append(d["chunk_data"])441 442            return splits443 444        except Exception as ex:445            logger.info(f"An exception occurred :: {ex}")446            traceback.print_exc()447            raise448