fix

pingcap · Mar 7, 2025 · 85e53dc · 85e53dc
1 parent 77f95f2
commit 85e53dc
Showing 1 changed file with 45 additions and 26 deletions.
diff --git a/core/tests/tidb-client-quickstart.ipynb b/core/tests/tidb-client-quickstart.ipynb
@@ -113,15 +113,19 @@
     "# Define your embedding model.\n",
     "text_embed = EmbeddingFunction(\"openai/text-embedding-3-small\")\n",
     "\n",
-    "class Chunk(TiDBModel, table=True):    \n",
+    "\n",
+    "class Chunk(TiDBModel, table=True):\n",
     "    __tablename__ = \"chunks\"\n",
-    "    __table_args__ = {'extend_existing': True}\n",
+    "    __table_args__ = {\"extend_existing\": True}\n",
     "\n",
     "    id: int = Field(primary_key=True)\n",
     "    text: str = Field()\n",
-    "    text_vec: Optional[Any] = text_embed.VectorField(source_field=\"text\")   # 👈 Define the vector field.\n",
+    "    text_vec: Optional[Any] = text_embed.VectorField(\n",
+    "        source_field=\"text\"\n",
+    "    )  # 👈 Define the vector field.\n",
     "    user_id: int = Field()\n",
     "\n",
+    "\n",
     "table = db.create_table(schema=Chunk)"
    ]
   },
@@ -153,12 +157,16 @@
     }
    ],
    "source": [
-    "table.insert(Chunk(id=1, text=\"The quick brown fox jumps over the lazy dog\", user_id=1),)\n",
-    "table.bulk_insert([\n",
-    "    Chunk(id=2, text=\"A quick brown dog runs in the park\", user_id=2),\n",
-    "    Chunk(id=3, text=\"The lazy fox sleeps under the tree\", user_id=2),\n",
-    "    Chunk(id=4, text=\"A dog and a fox play in the park\", user_id=3)\n",
-    "])\n",
+    "table.insert(\n",
+    "    Chunk(id=1, text=\"The quick brown fox jumps over the lazy dog\", user_id=1),\n",
+    ")\n",
+    "table.bulk_insert(\n",
+    "    [\n",
+    "        Chunk(id=2, text=\"A quick brown dog runs in the park\", user_id=2),\n",
+    "        Chunk(id=3, text=\"The lazy fox sleeps under the tree\", user_id=2),\n",
+    "        Chunk(id=4, text=\"A dog and a fox play in the park\", user_id=3),\n",
+    "    ]\n",
+    ")\n",
     "table.rows()"
    ]
   },
@@ -189,17 +197,13 @@
     }
    ],
    "source": [
-    "from autoflow.storage.tidb import DistanceMetric\n",
-    "\n",
     "chunks = (\n",
-    "    table.search(\"A quick fox in the park\")         # 👈 The query will be embedding automatically.\n",
-    "        # .distance_metric(metric=DistanceMetric.COSINE)\n",
-    "        # .num_candidate(20)\n",
-    "        .filter({\n",
-    "            \"user_id\": 2\n",
-    "        })\n",
-    "        .limit(2)\n",
-    "        .to_pydantic()\n",
+    "    table.search(\n",
+    "        \"A quick fox in the park\"\n",
+    "    )  # 👈 The query will be embedding automatically.\n",
+    "    .filter({\"user_id\": 2})\n",
+    "    .limit(2)\n",
+    "    .to_pydantic()\n",
     ")\n",
     "[(c.text, c.score) for c in chunks]"
    ]
@@ -245,10 +249,7 @@
    ],
    "source": [
     "chunks = table.query({\"user_id\": 1})\n",
-    "[\n",
-    "    (c.id,  c.text, c.user_id)\n",
-    "    for c in chunks\n",
-    "]"
+    "[(c.id, c.text, c.user_id) for c in chunks]"
    ]
   },
   {
@@ -284,7 +285,25 @@
    ]
   }
  ],
- "metadata": {},
- "nbformat": 5,
- "nbformat_minor": 9
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
 }