pgvector

-- Install the extension
CREATE EXTENSION vector;

-- Create table with vector column
CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    title TEXT,
    content TEXT,
    embedding vector(1536),  -- OpenAI embedding dimensions
    created_at TIMESTAMP DEFAULT now()
);

-- Create HNSW index for similarity search
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops);

-- Alternative: Create IVFFlat index
-- CREATE INDEX ON documents USING ivfflat (embedding vector_cosine_ops) WITH (lists = 1000);

-- Insert sample data with embeddings
INSERT INTO documents (title, content, embedding) VALUES
('AI Introduction', 'Basic concepts of artificial intelligence...', '[0.1, 0.2, 0.3, ...]'),
('Machine Learning', 'Overview of ML algorithms and techniques...', '[0.4, 0.5, 0.6, ...]'),
('Deep Learning', 'Neural networks and deep learning methods...', '[0.7, 0.8, 0.9, ...]');

-- Find most similar documents using cosine similarity
SELECT id, title, embedding <=> '[0.1, 0.2, 0.3, ...]' as distance
FROM documents
ORDER BY embedding <=> '[0.1, 0.2, 0.3, ...]'
LIMIT 5;

-- Find similar documents with threshold
SELECT id, title, content,
       1 - (embedding <=> '[0.1, 0.2, 0.3, ...]') as similarity
FROM documents
WHERE embedding <=> '[0.1, 0.2, 0.3, ...]' < 0.8  -- Less distance = more similar
ORDER BY similarity DESC
LIMIT 10;

-- Combine with traditional filters
SELECT id, title, similarity
FROM (
    SELECT id, title,
           1 - (embedding <=> '[0.1, 0.2, 0.3, ...]') as similarity
    FROM documents
    WHERE created_at >= '2024-01-01'
      AND title ILIKE '%machine learning%'
) t
WHERE similarity > 0.7
ORDER BY similarity DESC;

-- Use different similarity metrics
-- L2 distance (Euclidean)
SELECT id, title, embedding <-> '[0.1, 0.2, 0.3, ...]' as l2_distance
FROM documents
ORDER BY l2_distance
LIMIT 5;

-- Inner product
SELECT id, title, embedding <#> '[0.1, 0.2, 0.3, ...]' as inner_product
FROM documents
ORDER BY inner_product DESC  -- Higher is more similar for inner product
LIMIT 5;

import psycopg2
import numpy as np
from pgvector.psycopg2 import register_vector

# Connect to PostgreSQL with pgvector
conn = psycopg2.connect(
    host="localhost",
    database="vectordb", 
    user="postgres",
    password="password"
)

# Register vector type
register_vector(conn)

class VectorSearchService:
    def __init__(self, connection):
        self.conn = connection
        self.cursor = connection.cursor()
    
    def insert_document(self, title, content, embedding):
        """Insert document with vector embedding."""
        query = """
            INSERT INTO documents (title, content, embedding)
            VALUES (%s, %s, %s)
            RETURNING id;
        """
        self.cursor.execute(query, (title, content, embedding))
        doc_id = self.cursor.fetchone()[0]
        self.conn.commit()
        return doc_id
    
    def similarity_search(self, query_embedding, limit=10, threshold=0.8):
        """Find similar documents using cosine similarity."""
        query = """
            SELECT id, title, content,
                   1 - (embedding <=> %s) as similarity
            FROM documents
            WHERE embedding <=> %s < %s
            ORDER BY similarity DESC
            LIMIT %s;
        """
        self.cursor.execute(query, (
            query_embedding, 
            query_embedding, 
            1 - threshold,  # Convert similarity to distance
            limit
        ))
        return self.cursor.fetchall()
    
    def hybrid_search(self, query_embedding, text_query, limit=10):
        """Combine vector similarity with text search."""
        query = """
            SELECT id, title, content,
                   1 - (embedding <=> %s) as similarity,
                   ts_rank(to_tsvector('english', content), plainto_tsquery(%s)) as text_rank
            FROM documents
            WHERE to_tsvector('english', content) @@ plainto_tsquery(%s)
            ORDER BY (similarity * 0.7 + text_rank * 0.3) DESC
            LIMIT %s;
        """
        self.cursor.execute(query, (query_embedding, text_query, text_query, limit))
        return self.cursor.fetchall()
    
    def batch_insert(self, documents_with_embeddings):
        """Efficiently insert multiple documents."""
        query = """
            INSERT INTO documents (title, content, embedding)
            VALUES %s;
        """
        from psycopg2.extras import execute_values
        
        values = [(doc['title'], doc['content'], doc['embedding']) 
                 for doc in documents_with_embeddings]
        
        execute_values(self.cursor, query, values, page_size=1000)
        self.conn.commit()

# Example usage
vector_service = VectorSearchService(conn)

# Insert documents
embedding = np.random.random(1536).tolist()  # Example embedding
doc_id = vector_service.insert_document(
    "Vector Databases", 
    "Guide to vector databases and similarity search...",
    embedding
)

# Search for similar documents
query_embedding = np.random.random(1536).tolist()
similar_docs = vector_service.similarity_search(query_embedding, limit=5)

for doc_id, title, content, similarity in similar_docs:
    print(f"ID: {doc_id}, Title: {title}, Similarity: {similarity:.3f}")

# Hybrid search combining vector and text search
hybrid_results = vector_service.hybrid_search(
    query_embedding, 
    "machine learning algorithms",
    limit=5
)

conn.close()