chunklens 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.
- chunklens-0.1.0/.gitignore +41 -0
- chunklens-0.1.0/.python-version +1 -0
- chunklens-0.1.0/PKG-INFO +22 -0
- chunklens-0.1.0/README.md +5 -0
- chunklens-0.1.0/pyproject.toml +40 -0
- chunklens-0.1.0/src/chunklens/__init__.py +1 -0
- chunklens-0.1.0/src/chunklens/app.py +33 -0
- chunklens-0.1.0/src/chunklens/chroma_client.py +51 -0
- chunklens-0.1.0/src/chunklens/chroma_service.py +406 -0
- chunklens-0.1.0/src/chunklens/config.py +36 -0
- chunklens-0.1.0/src/chunklens/connection.py +38 -0
- chunklens-0.1.0/src/chunklens/deps.py +19 -0
- chunklens-0.1.0/src/chunklens/embedder_hints.py +66 -0
- chunklens-0.1.0/src/chunklens/embedders.py +151 -0
- chunklens-0.1.0/src/chunklens/routers/__init__.py +0 -0
- chunklens-0.1.0/src/chunklens/routers/collections.py +159 -0
- chunklens-0.1.0/src/chunklens/routers/connection.py +56 -0
- chunklens-0.1.0/src/chunklens/routers/embedders.py +26 -0
- chunklens-0.1.0/src/chunklens/routers/query.py +40 -0
- chunklens-0.1.0/src/chunklens/schemas.py +175 -0
- chunklens-0.1.0/src/chunklens/server.py +20 -0
- chunklens-0.1.0/src/chunklens/web/assets/Geist-Variable-Bj2R_7yk.woff2 +0 -0
- chunklens-0.1.0/src/chunklens/web/assets/GeistMono-Variable-Dispecij.woff2 +0 -0
- chunklens-0.1.0/src/chunklens/web/assets/index-BfnLujAD.js +40 -0
- chunklens-0.1.0/src/chunklens/web/assets/index-FTYKRLZj.css +1 -0
- chunklens-0.1.0/src/chunklens/web/index.html +13 -0
- chunklens-0.1.0/tests/conftest.py +50 -0
- chunklens-0.1.0/tests/test_api.py +268 -0
- chunklens-0.1.0/tests/test_chroma_service.py +225 -0
- chunklens-0.1.0/tests/test_config.py +36 -0
- chunklens-0.1.0/tests/test_connection.py +34 -0
- chunklens-0.1.0/tests/test_embedder_hints.py +44 -0
- chunklens-0.1.0/tests/test_embedder_hints_api.py +38 -0
- chunklens-0.1.0/tests/test_embedders.py +53 -0
- chunklens-0.1.0/tests/test_embedders_api.py +26 -0
- chunklens-0.1.0/tests/test_import_export.py +190 -0
- chunklens-0.1.0/tests/test_query_embedder.py +52 -0
- chunklens-0.1.0/tests/test_schemas.py +29 -0
- chunklens-0.1.0/tests/test_server.py +13 -0
- chunklens-0.1.0/tests/test_sources.py +63 -0
- chunklens-0.1.0/uv.lock +8243 -0
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
# ---- Python (backend) ----
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*.egg-info/
|
|
5
|
+
.eggs/
|
|
6
|
+
build/
|
|
7
|
+
dist/
|
|
8
|
+
.venv/
|
|
9
|
+
venv/
|
|
10
|
+
.env
|
|
11
|
+
.env.*
|
|
12
|
+
.pytest_cache/
|
|
13
|
+
.mypy_cache/
|
|
14
|
+
.ruff_cache/
|
|
15
|
+
|
|
16
|
+
# ---- Node (frontend) ----
|
|
17
|
+
node_modules/
|
|
18
|
+
frontend/dist/
|
|
19
|
+
*.tsbuildinfo
|
|
20
|
+
*.log
|
|
21
|
+
npm-debug.log*
|
|
22
|
+
pnpm-debug.log*
|
|
23
|
+
yarn-error.log*
|
|
24
|
+
|
|
25
|
+
# SPA build output served by the backend
|
|
26
|
+
backend/src/chunklens/web/
|
|
27
|
+
|
|
28
|
+
# ---- Chroma local data ----
|
|
29
|
+
.chroma/
|
|
30
|
+
.chroma-e2e/
|
|
31
|
+
chroma-data/
|
|
32
|
+
|
|
33
|
+
# ---- Test artifacts ----
|
|
34
|
+
test-results/
|
|
35
|
+
playwright-report/
|
|
36
|
+
frontend/playwright/.cache/
|
|
37
|
+
|
|
38
|
+
# ---- OS junk ----
|
|
39
|
+
.DS_Store
|
|
40
|
+
Thumbs.db
|
|
41
|
+
*.swp
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
3.12
|
chunklens-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: chunklens
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: The local-first inspector & retrieval debugger for ChromaDB.
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
Requires-Python: >=3.9
|
|
7
|
+
Requires-Dist: chromadb==1.5.9
|
|
8
|
+
Requires-Dist: cohere<6,>=5
|
|
9
|
+
Requires-Dist: fastapi<1,>=0.110
|
|
10
|
+
Requires-Dist: openai<3,>=2.44.0
|
|
11
|
+
Requires-Dist: pydantic<3,>=2.6
|
|
12
|
+
Requires-Dist: uvicorn[standard]<1,>=0.29
|
|
13
|
+
Requires-Dist: voyageai<1,>=0.2
|
|
14
|
+
Provides-Extra: local-embedders
|
|
15
|
+
Requires-Dist: sentence-transformers<6,>=2; extra == 'local-embedders'
|
|
16
|
+
Description-Content-Type: text/markdown
|
|
17
|
+
|
|
18
|
+
# chunklens (backend)
|
|
19
|
+
|
|
20
|
+
The FastAPI backend for ChunkLens - the local-first inspector & retrieval debugger for ChromaDB.
|
|
21
|
+
|
|
22
|
+
See the repository root `README.md` for the project quickstart.
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "chunklens"
|
|
7
|
+
dynamic = ["version"]
|
|
8
|
+
description = "The local-first inspector & retrieval debugger for ChromaDB."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = "MIT"
|
|
12
|
+
dependencies = [
|
|
13
|
+
"fastapi>=0.110,<1",
|
|
14
|
+
"uvicorn[standard]>=0.29,<1",
|
|
15
|
+
"chromadb==1.5.9",
|
|
16
|
+
"pydantic>=2.6,<3",
|
|
17
|
+
"openai>=2.44.0,<3",
|
|
18
|
+
"cohere>=5,<6",
|
|
19
|
+
"voyageai>=0.2,<1",
|
|
20
|
+
]
|
|
21
|
+
|
|
22
|
+
[project.optional-dependencies]
|
|
23
|
+
local-embedders = ["sentence-transformers>=2,<6"]
|
|
24
|
+
|
|
25
|
+
[dependency-groups]
|
|
26
|
+
dev = ["pytest>=8", "httpx>=0.27"]
|
|
27
|
+
|
|
28
|
+
[project.scripts]
|
|
29
|
+
chunklens = "chunklens.server:main"
|
|
30
|
+
|
|
31
|
+
[tool.hatch.version]
|
|
32
|
+
path = "src/chunklens/__init__.py"
|
|
33
|
+
|
|
34
|
+
[tool.hatch.build.targets.wheel]
|
|
35
|
+
# src/chunklens/web (the built SPA) is gitignored, but an explicit `packages` entry overrides
|
|
36
|
+
# hatchling's VCS exclusion, so the wheel AND sdist already ship web/ (verified). Do NOT add a
|
|
37
|
+
# force-include for web/: it collides with packages ("a second file is being added at the same
|
|
38
|
+
# path: chunklens/web/index.html"). scripts/build_release.py asserts the SPA is inside the wheel
|
|
39
|
+
# as the release-time guard, and requires web/ to be built first.
|
|
40
|
+
packages = ["src/chunklens"]
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.1.0"
|
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from pathlib import Path
|
|
4
|
+
|
|
5
|
+
from fastapi import FastAPI
|
|
6
|
+
from fastapi.staticfiles import StaticFiles
|
|
7
|
+
|
|
8
|
+
from . import __version__
|
|
9
|
+
from .routers import collections, connection, embedders, query
|
|
10
|
+
from .schemas import HealthResponse
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def create_app() -> FastAPI:
|
|
14
|
+
app = FastAPI(title="ChunkLens", version=__version__)
|
|
15
|
+
|
|
16
|
+
@app.get("/api/health", response_model=HealthResponse)
|
|
17
|
+
def health() -> HealthResponse:
|
|
18
|
+
return HealthResponse(status="ok", version=__version__)
|
|
19
|
+
|
|
20
|
+
app.include_router(collections.router)
|
|
21
|
+
app.include_router(query.router)
|
|
22
|
+
app.include_router(connection.router)
|
|
23
|
+
app.include_router(embedders.router)
|
|
24
|
+
|
|
25
|
+
# SPA build (present only after `npm run build`); mounted last so /api wins.
|
|
26
|
+
web_dir = Path(__file__).parent / "web"
|
|
27
|
+
if web_dir.is_dir():
|
|
28
|
+
app.mount("/", StaticFiles(directory=str(web_dir), html=True), name="web")
|
|
29
|
+
|
|
30
|
+
return app
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
app = create_app()
|
|
@@ -0,0 +1,51 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import chromadb
|
|
6
|
+
|
|
7
|
+
from .schemas import ConnectionConfig
|
|
8
|
+
|
|
9
|
+
DEFAULT_HOST = "localhost"
|
|
10
|
+
DEFAULT_PORT = 8000
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def make_client(
|
|
14
|
+
host: str = DEFAULT_HOST,
|
|
15
|
+
port: int = DEFAULT_PORT,
|
|
16
|
+
ssl: bool = False,
|
|
17
|
+
tenant: str = "default_tenant",
|
|
18
|
+
database: str = "default_database",
|
|
19
|
+
headers: Optional[dict] = None,
|
|
20
|
+
):
|
|
21
|
+
"""Build a real HTTP client against a running Chroma server."""
|
|
22
|
+
return chromadb.HttpClient(
|
|
23
|
+
host=host,
|
|
24
|
+
port=port,
|
|
25
|
+
ssl=ssl,
|
|
26
|
+
tenant=tenant,
|
|
27
|
+
database=database,
|
|
28
|
+
headers=headers or {},
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def heartbeat(client) -> int:
|
|
33
|
+
"""Return the server heartbeat (nanoseconds). Raises if unreachable."""
|
|
34
|
+
return client.heartbeat()
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def headers_for(cfg: ConnectionConfig) -> dict:
|
|
38
|
+
if cfg.auth_mode == "token" and cfg.token:
|
|
39
|
+
return {"Authorization": f"Bearer {cfg.token}"}
|
|
40
|
+
return {}
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def client_for_config(cfg: ConnectionConfig):
|
|
44
|
+
return make_client(
|
|
45
|
+
host=cfg.host,
|
|
46
|
+
port=cfg.port,
|
|
47
|
+
ssl=cfg.ssl,
|
|
48
|
+
tenant=cfg.tenant,
|
|
49
|
+
database=cfg.database,
|
|
50
|
+
headers=headers_for(cfg),
|
|
51
|
+
)
|
|
@@ -0,0 +1,406 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Any, Optional
|
|
4
|
+
|
|
5
|
+
from chromadb.errors import InvalidArgumentError, NotFoundError
|
|
6
|
+
|
|
7
|
+
from .schemas import (
|
|
8
|
+
CollectionDetails,
|
|
9
|
+
CollectionSummary,
|
|
10
|
+
ExportCollection,
|
|
11
|
+
ExportFile,
|
|
12
|
+
ExportRecord,
|
|
13
|
+
MetadataKeyInfo,
|
|
14
|
+
MetadataKeysResponse,
|
|
15
|
+
QueryHit,
|
|
16
|
+
QueryResult,
|
|
17
|
+
Record,
|
|
18
|
+
RecordsPage,
|
|
19
|
+
SourceInfo,
|
|
20
|
+
SourceList,
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def _name_of(c) -> str:
|
|
25
|
+
# chromadb has returned either Collection objects or bare names across versions.
|
|
26
|
+
return c if isinstance(c, str) else c.name
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def list_collections(client) -> list[CollectionSummary]:
|
|
30
|
+
summaries: list[CollectionSummary] = []
|
|
31
|
+
for c in client.list_collections():
|
|
32
|
+
name = _name_of(c)
|
|
33
|
+
col = client.get_collection(name)
|
|
34
|
+
summaries.append(CollectionSummary(name=name, count=col.count()))
|
|
35
|
+
return summaries
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def get_records(client, name: str, limit: int = 50, offset: int = 0, where: Optional[dict] = None) -> RecordsPage:
|
|
39
|
+
col = client.get_collection(name)
|
|
40
|
+
if where is None:
|
|
41
|
+
total = col.count()
|
|
42
|
+
else:
|
|
43
|
+
# ids-only fetch for the filtered count. (If chromadb rejects include=[],
|
|
44
|
+
# use include=["metadatas"] - ids are returned regardless of include.)
|
|
45
|
+
total = len((col.get(where=where, include=[]) or {}).get("ids") or [])
|
|
46
|
+
res = col.get(include=["documents", "metadatas"], limit=limit, offset=offset, where=where)
|
|
47
|
+
ids = res.get("ids") or []
|
|
48
|
+
docs = res.get("documents") or [None] * len(ids)
|
|
49
|
+
metas = res.get("metadatas") or [None] * len(ids)
|
|
50
|
+
items = [
|
|
51
|
+
Record(id=i, document=d, metadata=m)
|
|
52
|
+
for i, d, m in zip(ids, docs, metas)
|
|
53
|
+
]
|
|
54
|
+
return RecordsPage(items=items, limit=limit, offset=offset, total=total)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def query(
|
|
58
|
+
client,
|
|
59
|
+
name: str,
|
|
60
|
+
*,
|
|
61
|
+
query_text: Optional[str] = None,
|
|
62
|
+
query_embedding: Optional[list[float]] = None,
|
|
63
|
+
n_results: int = 10,
|
|
64
|
+
where: Optional[dict] = None,
|
|
65
|
+
where_document: Optional[dict] = None,
|
|
66
|
+
) -> QueryResult:
|
|
67
|
+
col = client.get_collection(name)
|
|
68
|
+
kwargs: dict[str, Any] = {
|
|
69
|
+
"n_results": n_results,
|
|
70
|
+
"include": ["documents", "metadatas", "distances"],
|
|
71
|
+
}
|
|
72
|
+
if query_embedding is not None:
|
|
73
|
+
kwargs["query_embeddings"] = [query_embedding]
|
|
74
|
+
elif query_text is not None:
|
|
75
|
+
kwargs["query_texts"] = [query_text]
|
|
76
|
+
else:
|
|
77
|
+
raise ValueError("Provide query_text or query_embedding")
|
|
78
|
+
if where:
|
|
79
|
+
kwargs["where"] = where
|
|
80
|
+
if where_document:
|
|
81
|
+
kwargs["where_document"] = where_document
|
|
82
|
+
|
|
83
|
+
try:
|
|
84
|
+
res = col.query(**kwargs)
|
|
85
|
+
except Exception as exc:
|
|
86
|
+
if "dimension" in str(exc).lower():
|
|
87
|
+
got = len(query_embedding) if query_embedding is not None else "?"
|
|
88
|
+
raise ValueError(
|
|
89
|
+
f"Query vector ({got}-dim) doesn't match this collection's embedding "
|
|
90
|
+
f"dimension; the chosen model likely differs from the one that created "
|
|
91
|
+
f"these vectors. ({exc})"
|
|
92
|
+
) from exc
|
|
93
|
+
raise
|
|
94
|
+
ids = (res.get("ids") or [[]])[0]
|
|
95
|
+
docs = (res.get("documents") or [[None] * len(ids)])[0]
|
|
96
|
+
metas = (res.get("metadatas") or [[None] * len(ids)])[0]
|
|
97
|
+
dists = (res.get("distances") or [[0.0] * len(ids)])[0]
|
|
98
|
+
hits = [
|
|
99
|
+
QueryHit(id=i, document=d, metadata=m, distance=float(dist))
|
|
100
|
+
for i, d, m, dist in zip(ids, docs, metas, dists)
|
|
101
|
+
]
|
|
102
|
+
return QueryResult(hits=hits)
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
class NotFound(Exception):
|
|
106
|
+
"""Collection or record does not exist."""
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
class Conflict(Exception):
|
|
110
|
+
"""A collection with the target name already exists."""
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
class InvalidName(Exception):
|
|
114
|
+
"""Collection name violates Chroma's naming rules."""
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
_RESERVED_PREFIXES = ("hnsw:", "chroma:")
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def _is_reserved(key: str) -> bool:
|
|
121
|
+
return any(key.startswith(p) for p in _RESERVED_PREFIXES)
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def _user_metadata(meta: Optional[dict]) -> dict:
|
|
125
|
+
return {k: v for k, v in (meta or {}).items() if not _is_reserved(k)}
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def collection_exists(client, name: str) -> bool:
|
|
129
|
+
return any(_name_of(c) == name for c in client.list_collections())
|
|
130
|
+
|
|
131
|
+
|
|
132
|
+
def _get(client, name: str):
|
|
133
|
+
try:
|
|
134
|
+
return client.get_collection(name)
|
|
135
|
+
except NotFoundError as exc:
|
|
136
|
+
raise NotFound(f"Collection {name!r} not found") from exc
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
def _peek_dim(col) -> Optional[int]:
|
|
140
|
+
got = col.get(limit=1, include=["embeddings"])
|
|
141
|
+
embs = got.get("embeddings")
|
|
142
|
+
if embs is not None and len(embs) > 0:
|
|
143
|
+
return len(embs[0])
|
|
144
|
+
return None
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def _ef_name(col, cfg: dict) -> str:
|
|
148
|
+
# Source of truth is the persisted config, not col._embedding_function:
|
|
149
|
+
# get_collection() re-attaches a DefaultEmbeddingFunction object even when a
|
|
150
|
+
# collection was created with embedding_function=None, so the in-memory attribute
|
|
151
|
+
# can't distinguish "none" after a re-fetch. configuration_json records it faithfully.
|
|
152
|
+
if "embedding_function" in cfg:
|
|
153
|
+
ef_cfg = cfg["embedding_function"]
|
|
154
|
+
if ef_cfg is None:
|
|
155
|
+
return "none"
|
|
156
|
+
if isinstance(ef_cfg, dict) and ef_cfg.get("name"):
|
|
157
|
+
return ef_cfg["name"]
|
|
158
|
+
if getattr(col, "_embedding_function", None) is None:
|
|
159
|
+
return "none"
|
|
160
|
+
return type(col._embedding_function).__name__
|
|
161
|
+
|
|
162
|
+
|
|
163
|
+
def get_collection_details(client, name: str) -> CollectionDetails:
|
|
164
|
+
col = _get(client, name)
|
|
165
|
+
cfg = col.configuration_json or {}
|
|
166
|
+
metric = (cfg.get("hnsw") or {}).get("space", "l2")
|
|
167
|
+
return CollectionDetails(
|
|
168
|
+
name=name,
|
|
169
|
+
count=col.count(),
|
|
170
|
+
dimensionality=_peek_dim(col),
|
|
171
|
+
distance_metric=metric,
|
|
172
|
+
embedding_function=_ef_name(col, cfg),
|
|
173
|
+
metadata=_user_metadata(col.metadata),
|
|
174
|
+
)
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
def create_collection(
|
|
178
|
+
client,
|
|
179
|
+
name: str,
|
|
180
|
+
*,
|
|
181
|
+
distance_metric: str = "l2",
|
|
182
|
+
embedding_function: str = "default",
|
|
183
|
+
metadata: Optional[dict] = None,
|
|
184
|
+
) -> CollectionDetails:
|
|
185
|
+
if collection_exists(client, name):
|
|
186
|
+
raise Conflict(f"Collection {name!r} already exists")
|
|
187
|
+
user_meta = _user_metadata(metadata)
|
|
188
|
+
kwargs: dict[str, Any] = {}
|
|
189
|
+
if embedding_function == "none":
|
|
190
|
+
# Passing configuration= alongside embedding_function=None makes Chroma persist
|
|
191
|
+
# a '{type: legacy}' EF marker (and re-attach a broken default EF on re-fetch).
|
|
192
|
+
# Setting the hnsw space via the legacy metadata form instead keeps
|
|
193
|
+
# embedding_function persisted cleanly as None, while configuration_json still
|
|
194
|
+
# reflects the chosen space - so get_collection_details reports "none".
|
|
195
|
+
kwargs["embedding_function"] = None
|
|
196
|
+
kwargs["metadata"] = {"hnsw:space": distance_metric, **user_meta}
|
|
197
|
+
else:
|
|
198
|
+
kwargs["configuration"] = {"hnsw": {"space": distance_metric}}
|
|
199
|
+
if user_meta:
|
|
200
|
+
kwargs["metadata"] = user_meta
|
|
201
|
+
try:
|
|
202
|
+
client.create_collection(name, **kwargs)
|
|
203
|
+
except InvalidArgumentError as exc:
|
|
204
|
+
raise InvalidName(str(exc)) from exc
|
|
205
|
+
return get_collection_details(client, name)
|
|
206
|
+
|
|
207
|
+
|
|
208
|
+
def _merge_collection_metadata(current: Optional[dict], new_user: Optional[dict]) -> dict:
|
|
209
|
+
"""Collection modify() replaces metadata, so keep reserved keys and swap user keys."""
|
|
210
|
+
reserved = {k: v for k, v in (current or {}).items() if _is_reserved(k)}
|
|
211
|
+
return {**reserved, **_user_metadata(new_user)}
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
def rename_collection(client, name: str, new_name: str) -> CollectionDetails:
|
|
215
|
+
col = _get(client, name)
|
|
216
|
+
if new_name != name and collection_exists(client, new_name):
|
|
217
|
+
raise Conflict(f"Collection {new_name!r} already exists")
|
|
218
|
+
try:
|
|
219
|
+
col.modify(name=new_name)
|
|
220
|
+
except InvalidArgumentError as exc:
|
|
221
|
+
raise InvalidName(str(exc)) from exc
|
|
222
|
+
return get_collection_details(client, new_name)
|
|
223
|
+
|
|
224
|
+
|
|
225
|
+
def update_collection_metadata(client, name: str, metadata: Optional[dict]) -> CollectionDetails:
|
|
226
|
+
col = _get(client, name)
|
|
227
|
+
merged = _merge_collection_metadata(col.metadata, metadata)
|
|
228
|
+
if merged: # Chroma rejects modify(metadata={}); empty result is a no-op
|
|
229
|
+
col.modify(metadata=merged)
|
|
230
|
+
return get_collection_details(client, name)
|
|
231
|
+
|
|
232
|
+
|
|
233
|
+
def delete_collection(client, name: str) -> None:
|
|
234
|
+
if not collection_exists(client, name):
|
|
235
|
+
raise NotFound(f"Collection {name!r} not found")
|
|
236
|
+
client.delete_collection(name)
|
|
237
|
+
|
|
238
|
+
|
|
239
|
+
def _replace_payload(old: dict, new: dict) -> dict:
|
|
240
|
+
"""Emulate replace on Chroma's merge-update: set new keys, null removed keys."""
|
|
241
|
+
payload = dict(new)
|
|
242
|
+
for key in old:
|
|
243
|
+
if key not in new:
|
|
244
|
+
payload[key] = None
|
|
245
|
+
return payload
|
|
246
|
+
|
|
247
|
+
|
|
248
|
+
def update_record_metadata(client, name: str, record_id: str, metadata: dict) -> Record:
|
|
249
|
+
col = _get(client, name)
|
|
250
|
+
existing = col.get(ids=[record_id], include=["documents", "metadatas"])
|
|
251
|
+
if not existing["ids"]:
|
|
252
|
+
raise NotFound(f"Record {record_id!r} not found in {name!r}")
|
|
253
|
+
old = (existing["metadatas"] or [None])[0] or {}
|
|
254
|
+
payload = _replace_payload(old, dict(metadata))
|
|
255
|
+
if payload: # Chroma rejects update(metadatas=[{}])
|
|
256
|
+
col.update(ids=[record_id], metadatas=[payload])
|
|
257
|
+
document = (existing["documents"] or [None])[0]
|
|
258
|
+
return Record(id=record_id, document=document, metadata=(dict(metadata) or None))
|
|
259
|
+
|
|
260
|
+
|
|
261
|
+
def _scalar_type(v) -> str:
|
|
262
|
+
if isinstance(v, bool): # bool is a subclass of int - check first
|
|
263
|
+
return "bool"
|
|
264
|
+
if isinstance(v, int):
|
|
265
|
+
return "int"
|
|
266
|
+
if isinstance(v, float):
|
|
267
|
+
return "float"
|
|
268
|
+
return "string"
|
|
269
|
+
|
|
270
|
+
|
|
271
|
+
def sample_metadata_keys(client, name: str, sample: int = 200) -> MetadataKeysResponse:
|
|
272
|
+
col = _get(client, name)
|
|
273
|
+
total = col.count()
|
|
274
|
+
res = col.get(limit=sample, include=["metadatas"])
|
|
275
|
+
metas = res.get("metadatas") or []
|
|
276
|
+
agg: dict[str, set] = {}
|
|
277
|
+
for m in metas:
|
|
278
|
+
for k, v in (m or {}).items():
|
|
279
|
+
if _is_reserved(k):
|
|
280
|
+
continue
|
|
281
|
+
agg.setdefault(k, set()).add(_scalar_type(v))
|
|
282
|
+
keys = [MetadataKeyInfo(key=k, types=sorted(t)) for k, t in sorted(agg.items())]
|
|
283
|
+
return MetadataKeysResponse(keys=keys, sampled=len(metas), total=total)
|
|
284
|
+
|
|
285
|
+
|
|
286
|
+
DEFAULT_SOURCE_SCAN_CAP = 10000
|
|
287
|
+
_SOURCE_BATCH = 1000
|
|
288
|
+
|
|
289
|
+
|
|
290
|
+
def list_sources(client, name: str, key: str, cap: int = DEFAULT_SOURCE_SCAN_CAP) -> SourceList:
|
|
291
|
+
col = client.get_collection(name)
|
|
292
|
+
total = col.count()
|
|
293
|
+
limit_total = min(total, cap)
|
|
294
|
+
counts: dict[str, int] = {}
|
|
295
|
+
scanned = 0
|
|
296
|
+
while scanned < limit_total:
|
|
297
|
+
batch = col.get(
|
|
298
|
+
include=["metadatas"],
|
|
299
|
+
limit=min(_SOURCE_BATCH, limit_total - scanned),
|
|
300
|
+
offset=scanned,
|
|
301
|
+
)
|
|
302
|
+
metas = batch.get("metadatas") or []
|
|
303
|
+
if not metas:
|
|
304
|
+
break
|
|
305
|
+
for m in metas:
|
|
306
|
+
raw = (m or {}).get(key)
|
|
307
|
+
value = "(none)" if raw is None else str(raw)
|
|
308
|
+
counts[value] = counts.get(value, 0) + 1
|
|
309
|
+
scanned += len(metas)
|
|
310
|
+
sources = [SourceInfo(value=v, count=c) for v, c in counts.items()]
|
|
311
|
+
sources.sort(key=lambda s: (-s.count, s.value))
|
|
312
|
+
return SourceList(key=key, sources=sources, scanned=scanned, total=total)
|
|
313
|
+
|
|
314
|
+
|
|
315
|
+
_ADD_BATCH = 500
|
|
316
|
+
|
|
317
|
+
|
|
318
|
+
def add_records(client, name: str, records: list[ExportRecord]) -> int:
|
|
319
|
+
col = _get(client, name)
|
|
320
|
+
ef = _ef_name(col, col.configuration_json or {})
|
|
321
|
+
have = [r.embedding is not None for r in records]
|
|
322
|
+
if any(have) and not all(have):
|
|
323
|
+
raise ValueError("Either every record includes an embedding, or none do.")
|
|
324
|
+
use_embeddings = bool(records) and all(have)
|
|
325
|
+
if ef == "none" and records and not use_embeddings:
|
|
326
|
+
raise ValueError(
|
|
327
|
+
"Collection has no embedding function, so every record must include an embedding."
|
|
328
|
+
)
|
|
329
|
+
if use_embeddings:
|
|
330
|
+
dims = {len(r.embedding) for r in records}
|
|
331
|
+
if len(dims) > 1:
|
|
332
|
+
raise ValueError(f"Inconsistent embedding dimensions: {sorted(dims)}")
|
|
333
|
+
added = 0
|
|
334
|
+
for i in range(0, len(records), _ADD_BATCH):
|
|
335
|
+
batch = records[i : i + _ADD_BATCH]
|
|
336
|
+
metadatas = [r.metadata or None for r in batch]
|
|
337
|
+
kwargs: dict[str, Any] = {
|
|
338
|
+
"ids": [r.id for r in batch],
|
|
339
|
+
"documents": [r.document for r in batch],
|
|
340
|
+
}
|
|
341
|
+
if any(m is not None for m in metadatas):
|
|
342
|
+
kwargs["metadatas"] = metadatas
|
|
343
|
+
if use_embeddings:
|
|
344
|
+
kwargs["embeddings"] = [r.embedding for r in batch]
|
|
345
|
+
col.add(**kwargs)
|
|
346
|
+
added += len(batch)
|
|
347
|
+
return added
|
|
348
|
+
|
|
349
|
+
|
|
350
|
+
_EXPORT_BATCH = 500
|
|
351
|
+
|
|
352
|
+
|
|
353
|
+
def export_collection(client, name: str, *, include_embeddings: bool) -> ExportFile:
|
|
354
|
+
details = get_collection_details(client, name) # raises NotFound
|
|
355
|
+
col = client.get_collection(name)
|
|
356
|
+
ef = "none" if details.embedding_function == "none" else "default"
|
|
357
|
+
want_emb = include_embeddings or ef == "none"
|
|
358
|
+
include = ["documents", "metadatas"] + (["embeddings"] if want_emb else [])
|
|
359
|
+
total = col.count()
|
|
360
|
+
records: list[ExportRecord] = []
|
|
361
|
+
for offset in range(0, max(total, 0), _EXPORT_BATCH) or [0]:
|
|
362
|
+
res = col.get(include=include, limit=_EXPORT_BATCH, offset=offset)
|
|
363
|
+
ids = res.get("ids") or []
|
|
364
|
+
docs = res.get("documents") or [None] * len(ids)
|
|
365
|
+
metas = res.get("metadatas") or [None] * len(ids)
|
|
366
|
+
embs = res.get("embeddings") if want_emb else None
|
|
367
|
+
for j, rid in enumerate(ids):
|
|
368
|
+
emb = None
|
|
369
|
+
if want_emb and embs is not None and j < len(embs) and embs[j] is not None:
|
|
370
|
+
emb = [float(x) for x in embs[j]]
|
|
371
|
+
records.append(ExportRecord(id=rid, document=docs[j], metadata=metas[j] or None, embedding=emb))
|
|
372
|
+
return ExportFile(
|
|
373
|
+
chunklens_export=1,
|
|
374
|
+
collection=ExportCollection(
|
|
375
|
+
name=details.name,
|
|
376
|
+
distance_metric=details.distance_metric,
|
|
377
|
+
embedding_function=ef,
|
|
378
|
+
metadata=details.metadata,
|
|
379
|
+
),
|
|
380
|
+
records=records,
|
|
381
|
+
)
|
|
382
|
+
|
|
383
|
+
|
|
384
|
+
_SUPPORTED_EXPORT_VERSION = 1
|
|
385
|
+
|
|
386
|
+
|
|
387
|
+
def import_collection(client, data: ExportFile, *, name_override: Optional[str] = None) -> CollectionDetails:
|
|
388
|
+
if data.chunklens_export != _SUPPORTED_EXPORT_VERSION:
|
|
389
|
+
raise ValueError(f"Unsupported export version {data.chunklens_export}")
|
|
390
|
+
name = name_override or data.collection.name
|
|
391
|
+
create_collection( # raises Conflict / InvalidName BEFORE anything is added
|
|
392
|
+
client,
|
|
393
|
+
name,
|
|
394
|
+
distance_metric=data.collection.distance_metric,
|
|
395
|
+
embedding_function=data.collection.embedding_function,
|
|
396
|
+
metadata=data.collection.metadata or None,
|
|
397
|
+
)
|
|
398
|
+
try:
|
|
399
|
+
add_records(client, name, data.records)
|
|
400
|
+
except Exception:
|
|
401
|
+
try:
|
|
402
|
+
client.delete_collection(name) # roll back our just-created collection
|
|
403
|
+
except Exception:
|
|
404
|
+
pass
|
|
405
|
+
raise
|
|
406
|
+
return get_collection_details(client, name)
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import os
|
|
5
|
+
import stat
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
|
|
8
|
+
from .schemas import ConnectionConfig
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def config_dir() -> Path:
|
|
12
|
+
return Path(os.environ.get("CHUNKLENS_HOME", str(Path.home() / ".chunklens")))
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def config_path() -> Path:
|
|
16
|
+
return config_dir() / "config.json"
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def load_config() -> ConnectionConfig:
|
|
20
|
+
path = config_path()
|
|
21
|
+
if not path.is_file():
|
|
22
|
+
return ConnectionConfig()
|
|
23
|
+
try:
|
|
24
|
+
return ConnectionConfig(**json.loads(path.read_text(encoding="utf-8")))
|
|
25
|
+
except Exception:
|
|
26
|
+
return ConnectionConfig()
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def save_config(cfg: ConnectionConfig) -> None:
|
|
30
|
+
config_dir().mkdir(parents=True, exist_ok=True)
|
|
31
|
+
path = config_path()
|
|
32
|
+
path.write_text(json.dumps(cfg.model_dump(), indent=2), encoding="utf-8")
|
|
33
|
+
try:
|
|
34
|
+
os.chmod(path, stat.S_IRUSR | stat.S_IWUSR) # 0600 (effective on POSIX)
|
|
35
|
+
except OSError:
|
|
36
|
+
pass
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
from . import chroma_client
|
|
6
|
+
from . import config
|
|
7
|
+
from .schemas import ConnectionConfig
|
|
8
|
+
|
|
9
|
+
_active: Optional[ConnectionConfig] = None
|
|
10
|
+
_client = None
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def get_active() -> ConnectionConfig:
|
|
14
|
+
global _active
|
|
15
|
+
if _active is None:
|
|
16
|
+
_active = config.load_config()
|
|
17
|
+
return _active
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def set_active(cfg: ConnectionConfig) -> None:
|
|
21
|
+
global _active, _client
|
|
22
|
+
config.save_config(cfg)
|
|
23
|
+
_active = cfg
|
|
24
|
+
_client = None # invalidate cached client
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def get_active_client():
|
|
28
|
+
global _client
|
|
29
|
+
if _client is None:
|
|
30
|
+
_client = chroma_client.client_for_config(get_active())
|
|
31
|
+
return _client
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def reset() -> None:
|
|
35
|
+
"""Test helper: clear in-memory cache (next access reloads from disk)."""
|
|
36
|
+
global _active, _client
|
|
37
|
+
_active = None
|
|
38
|
+
_client = None
|