Coverage for src/lilbee/data/ingest/code.py: 100%

1"""Code-file ingestion via tree-sitter chunking.""" 

2 

3from __future__ import annotations 

4 

5from pathlib import Path 

6 

7from lilbee.app.services import get_services 

8from lilbee.data.extract.chunk import enforce_chunk_limit 

9from lilbee.data.extract.code_chunker import CodeChunk, chunk_code 

10from lilbee.data.store import ChunkType 

11from lilbee.data.types import ChunkRecord 

12from lilbee.runtime.progress import DetailedProgressCallback, noop_callback 

13 

14 

15def ingest_code_sync( 

16 path: Path, 

17 source_name: str, 

18 on_progress: DetailedProgressCallback = noop_callback, 

19) -> list[ChunkRecord]: 

20 """Parse code with tree-sitter, chunk, embed, and return store-ready records.""" 

21 code_chunks: list[CodeChunk] = chunk_code(path, source_name=source_name) 

22 if not code_chunks: 

23 return [] 

24 

25 enforce_chunk_limit(len(code_chunks)) 

26 texts = [cc.chunk for cc in code_chunks] 

27 embedder = get_services().embedder 

28 vectors = embedder.embed_batch(texts, source=source_name, on_progress=on_progress) 

29 

30 return [ 

31 ChunkRecord( 

32 source=source_name, 

33 content_type="code", 

34 chunk_type=ChunkType.RAW, 

35 page_start=0, 

36 page_end=0, 

37 line_start=cc.line_start, 

38 line_end=cc.line_end, 

39 chunk=cc.chunk, 

40 chunk_index=cc.chunk_index, 

41 vector=vec, 

42 ) 

43 for cc, vec in zip(code_chunks, vectors, strict=True) 

44 ]