Python使用Oracle向量数据库实现文本检索系统

import array
import json
import uuid
from contextlib import contextmanager
from typing import List, Dict, Any

import numpy as np
import oracledb
from pydantic import BaseModel, validator

class OracleConfig(BaseModel):
    host: str
    port: int
    user: str
    password: str
    database: str

    @validator('host', 'user', 'password', 'database')
    def check_not_empty(cls, v):
        if not v:
            raise ValueError("Field cannot be empty")
        return v

class TextEmbeddingStore:
    def __init__(self, config: OracleConfig):
        self.pool = self._create_connection_pool(config)
        self.table_name = "text_embeddings"
        self._create_table()

    def _create_connection_pool(self, config: OracleConfig):
        return oracledb.create_pool(
            user=config.user,
            password=config.password,
            dsn=f"{config.host}:{config.port}/{config.database}",
            min=1,
            max=5,
            increment=1
        )

    @contextmanager
    def _get_cursor(self):
        conn = self.pool.acquire()
        conn.inputtypehandler = self._input_type_handler
        conn.outputtypehandler = self._output_type_handler
        cur = conn.cursor()
        try:
            yield cur
        finally:
            cur.close()
            conn.commit()
            conn.close()

    def _input_type_handler(self, cursor, value, arraysize):
        if isinstance(value, np.ndarray):
            return cursor.var(
                oracledb.DB_TYPE_VECTOR,
                arraysize=arraysize,
                inconverter=self._numpy_to_array
            )

    def _output_type_handler(self, cursor, metadata):
        if metadata.type_code is oracledb.DB_TYPE_VECTOR:
            return cursor.var(
                metadata.type_code,
                arraysize=cursor.arraysize,
                outconverter=self._array_to_numpy
            )

    def _numpy_to_array(self, value):
        return array.array('f', value)

    def _array_to_numpy(self, value):
        return np.array(value, dtype=np.float32)

    def _create_table(self):
        with self._get_cursor() as cur:
            cur.execute(f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    id VARCHAR2(100) PRIMARY KEY,
                    text CLOB NOT NULL,
                    metadata JSON,
                    embedding VECTOR NOT NULL
                )
            """)

    def add_texts(self, texts: List[str], embeddings: List[List[float]], metadata: List[Dict] = None):
        if metadata is None:
            metadata = [{} for _ in texts]
        
        values = [
            (str(uuid.uuid4()), text, json.dumps(meta), np.array(emb, dtype=np.float32))
            for text, emb, meta in zip(texts, embeddings, metadata)
        ]

        with self._get_cursor() as cur:
            cur.executemany(
                f"INSERT INTO {self.table_name} (id, text, metadata, embedding) VALUES (:1, :2, :3, :4)",
                values
            )

    def search_similar(self, query_vector: List[float], top_k: int = 5) -> List[Dict[str, Any]]:
        query_vector = np.array(query_vector, dtype=np.float32)
        with self._get_cursor() as cur:
            cur.execute(
                f"""
                SELECT id, text, metadata, vector_distance(embedding, :1) AS distance
                FROM {self.table_name}
                ORDER BY distance
                FETCH FIRST :2 ROWS ONLY
                """,
                [query_vector, top_k]
            )
            results = []
            for id, text, metadata, distance in cur:
                results.append({
                    "id": id,
                    "text": text,
                    "metadata": json.loads(metadata),
                    "distance": distance,
                    "similarity": 1 - distance
                })
        return results

# 使用示例
if __name__ == "__main__":
    config = OracleConfig(
        host="localhost",
        port=1521,
        user="your_username",
        password="your_password",
        database="your_database"
    )
    
    store = TextEmbeddingStore(config)
    
    # 添加文本和嵌入
    texts = ["Hello world", "Python programming", "Vector database"]
    embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
    store.add_texts(texts, embeddings)
    
    # 搜索相似文本
    query_vector = [0.2, 0.3, 0.4]
    results = store.search_similar(query_vector, top_k=2)
    
    for result in results:
        print(f"Text: {result['text']}")
        print(f"Similarity: {result['similarity']:.4f}")
        print("---")