Code Chunking

1from kern.agent import Agent
2from kern.knowledge.chunking.code import CodeChunking
3from kern.knowledge.knowledge import Knowledge
4from kern.knowledge.reader.text_reader import TextReader
5from kern.vectordb.pgvector import PgVector
6
7db_url = "postgresql+psycopg://ai:ai@localhost:5532/ai"
8
9knowledge = Knowledge(
10    vector_db=PgVector(table_name="python_code_chunking", db_url=db_url),
11)
12
13knowledge.insert(
14    url="https://raw.githubusercontent.com/kern-agi/kern/main/libs/kern/kern/session/workflow.py",
15    reader=TextReader(
16        chunking_strategy=CodeChunking(
17            tokenizer="gpt2",
18            chunk_size=500,
19            language="python",
20        ),
21    ),
22)
23
24agent = Agent(knowledge=knowledge, search_knowledge=True)
25agent.print_response("How does the Workflow class work?", markdown=True)

1from typing import Sequence
2
3from kern.agent import Agent
4from kern.knowledge.chunking.code import CodeChunking
5from kern.knowledge.knowledge import Knowledge
6from kern.knowledge.reader.text_reader import TextReader
7from kern.vectordb.pgvector import PgVector
8from chonkie.tokenizer import Tokenizer
9
10db_url = "postgresql+psycopg://ai:ai@localhost:5532/ai"
11
12
13class LineTokenizer(Tokenizer):
14    """Custom tokenizer that counts lines of code."""
15
16    def __init__(self):
17        self.vocab = []
18        self.token2id = {}
19
20    def __repr__(self) -> str:
21        return f"LineTokenizer(vocab_size={len(self.vocab)})"
22
23    def tokenize(self, text: str) -> Sequence[str]:
24        if not text:
25            return []
26        return text.split("\n")
27
28    def encode(self, text: str) -> Sequence[int]:
29        encoded = []
30        for token in self.tokenize(text):
31            if token not in self.token2id:
32                self.token2id[token] = len(self.vocab)
33                self.vocab.append(token)
34            encoded.append(self.token2id[token])
35        return encoded
36
37    def decode(self, tokens: Sequence[int]) -> str:
38        try:
39            return "\n".join([self.vocab[token] for token in tokens])
40        except Exception as e:
41            raise ValueError(
42                f"Decoding failed. Tokens: {tokens} not found in vocab."
43            ) from e
44
45    def count_tokens(self, text: str) -> int:
46        if not text:
47            return 0
48        return len(text.split("\n"))
49
50
51knowledge = Knowledge(
52    vector_db=PgVector(table_name="code_custom_tokenizer", db_url=db_url),
53)
54
55knowledge.insert(
56    url="https://raw.githubusercontent.com/kern-agi/kern/main/libs/kern/kern/session/workflow.py",
57    reader=TextReader(
58        chunking_strategy=CodeChunking(
59            tokenizer=LineTokenizer(),
60            chunk_size=500,
61            language="python",
62        ),
63    ),
64)
65
66agent = Agent(knowledge=knowledge, search_knowledge=True)
67agent.print_response("How does the Workflow class work?", markdown=True)

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

1uv venv --python 3.12
2.venv\Scripts\activate