Skip to content

Commit a7084e1

Browse files
Mini256Mini256
committed
fix: refine SDK title and add raw sql example
1 parent 85e53dc commit a7084e1

File tree

3 files changed

+124
-64
lines changed

3 files changed

+124
-64
lines changed

‎core/autoflow/storage/tidb/README.md‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# TiDB Vector SDK V2
1+
# TiDB Python SDK V2
22

33
A powerful Python SDK for vector storage and retrieval operations with TiDB.
44

@@ -11,7 +11,6 @@ A powerful Python SDK for vector storage and retrieval operations with TiDB.
1111

1212
```bash
1313
pip install autoflow-ai
14-
# TODO: move to `pip install tidb-vector`
1514
```
1615

1716
## Configuration

‎core/autoflow/storage/tidb/query.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
import numpy as np
1616
import pandas as pd
1717
from pydantic import BaseModel
18-
from sqlalchemy import asc, desc, Row, select
18+
from sqlalchemy import asc, desc, select
1919
from sqlalchemy.orm import Session
2020

2121
from autoflow.storage.tidb.constants import DistanceMetric, VectorDataType
@@ -92,7 +92,7 @@ def execute(self):
9292
def _execute(self):
9393
pass
9494

95-
def to_rows(self) -> Sequence[Row]:
95+
def to_rows(self) -> Sequence[Any]:
9696
return self.execute()
9797

9898
def to_pandas(self) -> pd.DataFrame:
@@ -130,7 +130,7 @@ def limit(self, k: int) -> "TiDBVectorQuery":
130130
self._limit = k
131131
return self
132132

133-
def _execute(self) -> Sequence[Row]:
133+
def _execute(self) -> Sequence[Any]:
134134
num_candidate = self._num_candidate if self._num_candidate else self._limit * 10
135135

136136
if self._vector_column is None:

‎core/tests/tidb-client-quickstart.ipynb‎

Lines changed: 120 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
"id": "b90e936b724dae40",
66
"metadata": {},
77
"source": [
8-
"# TiDB Vector SDK V2\n",
8+
"# TiDB Python SDK V2\n",
99
"\n",
1010
"A powerful Python SDK for vector storage and retrieval operations with TiDB.\n",
1111
"\n",
@@ -25,14 +25,14 @@
2525
},
2626
{
2727
"cell_type": "code",
28-
"execution_count": null,
2928
"id": "b81aa3ab",
3029
"metadata": {},
31-
"outputs": [],
3230
"source": [
3331
"%pip install autoflow-ai==0.0.1.dev10\n",
3432
"%pip install dotenv ipywidgets pymysql sqlmodel"
35-
]
33+
],
34+
"outputs": [],
35+
"execution_count": null
3636
},
3737
{
3838
"cell_type": "markdown",
@@ -41,7 +41,7 @@
4141
"source": [
4242
"#### Configure environment variable\n",
4343
"\n",
44-
"Go [tidbcloud.com](http://tidbcloud.com/) or using [tiup playground](https://docs.pingcap.com/tidb/stable/tiup-playground/) to create a free TiDB database cluster\n",
44+
"Go [tidbcloud.com](http://tidbcloud.com/) or using [tiup playground](https://docs.pingcap.com/tidb/stable/tiup-playground/) to create a free TiDB database cluster.\n",
4545
"\n",
4646
"Configuration can be provided through environment variables, or using `.env`:"
4747
]
@@ -78,17 +78,22 @@
7878
},
7979
{
8080
"cell_type": "code",
81-
"execution_count": 39,
8281
"id": "4995a54f311c4b1c",
83-
"metadata": {},
84-
"outputs": [],
82+
"metadata": {
83+
"ExecuteTime": {
84+
"end_time": "2025-03-10T02:09:23.700467Z",
85+
"start_time": "2025-03-10T02:09:06.319163Z"
86+
}
87+
},
8588
"source": [
8689
"import os\n",
8790
"from autoflow.storage.tidb import TiDBClient\n",
8891
"\n",
8992
"# Format: mysql+pymysql://<username>:<password>@<host>:4000/<database>\n",
9093
"db = TiDBClient.connect(os.getenv(\"DATABASE_URL\"))"
91-
]
94+
],
95+
"outputs": [],
96+
"execution_count": 3
9297
},
9398
{
9499
"cell_type": "markdown",
@@ -100,10 +105,13 @@
100105
},
101106
{
102107
"cell_type": "code",
103-
"execution_count": 40,
104108
"id": "bdddb9f0a005b74d",
105-
"metadata": {},
106-
"outputs": [],
109+
"metadata": {
110+
"ExecuteTime": {
111+
"end_time": "2025-03-10T02:09:37.844382Z",
112+
"start_time": "2025-03-10T02:09:37.823858Z"
113+
}
114+
},
107115
"source": [
108116
"from typing import Optional, Any\n",
109117
"from autoflow.storage.tidb.base import TiDBModel\n",
@@ -127,7 +135,9 @@
127135
"\n",
128136
"\n",
129137
"table = db.create_table(schema=Chunk)"
130-
]
138+
],
139+
"outputs": [],
140+
"execution_count": 4
131141
},
132142
{
133143
"cell_type": "markdown",
@@ -141,34 +151,39 @@
141151
},
142152
{
143153
"cell_type": "code",
144-
"execution_count": 46,
145154
"id": "baec9a5ae06231be",
146-
"metadata": {},
155+
"metadata": {
156+
"ExecuteTime": {
157+
"end_time": "2025-03-10T02:09:46.528201Z",
158+
"start_time": "2025-03-10T02:09:46.490530Z"
159+
}
160+
},
161+
"source": [
162+
"table.insert(\n",
163+
" Chunk(text=\"The quick brown fox jumps over the lazy dog\", user_id=1),\n",
164+
")\n",
165+
"table.bulk_insert(\n",
166+
" [\n",
167+
" Chunk(text=\"A quick brown dog runs in the park\", user_id=2),\n",
168+
" Chunk(text=\"The lazy fox sleeps under the tree\", user_id=2),\n",
169+
" Chunk(text=\"A dog and a fox play in the park\", user_id=3),\n",
170+
" ]\n",
171+
")\n",
172+
"table.rows()"
173+
],
147174
"outputs": [
148175
{
149176
"data": {
150177
"text/plain": [
151-
"4"
178+
"8"
152179
]
153180
},
154-
"execution_count": 46,
181+
"execution_count": 5,
155182
"metadata": {},
156183
"output_type": "execute_result"
157184
}
158185
],
159-
"source": [
160-
"table.insert(\n",
161-
" Chunk(id=1, text=\"The quick brown fox jumps over the lazy dog\", user_id=1),\n",
162-
")\n",
163-
"table.bulk_insert(\n",
164-
" [\n",
165-
" Chunk(id=2, text=\"A quick brown dog runs in the park\", user_id=2),\n",
166-
" Chunk(id=3, text=\"The lazy fox sleeps under the tree\", user_id=2),\n",
167-
" Chunk(id=4, text=\"A dog and a fox play in the park\", user_id=3),\n",
168-
" ]\n",
169-
")\n",
170-
"table.rows()"
171-
]
186+
"execution_count": 5
172187
},
173188
{
174189
"cell_type": "markdown",
@@ -180,22 +195,13 @@
180195
},
181196
{
182197
"cell_type": "code",
183-
"execution_count": 53,
184198
"id": "3c4313022f06bd3e",
185-
"metadata": {},
186-
"outputs": [
187-
{
188-
"data": {
189-
"text/plain": [
190-
"[('A quick brown dog runs in the park', 0.665493189763966),\n",
191-
" ('The lazy fox sleeps under the tree', 0.554631888866523)]"
192-
]
193-
},
194-
"execution_count": 53,
195-
"metadata": {},
196-
"output_type": "execute_result"
199+
"metadata": {
200+
"ExecuteTime": {
201+
"end_time": "2025-03-10T02:09:56.362012Z",
202+
"start_time": "2025-03-10T02:09:56.330090Z"
197203
}
198-
],
204+
},
199205
"source": [
200206
"chunks = (\n",
201207
" table.search(\n",
@@ -206,7 +212,21 @@
206212
" .to_pydantic()\n",
207213
")\n",
208214
"[(c.text, c.score) for c in chunks]"
209-
]
215+
],
216+
"outputs": [
217+
{
218+
"data": {
219+
"text/plain": [
220+
"[('A quick brown dog runs in the park', 0.665493189763966),\n",
221+
" ('A quick brown dog runs in the park', 0.665493189763966)]"
222+
]
223+
},
224+
"execution_count": 6,
225+
"metadata": {},
226+
"output_type": "execute_result"
227+
}
228+
],
229+
"execution_count": 6
210230
},
211231
{
212232
"cell_type": "markdown",
@@ -232,25 +252,61 @@
232252
},
233253
{
234254
"cell_type": "code",
235-
"execution_count": 49,
236255
"id": "ace02b45",
256+
"metadata": {
257+
"ExecuteTime": {
258+
"end_time": "2025-03-10T02:10:03.161292Z",
259+
"start_time": "2025-03-10T02:10:03.154627Z"
260+
}
261+
},
262+
"source": [
263+
"chunks = table.query({\"user_id\": 1})\n",
264+
"[(c.id, c.text, c.user_id) for c in chunks]"
265+
],
266+
"outputs": [
267+
{
268+
"data": {
269+
"text/plain": [
270+
"[(1, 'The quick brown fox jumps over the lazy dog', 1),\n",
271+
" (5, 'The quick brown fox jumps over the lazy dog', 1)]"
272+
]
273+
},
274+
"execution_count": 7,
275+
"metadata": {},
276+
"output_type": "execute_result"
277+
}
278+
],
279+
"execution_count": 7
280+
},
281+
{
237282
"metadata": {},
283+
"cell_type": "markdown",
284+
"source": "### Execute raw SQL",
285+
"id": "f8dd3bc68287ade1"
286+
},
287+
{
288+
"metadata": {
289+
"ExecuteTime": {
290+
"end_time": "2025-03-10T02:13:36.455402Z",
291+
"start_time": "2025-03-10T02:13:36.445640Z"
292+
}
293+
},
294+
"cell_type": "code",
295+
"source": "db.execute(\"SELECT COUNT(*) FROM chunks\")",
296+
"id": "124cc1a2463c474d",
238297
"outputs": [
239298
{
240299
"data": {
241300
"text/plain": [
242-
"[(1, 'The quick brown fox jumps over the lazy dog', 1)]"
301+
"{'success': True, 'result': [(0,)], 'error': None}"
243302
]
244303
},
245-
"execution_count": 49,
304+
"execution_count": 13,
246305
"metadata": {},
247306
"output_type": "execute_result"
248307
}
249308
],
250-
"source": [
251-
"chunks = table.query({\"user_id\": 1})\n",
252-
"[(c.id, c.text, c.user_id) for c in chunks]"
253-
]
309+
"execution_count": 13
254310
},
255311
{
256312
"cell_type": "markdown",
@@ -264,25 +320,30 @@
264320
},
265321
{
266322
"cell_type": "code",
267-
"execution_count": 45,
268323
"id": "cceb0bf0",
269-
"metadata": {},
324+
"metadata": {
325+
"ExecuteTime": {
326+
"end_time": "2025-03-10T02:13:11.093835Z",
327+
"start_time": "2025-03-10T02:13:11.024531Z"
328+
}
329+
},
330+
"source": [
331+
"table.truncate()\n",
332+
"table.rows()"
333+
],
270334
"outputs": [
271335
{
272336
"data": {
273337
"text/plain": [
274338
"0"
275339
]
276340
},
277-
"execution_count": 45,
341+
"execution_count": 10,
278342
"metadata": {},
279343
"output_type": "execute_result"
280344
}
281345
],
282-
"source": [
283-
"table.truncate()\n",
284-
"table.rows()"
285-
]
346+
"execution_count": 10
286347
}
287348
],
288349
"metadata": {

0 commit comments

Comments
 (0)