Code Chunking

Code chunking splits code based on its structure, leveraging Abstract Syntax Trees (ASTs) to create contextually relevant segments. It uses the Chonkie library to identify natural code boundaries like functions, classes, and blocks. Learn more about code chunking. This preserves code semantics better than fixed-size chunking by ensuring related code stays together in the same chunk, while splitting occurs at meaningful structural boundaries.

Code chunking supports several built-in tokenizers or a custom Tokenizer instance.

1. Create a Python file

   Code Chunking

   from agno.agent import Agent
   from agno.knowledge.chunking.code import CodeChunking
   from agno.knowledge.knowledge import Knowledge
   from agno.knowledge.reader.text_reader import TextReader
   from agno.vectordb.pgvector import PgVector
   
   db_url = "postgresql+psycopg://ai:ai@localhost:5532/ai"
   
   knowledge = Knowledge(
       vector_db=PgVector(table_name="python_code_chunking", db_url=db_url),
   )
   
   knowledge.insert(
       url="https://raw.githubusercontent.com/agno-agi/agno/main/libs/agno/agno/session/workflow.py",
       reader=TextReader(
           chunking_strategy=CodeChunking(
               tokenizer="gpt2",
               chunk_size=500,
               language="python",
           ),
       ),
   )
   
   agent = Agent(knowledge=knowledge, search_knowledge=True)
   agent.print_response("How does the Workflow class work?", markdown=True)
   

   from typing import Sequence
   
   from agno.agent import Agent
   from agno.knowledge.chunking.code import CodeChunking
   from agno.knowledge.knowledge import Knowledge
   from agno.knowledge.reader.text_reader import TextReader
   from agno.vectordb.pgvector import PgVector
   from chonkie.tokenizer import Tokenizer
   
   db_url = "postgresql+psycopg://ai:ai@localhost:5532/ai"
   
   
   class LineTokenizer(Tokenizer):
       """Custom tokenizer that counts lines of code."""
   
       def __init__(self):
           self.vocab = []
           self.token2id = {}
   
       def __repr__(self) -> str:
           return f"LineTokenizer(vocab_size={len(self.vocab)})"
   
       def tokenize(self, text: str) -> Sequence[str]:
           if not text:
               return []
           return text.split("\n")
   
       def encode(self, text: str) -> Sequence[int]:
           encoded = []
           for token in self.tokenize(text):
               if token not in self.token2id:
                   self.token2id[token] = len(self.vocab)
                   self.vocab.append(token)
               encoded.append(self.token2id[token])
           return encoded
   
       def decode(self, tokens: Sequence[int]) -> str:
           try:
               return "\n".join([self.vocab[token] for token in tokens])
           except Exception as e:
               raise ValueError(
                   f"Decoding failed. Tokens: {tokens} not found in vocab."
               ) from e
   
       def count_tokens(self, text: str) -> int:
           if not text:
               return 0
           return len(text.split("\n"))
   
   
   knowledge = Knowledge(
       vector_db=PgVector(table_name="code_custom_tokenizer", db_url=db_url),
   )
   
   knowledge.insert(
       url="https://raw.githubusercontent.com/agno-agi/agno/main/libs/agno/agno/session/workflow.py",
       reader=TextReader(
           chunking_strategy=CodeChunking(
               tokenizer=LineTokenizer(),
               chunk_size=500,
               language="python",
           ),
       ),
   )
   
   agent = Agent(knowledge=knowledge, search_knowledge=True)
   agent.print_response("How does the Workflow class work?", markdown=True)
   
2. Set up your virtual environment

   Mac

   uv venv --python 3.12
   source .venv/bin/activate
   

   uv venv --python 3.12
   .venv\Scripts\activate
   
3. Install dependencies

   uv pip install -U agno sqlalchemy psycopg pgvector "chonkie[code]" openai
   
4. Run PgVector

   docker run -d \
     -e POSTGRES_DB=ai \
     -e POSTGRES_USER=ai \
     -e POSTGRES_PASSWORD=ai \
     -e PGDATA=/var/lib/postgresql/data/pgdata \
     -v pgvolume:/var/lib/postgresql/data \
     -p 5532:5432 \
     --name pgvector \
     agno/pgvector:16
   
5. Run the script

   python code_chunking.py
   

Code Chunking Params

Parameter	Type	Default	Description
tokenizer	Union[str, TokenizerProtocol]	"character"	The tokenizer for measuring chunk sizes. Supports several built-in tokenizers or a custom Tokenizer instance.
chunk_size	int	2048	Maximum size of each chunk in tokens (based on the selected tokenizer).
language	Union[Literal["auto"], Any]	"auto"	The programming language to parse. Use "auto" for automatic detection or specify a tree-sitter language name (e.g., "python", "javascript", "go", "rust").
include_nodes	bool	False	Whether to include AST nodes. Note: Chonkie's base Chunk type does not store node information.
chunker_params	Optional[Dict[str, Any]]	None	Additional parameters to pass directly to Chonkie's CodeChunker.