diff --git a/src/store/src/Bridge/Postgres/Store.php b/src/store/src/Bridge/Postgres/Store.php index f6ed8bfd..cfb70578 100644 --- a/src/store/src/Bridge/Postgres/Store.php +++ b/src/store/src/Bridge/Postgres/Store.php @@ -92,7 +92,20 @@ public function add(VectorDocument ...$documents): void */ public function query(Vector $vector, array $options = []): array { + $where = null; + $maxScore = $options['maxScore'] ?? null; + if ($maxScore) { + $where = "WHERE ({$this->vectorFieldName} {$this->distance->getComparisonSign()} :embedding) <= :maxScore"; + } + + if ($options['where'] ?? false) { + if ($where) { + $where .= ' AND ('.$options['where'].')'; + } else { + $where = 'WHERE '.$options['where']; + } + } $sql = \sprintf(<<vectorFieldName, $this->distance->getComparisonSign(), $this->tableName, - null !== $maxScore ? "WHERE ({$this->vectorFieldName} {$this->distance->getComparisonSign()} :embedding) <= :maxScore" : '', + $where ?? '', $options['limit'] ?? 5, ); $statement = $this->connection->prepare($sql); $params = [ 'embedding' => $this->toPgvector($vector), + ...$options['params'] ?? [], ]; if (null !== $maxScore) { $params['maxScore'] = $maxScore;