cocoindex-code 0.2.39__tar.gz → 0.2.41__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.
- cocoindex_code-0.2.39/README.md → cocoindex_code-0.2.41/PKG-INFO +63 -2
- cocoindex_code-0.2.39/PKG-INFO → cocoindex_code-0.2.41/README.md +19 -46
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/pyproject.toml +7 -7
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/_daemon_paths.py +36 -10
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/_version.py +2 -2
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/cli.py +18 -2
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/client.py +18 -8
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/daemon.py +11 -3
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/file_walk.py +36 -2
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/grep.py +3 -1
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/indexer.py +3 -1
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/project.py +39 -14
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/query.py +46 -13
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/server.py +18 -14
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/settings.py +102 -2
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/shared.py +48 -2
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/.gitignore +0 -0
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/LICENSE +0 -0
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/__init__.py +0 -0
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/__main__.py +0 -0
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/chunking.py +0 -0
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/embedder_defaults.py +0 -0
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/embedder_params.py +0 -0
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/litellm_embedder.py +0 -0
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/protocol.py +0 -0
- {cocoindex_code-0.2.39 → cocoindex_code-0.2.41}/src/cocoindex_code/schema.py +0 -0
|
@@ -1,3 +1,47 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: cocoindex-code
|
|
3
|
+
Version: 0.2.41
|
|
4
|
+
Summary: MCP server for indexing and querying codebases using CocoIndex
|
|
5
|
+
Project-URL: Homepage, https://github.com/cocoindex-io/cocoindex-code
|
|
6
|
+
Project-URL: Repository, https://github.com/cocoindex-io/cocoindex-code
|
|
7
|
+
Project-URL: Issues, https://github.com/cocoindex-io/cocoindex-code/issues
|
|
8
|
+
License-Expression: Apache-2.0
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Keywords: cocoindex,codebase,indexing,mcp,vector-search
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
18
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
19
|
+
Requires-Python: >=3.11
|
|
20
|
+
Requires-Dist: cocoindex[litellm]<1.1.0,>=1.0.17
|
|
21
|
+
Requires-Dist: einops>=0.8.2
|
|
22
|
+
Requires-Dist: mcp<3,>=2.0.0
|
|
23
|
+
Requires-Dist: msgspec>=0.19.0
|
|
24
|
+
Requires-Dist: numpy>=1.24.0
|
|
25
|
+
Requires-Dist: pathspec<1.0,>=0.12.1
|
|
26
|
+
Requires-Dist: pydantic>=2.0.0
|
|
27
|
+
Requires-Dist: pyyaml>=6.0
|
|
28
|
+
Requires-Dist: questionary>=2.0.0
|
|
29
|
+
Requires-Dist: sqlite-vec>=0.1.0
|
|
30
|
+
Requires-Dist: typer>=0.9.0
|
|
31
|
+
Provides-Extra: dev
|
|
32
|
+
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.17; extra == 'dev'
|
|
33
|
+
Requires-Dist: mypy>=1.0.0; extra == 'dev'
|
|
34
|
+
Requires-Dist: prek>=0.1.0; extra == 'dev'
|
|
35
|
+
Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
|
|
36
|
+
Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
|
|
37
|
+
Requires-Dist: pytest>=7.0.0; extra == 'dev'
|
|
38
|
+
Requires-Dist: ruff>=0.1.0; extra == 'dev'
|
|
39
|
+
Provides-Extra: embeddings-local
|
|
40
|
+
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.17; extra == 'embeddings-local'
|
|
41
|
+
Provides-Extra: full
|
|
42
|
+
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.17; extra == 'full'
|
|
43
|
+
Description-Content-Type: text/markdown
|
|
44
|
+
|
|
1
45
|
<p align="center">
|
|
2
46
|
<img width="2428" alt="cocoindex code" src="https://github.com/user-attachments/assets/d05961b4-0b7b-42ea-834a-59c3c01717ca" />
|
|
3
47
|
</p>
|
|
@@ -253,6 +297,7 @@ The background daemon starts automatically on first use.
|
|
|
253
297
|
| `ccc mcp` | Run as MCP server in stdio mode |
|
|
254
298
|
| `ccc doctor` | Run diagnostics — checks settings, daemon, model, file matching, and index health |
|
|
255
299
|
| `ccc reset` | Delete index databases. `--all` also removes settings. `-f` skips confirmation. |
|
|
300
|
+
| `ccc version` | Print the CLI version |
|
|
256
301
|
| `ccc daemon status` | Show daemon version, uptime, and loaded projects |
|
|
257
302
|
| `ccc daemon restart` | Restart the background daemon |
|
|
258
303
|
| `ccc daemon stop` | Stop the daemon |
|
|
@@ -504,6 +549,8 @@ embedding:
|
|
|
504
549
|
model: Snowflake/snowflake-arctic-embed-xs
|
|
505
550
|
device: mps # optional: cpu, cuda, mps (auto-detected if omitted)
|
|
506
551
|
min_interval_ms: 300 # optional: pace LiteLLM embedding requests to reduce 429s; defaults to 5 for LiteLLM
|
|
552
|
+
mps_low_watermark_ratio: 0.4 # optional: PyTorch allocator soft limit
|
|
553
|
+
mps_high_watermark_ratio: 0.5 # optional: PyTorch allocator hard limit
|
|
507
554
|
|
|
508
555
|
# Optional extra kwargs passed to the embedder, separately for indexing vs query.
|
|
509
556
|
# `ccc init` auto-populates these for known models (e.g. Cohere, Voyage, Nvidia NIM,
|
|
@@ -518,11 +565,16 @@ envs: # extra environment variabl
|
|
|
518
565
|
|
|
519
566
|
daemon:
|
|
520
567
|
idle_timeout_minutes: 180 # optional: exit the daemon after this long without client activity (default 180, 0 = never)
|
|
568
|
+
keep_alive_with_mcp: true # optional: keep the daemon warm while an MCP client is connected (default true)
|
|
521
569
|
```
|
|
522
570
|
|
|
523
571
|
> **Note:** The daemon inherits your shell environment. If an API key (e.g. `OPENAI_API_KEY`) is already set as an environment variable, you don't need to duplicate it in `envs`. The `envs` field is only for values that aren't in your environment.
|
|
524
572
|
|
|
525
|
-
> **
|
|
573
|
+
> **Apple Silicon memory safety:** MPS SentenceTransformer calls use [CocoIndex's isolated GPU subprocess](https://github.com/cocoindex-io/cocoindex/blob/v1.0.18/python/cocoindex/_internal/runner.py), keeping the model loaded while isolating Metal allocations from the daemon. The low and high watermarks are ratios of PyTorch's recommended maximum working set; they default here to `0.4` and `0.5`. CocoIndex retries MPS out-of-memory failures with progressively smaller batches, and cocoindex-code [releases unused allocator cache](https://docs.pytorch.org/docs/stable/generated/torch.mps.empty_cache.html) after each index run. Explicit `COCOINDEX_RUN_GPU_IN_SUBPROCESS`, `PYTORCH_MPS_LOW_WATERMARK_RATIO`, and `PYTORCH_MPS_HIGH_WATERMARK_RATIO` environment variables take precedence over these defaults.
|
|
574
|
+
|
|
575
|
+
> **Indexing concurrency:** Multiple projects may prepare indexes concurrently, while CocoIndex serializes their GPU calls through its single MPS subprocess. A search waits only when its own project still needs the initial index.
|
|
576
|
+
|
|
577
|
+
> **Idle timeout:** the background daemon holds the embedding model in RAM, so it exits after `daemon.idle_timeout_minutes` without client activity and is restarted automatically on your next `ccc` command or MCP search. By default, a live MCP session sends periodic heartbeats so the daemon remains warm while your coding agent is connected. Set `daemon.keep_alive_with_mcp: false` to let the daemon idle-exit during long-lived MCP sessions and release the model between real requests. Set `idle_timeout_minutes: 0` to keep the daemon running forever.
|
|
526
578
|
|
|
527
579
|
> **Custom location:** set `COCOINDEX_CODE_DIR` to place `global_settings.yml` somewhere other than `~/.cocoindex_code/` — useful if you want the file to live alongside your projects (e.g. on a synced folder).
|
|
528
580
|
|
|
@@ -570,6 +622,8 @@ exclude_patterns:
|
|
|
570
622
|
- "**/dist"
|
|
571
623
|
# ...
|
|
572
624
|
|
|
625
|
+
max_file_size: 500KB # skip files larger than this (default: no limit)
|
|
626
|
+
|
|
573
627
|
language_overrides:
|
|
574
628
|
- ext: inc # treat .inc files as PHP
|
|
575
629
|
lang: php
|
|
@@ -581,7 +635,14 @@ chunkers:
|
|
|
581
635
|
|
|
582
636
|
> `.cocoindex_code/` is automatically added to `.gitignore` during init.
|
|
583
637
|
|
|
584
|
-
|
|
638
|
+
`max_file_size` keeps bundled or generated files out of the index without having
|
|
639
|
+
to enumerate them in `exclude_patterns`. It accepts a plain byte count
|
|
640
|
+
(`1048576`) or a size with a binary unit suffix: `B`, `KB`, `MB`, `GB`
|
|
641
|
+
(case-insensitive, so `500KB` and `500 kb` are the same). The limit is
|
|
642
|
+
inclusive, and omitting the key indexes files of any size. It applies wherever
|
|
643
|
+
project file matching does, so `ccc grep` skips the same files.
|
|
644
|
+
|
|
645
|
+
After editing `include_patterns`, `exclude_patterns`, `max_file_size`, or `language_overrides`:
|
|
585
646
|
|
|
586
647
|
- Run `ccc doctor` to preview which files match.
|
|
587
648
|
- Run `ccc index` or `ccc search --refresh ...` to update the existing index.
|
|
@@ -1,47 +1,3 @@
|
|
|
1
|
-
Metadata-Version: 2.4
|
|
2
|
-
Name: cocoindex-code
|
|
3
|
-
Version: 0.2.39
|
|
4
|
-
Summary: MCP server for indexing and querying codebases using CocoIndex
|
|
5
|
-
Project-URL: Homepage, https://github.com/cocoindex-io/cocoindex-code
|
|
6
|
-
Project-URL: Repository, https://github.com/cocoindex-io/cocoindex-code
|
|
7
|
-
Project-URL: Issues, https://github.com/cocoindex-io/cocoindex-code/issues
|
|
8
|
-
License-Expression: Apache-2.0
|
|
9
|
-
License-File: LICENSE
|
|
10
|
-
Keywords: cocoindex,codebase,indexing,mcp,vector-search
|
|
11
|
-
Classifier: Development Status :: 3 - Alpha
|
|
12
|
-
Classifier: Intended Audience :: Developers
|
|
13
|
-
Classifier: License :: OSI Approved :: Apache Software License
|
|
14
|
-
Classifier: Programming Language :: Python :: 3
|
|
15
|
-
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
-
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
-
Classifier: Programming Language :: Python :: 3.13
|
|
18
|
-
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
19
|
-
Requires-Python: >=3.11
|
|
20
|
-
Requires-Dist: cocoindex[litellm]<1.1.0,>=1.0.13
|
|
21
|
-
Requires-Dist: einops>=0.8.2
|
|
22
|
-
Requires-Dist: mcp>=1.0.0
|
|
23
|
-
Requires-Dist: msgspec>=0.19.0
|
|
24
|
-
Requires-Dist: numpy>=1.24.0
|
|
25
|
-
Requires-Dist: pathspec>=0.12.1
|
|
26
|
-
Requires-Dist: pydantic>=2.0.0
|
|
27
|
-
Requires-Dist: pyyaml>=6.0
|
|
28
|
-
Requires-Dist: questionary>=2.0.0
|
|
29
|
-
Requires-Dist: sqlite-vec>=0.1.0
|
|
30
|
-
Requires-Dist: typer>=0.9.0
|
|
31
|
-
Provides-Extra: dev
|
|
32
|
-
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.13; extra == 'dev'
|
|
33
|
-
Requires-Dist: mypy>=1.0.0; extra == 'dev'
|
|
34
|
-
Requires-Dist: prek>=0.1.0; extra == 'dev'
|
|
35
|
-
Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
|
|
36
|
-
Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
|
|
37
|
-
Requires-Dist: pytest>=7.0.0; extra == 'dev'
|
|
38
|
-
Requires-Dist: ruff>=0.1.0; extra == 'dev'
|
|
39
|
-
Provides-Extra: embeddings-local
|
|
40
|
-
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.13; extra == 'embeddings-local'
|
|
41
|
-
Provides-Extra: full
|
|
42
|
-
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.13; extra == 'full'
|
|
43
|
-
Description-Content-Type: text/markdown
|
|
44
|
-
|
|
45
1
|
<p align="center">
|
|
46
2
|
<img width="2428" alt="cocoindex code" src="https://github.com/user-attachments/assets/d05961b4-0b7b-42ea-834a-59c3c01717ca" />
|
|
47
3
|
</p>
|
|
@@ -297,6 +253,7 @@ The background daemon starts automatically on first use.
|
|
|
297
253
|
| `ccc mcp` | Run as MCP server in stdio mode |
|
|
298
254
|
| `ccc doctor` | Run diagnostics — checks settings, daemon, model, file matching, and index health |
|
|
299
255
|
| `ccc reset` | Delete index databases. `--all` also removes settings. `-f` skips confirmation. |
|
|
256
|
+
| `ccc version` | Print the CLI version |
|
|
300
257
|
| `ccc daemon status` | Show daemon version, uptime, and loaded projects |
|
|
301
258
|
| `ccc daemon restart` | Restart the background daemon |
|
|
302
259
|
| `ccc daemon stop` | Stop the daemon |
|
|
@@ -548,6 +505,8 @@ embedding:
|
|
|
548
505
|
model: Snowflake/snowflake-arctic-embed-xs
|
|
549
506
|
device: mps # optional: cpu, cuda, mps (auto-detected if omitted)
|
|
550
507
|
min_interval_ms: 300 # optional: pace LiteLLM embedding requests to reduce 429s; defaults to 5 for LiteLLM
|
|
508
|
+
mps_low_watermark_ratio: 0.4 # optional: PyTorch allocator soft limit
|
|
509
|
+
mps_high_watermark_ratio: 0.5 # optional: PyTorch allocator hard limit
|
|
551
510
|
|
|
552
511
|
# Optional extra kwargs passed to the embedder, separately for indexing vs query.
|
|
553
512
|
# `ccc init` auto-populates these for known models (e.g. Cohere, Voyage, Nvidia NIM,
|
|
@@ -562,11 +521,16 @@ envs: # extra environment variabl
|
|
|
562
521
|
|
|
563
522
|
daemon:
|
|
564
523
|
idle_timeout_minutes: 180 # optional: exit the daemon after this long without client activity (default 180, 0 = never)
|
|
524
|
+
keep_alive_with_mcp: true # optional: keep the daemon warm while an MCP client is connected (default true)
|
|
565
525
|
```
|
|
566
526
|
|
|
567
527
|
> **Note:** The daemon inherits your shell environment. If an API key (e.g. `OPENAI_API_KEY`) is already set as an environment variable, you don't need to duplicate it in `envs`. The `envs` field is only for values that aren't in your environment.
|
|
568
528
|
|
|
569
|
-
> **
|
|
529
|
+
> **Apple Silicon memory safety:** MPS SentenceTransformer calls use [CocoIndex's isolated GPU subprocess](https://github.com/cocoindex-io/cocoindex/blob/v1.0.18/python/cocoindex/_internal/runner.py), keeping the model loaded while isolating Metal allocations from the daemon. The low and high watermarks are ratios of PyTorch's recommended maximum working set; they default here to `0.4` and `0.5`. CocoIndex retries MPS out-of-memory failures with progressively smaller batches, and cocoindex-code [releases unused allocator cache](https://docs.pytorch.org/docs/stable/generated/torch.mps.empty_cache.html) after each index run. Explicit `COCOINDEX_RUN_GPU_IN_SUBPROCESS`, `PYTORCH_MPS_LOW_WATERMARK_RATIO`, and `PYTORCH_MPS_HIGH_WATERMARK_RATIO` environment variables take precedence over these defaults.
|
|
530
|
+
|
|
531
|
+
> **Indexing concurrency:** Multiple projects may prepare indexes concurrently, while CocoIndex serializes their GPU calls through its single MPS subprocess. A search waits only when its own project still needs the initial index.
|
|
532
|
+
|
|
533
|
+
> **Idle timeout:** the background daemon holds the embedding model in RAM, so it exits after `daemon.idle_timeout_minutes` without client activity and is restarted automatically on your next `ccc` command or MCP search. By default, a live MCP session sends periodic heartbeats so the daemon remains warm while your coding agent is connected. Set `daemon.keep_alive_with_mcp: false` to let the daemon idle-exit during long-lived MCP sessions and release the model between real requests. Set `idle_timeout_minutes: 0` to keep the daemon running forever.
|
|
570
534
|
|
|
571
535
|
> **Custom location:** set `COCOINDEX_CODE_DIR` to place `global_settings.yml` somewhere other than `~/.cocoindex_code/` — useful if you want the file to live alongside your projects (e.g. on a synced folder).
|
|
572
536
|
|
|
@@ -614,6 +578,8 @@ exclude_patterns:
|
|
|
614
578
|
- "**/dist"
|
|
615
579
|
# ...
|
|
616
580
|
|
|
581
|
+
max_file_size: 500KB # skip files larger than this (default: no limit)
|
|
582
|
+
|
|
617
583
|
language_overrides:
|
|
618
584
|
- ext: inc # treat .inc files as PHP
|
|
619
585
|
lang: php
|
|
@@ -625,7 +591,14 @@ chunkers:
|
|
|
625
591
|
|
|
626
592
|
> `.cocoindex_code/` is automatically added to `.gitignore` during init.
|
|
627
593
|
|
|
628
|
-
|
|
594
|
+
`max_file_size` keeps bundled or generated files out of the index without having
|
|
595
|
+
to enumerate them in `exclude_patterns`. It accepts a plain byte count
|
|
596
|
+
(`1048576`) or a size with a binary unit suffix: `B`, `KB`, `MB`, `GB`
|
|
597
|
+
(case-insensitive, so `500KB` and `500 kb` are the same). The limit is
|
|
598
|
+
inclusive, and omitting the key indexes files of any size. It applies wherever
|
|
599
|
+
project file matching does, so `ccc grep` skips the same files.
|
|
600
|
+
|
|
601
|
+
After editing `include_patterns`, `exclude_patterns`, `max_file_size`, or `language_overrides`:
|
|
629
602
|
|
|
630
603
|
- Run `ccc doctor` to preview which files match.
|
|
631
604
|
- Run `ccc index` or `ccc search --refresh ...` to update the existing index.
|
|
@@ -22,15 +22,15 @@ classifiers = [
|
|
|
22
22
|
]
|
|
23
23
|
|
|
24
24
|
dependencies = [
|
|
25
|
-
"mcp>=
|
|
26
|
-
"cocoindex[litellm]>=1.0.
|
|
25
|
+
"mcp>=2.0.0,<3",
|
|
26
|
+
"cocoindex[litellm]>=1.0.17,<1.1.0",
|
|
27
27
|
"sqlite-vec>=0.1.0",
|
|
28
28
|
"pydantic>=2.0.0",
|
|
29
29
|
"numpy>=1.24.0",
|
|
30
30
|
"einops>=0.8.2",
|
|
31
31
|
"typer>=0.9.0",
|
|
32
32
|
"msgspec>=0.19.0",
|
|
33
|
-
"pathspec>=0.12.1",
|
|
33
|
+
"pathspec>=0.12.1,<1.0",
|
|
34
34
|
"pyyaml>=6.0",
|
|
35
35
|
"questionary>=2.0.0",
|
|
36
36
|
]
|
|
@@ -39,7 +39,7 @@ dependencies = [
|
|
|
39
39
|
# `embeddings-local` is the primary feature extra: it pulls in
|
|
40
40
|
# `sentence-transformers` (via cocoindex) so local embeddings work without
|
|
41
41
|
# an API key.
|
|
42
|
-
embeddings-local = ["cocoindex[sentence-transformers]>=1.0.
|
|
42
|
+
embeddings-local = ["cocoindex[sentence-transformers]>=1.0.17,<1.1.0"]
|
|
43
43
|
# `full` is the umbrella "batteries-included" alias. Today it's just
|
|
44
44
|
# `embeddings-local`, but we expect to bundle more optional niceties under
|
|
45
45
|
# it over time — users who want everything can keep using `[full]` and pick
|
|
@@ -47,7 +47,7 @@ embeddings-local = ["cocoindex[sentence-transformers]>=1.0.13,<1.1.0"]
|
|
|
47
47
|
# `:full` image variant for consistency across install paths. Contents are
|
|
48
48
|
# inlined rather than self-referencing `cocoindex-code[embeddings-local]`
|
|
49
49
|
# to avoid resolver edge cases with older pip.
|
|
50
|
-
full = ["cocoindex[sentence-transformers]>=1.0.
|
|
50
|
+
full = ["cocoindex[sentence-transformers]>=1.0.17,<1.1.0"]
|
|
51
51
|
dev = [
|
|
52
52
|
"pytest>=7.0.0",
|
|
53
53
|
"pytest-asyncio>=0.21.0",
|
|
@@ -55,7 +55,7 @@ dev = [
|
|
|
55
55
|
"ruff>=0.1.0",
|
|
56
56
|
"mypy>=1.0.0",
|
|
57
57
|
"prek>=0.1.0",
|
|
58
|
-
"cocoindex[sentence-transformers]>=1.0.
|
|
58
|
+
"cocoindex[sentence-transformers]>=1.0.17,<1.1.0",
|
|
59
59
|
]
|
|
60
60
|
|
|
61
61
|
[project.scripts]
|
|
@@ -89,7 +89,7 @@ dev = [
|
|
|
89
89
|
"mypy>=1.0.0",
|
|
90
90
|
"prek>=0.1.0",
|
|
91
91
|
"types-pyyaml>=6.0.12.20250915",
|
|
92
|
-
"cocoindex[sentence-transformers]>=1.0.
|
|
92
|
+
"cocoindex[sentence-transformers]>=1.0.17,<1.1.0",
|
|
93
93
|
]
|
|
94
94
|
|
|
95
95
|
[tool.ruff]
|
|
@@ -6,15 +6,22 @@ can import these without pulling in the full daemon stack.
|
|
|
6
6
|
|
|
7
7
|
from __future__ import annotations
|
|
8
8
|
|
|
9
|
+
import hashlib
|
|
9
10
|
import json
|
|
10
11
|
import os
|
|
11
12
|
import sys
|
|
13
|
+
import tempfile
|
|
12
14
|
import time
|
|
13
15
|
from dataclasses import dataclass
|
|
14
16
|
from pathlib import Path
|
|
15
17
|
|
|
16
18
|
from .settings import user_settings_dir
|
|
17
19
|
|
|
20
|
+
# Max bytes in sockaddr_un.sun_path, including the NUL terminator: 104 on the
|
|
21
|
+
# BSDs (incl. macOS), 108 on Linux. Exceeding it makes bind() fail with
|
|
22
|
+
# "AF_UNIX path too long" — unlike ordinary files, which get PATH_MAX (~1024).
|
|
23
|
+
_SUN_PATH_MAX = 104 if sys.platform == "darwin" else 108
|
|
24
|
+
|
|
18
25
|
|
|
19
26
|
def daemon_runtime_dir() -> Path:
|
|
20
27
|
"""Return the directory that holds daemon runtime artifacts.
|
|
@@ -39,18 +46,37 @@ def connection_family() -> str:
|
|
|
39
46
|
return "AF_PIPE" if sys.platform == "win32" else "AF_UNIX"
|
|
40
47
|
|
|
41
48
|
|
|
49
|
+
def _runtime_dir_hash() -> str:
|
|
50
|
+
"""Stable short id for the current runtime dir.
|
|
51
|
+
|
|
52
|
+
Distinguishes daemon instances that differ only by
|
|
53
|
+
``COCOINDEX_CODE_RUNTIME_DIR`` / ``COCOINDEX_CODE_DIR``, so tests, users,
|
|
54
|
+
and containers never collide on one address.
|
|
55
|
+
"""
|
|
56
|
+
return hashlib.md5(str(daemon_runtime_dir()).encode()).hexdigest()[:12]
|
|
57
|
+
|
|
58
|
+
|
|
42
59
|
def daemon_socket_path() -> str:
|
|
43
|
-
"""Return the daemon socket/pipe address.
|
|
60
|
+
"""Return the daemon socket/pipe address.
|
|
61
|
+
|
|
62
|
+
Normally ``<runtime dir>/daemon.sock``. When that exceeds the platform's
|
|
63
|
+
``sun_path`` limit — a deep ``$HOME``, a sandbox, a container path — it
|
|
64
|
+
would fail at bind() with a message that reads like a broken install, so
|
|
65
|
+
fall back to a short path under the temp dir keyed by the runtime dir.
|
|
66
|
+
Client and daemon both resolve through here, so they agree either way.
|
|
67
|
+
"""
|
|
44
68
|
if sys.platform == "win32":
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
69
|
+
return rf"\\.\pipe\cocoindex_code_{_runtime_dir_hash()}"
|
|
70
|
+
|
|
71
|
+
path = daemon_runtime_dir() / "daemon.sock"
|
|
72
|
+
if len(str(path).encode()) < _SUN_PATH_MAX:
|
|
73
|
+
return str(path)
|
|
74
|
+
|
|
75
|
+
# gettempdir() honors $TMPDIR, which is per-user and private on macOS; on
|
|
76
|
+
# Linux it is usually shared /tmp, so include the uid to avoid a foreign
|
|
77
|
+
# socket sitting at our address.
|
|
78
|
+
short = Path(tempfile.gettempdir()) / f"ccc-{os.getuid()}-{_runtime_dir_hash()}.sock"
|
|
79
|
+
return str(short)
|
|
54
80
|
|
|
55
81
|
|
|
56
82
|
def daemon_pid_path() -> Path:
|
|
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
|
|
|
18
18
|
commit_id: str | None
|
|
19
19
|
__commit_id__: str | None
|
|
20
20
|
|
|
21
|
-
__version__ = version = '0.2.
|
|
22
|
-
__version_tuple__ = version_tuple = (0, 2,
|
|
21
|
+
__version__ = version = '0.2.41'
|
|
22
|
+
__version_tuple__ = version_tuple = (0, 2, 41)
|
|
23
23
|
|
|
24
24
|
__commit_id__ = commit_id = None
|
|
@@ -22,6 +22,7 @@ if TYPE_CHECKING:
|
|
|
22
22
|
SearchResponse,
|
|
23
23
|
)
|
|
24
24
|
|
|
25
|
+
from ._version import __version__
|
|
25
26
|
from .settings import (
|
|
26
27
|
DEFAULT_ST_MODEL,
|
|
27
28
|
EmbeddingSettings,
|
|
@@ -195,6 +196,15 @@ def print_index_stats(status: ProjectStatusResponse) -> None:
|
|
|
195
196
|
_typer.echo(f" {lang}: {count} chunks")
|
|
196
197
|
|
|
197
198
|
|
|
199
|
+
def _echo_search_text(text: str) -> None:
|
|
200
|
+
"""Echo result text, replacing characters unsupported by the console codec."""
|
|
201
|
+
try:
|
|
202
|
+
_typer.echo(text)
|
|
203
|
+
except UnicodeEncodeError as error:
|
|
204
|
+
safe_text = text.encode(error.encoding, errors="replace").decode(error.encoding)
|
|
205
|
+
_typer.echo(safe_text)
|
|
206
|
+
|
|
207
|
+
|
|
198
208
|
def print_search_results(response: SearchResponse) -> None:
|
|
199
209
|
"""Print formatted search results."""
|
|
200
210
|
if not response.success:
|
|
@@ -207,8 +217,8 @@ def print_search_results(response: SearchResponse) -> None:
|
|
|
207
217
|
|
|
208
218
|
for i, r in enumerate(response.results, 1):
|
|
209
219
|
_typer.echo(f"\n--- Result {i} (score: {r.score:.3f}) ---")
|
|
210
|
-
|
|
211
|
-
|
|
220
|
+
_echo_search_text(f"File: {r.file_path}:{r.start_line}-{r.end_line} [{r.language}]")
|
|
221
|
+
_echo_search_text(r.content)
|
|
212
222
|
|
|
213
223
|
|
|
214
224
|
def _run_index_with_progress(project_root: str) -> None:
|
|
@@ -1120,6 +1130,12 @@ def daemon_stop() -> None:
|
|
|
1120
1130
|
_typer.echo("Daemon stopped.")
|
|
1121
1131
|
|
|
1122
1132
|
|
|
1133
|
+
@app.command()
|
|
1134
|
+
def version() -> None:
|
|
1135
|
+
"""Print the CLI version."""
|
|
1136
|
+
_typer.echo(__version__)
|
|
1137
|
+
|
|
1138
|
+
|
|
1123
1139
|
@app.command("run-daemon", hidden=True)
|
|
1124
1140
|
def run_daemon_cmd() -> None:
|
|
1125
1141
|
"""Internal: run the daemon process."""
|
|
@@ -228,6 +228,19 @@ def _handle_vanished_daemon() -> None:
|
|
|
228
228
|
)
|
|
229
229
|
|
|
230
230
|
|
|
231
|
+
def _daemon_error(resp: ErrorResponse) -> RuntimeError:
|
|
232
|
+
"""Build the exception for an ``ErrorResponse`` from the daemon.
|
|
233
|
+
|
|
234
|
+
The daemon-side traceback is appended when present: without it the caller
|
|
235
|
+
only ever sees the client's own re-raise frames, which says nothing about
|
|
236
|
+
where the failure actually happened (issue #270).
|
|
237
|
+
"""
|
|
238
|
+
detail = f"Daemon error: {resp.message}"
|
|
239
|
+
if resp.traceback:
|
|
240
|
+
detail += f"\n{resp.traceback}"
|
|
241
|
+
return RuntimeError(detail)
|
|
242
|
+
|
|
243
|
+
|
|
231
244
|
class _HandshakeResult(NamedTuple):
|
|
232
245
|
conn: Connection
|
|
233
246
|
resp: HandshakeResponse
|
|
@@ -261,7 +274,7 @@ def _raw_connect_and_handshake() -> _HandshakeResult:
|
|
|
261
274
|
raise DaemonProtocolError(f"Undecodable handshake reply from daemon: {e}") from e
|
|
262
275
|
if isinstance(resp, ErrorResponse):
|
|
263
276
|
conn.close()
|
|
264
|
-
raise
|
|
277
|
+
raise _daemon_error(resp)
|
|
265
278
|
if not isinstance(resp, HandshakeResponse):
|
|
266
279
|
conn.close()
|
|
267
280
|
raise RuntimeError(f"Unexpected handshake response: {type(resp).__name__}")
|
|
@@ -340,7 +353,7 @@ def _send(req: Request) -> Response:
|
|
|
340
353
|
conn.close()
|
|
341
354
|
resp = decode_response(data)
|
|
342
355
|
if isinstance(resp, ErrorResponse):
|
|
343
|
-
raise
|
|
356
|
+
raise _daemon_error(resp)
|
|
344
357
|
return resp
|
|
345
358
|
|
|
346
359
|
|
|
@@ -366,7 +379,7 @@ def index(
|
|
|
366
379
|
raise RuntimeError("Connection to daemon lost during indexing")
|
|
367
380
|
resp = decode_response(data)
|
|
368
381
|
if isinstance(resp, ErrorResponse):
|
|
369
|
-
raise
|
|
382
|
+
raise _daemon_error(resp)
|
|
370
383
|
if isinstance(resp, IndexWaitingNotice):
|
|
371
384
|
if on_waiting is not None:
|
|
372
385
|
on_waiting()
|
|
@@ -419,7 +432,7 @@ def search(
|
|
|
419
432
|
raise RuntimeError("Connection to daemon lost during search")
|
|
420
433
|
resp = decode_response(data)
|
|
421
434
|
if isinstance(resp, ErrorResponse):
|
|
422
|
-
raise
|
|
435
|
+
raise _daemon_error(resp)
|
|
423
436
|
if isinstance(resp, IndexWaitingNotice):
|
|
424
437
|
if on_waiting is not None:
|
|
425
438
|
on_waiting()
|
|
@@ -496,10 +509,7 @@ def doctor(
|
|
|
496
509
|
raise RuntimeError("Connection to daemon lost during doctor checks")
|
|
497
510
|
resp = decode_response(data)
|
|
498
511
|
if isinstance(resp, ErrorResponse):
|
|
499
|
-
|
|
500
|
-
if resp.traceback:
|
|
501
|
-
detail += f"\n{resp.traceback}"
|
|
502
|
-
raise RuntimeError(detail)
|
|
512
|
+
raise _daemon_error(resp)
|
|
503
513
|
if isinstance(resp, DoctorResponse):
|
|
504
514
|
results.append(resp.result)
|
|
505
515
|
if on_result is not None:
|
|
@@ -73,7 +73,7 @@ from .settings import (
|
|
|
73
73
|
target_sqlite_db_path,
|
|
74
74
|
user_settings_path,
|
|
75
75
|
)
|
|
76
|
-
from .shared import Embedder, check_embedding, create_embedder
|
|
76
|
+
from .shared import Embedder, check_embedding, configure_mps_environment, create_embedder
|
|
77
77
|
|
|
78
78
|
logger = logging.getLogger(__name__)
|
|
79
79
|
|
|
@@ -135,6 +135,7 @@ class ProjectRegistry:
|
|
|
135
135
|
|
|
136
136
|
_projects: dict[str, Project]
|
|
137
137
|
_embedder: Embedder | None
|
|
138
|
+
_clear_mps_cache_after_index: bool
|
|
138
139
|
indexing_params: dict[str, Any]
|
|
139
140
|
query_params: dict[str, Any]
|
|
140
141
|
|
|
@@ -143,9 +144,11 @@ class ProjectRegistry:
|
|
|
143
144
|
embedder: Embedder | None,
|
|
144
145
|
indexing_params: dict[str, Any] | None = None,
|
|
145
146
|
query_params: dict[str, Any] | None = None,
|
|
147
|
+
clear_mps_cache_after_index: bool = False,
|
|
146
148
|
) -> None:
|
|
147
149
|
self._projects = {}
|
|
148
150
|
self._embedder = embedder
|
|
151
|
+
self._clear_mps_cache_after_index = clear_mps_cache_after_index
|
|
149
152
|
self.indexing_params = dict(indexing_params) if indexing_params else {}
|
|
150
153
|
self.query_params = dict(query_params) if query_params else {}
|
|
151
154
|
|
|
@@ -165,6 +168,7 @@ class ProjectRegistry:
|
|
|
165
168
|
indexing_params=self.indexing_params,
|
|
166
169
|
query_params=self.query_params,
|
|
167
170
|
chunker_registry=chunker_registry,
|
|
171
|
+
clear_mps_cache_after_index=self._clear_mps_cache_after_index,
|
|
168
172
|
)
|
|
169
173
|
self._projects[project_root] = project
|
|
170
174
|
return self._projects[project_root]
|
|
@@ -342,7 +346,8 @@ async def _search_with_wait(
|
|
|
342
346
|
offset=req.offset,
|
|
343
347
|
)
|
|
344
348
|
except Exception as e:
|
|
345
|
-
|
|
349
|
+
logger.exception("Error handling search request")
|
|
350
|
+
yield ErrorResponse(message=str(e), traceback=traceback.format_exc())
|
|
346
351
|
|
|
347
352
|
|
|
348
353
|
async def _handle_doctor(
|
|
@@ -601,7 +606,7 @@ async def _dispatch(
|
|
|
601
606
|
return ErrorResponse(message=f"Unknown request type: {type(req).__name__}")
|
|
602
607
|
except Exception as e:
|
|
603
608
|
logger.exception("Error dispatching request")
|
|
604
|
-
return ErrorResponse(message=str(e))
|
|
609
|
+
return ErrorResponse(message=str(e), traceback=traceback.format_exc())
|
|
605
610
|
|
|
606
611
|
|
|
607
612
|
# ---------------------------------------------------------------------------
|
|
@@ -636,6 +641,7 @@ def run_daemon(
|
|
|
636
641
|
embedder: Embedder | None
|
|
637
642
|
indexing_params: dict[str, Any] = {}
|
|
638
643
|
query_params: dict[str, Any] = {}
|
|
644
|
+
clear_mps_cache_after_index = False
|
|
639
645
|
handshake_warnings: list[str] = []
|
|
640
646
|
daemon_settings = DaemonSettings()
|
|
641
647
|
if user_settings_path().is_file():
|
|
@@ -644,6 +650,7 @@ def run_daemon(
|
|
|
644
650
|
settings_env_keys = list(user_settings.envs.keys())
|
|
645
651
|
for key, value in user_settings.envs.items():
|
|
646
652
|
os.environ[key] = value
|
|
653
|
+
clear_mps_cache_after_index = configure_mps_environment(user_settings.embedding)
|
|
647
654
|
# Resolve params BEFORE constructing the embedder so invalid configs
|
|
648
655
|
# fail fast without paying the model-load cost.
|
|
649
656
|
try:
|
|
@@ -684,6 +691,7 @@ def run_daemon(
|
|
|
684
691
|
embedder,
|
|
685
692
|
indexing_params=indexing_params,
|
|
686
693
|
query_params=query_params,
|
|
694
|
+
clear_mps_cache_after_index=clear_mps_cache_after_index,
|
|
687
695
|
)
|
|
688
696
|
|
|
689
697
|
sock_path = daemon_socket_path()
|
|
@@ -120,6 +120,33 @@ class GitignoreAwareMatcher(FilePathMatcher):
|
|
|
120
120
|
return self._delegate.is_file_included(path)
|
|
121
121
|
|
|
122
122
|
|
|
123
|
+
class SizeLimitedMatcher(FilePathMatcher):
|
|
124
|
+
"""Wraps another matcher and drops files larger than ``max_bytes``.
|
|
125
|
+
|
|
126
|
+
The matcher protocol only sees paths relative to the project root, so this
|
|
127
|
+
resolves each candidate against ``project_root`` to stat it. Files that
|
|
128
|
+
cannot be stat'd (a broken symlink, a race with a delete) are left to the
|
|
129
|
+
delegate rather than silently dropped, since a size limit should not be the
|
|
130
|
+
thing that decides an unreadable file's fate.
|
|
131
|
+
"""
|
|
132
|
+
|
|
133
|
+
def __init__(self, delegate: FilePathMatcher, project_root: Path, max_bytes: int) -> None:
|
|
134
|
+
self._delegate = delegate
|
|
135
|
+
self._root = project_root
|
|
136
|
+
self._max_bytes = max_bytes
|
|
137
|
+
|
|
138
|
+
def is_dir_included(self, path: PurePath) -> bool:
|
|
139
|
+
return self._delegate.is_dir_included(path)
|
|
140
|
+
|
|
141
|
+
def is_file_included(self, path: PurePath) -> bool:
|
|
142
|
+
if not self._delegate.is_file_included(path):
|
|
143
|
+
return False
|
|
144
|
+
try:
|
|
145
|
+
return (self._root / path).stat().st_size <= self._max_bytes
|
|
146
|
+
except OSError:
|
|
147
|
+
return True
|
|
148
|
+
|
|
149
|
+
|
|
123
150
|
def find_git_root(start: Path) -> Path | None:
|
|
124
151
|
"""Walk up from ``start`` to the nearest directory holding a ``.git`` entry — a
|
|
125
152
|
directory for a normal repo, or a *file* for a submodule or linked worktree.
|
|
@@ -140,14 +167,21 @@ def build_matcher(
|
|
|
140
167
|
project_root: Path,
|
|
141
168
|
included_patterns: list[str],
|
|
142
169
|
excluded_patterns: list[str],
|
|
170
|
+
max_file_size: int | None = None,
|
|
143
171
|
) -> FilePathMatcher:
|
|
144
172
|
"""Build the project's file matcher: include/exclude globs plus nested
|
|
145
|
-
``.gitignore`` awareness anchored at ``project_root
|
|
173
|
+
``.gitignore`` awareness anchored at ``project_root``, and an optional
|
|
174
|
+
``max_file_size`` cap in bytes."""
|
|
146
175
|
base_matcher = PatternFilePathMatcher(
|
|
147
176
|
included_patterns=included_patterns,
|
|
148
177
|
excluded_patterns=excluded_patterns,
|
|
149
178
|
)
|
|
150
|
-
|
|
179
|
+
matcher: FilePathMatcher = GitignoreAwareMatcher(
|
|
180
|
+
base_matcher, load_gitignore_spec(project_root), project_root
|
|
181
|
+
)
|
|
182
|
+
if max_file_size is not None:
|
|
183
|
+
matcher = SizeLimitedMatcher(matcher, project_root, max_file_size)
|
|
184
|
+
return matcher
|
|
151
185
|
|
|
152
186
|
|
|
153
187
|
def iter_included_files(
|
|
@@ -214,17 +214,19 @@ def _iter_targets(req: GrepRequest, compiler: _PatternCompiler) -> Iterator[_Tar
|
|
|
214
214
|
if project_root is not None:
|
|
215
215
|
ps = load_project_settings(project_root)
|
|
216
216
|
included, excluded = ps.include_patterns, ps.exclude_patterns
|
|
217
|
+
max_file_size = ps.max_file_size
|
|
217
218
|
ext_overrides = {f".{lo.ext}": lo.lang for lo in ps.language_overrides}
|
|
218
219
|
base = project_root
|
|
219
220
|
else:
|
|
220
221
|
included = list(DEFAULT_INCLUDED_PATTERNS)
|
|
221
222
|
excluded = list(DEFAULT_EXCLUDED_PATTERNS)
|
|
222
223
|
ext_overrides = {}
|
|
224
|
+
max_file_size = None
|
|
223
225
|
# Anchor at the enclosing git repo so grepping a subdirectory still honors the
|
|
224
226
|
# repo-root (and intervening) .gitignore; fall back to the target dir itself.
|
|
225
227
|
base = find_git_root(root) or root
|
|
226
228
|
|
|
227
|
-
matcher = build_matcher(base, included, excluded)
|
|
229
|
+
matcher = build_matcher(base, included, excluded, max_file_size)
|
|
228
230
|
path_filter = (
|
|
229
231
|
PatternFilePathMatcher(included_patterns=[req.path_glob]) if req.path_glob else None
|
|
230
232
|
)
|
|
@@ -110,7 +110,9 @@ async def indexer_main() -> None:
|
|
|
110
110
|
),
|
|
111
111
|
)
|
|
112
112
|
|
|
113
|
-
matcher = build_matcher(
|
|
113
|
+
matcher = build_matcher(
|
|
114
|
+
project_root, ps.include_patterns, ps.exclude_patterns, ps.max_file_size
|
|
115
|
+
)
|
|
114
116
|
|
|
115
117
|
files = localfs.walk_dir(
|
|
116
118
|
CODEBASE_DIR,
|
|
@@ -3,6 +3,7 @@
|
|
|
3
3
|
from __future__ import annotations
|
|
4
4
|
|
|
5
5
|
import asyncio
|
|
6
|
+
import logging
|
|
6
7
|
import sqlite3
|
|
7
8
|
from collections.abc import AsyncIterator, Callable
|
|
8
9
|
from pathlib import Path
|
|
@@ -39,15 +40,21 @@ from .shared import (
|
|
|
39
40
|
QUERY_EMBED_PARAMS,
|
|
40
41
|
SQLITE_DB,
|
|
41
42
|
Embedder,
|
|
43
|
+
clear_mps_allocator_cache,
|
|
42
44
|
)
|
|
43
45
|
|
|
46
|
+
logger = logging.getLogger(__name__)
|
|
47
|
+
|
|
44
48
|
|
|
45
49
|
class Project:
|
|
46
50
|
_env: coco.Environment
|
|
47
51
|
_app: coco.App[[], None]
|
|
48
52
|
_project_root: Path
|
|
49
53
|
_index_lock: asyncio.Lock
|
|
54
|
+
_clear_mps_cache_after_index: bool
|
|
50
55
|
_initial_index_done: asyncio.Event
|
|
56
|
+
_initial_index_task: asyncio.Task[None] | None
|
|
57
|
+
_initial_index_started: asyncio.Event | None
|
|
51
58
|
_indexing_stats: IndexingProgress | None = None
|
|
52
59
|
|
|
53
60
|
def close(self) -> None:
|
|
@@ -67,13 +74,15 @@ class Project:
|
|
|
67
74
|
on_progress: Callable[[IndexingProgress], None] | None = None,
|
|
68
75
|
on_started: asyncio.Event | None = None,
|
|
69
76
|
) -> None:
|
|
70
|
-
"""Acquire the
|
|
77
|
+
"""Acquire the project lock, run indexing, and release it.
|
|
71
78
|
|
|
72
|
-
If *on_started* is provided, it is set once the lock is
|
|
73
|
-
|
|
74
|
-
|
|
79
|
+
If *on_started* is provided, it is set once the project lock is
|
|
80
|
+
acquired. On completion (success or failure) ``_initial_index_done``
|
|
81
|
+
is set.
|
|
75
82
|
"""
|
|
76
83
|
async with self._index_lock:
|
|
84
|
+
if on_started is not None:
|
|
85
|
+
on_started.set()
|
|
77
86
|
self._indexing_stats = IndexingProgress(
|
|
78
87
|
num_execution_starts=0,
|
|
79
88
|
num_unchanged=0,
|
|
@@ -82,8 +91,6 @@ class Project:
|
|
|
82
91
|
num_reprocesses=0,
|
|
83
92
|
num_errors=0,
|
|
84
93
|
)
|
|
85
|
-
if on_started is not None:
|
|
86
|
-
on_started.set()
|
|
87
94
|
await self._run_index_inner(on_progress=on_progress)
|
|
88
95
|
|
|
89
96
|
async def _run_index_inner(
|
|
@@ -109,20 +116,32 @@ class Project:
|
|
|
109
116
|
on_progress(progress)
|
|
110
117
|
await asyncio.sleep(0.1)
|
|
111
118
|
finally:
|
|
112
|
-
|
|
113
|
-
|
|
119
|
+
try:
|
|
120
|
+
if self._clear_mps_cache_after_index:
|
|
121
|
+
try:
|
|
122
|
+
await clear_mps_allocator_cache()
|
|
123
|
+
except Exception:
|
|
124
|
+
logger.exception("Unable to clear the MPS allocator cache after indexing")
|
|
125
|
+
finally:
|
|
126
|
+
self._initial_index_done.set()
|
|
127
|
+
self._indexing_stats = None
|
|
114
128
|
|
|
115
129
|
async def ensure_indexing_started(self) -> None:
|
|
116
|
-
"""Kick off background indexing and wait until it
|
|
130
|
+
"""Kick off background indexing and wait until it is active or queued.
|
|
117
131
|
|
|
118
|
-
Returns once the indexing task holds the lock.
|
|
119
|
-
times — only the first call spawns a task; subsequent calls
|
|
120
|
-
immediately.
|
|
132
|
+
Returns once the indexing task holds the project lock. Safe to call
|
|
133
|
+
multiple times — only the first call spawns a task; subsequent calls
|
|
134
|
+
return immediately.
|
|
121
135
|
"""
|
|
122
136
|
if self._initial_index_done.is_set() or self._index_lock.locked():
|
|
123
137
|
return
|
|
124
|
-
|
|
125
|
-
|
|
138
|
+
if self._initial_index_task is None or self._initial_index_task.done():
|
|
139
|
+
self._initial_index_started = asyncio.Event()
|
|
140
|
+
self._initial_index_task = asyncio.create_task(
|
|
141
|
+
self.run_index(on_started=self._initial_index_started)
|
|
142
|
+
)
|
|
143
|
+
started = self._initial_index_started
|
|
144
|
+
assert started is not None
|
|
126
145
|
await started.wait()
|
|
127
146
|
|
|
128
147
|
async def stream_index(self) -> AsyncIterator[IndexStreamResponse]:
|
|
@@ -263,6 +282,7 @@ class Project:
|
|
|
263
282
|
indexing_params: dict[str, Any],
|
|
264
283
|
query_params: dict[str, Any],
|
|
265
284
|
chunker_registry: dict[str, ChunkerFn] | None = None,
|
|
285
|
+
clear_mps_cache_after_index: bool = False,
|
|
266
286
|
) -> Project:
|
|
267
287
|
"""Create a project with explicit embedder and per-call params.
|
|
268
288
|
|
|
@@ -281,6 +301,8 @@ class Project:
|
|
|
281
301
|
chunker_registry: Optional mapping of file suffix (e.g. ``".toml"``)
|
|
282
302
|
to a ``ChunkerFn``. When a suffix matches, the registered
|
|
283
303
|
chunker is called instead of the built-in splitter.
|
|
304
|
+
clear_mps_cache_after_index: Whether to release unused MPS allocator
|
|
305
|
+
memory in CocoIndex's GPU subprocess after each index run.
|
|
284
306
|
"""
|
|
285
307
|
settings_dir = project_root / ".cocoindex_code"
|
|
286
308
|
settings_dir.mkdir(parents=True, exist_ok=True)
|
|
@@ -315,5 +337,8 @@ class Project:
|
|
|
315
337
|
result._app = app
|
|
316
338
|
result._project_root = project_root
|
|
317
339
|
result._index_lock = asyncio.Lock()
|
|
340
|
+
result._clear_mps_cache_after_index = clear_mps_cache_after_index
|
|
318
341
|
result._initial_index_done = asyncio.Event()
|
|
342
|
+
result._initial_index_task = None
|
|
343
|
+
result._initial_index_started = None
|
|
319
344
|
return result
|
|
@@ -16,6 +16,33 @@ def _l2_to_score(distance: float) -> float:
|
|
|
16
16
|
return 1.0 - distance * distance / 2.0
|
|
17
17
|
|
|
18
18
|
|
|
19
|
+
def _checked(rows: list[tuple[Any, ...]], query_shape: str) -> list[tuple[Any, ...]]:
|
|
20
|
+
"""Return *rows*, failing loudly if any of them carries a NULL distance.
|
|
21
|
+
|
|
22
|
+
``code_chunks_vec`` is a vec0 virtual table whose ``distance`` is a *hidden*
|
|
23
|
+
column: sqlite-vec populates it only under the KNN query plan and yields
|
|
24
|
+
NULL for it on a plain full scan. A NULL therefore means the plan we asked
|
|
25
|
+
for is not the plan we got, which is a bug worth reporting rather than a
|
|
26
|
+
result worth ranking. Left unchecked, the NULL flows into ``_l2_to_score``
|
|
27
|
+
and surfaces as an unrelated-looking ``TypeError: unsupported operand
|
|
28
|
+
type(s) for *: 'NoneType' and 'NoneType'`` (issue #270).
|
|
29
|
+
|
|
30
|
+
Only KNN queries need this guard: ``_full_scan_query`` computes its own
|
|
31
|
+
``vec_distance_L2(...)``, which works under any plan and raises (never
|
|
32
|
+
returns NULL) on bad input.
|
|
33
|
+
"""
|
|
34
|
+
bad = next((row for row in rows if row[5] is None), None)
|
|
35
|
+
if bad is not None:
|
|
36
|
+
raise RuntimeError(
|
|
37
|
+
f"Vector index returned a row with no distance ({query_shape}, "
|
|
38
|
+
f"file_path={bad[0]!r}) — the sqlite-vec KNN query plan was not "
|
|
39
|
+
"used. Please report this at "
|
|
40
|
+
"https://github.com/cocoindex-io/cocoindex-code/issues along with "
|
|
41
|
+
"the output of `ccc doctor`."
|
|
42
|
+
)
|
|
43
|
+
return rows
|
|
44
|
+
|
|
45
|
+
|
|
19
46
|
def _knn_query(
|
|
20
47
|
conn: sqlite3.Connection,
|
|
21
48
|
embedding_bytes: bytes,
|
|
@@ -24,24 +51,30 @@ def _knn_query(
|
|
|
24
51
|
) -> list[tuple[Any, ...]]:
|
|
25
52
|
"""Run a vec0 KNN query, optionally constrained to a language partition."""
|
|
26
53
|
if language is not None:
|
|
27
|
-
return
|
|
54
|
+
return _checked(
|
|
55
|
+
conn.execute(
|
|
56
|
+
"""
|
|
57
|
+
SELECT file_path, language, content, start_line, end_line, distance
|
|
58
|
+
FROM code_chunks_vec
|
|
59
|
+
WHERE embedding MATCH ? AND k = ? AND language = ?
|
|
60
|
+
ORDER BY distance
|
|
61
|
+
""",
|
|
62
|
+
(embedding_bytes, k, language),
|
|
63
|
+
).fetchall(),
|
|
64
|
+
f"knn language={language!r}",
|
|
65
|
+
)
|
|
66
|
+
return _checked(
|
|
67
|
+
conn.execute(
|
|
28
68
|
"""
|
|
29
69
|
SELECT file_path, language, content, start_line, end_line, distance
|
|
30
70
|
FROM code_chunks_vec
|
|
31
|
-
WHERE embedding MATCH ? AND k = ?
|
|
71
|
+
WHERE embedding MATCH ? AND k = ?
|
|
32
72
|
ORDER BY distance
|
|
33
73
|
""",
|
|
34
|
-
(embedding_bytes, k
|
|
35
|
-
).fetchall()
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
SELECT file_path, language, content, start_line, end_line, distance
|
|
39
|
-
FROM code_chunks_vec
|
|
40
|
-
WHERE embedding MATCH ? AND k = ?
|
|
41
|
-
ORDER BY distance
|
|
42
|
-
""",
|
|
43
|
-
(embedding_bytes, k),
|
|
44
|
-
).fetchall()
|
|
74
|
+
(embedding_bytes, k),
|
|
75
|
+
).fetchall(),
|
|
76
|
+
"knn unfiltered",
|
|
77
|
+
)
|
|
45
78
|
|
|
46
79
|
|
|
47
80
|
def _full_scan_query(
|
|
@@ -14,9 +14,10 @@ import json
|
|
|
14
14
|
import os
|
|
15
15
|
from pathlib import Path
|
|
16
16
|
|
|
17
|
-
from mcp.server.
|
|
17
|
+
from mcp.server.mcpserver import MCPServer
|
|
18
18
|
from pydantic import BaseModel, Field
|
|
19
19
|
|
|
20
|
+
from ._version import __version__
|
|
20
21
|
from .settings import DaemonSettings, load_user_settings
|
|
21
22
|
|
|
22
23
|
_MCP_INSTRUCTIONS = (
|
|
@@ -58,9 +59,9 @@ class SearchResultModel(BaseModel):
|
|
|
58
59
|
# === Daemon-backed MCP server factory ===
|
|
59
60
|
|
|
60
61
|
|
|
61
|
-
def create_mcp_server(project_root: str) ->
|
|
62
|
+
def create_mcp_server(project_root: str) -> MCPServer:
|
|
62
63
|
"""Create a lightweight MCP server that delegates to the daemon."""
|
|
63
|
-
mcp =
|
|
64
|
+
mcp = MCPServer("cocoindex-code", instructions=_MCP_INSTRUCTIONS, version=__version__)
|
|
64
65
|
|
|
65
66
|
@mcp.tool(
|
|
66
67
|
name="search",
|
|
@@ -174,24 +175,27 @@ def heartbeat_interval_s(idle_timeout_minutes: int) -> int:
|
|
|
174
175
|
|
|
175
176
|
async def run_heartbeat_loop() -> None:
|
|
176
177
|
"""Periodically heartbeat the daemon so it never idle-exits under this
|
|
177
|
-
live MCP session
|
|
178
|
-
|
|
178
|
+
live MCP session when ``daemon.keep_alive_with_mcp`` is enabled. When the
|
|
179
|
+
option is disabled, the daemon may idle-exit and the next real request
|
|
180
|
+
restarts it transparently. When this process dies (even SIGKILL), the
|
|
181
|
+
heartbeats stop and the daemon exits on its normal idle timeout.
|
|
179
182
|
|
|
180
183
|
Reads the timeout from the same ``global_settings.yml`` the daemon reads
|
|
181
184
|
(dataclass default when the file is missing or invalid). Returns
|
|
182
|
-
immediately when the timeout is 0 (daemon never idle-exits)
|
|
183
|
-
heartbeat itself never starts or restarts a
|
|
184
|
-
``client.send_heartbeat``.
|
|
185
|
+
immediately when the timeout is 0 (daemon never idle-exits) or MCP
|
|
186
|
+
keep-alive is disabled. The heartbeat itself never starts or restarts a
|
|
187
|
+
daemon — see ``client.send_heartbeat``.
|
|
185
188
|
"""
|
|
186
|
-
from .client import send_heartbeat
|
|
187
|
-
|
|
188
189
|
try:
|
|
189
|
-
|
|
190
|
+
daemon_settings = load_user_settings().daemon
|
|
190
191
|
except (FileNotFoundError, ValueError):
|
|
191
|
-
|
|
192
|
-
|
|
192
|
+
daemon_settings = DaemonSettings()
|
|
193
|
+
timeout_minutes = daemon_settings.idle_timeout_minutes
|
|
194
|
+
if timeout_minutes <= 0 or not daemon_settings.keep_alive_with_mcp:
|
|
193
195
|
return
|
|
194
196
|
|
|
197
|
+
from .client import send_heartbeat
|
|
198
|
+
|
|
195
199
|
interval = heartbeat_interval_s(timeout_minutes)
|
|
196
200
|
loop = asyncio.get_event_loop()
|
|
197
201
|
while True:
|
|
@@ -201,7 +205,7 @@ async def run_heartbeat_loop() -> None:
|
|
|
201
205
|
|
|
202
206
|
|
|
203
207
|
# Keep the old `mcp` global for backward compatibility in __init__.py
|
|
204
|
-
mcp:
|
|
208
|
+
mcp: MCPServer | None = None
|
|
205
209
|
|
|
206
210
|
|
|
207
211
|
# === Backward-compatible entry point ===
|
|
@@ -98,18 +98,38 @@ class EmbeddingSettings:
|
|
|
98
98
|
provider: str = "litellm"
|
|
99
99
|
device: str | None = None
|
|
100
100
|
min_interval_ms: int | None = None
|
|
101
|
+
# PyTorch MPS allocator limits used by CocoIndex's isolated GPU runner.
|
|
102
|
+
mps_low_watermark_ratio: float = 0.4
|
|
103
|
+
mps_high_watermark_ratio: float = 0.5
|
|
101
104
|
# Extra kwargs spread into ``embedder.embed()`` during indexing/query.
|
|
102
105
|
# ``None`` means the user did not set the key; ``{}`` is an explicit empty
|
|
103
106
|
# dict (used to opt out of the legacy-bridge warning).
|
|
104
107
|
indexing_params: dict[str, Any] | None = None
|
|
105
108
|
query_params: dict[str, Any] | None = None
|
|
106
109
|
|
|
110
|
+
def __post_init__(self) -> None:
|
|
111
|
+
ratios = (
|
|
112
|
+
("mps_low_watermark_ratio", self.mps_low_watermark_ratio),
|
|
113
|
+
("mps_high_watermark_ratio", self.mps_high_watermark_ratio),
|
|
114
|
+
)
|
|
115
|
+
for name, value in ratios:
|
|
116
|
+
if isinstance(value, bool) or not 0 < value <= 1:
|
|
117
|
+
raise ValueError(f"embedding.{name} must be greater than 0 and at most 1")
|
|
118
|
+
|
|
119
|
+
if self.mps_low_watermark_ratio > self.mps_high_watermark_ratio:
|
|
120
|
+
raise ValueError(
|
|
121
|
+
"embedding.mps_low_watermark_ratio must be at most "
|
|
122
|
+
"embedding.mps_high_watermark_ratio"
|
|
123
|
+
)
|
|
124
|
+
|
|
107
125
|
|
|
108
126
|
@dataclass
|
|
109
127
|
class DaemonSettings:
|
|
110
128
|
# Minutes without client activity before the daemon exits (0 = never exit).
|
|
111
129
|
# Clients auto-restart the daemon on the next request, so exiting is cheap.
|
|
112
130
|
idle_timeout_minutes: int = 180
|
|
131
|
+
# Keep the daemon warm while a long-lived MCP client is connected.
|
|
132
|
+
keep_alive_with_mcp: bool = True
|
|
113
133
|
|
|
114
134
|
|
|
115
135
|
@dataclass
|
|
@@ -137,6 +157,8 @@ class ProjectSettings:
|
|
|
137
157
|
exclude_patterns: list[str] = field(default_factory=lambda: list(DEFAULT_EXCLUDED_PATTERNS))
|
|
138
158
|
language_overrides: list[LanguageOverride] = field(default_factory=list)
|
|
139
159
|
chunkers: list[ChunkerMapping] = field(default_factory=list)
|
|
160
|
+
#: Files larger than this many bytes are excluded. ``None`` means no limit.
|
|
161
|
+
max_file_size: int | None = None
|
|
140
162
|
|
|
141
163
|
|
|
142
164
|
# ---------------------------------------------------------------------------
|
|
@@ -430,6 +452,11 @@ def _embedding_settings_to_dict(embedding: EmbeddingSettings) -> dict[str, Any]:
|
|
|
430
452
|
d["device"] = embedding.device
|
|
431
453
|
if embedding.min_interval_ms is not None:
|
|
432
454
|
d["min_interval_ms"] = embedding.min_interval_ms
|
|
455
|
+
defaults = EmbeddingSettings(model=embedding.model)
|
|
456
|
+
if embedding.mps_low_watermark_ratio != defaults.mps_low_watermark_ratio:
|
|
457
|
+
d["mps_low_watermark_ratio"] = embedding.mps_low_watermark_ratio
|
|
458
|
+
if embedding.mps_high_watermark_ratio != defaults.mps_high_watermark_ratio:
|
|
459
|
+
d["mps_high_watermark_ratio"] = embedding.mps_high_watermark_ratio
|
|
433
460
|
if embedding.indexing_params is not None:
|
|
434
461
|
d["indexing_params"] = dict(embedding.indexing_params)
|
|
435
462
|
if embedding.query_params is not None:
|
|
@@ -441,8 +468,14 @@ def _user_settings_to_dict(settings: UserSettings) -> dict[str, Any]:
|
|
|
441
468
|
d: dict[str, Any] = {"embedding": _embedding_settings_to_dict(settings.embedding)}
|
|
442
469
|
if settings.envs:
|
|
443
470
|
d["envs"] = dict(settings.envs)
|
|
444
|
-
|
|
445
|
-
|
|
471
|
+
daemon_defaults = DaemonSettings()
|
|
472
|
+
daemon_dict: dict[str, Any] = {}
|
|
473
|
+
if settings.daemon.idle_timeout_minutes != daemon_defaults.idle_timeout_minutes:
|
|
474
|
+
daemon_dict["idle_timeout_minutes"] = settings.daemon.idle_timeout_minutes
|
|
475
|
+
if settings.daemon.keep_alive_with_mcp != daemon_defaults.keep_alive_with_mcp:
|
|
476
|
+
daemon_dict["keep_alive_with_mcp"] = settings.daemon.keep_alive_with_mcp
|
|
477
|
+
if daemon_dict:
|
|
478
|
+
d["daemon"] = daemon_dict
|
|
446
479
|
return d
|
|
447
480
|
|
|
448
481
|
|
|
@@ -458,6 +491,10 @@ def _user_settings_from_dict(d: dict[str, Any]) -> UserSettings:
|
|
|
458
491
|
emb_kwargs["device"] = emb_dict["device"]
|
|
459
492
|
if "min_interval_ms" in emb_dict:
|
|
460
493
|
emb_kwargs["min_interval_ms"] = emb_dict["min_interval_ms"]
|
|
494
|
+
if "mps_low_watermark_ratio" in emb_dict:
|
|
495
|
+
emb_kwargs["mps_low_watermark_ratio"] = float(emb_dict["mps_low_watermark_ratio"])
|
|
496
|
+
if "mps_high_watermark_ratio" in emb_dict:
|
|
497
|
+
emb_kwargs["mps_high_watermark_ratio"] = float(emb_dict["mps_high_watermark_ratio"])
|
|
461
498
|
# indexing_params / query_params: missing → None (dataclass default);
|
|
462
499
|
# present-but-null → {} (treat the same as an empty dict, since both mean
|
|
463
500
|
# "user acknowledged the key and wants no extra kwargs").
|
|
@@ -472,15 +509,66 @@ def _user_settings_from_dict(d: dict[str, Any]) -> UserSettings:
|
|
|
472
509
|
daemon_kwargs: dict[str, Any] = {}
|
|
473
510
|
if "idle_timeout_minutes" in daemon_dict:
|
|
474
511
|
daemon_kwargs["idle_timeout_minutes"] = int(daemon_dict["idle_timeout_minutes"])
|
|
512
|
+
if "keep_alive_with_mcp" in daemon_dict:
|
|
513
|
+
keep_alive_with_mcp = daemon_dict["keep_alive_with_mcp"]
|
|
514
|
+
if not isinstance(keep_alive_with_mcp, bool):
|
|
515
|
+
raise ValueError("daemon.keep_alive_with_mcp must be a boolean")
|
|
516
|
+
daemon_kwargs["keep_alive_with_mcp"] = keep_alive_with_mcp
|
|
475
517
|
daemon = DaemonSettings(**daemon_kwargs)
|
|
476
518
|
return UserSettings(embedding=embedding, envs=envs, daemon=daemon)
|
|
477
519
|
|
|
478
520
|
|
|
521
|
+
_SIZE_UNITS: dict[str, int] = {
|
|
522
|
+
"B": 1,
|
|
523
|
+
"KB": 1024,
|
|
524
|
+
"MB": 1024**2,
|
|
525
|
+
"GB": 1024**3,
|
|
526
|
+
}
|
|
527
|
+
|
|
528
|
+
|
|
529
|
+
def parse_file_size(value: Any) -> int:
|
|
530
|
+
"""Parse a ``max_file_size`` value into bytes.
|
|
531
|
+
|
|
532
|
+
Accepts a plain number of bytes (``1048576``) or a string with an optional
|
|
533
|
+
unit suffix (``500KB``, ``1.5 MB``, ``2048``). Units are binary (1KB = 1024
|
|
534
|
+
bytes) and case-insensitive.
|
|
535
|
+
"""
|
|
536
|
+
if isinstance(value, bool): # bool is an int subclass; reject it explicitly.
|
|
537
|
+
raise ValueError(f"invalid max_file_size: {value!r}")
|
|
538
|
+
if isinstance(value, int):
|
|
539
|
+
size = value
|
|
540
|
+
elif isinstance(value, float):
|
|
541
|
+
size = int(value)
|
|
542
|
+
elif isinstance(value, str):
|
|
543
|
+
text = value.strip().upper()
|
|
544
|
+
if not text:
|
|
545
|
+
raise ValueError("invalid max_file_size: empty value")
|
|
546
|
+
for suffix in ("GB", "MB", "KB", "B"):
|
|
547
|
+
if text.endswith(suffix):
|
|
548
|
+
number = text[: -len(suffix)].strip()
|
|
549
|
+
multiplier = _SIZE_UNITS[suffix]
|
|
550
|
+
break
|
|
551
|
+
else:
|
|
552
|
+
number, multiplier = text, 1
|
|
553
|
+
try:
|
|
554
|
+
size = int(float(number) * multiplier)
|
|
555
|
+
except ValueError as exc:
|
|
556
|
+
raise ValueError(f"invalid max_file_size: {value!r}") from exc
|
|
557
|
+
else:
|
|
558
|
+
raise ValueError(f"invalid max_file_size: {value!r}")
|
|
559
|
+
|
|
560
|
+
if size <= 0:
|
|
561
|
+
raise ValueError(f"max_file_size must be positive, got {value!r}")
|
|
562
|
+
return size
|
|
563
|
+
|
|
564
|
+
|
|
479
565
|
def _project_settings_to_dict(settings: ProjectSettings) -> dict[str, Any]:
|
|
480
566
|
d: dict[str, Any] = {
|
|
481
567
|
"include_patterns": settings.include_patterns,
|
|
482
568
|
"exclude_patterns": settings.exclude_patterns,
|
|
483
569
|
}
|
|
570
|
+
if settings.max_file_size is not None:
|
|
571
|
+
d["max_file_size"] = settings.max_file_size
|
|
484
572
|
if settings.language_overrides:
|
|
485
573
|
d["language_overrides"] = [
|
|
486
574
|
{"ext": lo.ext, "lang": lo.lang} for lo in settings.language_overrides
|
|
@@ -495,11 +583,14 @@ def _project_settings_from_dict(d: dict[str, Any]) -> ProjectSettings:
|
|
|
495
583
|
LanguageOverride(ext=lo["ext"], lang=lo["lang"]) for lo in d.get("language_overrides", [])
|
|
496
584
|
]
|
|
497
585
|
chunkers = [ChunkerMapping(ext=cm["ext"], module=cm["module"]) for cm in d.get("chunkers", [])]
|
|
586
|
+
raw_max_size = d.get("max_file_size")
|
|
587
|
+
max_file_size = None if raw_max_size is None else parse_file_size(raw_max_size)
|
|
498
588
|
return ProjectSettings(
|
|
499
589
|
include_patterns=d.get("include_patterns", list(DEFAULT_INCLUDED_PATTERNS)),
|
|
500
590
|
exclude_patterns=d.get("exclude_patterns", list(DEFAULT_EXCLUDED_PATTERNS)),
|
|
501
591
|
language_overrides=overrides,
|
|
502
592
|
chunkers=chunkers,
|
|
593
|
+
max_file_size=max_file_size,
|
|
503
594
|
)
|
|
504
595
|
|
|
505
596
|
|
|
@@ -573,6 +664,13 @@ _PARAMS_COMMENT_BY_PROVIDER: dict[str, str] = {
|
|
|
573
664
|
),
|
|
574
665
|
}
|
|
575
666
|
|
|
667
|
+
_MPS_SAFETY_COMMENT = (
|
|
668
|
+
" #\n"
|
|
669
|
+
" # Apple Silicon MPS allocator limits for CocoIndex's GPU subprocess.\n"
|
|
670
|
+
" # mps_low_watermark_ratio: 0.4\n"
|
|
671
|
+
" # mps_high_watermark_ratio: 0.5\n"
|
|
672
|
+
)
|
|
673
|
+
|
|
576
674
|
|
|
577
675
|
def save_initial_user_settings(
|
|
578
676
|
embedding: EmbeddingSettings,
|
|
@@ -595,6 +693,8 @@ def save_initial_user_settings(
|
|
|
595
693
|
sort_keys=False,
|
|
596
694
|
)
|
|
597
695
|
content = _INITIAL_HEADER + emb_block
|
|
696
|
+
if embedding.provider == "sentence-transformers":
|
|
697
|
+
content += _MPS_SAFETY_COMMENT
|
|
598
698
|
if not defaults_applied:
|
|
599
699
|
hint = _PARAMS_COMMENT_BY_PROVIDER.get(embedding.provider)
|
|
600
700
|
if hint is not None:
|
|
@@ -4,7 +4,9 @@ from __future__ import annotations
|
|
|
4
4
|
|
|
5
5
|
import importlib.util
|
|
6
6
|
import logging
|
|
7
|
+
import os
|
|
7
8
|
import pathlib
|
|
9
|
+
import sys
|
|
8
10
|
import traceback as _tb
|
|
9
11
|
from dataclasses import dataclass
|
|
10
12
|
from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, Union
|
|
@@ -26,7 +28,10 @@ SBERT_PREFIX = "sbert/"
|
|
|
26
28
|
DEFAULT_LITELLM_MIN_INTERVAL_MS = 5
|
|
27
29
|
|
|
28
30
|
# Type alias
|
|
29
|
-
Embedder = Union[
|
|
31
|
+
Embedder = Union[
|
|
32
|
+
"SentenceTransformerEmbedder",
|
|
33
|
+
"LiteLLMEmbedder",
|
|
34
|
+
]
|
|
30
35
|
|
|
31
36
|
# Context keys
|
|
32
37
|
EMBEDDER = coco.ContextKey[Embedder]("embedder", detect_change=True)
|
|
@@ -45,6 +50,46 @@ def is_sentence_transformers_installed() -> bool:
|
|
|
45
50
|
return importlib.util.find_spec("sentence_transformers") is not None
|
|
46
51
|
|
|
47
52
|
|
|
53
|
+
def configure_mps_environment(settings: EmbeddingSettings) -> bool:
|
|
54
|
+
"""Install MPS safety defaults before CocoIndex's first GPU invocation.
|
|
55
|
+
|
|
56
|
+
CocoIndex reads ``COCOINDEX_RUN_GPU_IN_SUBPROCESS`` lazily on the first
|
|
57
|
+
``coco.GPU`` call, while PyTorch reads the allocator watermarks when the
|
|
58
|
+
child initializes MPS. Explicit user environment variables always win.
|
|
59
|
+
|
|
60
|
+
Returns whether this embedding configuration uses the guarded MPS path.
|
|
61
|
+
"""
|
|
62
|
+
|
|
63
|
+
use_mps = settings.provider == "sentence-transformers" and (
|
|
64
|
+
settings.device == "mps" or (settings.device is None and sys.platform == "darwin")
|
|
65
|
+
)
|
|
66
|
+
if not use_mps:
|
|
67
|
+
return False
|
|
68
|
+
|
|
69
|
+
os.environ.setdefault("COCOINDEX_RUN_GPU_IN_SUBPROCESS", "1")
|
|
70
|
+
os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", str(settings.mps_low_watermark_ratio))
|
|
71
|
+
os.environ.setdefault(
|
|
72
|
+
"PYTORCH_MPS_HIGH_WATERMARK_RATIO", str(settings.mps_high_watermark_ratio)
|
|
73
|
+
)
|
|
74
|
+
return os.environ.get("COCOINDEX_RUN_GPU_IN_SUBPROCESS") == "1"
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
@coco.fn.as_async(runner=coco.GPU)
|
|
78
|
+
def clear_mps_allocator_cache() -> None:
|
|
79
|
+
"""Release unused MPS allocator cache inside CocoIndex's GPU child."""
|
|
80
|
+
|
|
81
|
+
import gc
|
|
82
|
+
|
|
83
|
+
import torch
|
|
84
|
+
|
|
85
|
+
if not torch.backends.mps.is_available():
|
|
86
|
+
return
|
|
87
|
+
torch.mps.synchronize()
|
|
88
|
+
gc.collect()
|
|
89
|
+
torch.mps.empty_cache()
|
|
90
|
+
torch.mps.synchronize()
|
|
91
|
+
|
|
92
|
+
|
|
48
93
|
class EmbeddingCheckResult(NamedTuple):
|
|
49
94
|
"""Outcome of a single embed-test call. See `check_embedding`.
|
|
50
95
|
|
|
@@ -100,6 +145,7 @@ def create_embedder(
|
|
|
100
145
|
only and the indexing default is supplied at the call site via
|
|
101
146
|
:data:`INDEXING_EMBED_PARAMS`.
|
|
102
147
|
"""
|
|
148
|
+
instance: Embedder
|
|
103
149
|
if settings.provider == "sentence-transformers":
|
|
104
150
|
from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder
|
|
105
151
|
|
|
@@ -108,7 +154,7 @@ def create_embedder(
|
|
|
108
154
|
if model_name.startswith(SBERT_PREFIX):
|
|
109
155
|
model_name = model_name[len(SBERT_PREFIX) :]
|
|
110
156
|
|
|
111
|
-
instance
|
|
157
|
+
instance = SentenceTransformerEmbedder(
|
|
112
158
|
model_name,
|
|
113
159
|
device=settings.device,
|
|
114
160
|
trust_remote_code=True,
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|