如何使用Pgvector和Python实现带过滤器的语义搜索

docker pull timescale/timescaledb-ha:pg16

docker run -d --name timescaledb -p 5432:5432 -e POSTGRES_PASSWORD=password timescale/timescaledb-ha:pg16

psql -d "postgres://<username>:<password>@<host>:<port>/<database-name>"

CREATE EXTENSION IF NOT EXISTS ai CASCADE;
CREATE EXTENSION IF NOT EXISTS vectorscale CASCADE;

\dx

postgres=# \dx
List of installed extensions
Name | Version | Schema | Description
---------------------+---------+------------+---------------------------------------------------------------------------------------
ai | 0.6.0 | ai | helper functions for ai workflows
plpgsql | 1.0 | pg_catalog | PL/pgSQL procedural language
plpython3u | 1.0 | pg_catalog | PL/Python3U untrusted procedural language
timescaledb | 2.17.2 | public | Enables scalable inserts and complex queries for time-series data (Community Edition)
timescaledb_toolkit | 1.19.0 | public | Library of analytical hyperfunctions, time-series pipelining, and other SQL utilities
vector | 0.8.0 | public | vector data type and ivfflat and hnsw access methods
vectorscale | 0.5.1 | public | pgvectorscale: Advanced indexing for vector data
(7 rows)

OPENAI_API_KEY="your_secret_openai_api_key"
DB_URL="postgres://postgres:password@127.0.0.1:5432/postgres"

$ pip install jupyterlab
$ jupyter lab

!pip install psycopg2-binary
!pip install python-dotenv

!pip install pgai

$ pgai vectorizer worker -d "postgres://postgres:password@127.0.0.1:5432/postgres"

import os
from dotenv import load_dotenv
load_dotenv()
OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
DB_URL = os.environ["DB_URL"]
import psycopg2
with psycopg2.connect(DB_URL) as conn:
    with conn.cursor() as cur:
        # pass the API key as a parameter to the query. don't use string manipulations

cur.execute("SELECT * FROM ai.openai_list_models(api_key=>%s) ORDER BY created DESC", (OPENAI_API_KEY,))
records = cur.fetchall()
print(records)

[
  {
    "review_text": "The hotel was a haven of comfort and luxury. Impeccably clean rooms, friendly staff, and a prime location made our stay unforgettable. The delicious breakfast spread was an added bonus. Highly recommend this gem!",
    "date": "2025-01-01",
    "category": "hotels",
    "name": "Yates, Gonzalez and Mack",
    "rating_score": 5,
    "location": {
      "city": "Cooktown",
      "country": "Canada"
    },
    "customer_location": {
      "city": "South Staceyburgh",
      "country": "Bahamas"
    },
    "bedrooms": 3,
    "bathrooms": 1
  },
  {
    "review_text": "This hotel is exceptional! The rooms are pristine and beautifully decorated. Staff are incredibly welcoming and attentive. The amenities, including the luxurious pool and spa, are top-notch. An unforgettable stay with stunning views and a delightful atmosphere. Highly recommend!",
    "date": "2025-01-02",
    "category": "hotels",
    "name": "Hill-Hatfield",
    "rating_score": 3,
    "location": {
      "city": "Rachelland",
      "country": "Benin"
    },
    "customer_location": {
      "city": "Dorothyview",
      "country": "Cayman Islands"
    },
    "bedrooms": 3,
    "bathrooms": 2
  },
  {
    "review_text": "The room was filthy, with stained sheets and a persistently foul odor. Staff was rude and unresponsive to complaints. The so-called \"breakfast\" was inedible. Worst hotel experience I've ever had. Not recommended at all.",
    "date": "2025-01-06",
    "category": "hotels",
    "name": "Rocha, Robinson and Ellis",
    "rating_score": 1,
    "location": {
      "city": "Holmesstad",
      "country": "Benin"
    },
    "customer_location": {
      "city": "New Joshua",
      "country": "Senegal"
    },
    "bedrooms": 2,
    "bathrooms": 3
  },
  {
    "review_text": "The hostel offers a convenient location with friendly staff and basic amenities. Rooms are clean but can be cramped. Communal areas are lively, though sometimes noisy. Good for budget travelers seeking a social atmosphere.",
    "date": "2025-01-11",
    "category": "hostels",
    "name": "Lane-Kidd",
    "rating_score": 2,
    "location": {
      "city": "North Sarahton",
      "country": "Portugal"
    },
    "customer_location": {
      "city": "Port Manuelburgh",
      "country": "Philippines"
    },
    "bedrooms": 5,
    "bathrooms": 3
  },
  {
    "review_text": "The hotel was a nightmare—musty room, stained sheets, and a bathroom that screamed for cleaning. Unhelpful staff and paper-thin walls made sleep impossible. The only thing five-star about this place is the regret. Don't stay here.",
    "date": "2025-01-11",
    "category": "hotels",
    "name": "Love Group",
    "rating_score": 5,
    "location": {
      "city": "Mackville",
      "country": "New Zealand"
    },
    "customer_location": {
      "city": "East Cole",
      "country": "Pitcairn Islands"
    },
    "bedrooms": 5,
    "bathrooms": 2
  },
  {
    "review_text": "These apartments are fantastic! The spacious, modern design and top-notch amenities create a luxurious living experience. The friendly staff and well-maintained facilities make it a pleasure to call this place home. Highly recommended!",
    "date": "2025-01-05",
    "category": "apartments",
    "name": "Jordan LLC",
    "rating_score": 1,
    "location": {
      "city": "Lake Meganview",
      "country": "Nepal"
    },
    "customer_location": {
      "city": "Williebury",
      "country": "Mexico"
    },
    "bedrooms": 5,
    "bathrooms": 2
  },
  {
    "review_text": "I had an amazing stay at the downtown hostel! The staff was incredibly friendly and helpful, the rooms were clean and cozy, and the communal areas were perfect for meeting fellow travelers. Highly recommend!",
    "date": "2025-01-08",
    "category": "hostels",
    "name": "Torres-Branch",
    "rating_score": 2,
    "location": {
      "city": "Davidsonchester",
      "country": "Romania"
    },
    "customer_location": {
      "city": "Wallaceberg",
      "country": "Chile"
    },
    "bedrooms": 5,
    "bathrooms": 2
  }
]

conn = psycopg2.connect(DB_URL)
conn.autocommit = True
with conn.cursor() as cur:
    cur.execute('''
CREATE TABLE IF NOT EXISTS hotel_reviews (
id BIGINT PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
review_text TEXT,
embedding vector(1536),
category TEXT,
rating_score INTEGER,
location_city TEXT,
location_country TEXT,
customer_city TEXT,
customer_country TEXT,
name TEXT,
date DATE,
bedrooms INTEGER,
bathrooms INTEGER
);
''')

with conn.cursor() as cur:
  
   for review in hotel_reviews:

       print("Inserting review:")
       print(review)

       cur.execute('''
       INSERT INTO hotel_reviews (
           review_text, category, rating_score,
           location_city, location_country,
           customer_city, customer_country, name,
           date, bedrooms, bathrooms
       )
       SELECT
           %s,
           %s,
           %s,
           %s,
           %s,
           %s,
           %s,
           %s,
           %s,
           %s,
           %s
       ''', (
           review['review_text'],
           review['category'],
           review['rating_score'],
           review['location']['city'],
           review['location']['country'],
           review['customer_location']['city'],
           review['customer_location']['country'],
           review['name'],
           review['date'],
           review['bedrooms'],
           review['bathrooms']
       ))

with conn.cursor() as cur:
    cur.execute('''
SELECT ai.create_vectorizer(
    'Reviews'::regclass,
    destination => 'hotel_reviews_embeddings',
    embedding => ai.embedding_openai('text-embedding-3-small', 1536),
    chunking => ai.chunking_recursive_character_text_splitter('review_text'),
    formatting => ai.formatting_python_template('$chunk - Category: $category'),
    indexing => ai.indexing_hsnw(min_rows => 50000, opclass => 'vector_l2_ops')
);
''')

with conn.cursor() as cur:
    cur.execute('''
SELECT ai.create_vectorizer(
    'Reviews'::regclass,
    destination => 'hotel_reviews_embeddings',
    embedding => ai.embedding_openai('text-embedding-3-small', 1536),
    chunking => ai.chunking_recursive_character_text_splitter('review_text'),
    formatting => ai.formatting_python_template('$chunk - Category: $category'),
    indexing => ai.indexing_diskann(min_rows => 100000, storage_layout => 'memory_optimized')
);
''')

def semantic_search(query_text, category=None, min_rating=None, max_rating=None,
                  min_bedrooms=None, max_bedrooms=None,
                  min_bathrooms=None, max_bathrooms=None,
                  location_city=None, location_country=None, limit=5):
   """
   Perform semantic search with optional filters on ratings, rooms, and location
   """
   with conn.cursor() as cur:
       # Build the SQL query with filters
       sql = '''
       WITH query_embedding AS (
           SELECT ai.openai_embed('text-embedding-3-small', %s, api_key=>%s) as embedding
       )
       SELECT review_text, category, rating_score,
              location_city, location_country,
              bedrooms, bathrooms,
              embedding <=> (SELECT embedding FROM query_embedding) as distance
       FROM hotel_reviews
       WHERE 1=1
       '''
       params = [query_text, OPENAI_API_KEY]
      
       if category:
           sql += ' AND category = %s'
           params.append(category)
      
       if min_rating:
           sql += ' AND rating_score >= %s'
           params.append(min_rating)
          
       if max_rating:
           sql += ' AND rating_score <= %s'
           params.append(max_rating)
          
       if min_bedrooms:
           sql += ' AND bedrooms >= %s'
           params.append(min_bedrooms)
          
       if max_bedrooms:
           sql += ' AND bedrooms <= %s'
           params.append(max_bedrooms)
          
       if min_bathrooms:
           sql += ' AND bathrooms >= %s'
           params.append(min_bathrooms)
          
       if max_bathrooms:
           sql += ' AND bathrooms <= %s'
           params.append(max_bathrooms)
          
       if location_city:
           sql += ' AND location_city = %s'
           params.append(location_city)
          
       if location_country:
           sql += ' AND location_country = %s'
           params.append(location_country)
          
       sql += ' ORDER BY embedding <-> (SELECT embedding FROM query_embedding) LIMIT %s'
       params.append(limit)
      
       cur.execute(sql, params)
       results = cur.fetchall()
      
       return [
           {
               'review_text': r[0],
               'category': r[1],
               'rating_score': r[2],
               'location_city': r[3],
               'location_country': r[4],
               'bedrooms': r[5],
               'bathrooms': r[6],
               'similarity_score': 1 - r[7]  # Convert distance to similarity
           }
           for r in results
       ]

results = semantic_search(
   "private pool",
   category="villas",   
   limit=3
)

for i, result in enumerate(results, 1):
   print(f"\nResult {i}:")
   print(f"Review: {result['review_text']}")
   print(f"Category: {result['category']}")
   print(f"Rating: {result['rating_score']}")
   print(f"Location: {result['location_city']}, {result['location_country']}")
   print(f"Rooms: {result['bedrooms']} bedrooms, {result['bathrooms']} bathrooms")
   print(f"Similarity Score: {result['similarity_score']:.2f}")

Result 1:
Review: This villa is a stunning oasis of tranquility, offering breathtaking views, elegant interiors, and a serene private pool. It's the perfect blend of luxury and comfort, providing an unforgettable escape filled with beauty and relaxation.
Category: villas
Rating: 4
Location: New Craig, Christmas Island
Rooms: 3 bedrooms, 1 bathrooms
Similarity Score: 0.39

Result 2:
Review: This villa is a paradise! Stunning views, luxurious amenities, and impeccable service. The infinity pool and lush gardens create a serene escape. Perfect for relaxation and rejuvenation. A true gem for anyone seeking an unforgettable retreat!
Category: villas
Rating: 2
Location: Port Ashleyhaven, Reunion
Rooms: 3 bedrooms, 3 bathrooms
Similarity Score: 0.35

Result 3:
Review: The villas exude luxury and tranquility, with stunning ocean views, plush interiors, and a private infinity pool. Impeccable service and lush gardens make this a dreamy getaway. Truly a paradise retreat!
Category: villas
Rating: 5
Location: Hoffmanfort, South Africa
Rooms: 5 bedrooms, 3 bathrooms
Similarity Score: 0.35