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.
Files changed (37) hide show
  1. {vortexa-0.2.0 → vortexa-0.3.1}/PKG-INFO +93 -7
  2. {vortexa-0.2.0 → vortexa-0.3.1}/README.md +92 -6
  3. {vortexa-0.2.0 → vortexa-0.3.1}/pyproject.toml +1 -1
  4. vortexa-0.3.1/src/vortexa/core/context_engine.py +250 -0
  5. vortexa-0.3.1/src/vortexa/core/graph.py +246 -0
  6. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/indexer.py +4 -3
  7. vortexa-0.3.1/src/vortexa/core/inference.py +248 -0
  8. vortexa-0.3.1/src/vortexa/core/lf4_v4_model.py +258 -0
  9. vortexa-0.3.1/src/vortexa/core/v4_embedder.py +150 -0
  10. vortexa-0.3.1/src/vortexa/core/vortex_score.py +208 -0
  11. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/interfaces/cli.py +63 -1
  12. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/PKG-INFO +93 -7
  13. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/SOURCES.txt +6 -0
  14. {vortexa-0.2.0 → vortexa-0.3.1}/LICENSE +0 -0
  15. {vortexa-0.2.0 → vortexa-0.3.1}/setup.cfg +0 -0
  16. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/__init__.py +0 -0
  17. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/__init__.py +0 -0
  18. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/chunking.py +0 -0
  19. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/embedding.py +0 -0
  20. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/language.py +0 -0
  21. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/lf4_model.py +0 -0
  22. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/core/types.py +0 -0
  23. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/interfaces/__init__.py +0 -0
  24. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/interfaces/mcp_server.py +0 -0
  25. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/interfaces/watcher.py +0 -0
  26. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/search/__init__.py +0 -0
  27. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/search/ranking.py +0 -0
  28. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/search/search.py +0 -0
  29. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/search/tokens.py +0 -0
  30. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/storage/__init__.py +0 -0
  31. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/storage/bm25.py +0 -0
  32. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/storage/vector_store.py +0 -0
  33. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa/storage/walker.py +0 -0
  34. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/dependency_links.txt +0 -0
  35. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/entry_points.txt +0 -0
  36. {vortexa-0.2.0 → vortexa-0.3.1}/src/vortexa.egg-info/requires.txt +0 -0
  37. {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.2.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 static or learned embeddings (Model2Vec / SentenceTransformers)
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 models (Model2Vec, SentenceTransformers)
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/Vortex-Embed-4.7M` |
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 static or learned embeddings (Model2Vec / SentenceTransformers)
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 models (Model2Vec, SentenceTransformers)
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/Vortex-Embed-4.7M` |
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 |
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "vortexa"
7
- version = "0.2.0"
7
+ version = "0.3.1"
8
8
  description = "Codebase indexing and semantic search engine"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10"
@@ -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)