Skip to content

Commit bae0a94

Browse files
ceviancevian
committed
add constants for accessing search results
1 parent e47ef81 commit bae0a94

4 files changed

Lines changed: 182 additions & 13 deletions

File tree

‎README.md‎

Lines changed: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -78,24 +78,57 @@ Now you can query for similar items:
7878
await vec.search([1.0, 9.0])
7979
```
8080

81-
[<Record id=UUID('393ef5c7-1587-468e-870c-b9a6a6716321') metadata='{"action": "jump", "animal": "fox"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>,
82-
<Record id=UUID('7bec0df8-463f-4ea9-8d50-197906f9237e') metadata='{"animal": "fox"}' contents='the brown fox' embedding=array([1. , 1.3], dtype=float32) distance=0.14489260377438218>]
81+
[<Record id=UUID('0c2aceca-9375-4c04-b5b4-01db19eea3c8') metadata='{"action": "jump", "animal": "fox"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>,
82+
<Record id=UUID('4954646e-2677-49c3-91e8-fdc38b8eee75') metadata='{"animal": "fox"}' contents='the brown fox' embedding=array([1. , 1.3], dtype=float32) distance=0.14489260377438218>]
8383

8484
You can specify the number of records to return.
8585

8686
``` python
8787
await vec.search([1.0, 9.0], k=1)
8888
```
8989

90-
[<Record id=UUID('393ef5c7-1587-468e-870c-b9a6a6716321') metadata='{"action": "jump", "animal": "fox"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>]
90+
[<Record id=UUID('0c2aceca-9375-4c04-b5b4-01db19eea3c8') metadata='{"action": "jump", "animal": "fox"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>]
9191

9292
You can also specify a filter on the metadata as a simple dictionary
9393

9494
``` python
9595
await vec.search([1.0, 9.0], k=1, filter={"action": "jump"})
9696
```
9797

98-
[<Record id=UUID('393ef5c7-1587-468e-870c-b9a6a6716321') metadata='{"action": "jump", "animal": "fox"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>]
98+
[<Record id=UUID('0c2aceca-9375-4c04-b5b4-01db19eea3c8') metadata='{"action": "jump", "animal": "fox"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>]
99+
100+
You can access the fields as follows
101+
102+
``` python
103+
records = await vec.search([1.0, 9.0], k=1, filter={"action": "jump"})
104+
records[0][client.SEARCH_RESULT_ID_IDX]
105+
```
106+
107+
UUID('0c2aceca-9375-4c04-b5b4-01db19eea3c8')
108+
109+
``` python
110+
records[0][client.SEARCH_RESULT_METADATA_IDX]
111+
```
112+
113+
'{"action": "jump", "animal": "fox"}'
114+
115+
``` python
116+
records[0][client.SEARCH_RESULT_CONTENTS_IDX]
117+
```
118+
119+
'jumped over the'
120+
121+
``` python
122+
records[0][client.SEARCH_RESULT_EMBEDDING_IDX]
123+
```
124+
125+
array([ 1. , 10.8], dtype=float32)
126+
127+
``` python
128+
records[0][client.SEARCH_RESULT_DISTANCE_IDX]
129+
```
130+
131+
0.00016793422934946456
99132

100133
## Advanced Usage
101134

‎nbs/00_vector.ipynb‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,21 @@
7070
"import json "
7171
]
7272
},
73+
{
74+
"cell_type": "code",
75+
"execution_count": null,
76+
"metadata": {},
77+
"outputs": [],
78+
"source": [
79+
"#| export\n",
80+
"\n",
81+
"SEARCH_RESULT_ID_IDX = 0\n",
82+
"SEARCH_RESULT_METADATA_IDX = 1\n",
83+
"SEARCH_RESULT_CONTENTS_IDX = 2\n",
84+
"SEARCH_RESULT_EMBEDDING_IDX = 3\n",
85+
"SEARCH_RESULT_DISTANCE_IDX = 4"
86+
]
87+
},
7388
{
7489
"cell_type": "code",
7590
"execution_count": null,
@@ -1100,7 +1115,12 @@
11001115
" ])\n",
11011116
" assert False\n",
11021117
"except BaseException as e:\n",
1103-
" pass"
1118+
" pass\n",
1119+
"\n",
1120+
"rec = vec.search([1.0, 2.0], filter={\"key_1\":\"val_1\", \"key_2\":\"val_2\"})\n",
1121+
"assert rec[0][SEARCH_RESULT_CONTENTS_IDX] == 'the brown fox'\n",
1122+
"assert rec[0][SEARCH_RESULT_METADATA_IDX] == {'key_1': 'val_1', 'key_2': 'val_2'}\n",
1123+
"assert rec[0][SEARCH_RESULT_DISTANCE_IDX] == 0.0009438353921149556"
11041124
]
11051125
},
11061126
{

‎nbs/index.ipynb‎

Lines changed: 112 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -187,8 +187,8 @@
187187
{
188188
"data": {
189189
"text/plain": [
190-
"[<Record id=UUID('393ef5c7-1587-468e-870c-b9a6a6716321') metadata='{\"action\": \"jump\", \"animal\": \"fox\"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>,\n",
191-
" <Record id=UUID('7bec0df8-463f-4ea9-8d50-197906f9237e') metadata='{\"animal\": \"fox\"}' contents='the brown fox' embedding=array([1. , 1.3], dtype=float32) distance=0.14489260377438218>]"
190+
"[<Record id=UUID('0c2aceca-9375-4c04-b5b4-01db19eea3c8') metadata='{\"action\": \"jump\", \"animal\": \"fox\"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>,\n",
191+
" <Record id=UUID('4954646e-2677-49c3-91e8-fdc38b8eee75') metadata='{\"animal\": \"fox\"}' contents='the brown fox' embedding=array([1. , 1.3], dtype=float32) distance=0.14489260377438218>]"
192192
]
193193
},
194194
"execution_count": null,
@@ -215,7 +215,7 @@
215215
{
216216
"data": {
217217
"text/plain": [
218-
"[<Record id=UUID('393ef5c7-1587-468e-870c-b9a6a6716321') metadata='{\"action\": \"jump\", \"animal\": \"fox\"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>]"
218+
"[<Record id=UUID('0c2aceca-9375-4c04-b5b4-01db19eea3c8') metadata='{\"action\": \"jump\", \"animal\": \"fox\"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>]"
219219
]
220220
},
221221
"execution_count": null,
@@ -242,7 +242,7 @@
242242
{
243243
"data": {
244244
"text/plain": [
245-
"[<Record id=UUID('393ef5c7-1587-468e-870c-b9a6a6716321') metadata='{\"action\": \"jump\", \"animal\": \"fox\"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>]"
245+
"[<Record id=UUID('0c2aceca-9375-4c04-b5b4-01db19eea3c8') metadata='{\"action\": \"jump\", \"animal\": \"fox\"}' contents='jumped over the' embedding=array([ 1. , 10.8], dtype=float32) distance=0.00016793422934946456>]"
246246
]
247247
},
248248
"execution_count": null,
@@ -254,6 +254,114 @@
254254
"await vec.search([1.0, 9.0], k=1, filter={\"action\": \"jump\"})"
255255
]
256256
},
257+
{
258+
"cell_type": "markdown",
259+
"metadata": {},
260+
"source": [
261+
"You can access the fields as follows"
262+
]
263+
},
264+
{
265+
"cell_type": "code",
266+
"execution_count": null,
267+
"metadata": {},
268+
"outputs": [
269+
{
270+
"data": {
271+
"text/plain": [
272+
"UUID('0c2aceca-9375-4c04-b5b4-01db19eea3c8')"
273+
]
274+
},
275+
"execution_count": null,
276+
"metadata": {},
277+
"output_type": "execute_result"
278+
}
279+
],
280+
"source": [
281+
"records = await vec.search([1.0, 9.0], k=1, filter={\"action\": \"jump\"})\n",
282+
"records[0][client.SEARCH_RESULT_ID_IDX]"
283+
]
284+
},
285+
{
286+
"cell_type": "code",
287+
"execution_count": null,
288+
"metadata": {},
289+
"outputs": [
290+
{
291+
"data": {
292+
"text/plain": [
293+
"'{\"action\": \"jump\", \"animal\": \"fox\"}'"
294+
]
295+
},
296+
"execution_count": null,
297+
"metadata": {},
298+
"output_type": "execute_result"
299+
}
300+
],
301+
"source": [
302+
"records[0][client.SEARCH_RESULT_METADATA_IDX]"
303+
]
304+
},
305+
{
306+
"cell_type": "code",
307+
"execution_count": null,
308+
"metadata": {},
309+
"outputs": [
310+
{
311+
"data": {
312+
"text/plain": [
313+
"'jumped over the'"
314+
]
315+
},
316+
"execution_count": null,
317+
"metadata": {},
318+
"output_type": "execute_result"
319+
}
320+
],
321+
"source": [
322+
"records[0][client.SEARCH_RESULT_CONTENTS_IDX]"
323+
]
324+
},
325+
{
326+
"cell_type": "code",
327+
"execution_count": null,
328+
"metadata": {},
329+
"outputs": [
330+
{
331+
"data": {
332+
"text/plain": [
333+
"array([ 1. , 10.8], dtype=float32)"
334+
]
335+
},
336+
"execution_count": null,
337+
"metadata": {},
338+
"output_type": "execute_result"
339+
}
340+
],
341+
"source": [
342+
"records[0][client.SEARCH_RESULT_EMBEDDING_IDX]"
343+
]
344+
},
345+
{
346+
"cell_type": "code",
347+
"execution_count": null,
348+
"metadata": {},
349+
"outputs": [
350+
{
351+
"data": {
352+
"text/plain": [
353+
"0.00016793422934946456"
354+
]
355+
},
356+
"execution_count": null,
357+
"metadata": {},
358+
"output_type": "execute_result"
359+
}
360+
],
361+
"source": [
362+
"records[0][client.SEARCH_RESULT_DISTANCE_IDX]"
363+
]
364+
},
257365
{
258366
"cell_type": "markdown",
259367
"metadata": {},

‎timescale_vector/client.py‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_vector.ipynb.
22

33
# %% auto 0
4-
__all__ = ['QueryBuilder', 'Async', 'Sync']
4+
__all__ = ['SEARCH_RESULT_ID_IDX', 'SEARCH_RESULT_METADATA_IDX', 'SEARCH_RESULT_CONTENTS_IDX', 'SEARCH_RESULT_EMBEDDING_IDX',
5+
'SEARCH_RESULT_DISTANCE_IDX', 'QueryBuilder', 'Async', 'Sync']
56

67
# %% ../nbs/00_vector.ipynb 6
78
import asyncpg
@@ -11,6 +12,13 @@
1112
import json
1213

1314
# %% ../nbs/00_vector.ipynb 7
15+
SEARCH_RESULT_ID_IDX = 0
16+
SEARCH_RESULT_METADATA_IDX = 1
17+
SEARCH_RESULT_CONTENTS_IDX = 2
18+
SEARCH_RESULT_EMBEDDING_IDX = 3
19+
SEARCH_RESULT_DISTANCE_IDX = 4
20+
21+
# %% ../nbs/00_vector.ipynb 8
1422
class QueryBuilder:
1523
def __init__(
1624
self,
@@ -166,7 +174,7 @@ def search_query(self, query_embedding: List[float], k: int=10, filter: Optional
166174
'''.format(distance=distance, where=where, table_name=self._quote_ident(self.table_name), k=k)
167175
return (query, params)
168176

169-
# %% ../nbs/00_vector.ipynb 10
177+
# %% ../nbs/00_vector.ipynb 11
170178
class Async(QueryBuilder):
171179
def __init__(
172180
self,
@@ -298,15 +306,15 @@ async def search(self,
298306
async with await self.connect() as pool:
299307
return await pool.fetch(query, *params)
300308

301-
# %% ../nbs/00_vector.ipynb 18
309+
# %% ../nbs/00_vector.ipynb 19
302310
import psycopg2.pool
303311
from contextlib import contextmanager
304312
import psycopg2.extras
305313
import pgvector.psycopg2
306314
import numpy as np
307315
import re
308316

309-
# %% ../nbs/00_vector.ipynb 19
317+
# %% ../nbs/00_vector.ipynb 20
310318
class Sync:
311319
translated_queries = {}
312320

0 commit comments

Comments
 (0)