pgvector cosine similarity search with SQLAlchemy

<p>Storing OpenAI embedding vectors in PostgreSQL with pgvector. Need to query for the N most semantically similar traces given a query embedding. Query must filter by tags and status before vector ranking to avoid full-table scans.</p>

<p>Use pgvector's cosine distance operator with SQLAlchemy and filter before ranking:</p> <div class="highlight"><pre><span></span><code><span class="kn">from</span><span class="w"> </span><span class="nn">pgvector.sqlalchemy</span><span class="w"> </span><span class="kn">import</span> <span class="n">Vector</span> <span class="kn">from</span><span class="w"> </span><span class="nn">sqlalchemy</span><span class="w"> </span><span class="kn">import</span> <span class="n">select</span><span class="p">,</span> <span class="n">func</span><span class="p">,</span> <span class="n">and_</span><span class="p">,</span> <span class="n">Float</span> <span class="kn">from</span><span class="w"> </span><span class="nn">sqlalchemy.ext.asyncio</span><span class="w"> </span><span class="kn">import</span> <span class="n">AsyncSession</span> <span class="c1"># Model definition</span> <span class="k">class</span><span class="w"> </span><span class="nc">Trace</span><span class="p">(</span><span class="n">Base</span><span class="p">):</span> <span class="n">__tablename__</span> <span class="o">=</span> <span class="s1">'traces'</span> <span class="n">embedding</span><span class="p">:</span> <span class="n">Mapped</span><span class="p">[</span><span class="n">Optional</span><span class="p">[</span><span class="nb">list</span><span class="p">[</span><span class="nb">float</span><span class="p">]]]</span> <span class="o">=</span> <span class="n">mapped_column</span><span class="p">(</span> <span class="n">Vector</span><span class="p">(</span><span class="mi">1536</span><span class="p">),</span> <span class="n">nullable</span><span class="o">=</span><span class="kc">True</span> <span class="p">)</span> <span class="c1"># Search function</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">semantic_search</span><span class="p">(</span> <span class="n">session</span><span class="p">:</span> <span class="n">AsyncSession</span><span class="p">,</span> <span class="n">query_embedding</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="nb">float</span><span class="p">],</span> <span class="n">tags</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">limit</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">20</span><span class="p">,</span> <span class="n">ann_limit</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">100</span><span class="p">,</span> <span class="c1"># Over-fetch for re-ranking</span> <span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">list</span><span class="p">[</span><span class="n">Trace</span><span class="p">]:</span> <span class="c1"># Cosine distance (1 - cosine_similarity), lower is more similar</span> <span class="n">cosine_dist</span> <span class="o">=</span> <span class="n">Trace</span><span class="o">.</span><span class="n">embedding</span><span class="o">.</span><span class="n">cosine_distance</span><span class="p">(</span><span class="n">query_embedding</span><span class="p">)</span> <span class="n">stmt</span> <span class="o">=</span> <span class="p">(</span> <span class="n">select</span><span class="p">(</span> <span class="n">Trace</span><span class="p">,</span> <span class="n">cosine_dist</span><span class="o">.</span><span class="n">label</span><span class="p">(</span><span class="s1">'similarity_distance'</span><span class="p">),</span> <span class="p">)</span> <span class="o">.</span><span class="n">where</span><span class="p">(</span> <span class="n">and_</span><span class="p">(</span> <span class="n">Trace</span><span class="o">.</span><span class="n">status</span> <span class="o">==</span> <span class="s1">'validated'</span><span class="p">,</span> <span class="n">Trace</span><span class="o">.</span><span class="n">embedding</span><span class="o">.</span><span class="n">is_not</span><span class="p">(</span><span class="kc">None</span><span class="p">),</span> <span class="c1"># Only embedded traces</span> <span class="p">)</span> <span class="p">)</span> <span class="o">.</span><span class="n">order_by</span><span class="p">(</span><span class="n">cosine_dist</span><span class="p">)</span> <span class="c1"># Ascending: smaller distance = more similar</span> <span class="o">.</span><span class="n">limit</span><span class="p">(</span><span class="n">ann_limit</span><span class="p">)</span> <span class="c1"># Over-fetch for re-ranking by trust score</span> <span class="p">)</span> <span class="c1"># Optional tag filter</span> <span class="k">if</span> <span class="n">tags</span><span class="p">:</span> <span class="n">stmt</span> <span class="o">=</span> <span class="n">stmt</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">Trace</span><span class="o">.</span><span class="n">tags</span><span class="p">)</span><span class="o">.</span><span class="n">where</span><span class="p">(</span> <span class="n">Tag</span><span class="o">.</span><span class="n">name</span><span class="o">.</span><span class="n">in_</span><span class="p">(</span><span class="n">tags</span><span class="p">)</span> <span class="p">)</span><span class="o">.</span><span class="n">group_by</span><span class="p">(</span><span class="n">Trace</span><span class="o">.</span><span class="n">id</span><span class="p">)</span><span class="o">.</span><span class="n">having</span><span class="p">(</span> <span class="n">func</span><span class="o">.</span><span class="n">count</span><span class="p">(</span><span class="n">Tag</span><span class="o">.</span><span class="n">id</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="p">)</span> <span class="n">result</span> <span class="o">=</span> <span class="k">await</span> <span class="n">session</span><span class="o">.</span><span class="n">execute</span><span class="p">(</span><span class="n">stmt</span><span class="p">)</span> <span class="n">rows</span> <span class="o">=</span> <span class="n">result</span><span class="o">.</span><span class="n">all</span><span class="p">()</span> <span class="c1"># Re-rank by combining similarity and trust score</span> <span class="k">def</span><span class="w"> </span><span class="nf">combined_score</span><span class="p">(</span><span class="n">row</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">float</span><span class="p">:</span> <span class="n">similarity</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">row</span><span class="o">.</span><span class="n">similarity_distance</span> <span class="c1"># Convert distance to similarity</span> <span class="k">return</span> <span class="mf">0.7</span> <span class="o">*</span> <span class="n">similarity</span> <span class="o">+</span> <span class="mf">0.3</span> <span class="o">*</span> <span class="n">row</span><span class="o">.</span><span class="n">Trace</span><span class="o">.</span><span class="n">trust_score</span> <span class="n">ranked</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">rows</span><span class="p">,</span> <span class="n">key</span><span class="o">=</span><span class="n">combined_score</span><span class="p">,</span> <span class="n">reverse</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="k">return</span> <span class="p">[</span><span class="n">row</span><span class="o">.</span><span class="n">Trace</span> <span class="k">for</span> <span class="n">row</span> <span class="ow">in</span> <span class="n">ranked</span><span class="p">[:</span><span class="n">limit</span><span class="p">]]</span> <span class="c1"># HNSW index for fast approximate nearest neighbor</span> <span class="c1"># CREATE INDEX ON traces USING hnsw (embedding vector_cosine_ops)</span> <span class="c1"># WITH (m = 16, ef_construction = 64);</span> </code></pre></div> <p>Over-fetch (<code>ann_limit=100</code>) then re-rank allows combining vector similarity with domain-specific scores (trust, recency). HNSW index makes vector search O(log N) instead of O(N). <code>cosine_distance</code> returns values in [0, 2]; 0 = identical, 2 = opposite.</p>