langchain-shannonbase 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- langchain_shannonbase-0.1.0/.github/workflows/ci.yml +18 -0
- langchain_shannonbase-0.1.0/.github/workflows/publish.yml +17 -0
- langchain_shannonbase-0.1.0/.gitignore +15 -0
- langchain_shannonbase-0.1.0/CHANGELOG.md +11 -0
- langchain_shannonbase-0.1.0/LICENSE +21 -0
- langchain_shannonbase-0.1.0/PKG-INFO +127 -0
- langchain_shannonbase-0.1.0/README.md +101 -0
- langchain_shannonbase-0.1.0/examples/basic.py +34 -0
- langchain_shannonbase-0.1.0/pyproject.toml +37 -0
- langchain_shannonbase-0.1.0/src/langchain_shannonbase/__init__.py +8 -0
- langchain_shannonbase-0.1.0/src/langchain_shannonbase/_sql.py +67 -0
- langchain_shannonbase-0.1.0/src/langchain_shannonbase/_store.py +150 -0
- langchain_shannonbase-0.1.0/src/langchain_shannonbase/vectorstores.py +132 -0
- langchain_shannonbase-0.1.0/tests/test_integration.py +57 -0
- langchain_shannonbase-0.1.0/tests/test_sql.py +40 -0
- langchain_shannonbase-0.1.0/tests/test_standard.py +18 -0
- langchain_shannonbase-0.1.0/tests/test_vectorstore.py +81 -0
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
name: CI
|
|
2
|
+
on:
|
|
3
|
+
push:
|
|
4
|
+
branches: [main]
|
|
5
|
+
pull_request:
|
|
6
|
+
jobs:
|
|
7
|
+
test:
|
|
8
|
+
runs-on: ubuntu-latest
|
|
9
|
+
strategy:
|
|
10
|
+
matrix:
|
|
11
|
+
python-version: ["3.9", "3.11", "3.12"]
|
|
12
|
+
steps:
|
|
13
|
+
- uses: actions/checkout@v4
|
|
14
|
+
- uses: actions/setup-python@v5
|
|
15
|
+
with:
|
|
16
|
+
python-version: ${{ matrix.python-version }}
|
|
17
|
+
- run: pip install -e ".[dev]"
|
|
18
|
+
- run: pytest -q
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
on:
|
|
3
|
+
release:
|
|
4
|
+
types: [published]
|
|
5
|
+
jobs:
|
|
6
|
+
build-and-publish:
|
|
7
|
+
runs-on: ubuntu-latest
|
|
8
|
+
environment: pypi
|
|
9
|
+
permissions:
|
|
10
|
+
id-token: write
|
|
11
|
+
steps:
|
|
12
|
+
- uses: actions/checkout@v4
|
|
13
|
+
- uses: actions/setup-python@v5
|
|
14
|
+
with:
|
|
15
|
+
python-version: "3.12"
|
|
16
|
+
- run: pip install build && python -m build
|
|
17
|
+
- uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
# Changelog
|
|
2
|
+
|
|
3
|
+
## 0.1.0
|
|
4
|
+
- Initial release.
|
|
5
|
+
- `ShannonBaseVectorStore` implementing LangChain's VectorStore interface on
|
|
6
|
+
MySQL 9's native VECTOR type (ShannonBase / MySQL 9 / HeatWave).
|
|
7
|
+
- add_texts, similarity_search(_with_score / _by_vector), get_by_ids, delete, from_texts.
|
|
8
|
+
- cosine / dot / euclidean metrics.
|
|
9
|
+
- Offline InMemoryStore for deterministic tests; gated live integration test.
|
|
10
|
+
- Passes LangChain's standard VectorStore integration suite (langchain-tests),
|
|
11
|
+
run offline against the in-memory backend.
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Apoorva Verma
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,127 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: langchain-shannonbase
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: LangChain VectorStore for MySQL 9's native VECTOR type — works with ShannonBase, self-hosted MySQL, and MySQL HeatWave.
|
|
5
|
+
Project-URL: Homepage, https://github.com/apoorva-01/langchain-shannonbase
|
|
6
|
+
Project-URL: Issues, https://github.com/apoorva-01/langchain-shannonbase/issues
|
|
7
|
+
Author: Apoorva Verma
|
|
8
|
+
License: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Keywords: embeddings,heatwave,langchain,mysql,rag,shannonbase,vector-search,vectorstore
|
|
11
|
+
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Topic :: Database
|
|
16
|
+
Requires-Python: >=3.9
|
|
17
|
+
Requires-Dist: langchain-core>=0.3
|
|
18
|
+
Provides-Extra: dev
|
|
19
|
+
Requires-Dist: langchain-tests>=0.3; extra == 'dev'
|
|
20
|
+
Requires-Dist: mysql-connector-python>=8.3; extra == 'dev'
|
|
21
|
+
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
|
|
22
|
+
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
23
|
+
Provides-Extra: mysql
|
|
24
|
+
Requires-Dist: mysql-connector-python>=8.3; extra == 'mysql'
|
|
25
|
+
Description-Content-Type: text/markdown
|
|
26
|
+
|
|
27
|
+
# langchain-shannonbase
|
|
28
|
+
|
|
29
|
+
A [LangChain](https://python.langchain.com) `VectorStore` for **MySQL 9's native
|
|
30
|
+
`VECTOR` type** — so you can do RAG on a database you already run.
|
|
31
|
+
|
|
32
|
+
Works with **[ShannonBase](https://github.com/Shannon-Data/ShannonBase)** (the
|
|
33
|
+
open-source MySQL-for-AI), **self-hosted MySQL 9**, and **MySQL HeatWave** — they
|
|
34
|
+
all share the same `VECTOR` / `STRING_TO_VECTOR` / `DISTANCE` surface.
|
|
35
|
+
|
|
36
|
+
## Why this exists
|
|
37
|
+
|
|
38
|
+
If your data already lives in MySQL, your options for LangChain vector search were
|
|
39
|
+
thin: the only MySQL `VectorStore` is locked to Google Cloud SQL, and ShannonBase's
|
|
40
|
+
LangChain integration was on its wishlist but unbuilt. This fills that gap — a
|
|
41
|
+
plain, self-hostable adapter that plugs MySQL 9 into the LangChain ecosystem, no
|
|
42
|
+
separate vector database required.
|
|
43
|
+
|
|
44
|
+
## Install
|
|
45
|
+
|
|
46
|
+
```bash
|
|
47
|
+
pip install "langchain-shannonbase[mysql]"
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
## Use
|
|
51
|
+
|
|
52
|
+
```python
|
|
53
|
+
from langchain_openai import OpenAIEmbeddings
|
|
54
|
+
from langchain_shannonbase import ShannonBaseVectorStore
|
|
55
|
+
|
|
56
|
+
store = ShannonBaseVectorStore(
|
|
57
|
+
embedding=OpenAIEmbeddings(model="text-embedding-3-small"),
|
|
58
|
+
table="documents",
|
|
59
|
+
host="127.0.0.1", port=3306, user="root", password="", database="rag",
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
store.add_texts(
|
|
63
|
+
["Refunds are accepted within 30 days.", "Free shipping over $50."],
|
|
64
|
+
metadatas=[{"topic": "refunds"}, {"topic": "shipping"}],
|
|
65
|
+
)
|
|
66
|
+
|
|
67
|
+
# It's a normal LangChain vector store — use it directly or as a retriever:
|
|
68
|
+
docs = store.similarity_search("return policy?", k=2)
|
|
69
|
+
retriever = store.as_retriever(search_kwargs={"k": 3})
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
Because it implements LangChain's `VectorStore` interface, it drops into any
|
|
73
|
+
LangChain chain, retriever, or RAG pipeline unchanged.
|
|
74
|
+
|
|
75
|
+
## How it works
|
|
76
|
+
|
|
77
|
+
Under the hood it uses MySQL 9's native vector features — no extensions:
|
|
78
|
+
|
|
79
|
+
```sql
|
|
80
|
+
CREATE TABLE documents (
|
|
81
|
+
id VARCHAR(36) PRIMARY KEY,
|
|
82
|
+
content TEXT, metadata JSON,
|
|
83
|
+
embedding VECTOR(1536)
|
|
84
|
+
);
|
|
85
|
+
-- inserts go through STRING_TO_VECTOR('[...]')
|
|
86
|
+
-- search: ORDER BY DISTANCE(embedding, STRING_TO_VECTOR('[...]'), 'COSINE')
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
Similarity search reads back the nearest rows by cosine distance and returns them
|
|
90
|
+
as LangChain `Document`s with a score (`1 - distance`).
|
|
91
|
+
|
|
92
|
+
## API
|
|
93
|
+
|
|
94
|
+
| Method | Does |
|
|
95
|
+
|--------|------|
|
|
96
|
+
| `add_texts(texts, metadatas, ids)` | embed + upsert, returns ids |
|
|
97
|
+
| `similarity_search(query, k)` | top-k `Document`s |
|
|
98
|
+
| `similarity_search_with_score(query, k)` | with cosine similarity scores |
|
|
99
|
+
| `similarity_search_by_vector(embedding, k)` | search with a raw vector |
|
|
100
|
+
| `delete(ids)` | remove by id |
|
|
101
|
+
| `from_texts(texts, embedding, ...)` | build a store in one call |
|
|
102
|
+
| `metric=` | `"cosine"` (default), `"dot"`, `"euclidean"` |
|
|
103
|
+
|
|
104
|
+
## Testing
|
|
105
|
+
|
|
106
|
+
The core logic is unit-tested offline via an in-memory backend (no database
|
|
107
|
+
needed — `pytest`). A live round-trip test runs against a real instance when you
|
|
108
|
+
set the connection env vars:
|
|
109
|
+
|
|
110
|
+
```bash
|
|
111
|
+
export SB_HOST=127.0.0.1 SB_USER=root SB_PASSWORD=... SB_DATABASE=test
|
|
112
|
+
pytest tests/test_integration.py
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
> Local dev tip: run [ShannonBase](https://github.com/Shannon-Data/ShannonBase) to
|
|
116
|
+
> get MySQL-9 vector features without a HeatWave subscription.
|
|
117
|
+
|
|
118
|
+
## Requirements
|
|
119
|
+
|
|
120
|
+
- Python 3.9+
|
|
121
|
+
- A MySQL-9-compatible database with the `VECTOR` type (ShannonBase, MySQL 9, or
|
|
122
|
+
HeatWave)
|
|
123
|
+
- `mysql-connector-python` (installed via the `[mysql]` extra)
|
|
124
|
+
|
|
125
|
+
## License
|
|
126
|
+
|
|
127
|
+
MIT © Apoorva Verma
|
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
# langchain-shannonbase
|
|
2
|
+
|
|
3
|
+
A [LangChain](https://python.langchain.com) `VectorStore` for **MySQL 9's native
|
|
4
|
+
`VECTOR` type** — so you can do RAG on a database you already run.
|
|
5
|
+
|
|
6
|
+
Works with **[ShannonBase](https://github.com/Shannon-Data/ShannonBase)** (the
|
|
7
|
+
open-source MySQL-for-AI), **self-hosted MySQL 9**, and **MySQL HeatWave** — they
|
|
8
|
+
all share the same `VECTOR` / `STRING_TO_VECTOR` / `DISTANCE` surface.
|
|
9
|
+
|
|
10
|
+
## Why this exists
|
|
11
|
+
|
|
12
|
+
If your data already lives in MySQL, your options for LangChain vector search were
|
|
13
|
+
thin: the only MySQL `VectorStore` is locked to Google Cloud SQL, and ShannonBase's
|
|
14
|
+
LangChain integration was on its wishlist but unbuilt. This fills that gap — a
|
|
15
|
+
plain, self-hostable adapter that plugs MySQL 9 into the LangChain ecosystem, no
|
|
16
|
+
separate vector database required.
|
|
17
|
+
|
|
18
|
+
## Install
|
|
19
|
+
|
|
20
|
+
```bash
|
|
21
|
+
pip install "langchain-shannonbase[mysql]"
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
## Use
|
|
25
|
+
|
|
26
|
+
```python
|
|
27
|
+
from langchain_openai import OpenAIEmbeddings
|
|
28
|
+
from langchain_shannonbase import ShannonBaseVectorStore
|
|
29
|
+
|
|
30
|
+
store = ShannonBaseVectorStore(
|
|
31
|
+
embedding=OpenAIEmbeddings(model="text-embedding-3-small"),
|
|
32
|
+
table="documents",
|
|
33
|
+
host="127.0.0.1", port=3306, user="root", password="", database="rag",
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
store.add_texts(
|
|
37
|
+
["Refunds are accepted within 30 days.", "Free shipping over $50."],
|
|
38
|
+
metadatas=[{"topic": "refunds"}, {"topic": "shipping"}],
|
|
39
|
+
)
|
|
40
|
+
|
|
41
|
+
# It's a normal LangChain vector store — use it directly or as a retriever:
|
|
42
|
+
docs = store.similarity_search("return policy?", k=2)
|
|
43
|
+
retriever = store.as_retriever(search_kwargs={"k": 3})
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
Because it implements LangChain's `VectorStore` interface, it drops into any
|
|
47
|
+
LangChain chain, retriever, or RAG pipeline unchanged.
|
|
48
|
+
|
|
49
|
+
## How it works
|
|
50
|
+
|
|
51
|
+
Under the hood it uses MySQL 9's native vector features — no extensions:
|
|
52
|
+
|
|
53
|
+
```sql
|
|
54
|
+
CREATE TABLE documents (
|
|
55
|
+
id VARCHAR(36) PRIMARY KEY,
|
|
56
|
+
content TEXT, metadata JSON,
|
|
57
|
+
embedding VECTOR(1536)
|
|
58
|
+
);
|
|
59
|
+
-- inserts go through STRING_TO_VECTOR('[...]')
|
|
60
|
+
-- search: ORDER BY DISTANCE(embedding, STRING_TO_VECTOR('[...]'), 'COSINE')
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
Similarity search reads back the nearest rows by cosine distance and returns them
|
|
64
|
+
as LangChain `Document`s with a score (`1 - distance`).
|
|
65
|
+
|
|
66
|
+
## API
|
|
67
|
+
|
|
68
|
+
| Method | Does |
|
|
69
|
+
|--------|------|
|
|
70
|
+
| `add_texts(texts, metadatas, ids)` | embed + upsert, returns ids |
|
|
71
|
+
| `similarity_search(query, k)` | top-k `Document`s |
|
|
72
|
+
| `similarity_search_with_score(query, k)` | with cosine similarity scores |
|
|
73
|
+
| `similarity_search_by_vector(embedding, k)` | search with a raw vector |
|
|
74
|
+
| `delete(ids)` | remove by id |
|
|
75
|
+
| `from_texts(texts, embedding, ...)` | build a store in one call |
|
|
76
|
+
| `metric=` | `"cosine"` (default), `"dot"`, `"euclidean"` |
|
|
77
|
+
|
|
78
|
+
## Testing
|
|
79
|
+
|
|
80
|
+
The core logic is unit-tested offline via an in-memory backend (no database
|
|
81
|
+
needed — `pytest`). A live round-trip test runs against a real instance when you
|
|
82
|
+
set the connection env vars:
|
|
83
|
+
|
|
84
|
+
```bash
|
|
85
|
+
export SB_HOST=127.0.0.1 SB_USER=root SB_PASSWORD=... SB_DATABASE=test
|
|
86
|
+
pytest tests/test_integration.py
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
> Local dev tip: run [ShannonBase](https://github.com/Shannon-Data/ShannonBase) to
|
|
90
|
+
> get MySQL-9 vector features without a HeatWave subscription.
|
|
91
|
+
|
|
92
|
+
## Requirements
|
|
93
|
+
|
|
94
|
+
- Python 3.9+
|
|
95
|
+
- A MySQL-9-compatible database with the `VECTOR` type (ShannonBase, MySQL 9, or
|
|
96
|
+
HeatWave)
|
|
97
|
+
- `mysql-connector-python` (installed via the `[mysql]` extra)
|
|
98
|
+
|
|
99
|
+
## License
|
|
100
|
+
|
|
101
|
+
MIT © Apoorva Verma
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
"""Minimal RAG-style usage against a real ShannonBase / MySQL 9 / HeatWave DB.
|
|
2
|
+
|
|
3
|
+
pip install 'langchain-shannonbase[mysql]' langchain-openai
|
|
4
|
+
export OPENAI_API_KEY=...
|
|
5
|
+
python examples/basic.py
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from langchain_openai import OpenAIEmbeddings
|
|
9
|
+
|
|
10
|
+
from langchain_shannonbase import ShannonBaseVectorStore
|
|
11
|
+
|
|
12
|
+
store = ShannonBaseVectorStore(
|
|
13
|
+
embedding=OpenAIEmbeddings(model="text-embedding-3-small"),
|
|
14
|
+
table="documents",
|
|
15
|
+
host="127.0.0.1",
|
|
16
|
+
port=3306,
|
|
17
|
+
user="root",
|
|
18
|
+
password="",
|
|
19
|
+
database="rag",
|
|
20
|
+
)
|
|
21
|
+
|
|
22
|
+
store.add_texts(
|
|
23
|
+
[
|
|
24
|
+
"Our refund policy allows returns within 30 days of purchase.",
|
|
25
|
+
"Support is available Monday to Friday, 9am to 6pm.",
|
|
26
|
+
"Free shipping on orders over $50.",
|
|
27
|
+
],
|
|
28
|
+
metadatas=[{"topic": "refunds"}, {"topic": "support"}, {"topic": "shipping"}],
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
# Use it like any LangChain vector store:
|
|
32
|
+
retriever = store.as_retriever(search_kwargs={"k": 2})
|
|
33
|
+
for doc in retriever.invoke("how long do I have to return something?"):
|
|
34
|
+
print(f"- {doc.page_content} ({doc.metadata})")
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "langchain-shannonbase"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "LangChain VectorStore for MySQL 9's native VECTOR type — works with ShannonBase, self-hosted MySQL, and MySQL HeatWave."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [{ name = "Apoorva Verma" }]
|
|
13
|
+
keywords = ["langchain", "vectorstore", "mysql", "shannonbase", "heatwave", "rag", "vector-search", "embeddings"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Development Status :: 4 - Beta",
|
|
16
|
+
"Intended Audience :: Developers",
|
|
17
|
+
"License :: OSI Approved :: MIT License",
|
|
18
|
+
"Programming Language :: Python :: 3",
|
|
19
|
+
"Topic :: Database",
|
|
20
|
+
]
|
|
21
|
+
dependencies = [
|
|
22
|
+
"langchain-core>=0.3",
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
[project.optional-dependencies]
|
|
26
|
+
mysql = ["mysql-connector-python>=8.3"]
|
|
27
|
+
dev = ["pytest>=7.0", "pytest-asyncio>=0.23", "mysql-connector-python>=8.3", "langchain-tests>=0.3"]
|
|
28
|
+
|
|
29
|
+
[project.urls]
|
|
30
|
+
Homepage = "https://github.com/apoorva-01/langchain-shannonbase"
|
|
31
|
+
Issues = "https://github.com/apoorva-01/langchain-shannonbase/issues"
|
|
32
|
+
|
|
33
|
+
[tool.hatch.build.targets.wheel]
|
|
34
|
+
packages = ["src/langchain_shannonbase"]
|
|
35
|
+
|
|
36
|
+
[tool.pytest.ini_options]
|
|
37
|
+
asyncio_mode = "auto"
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
"""langchain-shannonbase — a LangChain VectorStore for MySQL 9's VECTOR type."""
|
|
2
|
+
|
|
3
|
+
__version__ = "0.1.0"
|
|
4
|
+
|
|
5
|
+
from ._store import InMemoryStore, MySQLStore
|
|
6
|
+
from .vectorstores import ShannonBaseVectorStore
|
|
7
|
+
|
|
8
|
+
__all__ = ["ShannonBaseVectorStore", "MySQLStore", "InMemoryStore", "__version__"]
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
"""SQL for MySQL 9's native VECTOR type (ShannonBase / MySQL / HeatWave).
|
|
2
|
+
|
|
3
|
+
Kept as pure string builders so they can be unit-tested without a database.
|
|
4
|
+
Embeddings are passed to MySQL as a JSON array string via STRING_TO_VECTOR(),
|
|
5
|
+
and cosine distance comes from the native DISTANCE(a, b, 'COSINE') function.
|
|
6
|
+
similarity = 1 - distance.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import json
|
|
12
|
+
from typing import List
|
|
13
|
+
|
|
14
|
+
# Distance metric -> MySQL DISTANCE() metric name.
|
|
15
|
+
_METRICS = {"cosine": "COSINE", "dot": "DOT", "euclidean": "EUCLIDEAN"}
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def vector_literal(embedding: List[float]) -> str:
|
|
19
|
+
"""Serialize a vector to the JSON-array string STRING_TO_VECTOR expects."""
|
|
20
|
+
return json.dumps([float(x) for x in embedding])
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def create_table_sql(table: str, dim: int) -> str:
|
|
24
|
+
return (
|
|
25
|
+
f"CREATE TABLE IF NOT EXISTS `{table}` ("
|
|
26
|
+
" id VARCHAR(36) PRIMARY KEY,"
|
|
27
|
+
" content TEXT,"
|
|
28
|
+
" metadata JSON,"
|
|
29
|
+
f" embedding VECTOR({dim})"
|
|
30
|
+
")"
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def insert_sql(table: str) -> str:
|
|
35
|
+
"""Upsert one row. Embedding bind param is wrapped in STRING_TO_VECTOR()."""
|
|
36
|
+
return (
|
|
37
|
+
f"INSERT INTO `{table}` (id, content, metadata, embedding) "
|
|
38
|
+
"VALUES (%s, %s, %s, STRING_TO_VECTOR(%s)) "
|
|
39
|
+
"AS new ON DUPLICATE KEY UPDATE "
|
|
40
|
+
"content = new.content, metadata = new.metadata, embedding = new.embedding"
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def search_sql(table: str, metric: str = "cosine") -> str:
|
|
45
|
+
m = _METRICS.get(metric)
|
|
46
|
+
if m is None:
|
|
47
|
+
raise ValueError(f"unknown metric {metric!r}; use one of {list(_METRICS)}")
|
|
48
|
+
return (
|
|
49
|
+
"SELECT id, content, metadata, "
|
|
50
|
+
f"DISTANCE(embedding, STRING_TO_VECTOR(%s), '{m}') AS dist "
|
|
51
|
+
f"FROM `{table}` ORDER BY dist ASC LIMIT %s"
|
|
52
|
+
)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def delete_sql(table: str, n_ids: int) -> str:
|
|
56
|
+
placeholders = ", ".join(["%s"] * n_ids)
|
|
57
|
+
return f"DELETE FROM `{table}` WHERE id IN ({placeholders})"
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def select_by_ids_sql(table: str, n_ids: int) -> str:
|
|
61
|
+
placeholders = ", ".join(["%s"] * n_ids)
|
|
62
|
+
return f"SELECT id, content, metadata FROM `{table}` WHERE id IN ({placeholders})"
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def distance_to_score(distance: float) -> float:
|
|
66
|
+
"""Convert a cosine distance to a [0, 1]-ish similarity score."""
|
|
67
|
+
return 1.0 - float(distance)
|
|
@@ -0,0 +1,150 @@
|
|
|
1
|
+
"""Storage backends behind the vector store.
|
|
2
|
+
|
|
3
|
+
MySQLStore talks to a real ShannonBase / MySQL 9 / HeatWave instance. InMemoryStore
|
|
4
|
+
emulates the same behavior in pure Python (cosine over stored vectors) so the
|
|
5
|
+
vector store can be fully unit-tested offline — same pattern that keeps the tests
|
|
6
|
+
deterministic and CI-friendly without provisioning a database.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import json
|
|
12
|
+
import math
|
|
13
|
+
from dataclasses import dataclass
|
|
14
|
+
from typing import List, Optional, Protocol, Tuple
|
|
15
|
+
|
|
16
|
+
from . import _sql
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@dataclass
|
|
20
|
+
class Row:
|
|
21
|
+
id: str
|
|
22
|
+
content: str
|
|
23
|
+
metadata: dict
|
|
24
|
+
distance: float
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
class Store(Protocol):
|
|
28
|
+
def ensure_table(self, dim: int) -> None: ...
|
|
29
|
+
def upsert(self, rows: List[Tuple[str, str, dict, List[float]]]) -> None: ...
|
|
30
|
+
def search(self, embedding: List[float], k: int, metric: str) -> List[Row]: ...
|
|
31
|
+
def get(self, ids: List[str]) -> List[Row]: ...
|
|
32
|
+
def delete(self, ids: List[str]) -> None: ...
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class MySQLStore:
|
|
36
|
+
"""Real backend. Requires: pip install 'langchain-shannonbase[mysql]'."""
|
|
37
|
+
|
|
38
|
+
def __init__(self, table: str, **connection_kwargs):
|
|
39
|
+
import mysql.connector # noqa: F401 (lazy; validates the extra is installed)
|
|
40
|
+
self.table = table
|
|
41
|
+
self._conn_kwargs = connection_kwargs
|
|
42
|
+
|
|
43
|
+
def _connect(self):
|
|
44
|
+
import mysql.connector
|
|
45
|
+
return mysql.connector.connect(**self._conn_kwargs)
|
|
46
|
+
|
|
47
|
+
def ensure_table(self, dim: int) -> None:
|
|
48
|
+
conn = self._connect()
|
|
49
|
+
try:
|
|
50
|
+
conn.cursor().execute(_sql.create_table_sql(self.table, dim))
|
|
51
|
+
conn.commit()
|
|
52
|
+
finally:
|
|
53
|
+
conn.close()
|
|
54
|
+
|
|
55
|
+
def upsert(self, rows):
|
|
56
|
+
conn = self._connect()
|
|
57
|
+
try:
|
|
58
|
+
cur = conn.cursor()
|
|
59
|
+
stmt = _sql.insert_sql(self.table)
|
|
60
|
+
cur.executemany(
|
|
61
|
+
stmt,
|
|
62
|
+
[(rid, content, json.dumps(meta), _sql.vector_literal(emb))
|
|
63
|
+
for rid, content, meta, emb in rows],
|
|
64
|
+
)
|
|
65
|
+
conn.commit()
|
|
66
|
+
finally:
|
|
67
|
+
conn.close()
|
|
68
|
+
|
|
69
|
+
def search(self, embedding, k, metric):
|
|
70
|
+
conn = self._connect()
|
|
71
|
+
try:
|
|
72
|
+
cur = conn.cursor()
|
|
73
|
+
cur.execute(_sql.search_sql(self.table, metric),
|
|
74
|
+
(_sql.vector_literal(embedding), k))
|
|
75
|
+
out = []
|
|
76
|
+
for rid, content, meta, dist in cur.fetchall():
|
|
77
|
+
md = meta if isinstance(meta, dict) else json.loads(meta or "{}")
|
|
78
|
+
out.append(Row(rid, content, md, float(dist)))
|
|
79
|
+
return out
|
|
80
|
+
finally:
|
|
81
|
+
conn.close()
|
|
82
|
+
|
|
83
|
+
def get(self, ids):
|
|
84
|
+
if not ids:
|
|
85
|
+
return []
|
|
86
|
+
conn = self._connect()
|
|
87
|
+
try:
|
|
88
|
+
cur = conn.cursor()
|
|
89
|
+
cur.execute(_sql.select_by_ids_sql(self.table, len(ids)), tuple(ids))
|
|
90
|
+
out = []
|
|
91
|
+
for rid, content, meta in cur.fetchall():
|
|
92
|
+
md = meta if isinstance(meta, dict) else json.loads(meta or "{}")
|
|
93
|
+
out.append(Row(rid, content, md, 0.0))
|
|
94
|
+
return out
|
|
95
|
+
finally:
|
|
96
|
+
conn.close()
|
|
97
|
+
|
|
98
|
+
def delete(self, ids):
|
|
99
|
+
if not ids:
|
|
100
|
+
return
|
|
101
|
+
conn = self._connect()
|
|
102
|
+
try:
|
|
103
|
+
cur = conn.cursor()
|
|
104
|
+
cur.execute(_sql.delete_sql(self.table, len(ids)), tuple(ids))
|
|
105
|
+
conn.commit()
|
|
106
|
+
finally:
|
|
107
|
+
conn.close()
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
class InMemoryStore:
|
|
111
|
+
"""Deterministic offline backend that mirrors MySQLStore's cosine behavior."""
|
|
112
|
+
|
|
113
|
+
def __init__(self):
|
|
114
|
+
self._rows: dict[str, Tuple[str, dict, List[float]]] = {}
|
|
115
|
+
|
|
116
|
+
def ensure_table(self, dim: int) -> None:
|
|
117
|
+
pass
|
|
118
|
+
|
|
119
|
+
def upsert(self, rows):
|
|
120
|
+
for rid, content, meta, emb in rows:
|
|
121
|
+
self._rows[rid] = (content, dict(meta or {}), list(emb))
|
|
122
|
+
|
|
123
|
+
def search(self, embedding, k, metric):
|
|
124
|
+
scored = []
|
|
125
|
+
for rid, (content, meta, emb) in self._rows.items():
|
|
126
|
+
scored.append(Row(rid, content, meta, _cosine_distance(embedding, emb)))
|
|
127
|
+
scored.sort(key=lambda r: r.distance)
|
|
128
|
+
return scored[:k]
|
|
129
|
+
|
|
130
|
+
def get(self, ids):
|
|
131
|
+
out = []
|
|
132
|
+
for i in ids:
|
|
133
|
+
row = self._rows.get(i)
|
|
134
|
+
if row is not None:
|
|
135
|
+
content, meta, _ = row
|
|
136
|
+
out.append(Row(i, content, dict(meta), 0.0))
|
|
137
|
+
return out
|
|
138
|
+
|
|
139
|
+
def delete(self, ids):
|
|
140
|
+
for i in ids:
|
|
141
|
+
self._rows.pop(i, None)
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
def _cosine_distance(a: List[float], b: List[float]) -> float:
|
|
145
|
+
dot = sum(x * y for x, y in zip(a, b))
|
|
146
|
+
na = math.sqrt(sum(x * x for x in a))
|
|
147
|
+
nb = math.sqrt(sum(y * y for y in b))
|
|
148
|
+
if na == 0 or nb == 0:
|
|
149
|
+
return 1.0
|
|
150
|
+
return 1.0 - dot / (na * nb)
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
"""ShannonBaseVectorStore — a LangChain VectorStore backed by MySQL 9's native
|
|
2
|
+
VECTOR type. Works with ShannonBase, self-hosted MySQL 9, and MySQL HeatWave,
|
|
3
|
+
which all share the same VECTOR / STRING_TO_VECTOR / DISTANCE surface.
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
import uuid
|
|
9
|
+
from typing import Any, Iterable, List, Optional, Tuple
|
|
10
|
+
|
|
11
|
+
from langchain_core.documents import Document
|
|
12
|
+
from langchain_core.embeddings import Embeddings
|
|
13
|
+
from langchain_core.vectorstores import VectorStore
|
|
14
|
+
|
|
15
|
+
from ._store import InMemoryStore, MySQLStore, Store
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class ShannonBaseVectorStore(VectorStore):
|
|
19
|
+
"""Store and query embeddings in a MySQL-9-compatible database.
|
|
20
|
+
|
|
21
|
+
Typical use points it at a real database:
|
|
22
|
+
|
|
23
|
+
from langchain_shannonbase import ShannonBaseVectorStore
|
|
24
|
+
store = ShannonBaseVectorStore(
|
|
25
|
+
embedding=my_embeddings,
|
|
26
|
+
table="documents",
|
|
27
|
+
host="127.0.0.1", user="root", password="...", database="rag",
|
|
28
|
+
)
|
|
29
|
+
store.add_texts(["hello world"], metadatas=[{"src": "demo"}])
|
|
30
|
+
docs = store.similarity_search("greeting", k=3)
|
|
31
|
+
|
|
32
|
+
Pass `store=InMemoryStore()` instead of connection kwargs for offline tests.
|
|
33
|
+
"""
|
|
34
|
+
|
|
35
|
+
def __init__(
|
|
36
|
+
self,
|
|
37
|
+
embedding: Embeddings,
|
|
38
|
+
table: str = "langchain_vectors",
|
|
39
|
+
metric: str = "cosine",
|
|
40
|
+
store: Optional[Store] = None,
|
|
41
|
+
**connection_kwargs: Any,
|
|
42
|
+
):
|
|
43
|
+
self._embedding = embedding
|
|
44
|
+
self.table = table
|
|
45
|
+
self.metric = metric
|
|
46
|
+
self._store: Store = store if store is not None else MySQLStore(table, **connection_kwargs)
|
|
47
|
+
self._dim: Optional[int] = None
|
|
48
|
+
|
|
49
|
+
@property
|
|
50
|
+
def embeddings(self) -> Embeddings:
|
|
51
|
+
return self._embedding
|
|
52
|
+
|
|
53
|
+
def _ensure_dim(self, dim: int) -> None:
|
|
54
|
+
if self._dim is None:
|
|
55
|
+
self._dim = dim
|
|
56
|
+
self._store.ensure_table(dim)
|
|
57
|
+
|
|
58
|
+
def add_texts(
|
|
59
|
+
self,
|
|
60
|
+
texts: Iterable[str],
|
|
61
|
+
metadatas: Optional[List[dict]] = None,
|
|
62
|
+
ids: Optional[List[str]] = None,
|
|
63
|
+
**kwargs: Any,
|
|
64
|
+
) -> List[str]:
|
|
65
|
+
texts = list(texts)
|
|
66
|
+
if not texts:
|
|
67
|
+
return []
|
|
68
|
+
vectors = self._embedding.embed_documents(texts)
|
|
69
|
+
self._ensure_dim(len(vectors[0]))
|
|
70
|
+
metadatas = metadatas or [{} for _ in texts]
|
|
71
|
+
ids = ids or [str(uuid.uuid4()) for _ in texts]
|
|
72
|
+
rows: List[Tuple[str, str, dict, List[float]]] = [
|
|
73
|
+
(ids[i], texts[i], metadatas[i], vectors[i]) for i in range(len(texts))
|
|
74
|
+
]
|
|
75
|
+
self._store.upsert(rows)
|
|
76
|
+
return ids
|
|
77
|
+
|
|
78
|
+
def similarity_search(self, query: str, k: int = 4, **kwargs: Any) -> List[Document]:
|
|
79
|
+
return [doc for doc, _ in self.similarity_search_with_score(query, k, **kwargs)]
|
|
80
|
+
|
|
81
|
+
def similarity_search_with_score(
|
|
82
|
+
self, query: str, k: int = 4, **kwargs: Any
|
|
83
|
+
) -> List[Tuple[Document, float]]:
|
|
84
|
+
vector = self._embedding.embed_query(query)
|
|
85
|
+
return self.similarity_search_by_vector_with_score(vector, k, **kwargs)
|
|
86
|
+
|
|
87
|
+
def similarity_search_by_vector(
|
|
88
|
+
self, embedding: List[float], k: int = 4, **kwargs: Any
|
|
89
|
+
) -> List[Document]:
|
|
90
|
+
return [doc for doc, _ in self.similarity_search_by_vector_with_score(embedding, k)]
|
|
91
|
+
|
|
92
|
+
def similarity_search_by_vector_with_score(
|
|
93
|
+
self, embedding: List[float], k: int = 4, **kwargs: Any
|
|
94
|
+
) -> List[Tuple[Document, float]]:
|
|
95
|
+
rows = self._store.search(embedding, k, self.metric)
|
|
96
|
+
# DISTANCE is smaller-is-closer; expose it as a score (1 - distance).
|
|
97
|
+
return [
|
|
98
|
+
(Document(id=r.id, page_content=r.content, metadata=dict(r.metadata)),
|
|
99
|
+
1.0 - r.distance)
|
|
100
|
+
for r in rows
|
|
101
|
+
]
|
|
102
|
+
|
|
103
|
+
def get_by_ids(self, ids: Iterable[str]) -> List[Document]:
|
|
104
|
+
ids = list(ids)
|
|
105
|
+
found = {r.id: r for r in self._store.get(ids)}
|
|
106
|
+
return [
|
|
107
|
+
Document(id=i, page_content=found[i].content, metadata=dict(found[i].metadata))
|
|
108
|
+
for i in ids
|
|
109
|
+
if i in found
|
|
110
|
+
]
|
|
111
|
+
|
|
112
|
+
def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:
|
|
113
|
+
if not ids:
|
|
114
|
+
return False
|
|
115
|
+
self._store.delete(ids)
|
|
116
|
+
return True
|
|
117
|
+
|
|
118
|
+
@classmethod
|
|
119
|
+
def from_texts(
|
|
120
|
+
cls,
|
|
121
|
+
texts: List[str],
|
|
122
|
+
embedding: Embeddings,
|
|
123
|
+
metadatas: Optional[List[dict]] = None,
|
|
124
|
+
table: str = "langchain_vectors",
|
|
125
|
+
metric: str = "cosine",
|
|
126
|
+
store: Optional[Store] = None,
|
|
127
|
+
ids: Optional[List[str]] = None,
|
|
128
|
+
**connection_kwargs: Any,
|
|
129
|
+
) -> "ShannonBaseVectorStore":
|
|
130
|
+
vs = cls(embedding=embedding, table=table, metric=metric, store=store, **connection_kwargs)
|
|
131
|
+
vs.add_texts(texts, metadatas=metadatas, ids=ids)
|
|
132
|
+
return vs
|
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
"""Live integration test against a real MySQL 9 / ShannonBase / HeatWave instance.
|
|
2
|
+
|
|
3
|
+
Skipped unless SHANNONBASE_TEST_DSN-style env vars are set, so CI stays green
|
|
4
|
+
without provisioning a database. To run locally against a ShannonBase or MySQL 9:
|
|
5
|
+
|
|
6
|
+
export SB_HOST=127.0.0.1 SB_PORT=3306 SB_USER=root SB_PASSWORD=... SB_DATABASE=test
|
|
7
|
+
pytest tests/test_integration.py -v
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
import os
|
|
11
|
+
|
|
12
|
+
import pytest
|
|
13
|
+
|
|
14
|
+
pytestmark = pytest.mark.skipif(
|
|
15
|
+
not os.environ.get("SB_HOST"),
|
|
16
|
+
reason="set SB_HOST/SB_USER/SB_PASSWORD/SB_DATABASE to run live DB tests",
|
|
17
|
+
)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def _conn_kwargs():
|
|
21
|
+
return dict(
|
|
22
|
+
host=os.environ["SB_HOST"],
|
|
23
|
+
port=int(os.environ.get("SB_PORT", 3306)),
|
|
24
|
+
user=os.environ["SB_USER"],
|
|
25
|
+
password=os.environ.get("SB_PASSWORD", ""),
|
|
26
|
+
database=os.environ["SB_DATABASE"],
|
|
27
|
+
)
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def test_roundtrip_against_real_db():
|
|
31
|
+
import hashlib
|
|
32
|
+
import re
|
|
33
|
+
|
|
34
|
+
from langchain_core.embeddings import Embeddings
|
|
35
|
+
|
|
36
|
+
from langchain_shannonbase import ShannonBaseVectorStore
|
|
37
|
+
|
|
38
|
+
tok = re.compile(r"[a-z0-9]+")
|
|
39
|
+
|
|
40
|
+
class E(Embeddings):
|
|
41
|
+
def _e(self, t):
|
|
42
|
+
v = [0.0] * 64
|
|
43
|
+
for w in tok.findall(t.lower()):
|
|
44
|
+
v[int(hashlib.md5(w.encode()).hexdigest(), 16) % 64] += 1.0
|
|
45
|
+
return v
|
|
46
|
+
|
|
47
|
+
def embed_documents(self, x):
|
|
48
|
+
return [self._e(t) for t in x]
|
|
49
|
+
|
|
50
|
+
def embed_query(self, t):
|
|
51
|
+
return self._e(t)
|
|
52
|
+
|
|
53
|
+
vs = ShannonBaseVectorStore(embedding=E(), table="lc_sb_itest", **_conn_kwargs())
|
|
54
|
+
vs.add_texts(["reset my password", "cancel my subscription"], ids=["a", "b"])
|
|
55
|
+
docs = vs.similarity_search("how do I reset my password", k=1)
|
|
56
|
+
assert docs[0].page_content == "reset my password"
|
|
57
|
+
vs.delete(["a", "b"])
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
import json
|
|
2
|
+
|
|
3
|
+
import pytest
|
|
4
|
+
|
|
5
|
+
from langchain_shannonbase import _sql
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def test_vector_literal_is_json_array():
|
|
9
|
+
assert json.loads(_sql.vector_literal([1, 2.5, 3])) == [1.0, 2.5, 3.0]
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def test_create_table_has_vector_dim():
|
|
13
|
+
sql = _sql.create_table_sql("docs", 384)
|
|
14
|
+
assert "VECTOR(384)" in sql
|
|
15
|
+
assert "`docs`" in sql
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def test_insert_wraps_embedding_in_string_to_vector():
|
|
19
|
+
assert "STRING_TO_VECTOR(%s)" in _sql.insert_sql("docs")
|
|
20
|
+
assert "ON DUPLICATE KEY UPDATE" in _sql.insert_sql("docs")
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def test_search_uses_distance_and_metric():
|
|
24
|
+
sql = _sql.search_sql("docs", "cosine")
|
|
25
|
+
assert "DISTANCE(embedding, STRING_TO_VECTOR(%s), 'COSINE')" in sql
|
|
26
|
+
assert "ORDER BY dist ASC LIMIT %s" in sql
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def test_search_rejects_unknown_metric():
|
|
30
|
+
with pytest.raises(ValueError):
|
|
31
|
+
_sql.search_sql("docs", "nope")
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def test_delete_has_one_placeholder_per_id():
|
|
35
|
+
assert _sql.delete_sql("docs", 3).count("%s") == 3
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def test_distance_to_score():
|
|
39
|
+
assert _sql.distance_to_score(0.0) == 1.0
|
|
40
|
+
assert _sql.distance_to_score(0.25) == 0.75
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
"""LangChain's standard VectorStore integration suite, run offline.
|
|
2
|
+
|
|
3
|
+
The suite exercises add/get_by_ids/delete/search against a real store. We point it
|
|
4
|
+
at InMemoryStore so it runs in CI without provisioning a database — same behavior
|
|
5
|
+
as the MySQL backend, just in-process.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import pytest
|
|
9
|
+
from langchain_core.vectorstores import VectorStore
|
|
10
|
+
from langchain_tests.integration_tests import VectorStoreIntegrationTests
|
|
11
|
+
|
|
12
|
+
from langchain_shannonbase import InMemoryStore, ShannonBaseVectorStore
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class TestShannonBaseVectorStore(VectorStoreIntegrationTests):
|
|
16
|
+
@pytest.fixture()
|
|
17
|
+
def vectorstore(self) -> VectorStore:
|
|
18
|
+
return ShannonBaseVectorStore(embedding=self.get_embeddings(), store=InMemoryStore())
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
"""End-to-end tests of ShannonBaseVectorStore against the offline InMemoryStore.
|
|
2
|
+
|
|
3
|
+
A deterministic bag-of-words embedder stands in for a real embedding model so the
|
|
4
|
+
semantics (paraphrases rank closer than unrelated text) are reproducible offline.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import hashlib
|
|
8
|
+
import re
|
|
9
|
+
|
|
10
|
+
from langchain_core.embeddings import Embeddings
|
|
11
|
+
|
|
12
|
+
from langchain_shannonbase import InMemoryStore, ShannonBaseVectorStore
|
|
13
|
+
|
|
14
|
+
_TOKEN = re.compile(r"[a-z0-9]+")
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class HashEmbeddings(Embeddings):
|
|
18
|
+
dim = 64
|
|
19
|
+
|
|
20
|
+
def _embed(self, text: str):
|
|
21
|
+
vec = [0.0] * self.dim
|
|
22
|
+
for tok in _TOKEN.findall(text.lower()):
|
|
23
|
+
vec[int(hashlib.md5(tok.encode()).hexdigest(), 16) % self.dim] += 1.0
|
|
24
|
+
return vec
|
|
25
|
+
|
|
26
|
+
def embed_documents(self, texts):
|
|
27
|
+
return [self._embed(t) for t in texts]
|
|
28
|
+
|
|
29
|
+
def embed_query(self, text):
|
|
30
|
+
return self._embed(text)
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def _store():
|
|
34
|
+
return ShannonBaseVectorStore(embedding=HashEmbeddings(), store=InMemoryStore())
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_add_and_search_returns_relevant_doc():
|
|
38
|
+
vs = _store()
|
|
39
|
+
vs.add_texts(
|
|
40
|
+
["the cat sat on the mat", "quarterly financial report", "a dog in the park"],
|
|
41
|
+
metadatas=[{"src": "a"}, {"src": "b"}, {"src": "c"}],
|
|
42
|
+
)
|
|
43
|
+
docs = vs.similarity_search("cat on a mat", k=1)
|
|
44
|
+
assert len(docs) == 1
|
|
45
|
+
assert "cat" in docs[0].page_content
|
|
46
|
+
assert docs[0].metadata["src"] == "a"
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def test_add_texts_returns_ids_and_respects_custom_ids():
|
|
50
|
+
vs = _store()
|
|
51
|
+
ids = vs.add_texts(["hello"], ids=["fixed-id"])
|
|
52
|
+
assert ids == ["fixed-id"]
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def test_similarity_search_with_score_orders_by_similarity():
|
|
56
|
+
vs = _store()
|
|
57
|
+
vs.add_texts(["reset my password", "cancel my subscription"])
|
|
58
|
+
results = vs.similarity_search_with_score("how do I reset my password", k=2)
|
|
59
|
+
assert results[0][0].page_content == "reset my password"
|
|
60
|
+
assert results[0][1] >= results[1][1] # higher score first
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def test_k_limits_results():
|
|
64
|
+
vs = _store()
|
|
65
|
+
vs.add_texts([f"doc number {i}" for i in range(10)])
|
|
66
|
+
assert len(vs.similarity_search("doc", k=3)) == 3
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def test_delete_removes_document():
|
|
70
|
+
vs = _store()
|
|
71
|
+
vs.add_texts(["keep me", "delete me"], ids=["keep", "del"])
|
|
72
|
+
assert vs.delete(["del"]) is True
|
|
73
|
+
remaining = vs.similarity_search("me", k=10)
|
|
74
|
+
assert all(d.id != "del" for d in remaining)
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def test_from_texts_classmethod():
|
|
78
|
+
vs = ShannonBaseVectorStore.from_texts(
|
|
79
|
+
["alpha", "beta"], embedding=HashEmbeddings(), store=InMemoryStore()
|
|
80
|
+
)
|
|
81
|
+
assert len(vs.similarity_search("alpha", k=2)) == 2
|