vortexa 0.2.0__tar.gz → 0.3.1__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.
- {vortexa-0.2.0 → vortexa-0.3.1}/PKG-INFO +93 -7
- {vortexa-0.2.0 → vortexa-0.3.1}/README.md +92 -6
- {vortexa-0.2.0 → vortexa-0.3.1}/pyproject.toml +1 -1
- vortexa-0.3.1/src/vortexa/core/context_engine.py +250 -0
- vortexa-0.3.1/src/vortexa/core/graph.py +246 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/indexer.py +4 -3
- vortexa-0.3.1/src/vortexa/core/inference.py +248 -0
- vortexa-0.3.1/src/vortexa/core/lf4_v4_model.py +258 -0
- vortexa-0.3.1/src/vortexa/core/v4_embedder.py +150 -0
- vortexa-0.3.1/src/vortexa/core/vortex_score.py +208 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/interfaces/cli.py +63 -1
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/PKG-INFO +93 -7
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/SOURCES.txt +6 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/LICENSE +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/setup.cfg +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/__init__.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/__init__.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/chunking.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/embedding.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/language.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/lf4_model.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/types.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/interfaces/__init__.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/interfaces/mcp_server.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/interfaces/watcher.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/search/__init__.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/search/ranking.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/search/search.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/search/tokens.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/storage/__init__.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/storage/bm25.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/storage/vector_store.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/storage/walker.py +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/dependency_links.txt +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/entry_points.txt +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/requires.txt +0 -0
- {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: vortexa
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.1
|
|
4
4
|
Summary: Codebase indexing and semantic search engine
|
|
5
5
|
Author-email: VortexAI <koulabhay25@gmail.com>
|
|
6
6
|
License-Expression: Apache-2.0
|
|
@@ -78,7 +78,7 @@ _Dense + sparse hybrid retrieval · AST-aware chunking · LMDB persistence · MC
|
|
|
78
78
|
|
|
79
79
|
vortexa is a standalone **codebase indexing and semantic search engine** designed for AI agents and developers. It builds a persistent, hybrid search index over source code using:
|
|
80
80
|
|
|
81
|
-
- **Dense retrieval** via
|
|
81
|
+
- **Dense retrieval** via VortexEmbedderV4 (on-the-fly LF4 4-bit dequant, SIF+PC, Matryoshka) or static embeddings (Model2Vec / SentenceTransformers)
|
|
82
82
|
- **Sparse retrieval** via BM25 keyword scoring
|
|
83
83
|
- **AST-aware chunking** that respects function and class boundaries via tree-sitter
|
|
84
84
|
- **LMDB-backed storage** for fast, persistent vector and chunk storage
|
|
@@ -311,6 +311,7 @@ Useful flags:
|
|
|
311
311
|
| `--force` | Force a full re-index before searching. |
|
|
312
312
|
| `--no-index` | Search the existing index only. |
|
|
313
313
|
| `--plain` | Print human-readable results instead of JSON. |
|
|
314
|
+
| `--model` | Embedding model ID or alias (`mini`, `nano`). Default: `mini`. |
|
|
314
315
|
|
|
315
316
|
By default CLI output is JSON:
|
|
316
317
|
|
|
@@ -336,14 +337,99 @@ vortexa-serve
|
|
|
336
337
|
|
|
337
338
|
---
|
|
338
339
|
|
|
340
|
+
## Model Configuration
|
|
341
|
+
|
|
342
|
+
vortexa supports configurable embedding models. The default is `mini` (`VTXAI/vtx-embed-7M`).
|
|
343
|
+
|
|
344
|
+
### Available Models
|
|
345
|
+
|
|
346
|
+
| Alias | Model ID | Description |
|
|
347
|
+
|-------|----------|-------------|
|
|
348
|
+
| `mini` | `VTXAI/vtx-embed-7M` | Default. 7M-parameter Vortex-Embed v4.5 with LF4 4-bit dequant, SIF+PC, Matryoshka. |
|
|
349
|
+
| `nano` | `VTXAI/vtx-embed-1M` | Lightweight 1M-parameter variant. Lower RAM, lower dim. |
|
|
350
|
+
|
|
351
|
+
### CLI Usage
|
|
352
|
+
|
|
353
|
+
```bash
|
|
354
|
+
# Use the nano model (smaller, faster)
|
|
355
|
+
vortexa -q "authentication" --model nano /path/to/project
|
|
356
|
+
|
|
357
|
+
# Use mini explicitly (same as default)
|
|
358
|
+
vortexa -q "authentication" --model mini /path/to/project
|
|
359
|
+
|
|
360
|
+
# Use a custom HuggingFace model ID
|
|
361
|
+
vortexa -q "authentication" --model VTXAI/vtx-embed-1M /path/to/project
|
|
362
|
+
```
|
|
363
|
+
|
|
364
|
+
### Python API
|
|
365
|
+
|
|
366
|
+
```python
|
|
367
|
+
from vortexa.core.indexer import CodebaseIndexer
|
|
368
|
+
|
|
369
|
+
# Use nano model
|
|
370
|
+
indexer = CodebaseIndexer(root="/path/to/project", model_id="VTXAI/vtx-embed-1M")
|
|
371
|
+
|
|
372
|
+
# Use mini model (default)
|
|
373
|
+
indexer = CodebaseIndexer(root="/path/to/project", model_id="VTXAI/vtx-embed-7M")
|
|
374
|
+
```
|
|
375
|
+
|
|
376
|
+
### Alternative Embedders
|
|
377
|
+
|
|
378
|
+
The `embedding.py` module also provides alternative embedders that can be passed directly:
|
|
379
|
+
|
|
380
|
+
- `Model2VecEmbedder` — static embeddings via Model2Vec (`AI4free/JARVIS-tool-search-v1`)
|
|
381
|
+
- `SentenceTransformerEmbedder` — Transformer-based dense embeddings (`all-MiniLM-L6-v2`)
|
|
382
|
+
- `LF4Embedder` — 4-bit quantized static embeddings (`VTXAI/Vortex-Embed-4.7M`)
|
|
383
|
+
|
|
384
|
+
---
|
|
385
|
+
|
|
386
|
+
## Inference API
|
|
387
|
+
|
|
388
|
+
vortexa can be used as a standalone embedding inference engine for VTXAI models,
|
|
389
|
+
similar to how sentence-transformers works but purpose-built for Vortex-Embed models.
|
|
390
|
+
|
|
391
|
+
### Python API
|
|
392
|
+
|
|
393
|
+
```python
|
|
394
|
+
from vortexa.core.inference import embed
|
|
395
|
+
|
|
396
|
+
# Encode a single string (default: mini model)
|
|
397
|
+
vec = embed("India is a diverse country")
|
|
398
|
+
|
|
399
|
+
# Encode multiple strings with the nano model
|
|
400
|
+
vecs = embed(["Indian cricket team is strong", "Chennai is a major city"], model="nano")
|
|
401
|
+
|
|
402
|
+
# Use any HuggingFace model ID
|
|
403
|
+
vecs = embed(["Indian agriculture output"], model="VTXAI/vtx-embed-7M")
|
|
404
|
+
|
|
405
|
+
# Shape: (1, D) for single text, (N, D) for batch
|
|
406
|
+
print(vec.shape)
|
|
407
|
+
```
|
|
408
|
+
|
|
409
|
+
### CLI
|
|
410
|
+
|
|
411
|
+
```bash
|
|
412
|
+
# Encode text with the mini model (default)
|
|
413
|
+
vortexa embed "India is a diverse country"
|
|
414
|
+
|
|
415
|
+
# Encode with the nano model
|
|
416
|
+
vortexa embed "Indian cricket team is strong" --model nano
|
|
417
|
+
|
|
418
|
+
# Encode multiple strings
|
|
419
|
+
vortexa embed "India has 28 states" "Chennai is in Tamil Nadu" --model nano
|
|
420
|
+
|
|
421
|
+
# Use a custom model ID
|
|
422
|
+
vortexa embed "Indian monsoon patterns" --model VTXAI/vtx-embed-7M
|
|
423
|
+
```
|
|
424
|
+
|
|
425
|
+
---
|
|
426
|
+
|
|
339
427
|
<div align="center">
|
|
340
428
|
|
|
341
429
|
## MCP Server
|
|
342
430
|
|
|
343
431
|
</div>
|
|
344
432
|
|
|
345
|
-
vortexa ships with a built-in **MCP (Model Context Protocol) server** that exposes codebase search as a single `search` tool. Start it with:
|
|
346
|
-
|
|
347
433
|
```bash
|
|
348
434
|
# Auto-indexes current directory, serves on stdio
|
|
349
435
|
python -m vortexa.interfaces.mcp_server
|
|
@@ -398,7 +484,7 @@ vortexa/
|
|
|
398
484
|
├── core/
|
|
399
485
|
│ ├── indexer.py # CodebaseIndexer — main orchestrator
|
|
400
486
|
│ ├── chunking.py # AST-aware (tree-sitter) + line-based chunking
|
|
401
|
-
│ ├── embedding.py # Embedding
|
|
487
|
+
│ ├── embedding.py # Embedding model wrappers (Model2Vec, SentenceTransformers, LF4)
|
|
402
488
|
│ ├── language.py # Language detection & file extension mapping
|
|
403
489
|
│ └── types.py # Shared types (Chunk, ChunkConfig, IndexStats, SearchResult, ...)
|
|
404
490
|
├── storage/
|
|
@@ -459,7 +545,7 @@ graph LR
|
|
|
459
545
|
D -->|No| F[Line-based Splitter<br/>Configurable size/overlap]
|
|
460
546
|
E --> G[Chunk Set]
|
|
461
547
|
F --> G
|
|
462
|
-
G --> H[Embedding Model<br/>Model2Vec / SentenceTransformer]
|
|
548
|
+
G --> H[Embedding Model<br/>VortexEmbedderV4 (LF4) / Model2Vec / SentenceTransformer]
|
|
463
549
|
G --> I[BM25 Tokenizer]
|
|
464
550
|
H --> J[(LMDB Vector Store)]
|
|
465
551
|
I --> K[(BM25 Index)]
|
|
@@ -531,7 +617,7 @@ graph TD
|
|
|
531
617
|
| `bm25s` | Yes | Fast BM25 keyword index and persistence |
|
|
532
618
|
| `pathspec` | Yes | `.gitignore` pattern matching in file walker |
|
|
533
619
|
| `model2vec` | Optional | Alternative static embeddings |
|
|
534
|
-
| `huggingface-hub` | Yes (default model) | Loading `VTXAI/
|
|
620
|
+
| `huggingface-hub` | Yes (default model) | Loading `VTXAI/vtx-embed-7M` (mini) or `VTXAI/vtx-embed-1M` (nano) |
|
|
535
621
|
| `tokenizers` | Yes (default model) | HF tokenizer for embedding model |
|
|
536
622
|
| `safetensors` | Yes (default model) | Safe tensor loading for 4-bit weights |
|
|
537
623
|
| `sentence-transformers` | Optional | Transformer-based dense embeddings |
|
|
@@ -42,7 +42,7 @@ _Dense + sparse hybrid retrieval · AST-aware chunking · LMDB persistence · MC
|
|
|
42
42
|
|
|
43
43
|
vortexa is a standalone **codebase indexing and semantic search engine** designed for AI agents and developers. It builds a persistent, hybrid search index over source code using:
|
|
44
44
|
|
|
45
|
-
- **Dense retrieval** via
|
|
45
|
+
- **Dense retrieval** via VortexEmbedderV4 (on-the-fly LF4 4-bit dequant, SIF+PC, Matryoshka) or static embeddings (Model2Vec / SentenceTransformers)
|
|
46
46
|
- **Sparse retrieval** via BM25 keyword scoring
|
|
47
47
|
- **AST-aware chunking** that respects function and class boundaries via tree-sitter
|
|
48
48
|
- **LMDB-backed storage** for fast, persistent vector and chunk storage
|
|
@@ -275,6 +275,7 @@ Useful flags:
|
|
|
275
275
|
| `--force` | Force a full re-index before searching. |
|
|
276
276
|
| `--no-index` | Search the existing index only. |
|
|
277
277
|
| `--plain` | Print human-readable results instead of JSON. |
|
|
278
|
+
| `--model` | Embedding model ID or alias (`mini`, `nano`). Default: `mini`. |
|
|
278
279
|
|
|
279
280
|
By default CLI output is JSON:
|
|
280
281
|
|
|
@@ -300,14 +301,99 @@ vortexa-serve
|
|
|
300
301
|
|
|
301
302
|
---
|
|
302
303
|
|
|
304
|
+
## Model Configuration
|
|
305
|
+
|
|
306
|
+
vortexa supports configurable embedding models. The default is `mini` (`VTXAI/vtx-embed-7M`).
|
|
307
|
+
|
|
308
|
+
### Available Models
|
|
309
|
+
|
|
310
|
+
| Alias | Model ID | Description |
|
|
311
|
+
|-------|----------|-------------|
|
|
312
|
+
| `mini` | `VTXAI/vtx-embed-7M` | Default. 7M-parameter Vortex-Embed v4.5 with LF4 4-bit dequant, SIF+PC, Matryoshka. |
|
|
313
|
+
| `nano` | `VTXAI/vtx-embed-1M` | Lightweight 1M-parameter variant. Lower RAM, lower dim. |
|
|
314
|
+
|
|
315
|
+
### CLI Usage
|
|
316
|
+
|
|
317
|
+
```bash
|
|
318
|
+
# Use the nano model (smaller, faster)
|
|
319
|
+
vortexa -q "authentication" --model nano /path/to/project
|
|
320
|
+
|
|
321
|
+
# Use mini explicitly (same as default)
|
|
322
|
+
vortexa -q "authentication" --model mini /path/to/project
|
|
323
|
+
|
|
324
|
+
# Use a custom HuggingFace model ID
|
|
325
|
+
vortexa -q "authentication" --model VTXAI/vtx-embed-1M /path/to/project
|
|
326
|
+
```
|
|
327
|
+
|
|
328
|
+
### Python API
|
|
329
|
+
|
|
330
|
+
```python
|
|
331
|
+
from vortexa.core.indexer import CodebaseIndexer
|
|
332
|
+
|
|
333
|
+
# Use nano model
|
|
334
|
+
indexer = CodebaseIndexer(root="/path/to/project", model_id="VTXAI/vtx-embed-1M")
|
|
335
|
+
|
|
336
|
+
# Use mini model (default)
|
|
337
|
+
indexer = CodebaseIndexer(root="/path/to/project", model_id="VTXAI/vtx-embed-7M")
|
|
338
|
+
```
|
|
339
|
+
|
|
340
|
+
### Alternative Embedders
|
|
341
|
+
|
|
342
|
+
The `embedding.py` module also provides alternative embedders that can be passed directly:
|
|
343
|
+
|
|
344
|
+
- `Model2VecEmbedder` — static embeddings via Model2Vec (`AI4free/JARVIS-tool-search-v1`)
|
|
345
|
+
- `SentenceTransformerEmbedder` — Transformer-based dense embeddings (`all-MiniLM-L6-v2`)
|
|
346
|
+
- `LF4Embedder` — 4-bit quantized static embeddings (`VTXAI/Vortex-Embed-4.7M`)
|
|
347
|
+
|
|
348
|
+
---
|
|
349
|
+
|
|
350
|
+
## Inference API
|
|
351
|
+
|
|
352
|
+
vortexa can be used as a standalone embedding inference engine for VTXAI models,
|
|
353
|
+
similar to how sentence-transformers works but purpose-built for Vortex-Embed models.
|
|
354
|
+
|
|
355
|
+
### Python API
|
|
356
|
+
|
|
357
|
+
```python
|
|
358
|
+
from vortexa.core.inference import embed
|
|
359
|
+
|
|
360
|
+
# Encode a single string (default: mini model)
|
|
361
|
+
vec = embed("India is a diverse country")
|
|
362
|
+
|
|
363
|
+
# Encode multiple strings with the nano model
|
|
364
|
+
vecs = embed(["Indian cricket team is strong", "Chennai is a major city"], model="nano")
|
|
365
|
+
|
|
366
|
+
# Use any HuggingFace model ID
|
|
367
|
+
vecs = embed(["Indian agriculture output"], model="VTXAI/vtx-embed-7M")
|
|
368
|
+
|
|
369
|
+
# Shape: (1, D) for single text, (N, D) for batch
|
|
370
|
+
print(vec.shape)
|
|
371
|
+
```
|
|
372
|
+
|
|
373
|
+
### CLI
|
|
374
|
+
|
|
375
|
+
```bash
|
|
376
|
+
# Encode text with the mini model (default)
|
|
377
|
+
vortexa embed "India is a diverse country"
|
|
378
|
+
|
|
379
|
+
# Encode with the nano model
|
|
380
|
+
vortexa embed "Indian cricket team is strong" --model nano
|
|
381
|
+
|
|
382
|
+
# Encode multiple strings
|
|
383
|
+
vortexa embed "India has 28 states" "Chennai is in Tamil Nadu" --model nano
|
|
384
|
+
|
|
385
|
+
# Use a custom model ID
|
|
386
|
+
vortexa embed "Indian monsoon patterns" --model VTXAI/vtx-embed-7M
|
|
387
|
+
```
|
|
388
|
+
|
|
389
|
+
---
|
|
390
|
+
|
|
303
391
|
<div align="center">
|
|
304
392
|
|
|
305
393
|
## MCP Server
|
|
306
394
|
|
|
307
395
|
</div>
|
|
308
396
|
|
|
309
|
-
vortexa ships with a built-in **MCP (Model Context Protocol) server** that exposes codebase search as a single `search` tool. Start it with:
|
|
310
|
-
|
|
311
397
|
```bash
|
|
312
398
|
# Auto-indexes current directory, serves on stdio
|
|
313
399
|
python -m vortexa.interfaces.mcp_server
|
|
@@ -362,7 +448,7 @@ vortexa/
|
|
|
362
448
|
├── core/
|
|
363
449
|
│ ├── indexer.py # CodebaseIndexer — main orchestrator
|
|
364
450
|
│ ├── chunking.py # AST-aware (tree-sitter) + line-based chunking
|
|
365
|
-
│ ├── embedding.py # Embedding
|
|
451
|
+
│ ├── embedding.py # Embedding model wrappers (Model2Vec, SentenceTransformers, LF4)
|
|
366
452
|
│ ├── language.py # Language detection & file extension mapping
|
|
367
453
|
│ └── types.py # Shared types (Chunk, ChunkConfig, IndexStats, SearchResult, ...)
|
|
368
454
|
├── storage/
|
|
@@ -423,7 +509,7 @@ graph LR
|
|
|
423
509
|
D -->|No| F[Line-based Splitter<br/>Configurable size/overlap]
|
|
424
510
|
E --> G[Chunk Set]
|
|
425
511
|
F --> G
|
|
426
|
-
G --> H[Embedding Model<br/>Model2Vec / SentenceTransformer]
|
|
512
|
+
G --> H[Embedding Model<br/>VortexEmbedderV4 (LF4) / Model2Vec / SentenceTransformer]
|
|
427
513
|
G --> I[BM25 Tokenizer]
|
|
428
514
|
H --> J[(LMDB Vector Store)]
|
|
429
515
|
I --> K[(BM25 Index)]
|
|
@@ -495,7 +581,7 @@ graph TD
|
|
|
495
581
|
| `bm25s` | Yes | Fast BM25 keyword index and persistence |
|
|
496
582
|
| `pathspec` | Yes | `.gitignore` pattern matching in file walker |
|
|
497
583
|
| `model2vec` | Optional | Alternative static embeddings |
|
|
498
|
-
| `huggingface-hub` | Yes (default model) | Loading `VTXAI/
|
|
584
|
+
| `huggingface-hub` | Yes (default model) | Loading `VTXAI/vtx-embed-7M` (mini) or `VTXAI/vtx-embed-1M` (nano) |
|
|
499
585
|
| `tokenizers` | Yes (default model) | HF tokenizer for embedding model |
|
|
500
586
|
| `safetensors` | Yes (default model) | Safe tensor loading for 4-bit weights |
|
|
501
587
|
| `sentence-transformers` | Optional | Transformer-based dense embeddings |
|
|
@@ -0,0 +1,250 @@
|
|
|
1
|
+
"""VortexA v2 — integrated context engine.
|
|
2
|
+
|
|
3
|
+
Combines:
|
|
4
|
+
- VortexEmbedderV4 (4-bit dense retrieval, SIF+PC, on-the-fly LF4 dequant)
|
|
5
|
+
- RepoGraph (knowledge graph)
|
|
6
|
+
- VortexScore (multi-signal ranking)
|
|
7
|
+
- Context expansion (related files/tests)
|
|
8
|
+
|
|
9
|
+
The main entry point is VortexContextEngine, which takes a query and
|
|
10
|
+
returns a "context pack": the top files plus related tests, imports,
|
|
11
|
+
callers, and callees.
|
|
12
|
+
"""
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
import json
|
|
16
|
+
import logging
|
|
17
|
+
import math
|
|
18
|
+
import time
|
|
19
|
+
from dataclasses import dataclass, field
|
|
20
|
+
from pathlib import Path
|
|
21
|
+
from typing import Dict, List, Optional, Set, Tuple
|
|
22
|
+
|
|
23
|
+
import numpy as np
|
|
24
|
+
|
|
25
|
+
from vortexa.core.chunking import chunk_source
|
|
26
|
+
from vortexa.core.embedding import Embedder
|
|
27
|
+
from vortexa.core.graph import RepoGraph, RepoGraphBuilder
|
|
28
|
+
from vortexa.core.language import detect_language, get_extensions
|
|
29
|
+
from vortexa.core.v4_embedder import VortexEmbedderV4
|
|
30
|
+
from vortexa.core.vortex_score import VortexScoreWeights, vortex_score, tokenize
|
|
31
|
+
|
|
32
|
+
logger = logging.getLogger(__name__)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
@dataclass
|
|
36
|
+
class ContextFile:
|
|
37
|
+
"""A file in the context pack with its score and reason."""
|
|
38
|
+
path: str
|
|
39
|
+
score: float
|
|
40
|
+
components: Dict[str, float] = field(default_factory=dict)
|
|
41
|
+
snippet: str = ""
|
|
42
|
+
chunk_ids: List[str] = field(default_factory=list)
|
|
43
|
+
relations: Dict[str, List[str]] = field(default_factory=dict) # {relation: [related_paths]}
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
@dataclass
|
|
47
|
+
class ContextPack:
|
|
48
|
+
"""A full context pack returned by VortexA v2."""
|
|
49
|
+
query: str
|
|
50
|
+
primary: List[ContextFile] # top files
|
|
51
|
+
related: List[ContextFile] # tests, imports, callers, callees
|
|
52
|
+
graph_hits: List[str] = field(default_factory=list) # files hit via graph
|
|
53
|
+
elapsed_ms: float = 0.0
|
|
54
|
+
n_chunks_scanned: int = 0
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
class VortexContextEngine:
|
|
58
|
+
"""The V2 context engine.
|
|
59
|
+
|
|
60
|
+
Workflow:
|
|
61
|
+
1. Encode query with V3 embedder
|
|
62
|
+
2. Dense search against chunk index
|
|
63
|
+
3. Re-rank with Vortex Score (multi-signal)
|
|
64
|
+
4. Expand to related files via graph
|
|
65
|
+
5. Return ContextPack
|
|
66
|
+
|
|
67
|
+
Args:
|
|
68
|
+
embedder: VortexEmbedderV4 (or compatible Embedder)
|
|
69
|
+
chunks: List[dict] with 'path', 'content', 'start', 'end', 'chunk_id'
|
|
70
|
+
chunk_embeddings: np.ndarray (N, dim)
|
|
71
|
+
graph: Optional RepoGraph
|
|
72
|
+
file_imports: Optional dict {path: set of imported paths}
|
|
73
|
+
file_adjacency: Optional dict {path: set of directly related paths}
|
|
74
|
+
weights: VortexScoreWeights
|
|
75
|
+
"""
|
|
76
|
+
|
|
77
|
+
def __init__(
|
|
78
|
+
self,
|
|
79
|
+
embedder: Embedder,
|
|
80
|
+
chunks: List[dict],
|
|
81
|
+
chunk_embeddings: np.ndarray,
|
|
82
|
+
*,
|
|
83
|
+
graph: Optional[RepoGraph] = None,
|
|
84
|
+
file_imports: Optional[Dict[str, Set[str]]] = None,
|
|
85
|
+
file_adjacency: Optional[Dict[str, Set[str]]] = None,
|
|
86
|
+
file_idf: Optional[Dict[str, Dict[str, float]]] = None,
|
|
87
|
+
weights: Optional[VortexScoreWeights] = None,
|
|
88
|
+
):
|
|
89
|
+
self.embedder = embedder
|
|
90
|
+
self.chunks = chunks
|
|
91
|
+
self.chunk_embeddings = chunk_embeddings
|
|
92
|
+
self.graph = graph
|
|
93
|
+
self.file_imports = file_imports or {}
|
|
94
|
+
self.file_adjacency = file_adjacency or {}
|
|
95
|
+
# file_idf: {file_path: {term: idf_weight}}
|
|
96
|
+
self.file_idf = file_idf or {}
|
|
97
|
+
self.weights = weights or VortexScoreWeights()
|
|
98
|
+
# Build path -> chunks index
|
|
99
|
+
self._path_to_chunks: Dict[str, List[int]] = {}
|
|
100
|
+
for i, c in enumerate(chunks):
|
|
101
|
+
self._path_to_chunks.setdefault(c["path"], []).append(i)
|
|
102
|
+
# Build path -> symbols index
|
|
103
|
+
self._path_to_symbols: Dict[str, Set[str]] = {}
|
|
104
|
+
if self.graph is not None:
|
|
105
|
+
for nid, node in self.graph.nodes.items():
|
|
106
|
+
if node.kind != "file":
|
|
107
|
+
self._path_to_symbols.setdefault(node.path, set()).add(nid)
|
|
108
|
+
|
|
109
|
+
def search(
|
|
110
|
+
self,
|
|
111
|
+
query: str,
|
|
112
|
+
top_k_dense: int = 50,
|
|
113
|
+
top_k_final: int = 10,
|
|
114
|
+
expand_related: bool = True,
|
|
115
|
+
max_related: int = 8,
|
|
116
|
+
) -> ContextPack:
|
|
117
|
+
"""Search and return a ContextPack."""
|
|
118
|
+
t0 = time.perf_counter()
|
|
119
|
+
# 1. Dense retrieval
|
|
120
|
+
q_emb = self.embedder.embed(query)
|
|
121
|
+
scores = self.chunk_embeddings @ q_emb
|
|
122
|
+
top_idx = np.argsort(scores)[::-1][:top_k_dense]
|
|
123
|
+
|
|
124
|
+
# 2. Aggregate to file-level (best chunk per file)
|
|
125
|
+
file_best: Dict[str, Tuple[float, int]] = {}
|
|
126
|
+
for idx in top_idx:
|
|
127
|
+
path = self.chunks[idx]["path"]
|
|
128
|
+
s = float(scores[idx])
|
|
129
|
+
if path not in file_best or s > file_best[path][0]:
|
|
130
|
+
file_best[path] = (s, idx)
|
|
131
|
+
|
|
132
|
+
# 3. Build per-file file_idf (term frequency in this file)
|
|
133
|
+
# 4. Get seed files for graph expansion
|
|
134
|
+
seed_files = set(file_best.keys())
|
|
135
|
+
|
|
136
|
+
# 5. Re-rank with Vortex Score
|
|
137
|
+
scored: List[Tuple[str, float, dict, int]] = [] # (path, score, components, chunk_idx)
|
|
138
|
+
for path, (emb_score, chunk_idx) in file_best.items():
|
|
139
|
+
file_syms = self._path_to_symbols.get(path, set())
|
|
140
|
+
chunk_text = self.chunks[chunk_idx].get("content", "")
|
|
141
|
+
file_idf = self.file_idf.get(path, {})
|
|
142
|
+
v_score, comps = vortex_score(
|
|
143
|
+
emb_score, query, path,
|
|
144
|
+
file_symbols=file_syms, chunk_text=chunk_text,
|
|
145
|
+
file_idf=file_idf, seed_files=seed_files,
|
|
146
|
+
adjacency=self.file_adjacency, imports=self.file_imports,
|
|
147
|
+
weights=self.weights,
|
|
148
|
+
)
|
|
149
|
+
scored.append((path, v_score, comps, chunk_idx))
|
|
150
|
+
# Sort by Vortex score
|
|
151
|
+
scored.sort(key=lambda x: -x[1])
|
|
152
|
+
|
|
153
|
+
# 6. Build top-K primary files
|
|
154
|
+
primary = []
|
|
155
|
+
for path, vscore, comps, chunk_idx in scored[:top_k_final]:
|
|
156
|
+
cf = ContextFile(
|
|
157
|
+
path=path, score=vscore, components=comps,
|
|
158
|
+
snippet=self.chunks[chunk_idx].get("content", "")[:500],
|
|
159
|
+
chunk_ids=[self.chunks[chunk_idx].get("chunk_id", "")],
|
|
160
|
+
)
|
|
161
|
+
primary.append(cf)
|
|
162
|
+
|
|
163
|
+
# 7. Expand to related files
|
|
164
|
+
related = []
|
|
165
|
+
if expand_related:
|
|
166
|
+
related_paths: Set[str] = set()
|
|
167
|
+
# Related via graph (1-hop neighbors)
|
|
168
|
+
if self.graph is not None:
|
|
169
|
+
for cf in primary:
|
|
170
|
+
nid = f"file:{cf.path}"
|
|
171
|
+
if nid in self.graph.nodes:
|
|
172
|
+
for nb in self.graph.neighbors(nid)[:5]:
|
|
173
|
+
nb_path = nb.split("file:")[-1] if nb.startswith("file:") else None
|
|
174
|
+
if nb_path and nb_path not in seed_files and nb_path in self._path_to_chunks:
|
|
175
|
+
related_paths.add(nb_path)
|
|
176
|
+
# Related via adjacency (imports)
|
|
177
|
+
for cf in primary:
|
|
178
|
+
for nb in list(self.file_adjacency.get(cf.path, set()))[:3]:
|
|
179
|
+
if nb not in seed_files and nb in self._path_to_chunks:
|
|
180
|
+
related_paths.add(nb)
|
|
181
|
+
# Related via test files (heuristic: file path contains "test" and shares parent dir)
|
|
182
|
+
for cf in primary:
|
|
183
|
+
base = cf.path.replace(".py", "")
|
|
184
|
+
test_candidates = [
|
|
185
|
+
f"tests/test_{base.split('/')[-1]}.py",
|
|
186
|
+
f"test_{base.split('/')[-1]}.py",
|
|
187
|
+
f"{base}_test.py",
|
|
188
|
+
]
|
|
189
|
+
for tc in test_candidates:
|
|
190
|
+
if tc in self._path_to_chunks and tc not in seed_files:
|
|
191
|
+
related_paths.add(tc)
|
|
192
|
+
# Score the related files
|
|
193
|
+
related_scored = []
|
|
194
|
+
for path in related_paths:
|
|
195
|
+
# Use a lower-priority scoring (no emb contribution from seed)
|
|
196
|
+
# Find best chunk for this file
|
|
197
|
+
chunk_indices = self._path_to_chunks.get(path, [])
|
|
198
|
+
if not chunk_indices:
|
|
199
|
+
continue
|
|
200
|
+
# Use first chunk for scoring context
|
|
201
|
+
chunk_idx = chunk_indices[0]
|
|
202
|
+
v_score, comps = vortex_score(
|
|
203
|
+
0.0, query, path,
|
|
204
|
+
file_symbols=self._path_to_symbols.get(path, set()),
|
|
205
|
+
chunk_text=self.chunks[chunk_idx].get("content", ""),
|
|
206
|
+
file_idf=self.file_idf.get(path, {}),
|
|
207
|
+
seed_files=set(cf.path for cf in primary),
|
|
208
|
+
adjacency=self.file_adjacency, imports=self.file_imports,
|
|
209
|
+
weights=self.weights,
|
|
210
|
+
)
|
|
211
|
+
related_scored.append((path, v_score, comps, chunk_idx))
|
|
212
|
+
related_scored.sort(key=lambda x: -x[1])
|
|
213
|
+
for path, vscore, comps, chunk_idx in related_scored[:max_related]:
|
|
214
|
+
related.append(ContextFile(
|
|
215
|
+
path=path, score=vscore, components=comps,
|
|
216
|
+
snippet=self.chunks[chunk_idx].get("content", "")[:500],
|
|
217
|
+
chunk_ids=[self.chunks[chunk_idx].get("chunk_id", "")],
|
|
218
|
+
))
|
|
219
|
+
|
|
220
|
+
elapsed = (time.perf_counter() - t0) * 1000
|
|
221
|
+
return ContextPack(
|
|
222
|
+
query=query,
|
|
223
|
+
primary=primary,
|
|
224
|
+
related=related,
|
|
225
|
+
graph_hits=[cf.path for cf in related if cf.path not in seed_files],
|
|
226
|
+
elapsed_ms=elapsed,
|
|
227
|
+
n_chunks_scanned=len(top_idx),
|
|
228
|
+
)
|
|
229
|
+
|
|
230
|
+
def to_text(self, pack: ContextPack) -> str:
|
|
231
|
+
"""Format a context pack as text for an LLM."""
|
|
232
|
+
lines = [f"# Context for: {pack.query}", ""]
|
|
233
|
+
lines.append("## Primary files")
|
|
234
|
+
for cf in pack.primary:
|
|
235
|
+
lines.append(f"\n### {cf.path} (score: {cf.score:.3f})")
|
|
236
|
+
if cf.components:
|
|
237
|
+
comp_str = ", ".join(f"{k}={v:.2f}" for k, v in cf.components.items() if v > 0)
|
|
238
|
+
if comp_str:
|
|
239
|
+
lines.append(f" Signals: {comp_str}")
|
|
240
|
+
lines.append("```")
|
|
241
|
+
lines.append(cf.snippet)
|
|
242
|
+
lines.append("```")
|
|
243
|
+
if pack.related:
|
|
244
|
+
lines.append("\n## Related files (context expansion)")
|
|
245
|
+
for cf in pack.related:
|
|
246
|
+
lines.append(f"\n### {cf.path} (score: {cf.score:.3f})")
|
|
247
|
+
lines.append("```")
|
|
248
|
+
lines.append(cf.snippet[:300])
|
|
249
|
+
lines.append("```")
|
|
250
|
+
return "\n".join(lines)
|