Python

conn = psycopg2.connect(database="postgres", user="vbadmin", password="Vbase@123", host="127.0.0.1", port="5432")
conn.close()

conn = psycopg2.connect(
                        host="xxx.xxx.xxx.xxa,xxx.xxx.xxx.xxb",
                        port=5432,
                        database='postgres',
                        user='vbadmin',
                        password='Vbase@123',
                        target_session_attrs='primary'
                        )

import string
import ctypes
import decimal
import platform
import unittest
import random
import testutils
import ast
import csv
import datetime
import io 
from io import StringIO
import json
import struct
import numpy
from typing import Any
# 注意修改ConnectingTestCase中的数据库连接信息
from testutils import ConnectingTestCase, restore_types

import psycopg2
from psycopg2.extensions import FloatVector

class FloatVectorTest(ConnectingTestCase):

    def test_floatvector_dql(self):
        curs = self.conn.cursor()
        try:
            curs.execute("select '[1,2,3]'::floatvector(3);") 
            result = curs.fetchone()[0]
            print(result.tolist())
            print(FloatVector([1,2,3]).tolist())
            print(str(FloatVector([1,2,3])))

            self.assertEqual(result.tolist(),FloatVector([1,2,3]).tolist())
            self.assertEqual(str(result),str(FloatVector([1,2,3])))

            curs.execute("select '[0]'::floatvector(1);") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([0]).tolist())
            self.assertEqual(str(result),str(FloatVector([0])))

            curs.execute("select '[1.00001,0.001234]'::floatvector(2);") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([1.00001,0.001234]).tolist())
            self.assertEqual(str(result),str(FloatVector([1.00001,0.001234])))

            curs.execute("SELECT ARRAY[1,1,1]::floatvector;") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([1,1,1]).tolist())
            self.assertEqual(str(result),str(FloatVector([1,1,1])))

            curs.execute("SELECT '[1,1,1]'::floatvector;") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([1,1,1]).tolist())
            self.assertEqual(str(result),str(FloatVector([1,1,1])))

            curs.execute("SELECT '[1.0]'::text::floatvector AS vector_example;") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([1.0]).tolist())
            self.assertEqual(str(result),str(FloatVector([1.0])))

            curs.execute("SELECT '[1.0]'::char(5)::floatvector AS vector_example;") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([1.0]).tolist())
            self.assertEqual(str(result),str(FloatVector([1.0])))

            curs.execute("SELECT '[1.0]'::nchar(5)::floatvector AS vector_example;") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([1.0]).tolist())
            self.assertEqual(str(result),str(FloatVector([1.0])))

            curs.execute("SELECT '[1.0]'::varchar(5)::floatvector AS vector_example;") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([1.0]).tolist())
            self.assertEqual(str(result),str(FloatVector([1.0])))

            curs.execute("SELECT '[1.0]'::nvarchar(5)::floatvector AS vector_example;") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([1.0]).tolist())
            self.assertEqual(str(result),str(FloatVector([1.0])))

            curs.execute("SELECT '[1.0]'::clob::floatvector AS vector_example;") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([1.0]).tolist())
            self.assertEqual(str(result),str(FloatVector([1.0])))

            curs.execute("SELECT '[1.0]'::bpchar(5)::floatvector AS vector_example;") 
            result = curs.fetchone()[0]
            print(str(result))
            print("------------------")
            self.assertEqual(result.tolist(),FloatVector([1.0]).tolist())
            self.assertEqual(str(result),str(FloatVector([1.0])))

            print("query result: ",result)
            self.conn.commit()
        except Exception:
            self.conn.rollback()
            raise
        finally:
            curs.close()

pip install psycopg2-binary #必须
pip install sqlalchemy #如果要使用SQLAlchemy集成
pip install numpy #可选，支持直接映射numpy到向量类型
pip install pyvector-vexdb.xxx.whl

import numpy as np
from vexdb.psycopg2 import register_vector
import psycopg2
from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor
from psycopg2.pool import ThreadedConnectionPool

# 使用dsn
conn = psycopg2.connect("dbname=postgres user=test password=Test@1234 host=localhost port=5432")
# 使用关键字
# conn = psycopg2.connect(
#     dbname="postgres",
#     user="test",
#     password="Test@1234",
#     host="localhost",
#     port=5432)
conn.autocommit = True

cur = conn.cursor()
cur.execute('DROP TABLE IF EXISTS psycopg2_items')
cur.execute('CREATE TABLE psycopg2_items (id bigserial PRIMARY KEY, embedding floatvector(3))')

# 注册到连接
register_vector(conn)
# 注册到游标
# register_vector(cur)

class TestPsycopg2:
    def setup_method(self, test_method):
        cur.execute('DELETE FROM psycopg2_items')

    def test_vector(self):
        embedding = np.array([1.5, 2, 3])
        cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,))

        cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id')
        res = cur.fetchall()
        assert np.array_equal(res[0][0], embedding)
        assert res[0][0].dtype == np.float32
        assert res[1][0] is None

    def test_query(self):
        pass

    def test_cursor_factory(self):
        for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]:
            conn = psycopg2.connect(host='localhost', port=5432, database='postgres', user='test', password='Test@1234')
            cur = conn.cursor(cursor_factory=cursor_factory)
            register_vector(cur, globally=False)
            conn.close()

    def test_cursor_factory_connection(self):
        for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]:
            conn = psycopg2.connect(host='localhost', port=5432, database='postgres', user='test', password='Test@1234', cursor_factory=cursor_factory)
            register_vector(conn, globally=False)
            conn.close()

    def test_pool(self):
        pool = ThreadedConnectionPool(1, 1, host='localhost', port=5432, database='postgres', user='test', password='Test@1234')

        conn = pool.getconn()
        try:
            # use globally=True for apps to ensure registered with all connections
            register_vector(conn, globally=False)
        finally:
            pool.putconn(conn)

        conn = pool.getconn()
        try:
            cur = conn.cursor()
            cur.execute("SELECT '[1,2,3]'::floatvector")
            res = cur.fetchone()
            assert np.array_equal(res[0], np.array([1, 2, 3]))
        finally:
            pool.putconn(conn)

        pool.closeall()

from sqlalchemy import create_engine, Column, Integer, Text, func
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.sql import text
from vexdb.sqlalchemy import FloatVector
import numpy as np
from typing import List, Optional, Tuple
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 数据库配置
#DATABASE_URL = "postgresql+psycopg2://test:Test%401234@172.16.100.97:5432/postgres"
DATABASE_URL = "postgresql+psycopg2://postgres:12345Aa*@172.16.97.114:5432/postgres" 

# 创建基类
Base = declarative_base()

class Product(Base):
    """手机产品模型"""
    __tablename__ = 'products'
    
    id = Column(Integer, primary_key=True)
    name = Column(Text, nullable=False)
    color = Column(Text, nullable=False)  # 颜色枚举值：'black', 'white', 'red', 'blue', 'gold'
    description = Column(Text, nullable=False)
    features = Column(FloatVector(512), nullable=False)  # 512维特征向量
    
    def __repr__(self):
        return f"<Product(id={self.id}, name='{self.name}', color='{self.color}')>"


class ProductVectorSearch:
    """手机产品向量检索服务"""
    
    def __init__(self, database_url: str):
        self.engine = create_engine(database_url)
        self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)
        
        # 创建表
        self._create_tables_and_extensions()
    
    def _create_tables_and_extensions(self):
        """创建表"""
        # 清空表，便于重新创建
        Base.metadata.drop_all(bind=self.engine)  
        # 创建所有表
        Base.metadata.create_all(bind=self.engine)
        
        # 创建向量索引以提升检索性能
        with self.engine.connect() as conn:
            try:
                conn.execute(text("""
                    CREATE INDEX IF NOT EXISTS products_features_idx 
                    ON products USING ivfflat (features floatvector_cosine_ops) 
                    WITH (ivf_nlist = 100)
                """))
                conn.commit()
                logger.info("向量索引创建成功")
            except Exception as e:
                logger.warning(f"创建向量索引失败: {e}")
    
    def add_product(self, id: int, name: str, color: str, description: str, features: np.ndarray) -> Product:
        """添加新产品"""
        with self.SessionLocal() as session:
            product = Product(
                id=id,
                name=name,
                color=color,
                description=description,
                features=features.tolist()  # 转换为列表格式
            )
            session.add(product)
            session.commit()
            session.refresh(product)
            return product
    
    def vector_search_with_color_filter(
        self, 
        query_vector: np.ndarray, 
        color_filter: Optional[str] = None,
        colors_filter: Optional[List[str]] = None,
        limit: int = 10,
        similarity_threshold: float = 0.5
    ) -> List[Tuple[Product, float]]:
        """
        使用向量检索结合颜色过滤的联合查询
        
        Args:
            query_vector: 查询向量 (512维)
            color_filter: 单个颜色过滤条件
            colors_filter: 多个颜色过滤条件列表
            limit: 返回结果数量限制
            similarity_threshold: 相似度阈值 (0-1)
        
        Returns:
            List[Tuple[Product, float]]: 产品和相似度分数的元组列表
        """
        with self.SessionLocal() as session:
            # 构建基础查询，计算余弦相似度
            query = session.query(
                Product,
                (1 - Product.features.cosine_distance(query_vector)).label('similarity')
            )
            
            # 添加颜色过滤条件
            if color_filter:
                query = query.filter(Product.color == color_filter)
            elif colors_filter:
                query = query.filter(Product.color.in_(colors_filter))
            
            # 添加相似度阈值过滤
            query = query.filter(
                (1 - Product.features.cosine_distance(query_vector)) >= similarity_threshold
            )
            
            # 按相似度降序排列并限制结果数量
            results = query.order_by(
                (1 - Product.features.cosine_distance(query_vector)).desc()
            ).limit(limit).all()
            
            return [(product, float(similarity)) for product, similarity in results]
    
    def hybrid_search(
        self,
        query_vector: np.ndarray,
        color_filter: Optional[str] = None,
        price_range: Optional[Tuple[float, float]] = None,
        limit: int = 10
    ) -> List[Tuple[Product, float]]:
        """
        混合搜索：向量相似度 + 多重标量字段过滤
        """
        with self.SessionLocal() as session:
            query = session.query(
                Product,
                (1 - Product.features.cosine_distance(query_vector)).label('similarity')
            )
            
            # 颜色过滤
            if color_filter:
                query = query.filter(Product.color == color_filter)
            
            # 可以扩展更多过滤条件，比如价格范围
            # if price_range:
            #     query = query.filter(Product.price.between(price_range[0], price_range[1]))
            
            results = query.order_by(
                (1 - Product.features.cosine_distance(query_vector)).desc()
            ).limit(limit).all()
            
            return [(product, float(similarity)) for product, similarity in results]
    
    def get_similar_products_by_description(
        self,
        description: str,
        embedding_function,  # 文本转向量的函数
        color_filter: Optional[str] = None,
        limit: int = 5
    ) -> List[Tuple[Product, float]]:
        """
        根据描述文本查找相似产品
        
        Args:
            description: 产品描述文本
            embedding_function: 将文本转换为向量的函数
            color_filter: 颜色过滤
            limit: 结果数量限制
        """
        # 将描述转换为向量
        query_vector = embedding_function(description)
        
        return self.vector_search_with_color_filter(
            query_vector=query_vector,
            color_filter=color_filter,
            limit=limit
        )
    
    def advanced_search(
        self,
        query_vector: np.ndarray,
        filters: dict,
        limit: int = 10
    ) -> List[Tuple[Product, float]]:
        """
        高级搜索：支持动态过滤条件
        
        Args:
            query_vector: 查询向量
            filters: 过滤条件字典，如 {'color': 'black', 'colors': ['black', 'white']}
            limit: 结果数量限制
        """
        with self.SessionLocal() as session:
            query = session.query(
                Product,
                (1 - Product.features.cosine_distance(query_vector)).label('similarity')
            )
            
            # 动态添加过滤条件
            if 'color' in filters:
                query = query.filter(Product.color == filters['color'])
            
            if 'colors' in filters:
                query = query.filter(Product.color.in_(filters['colors']))
            
            if 'name_contains' in filters:
                query = query.filter(Product.name.ilike(f"%{filters['name_contains']}%"))
            
            results = query.order_by(
                (1 - Product.features.cosine_distance(query_vector)).desc()
            ).limit(limit).all()
            
            return [(product, float(similarity)) for product, similarity in results]


# 使用示例
def example_usage():
    """使用示例"""
    
    # 初始化搜索服务
    search_service = ProductVectorSearch(DATABASE_URL)
    
    # 模拟添加一些手机产品数据
    sample_products = [
        {
            "id": 1,
            "name": "iPhone 15 Pro",
            "color": "black",
            "description": "高端智能手机，配备A17 Pro芯片，钛金属机身，三摄系统",
            "features": np.random.rand(512).astype(np.float32)  # 实际应用中这里是真实的嵌入向量
        },
        {
            "id": 2,
            "name": "Samsung Galaxy S24",
            "color": "white",
            "description": "Android旗舰手机，骁龙8 Gen3处理器，AI拍照功能",
            "features": np.random.rand(512).astype(np.float32)
        },
        {
            "id": 3,
            "name": "小米14 Pro",
            "color": "blue",
            "description": "徕卡影像系统，骁龙8 Gen3，120W快充",
            "features": np.random.rand(512).astype(np.float32)
        }
    ]
    
    # 添加产品到数据库
    for product_data in sample_products:
        search_service.add_product(**product_data)
    
    # 示例1: 基于向量查询，过滤黑色手机
    query_vector = np.random.rand(512).astype(np.float32)
    results = search_service.vector_search_with_color_filter(
        query_vector=query_vector,
        color_filter="black",
        limit=5
    )
    
    print("基于向量查询，过滤黑色手机:")
    for product, similarity in results:
        print(f"  {product.name} (相似度: {similarity:.3f})")
    
    # 示例2: 多颜色过滤
    results = search_service.vector_search_with_color_filter(
        query_vector=query_vector,
        colors_filter=["black", "white"],
        limit=5
    )
    
    print("\n黑色或白色手机搜索结果:")
    for product, similarity in results:
        print(f"  {product.name} - {product.color} (相似度: {similarity:.3f})")
    
    # 示例3: 高级搜索
    results = search_service.advanced_search(
        query_vector=query_vector,
        filters={
            "colors": ["blue", "white"],
            "name_contains": "Pro"
        },
        limit=3
    )
    
    print("\n高级搜索结果 (蓝色或白色 + 名称包含Pro):")
    for product, similarity in results:
        print(f"  {product.name} - {product.color} (相似度: {similarity:.3f})")


if __name__ == "__main__":
    # 运行示例
    example_usage()