Custom Chunking

from pathlib import Path
from typing import List

from agno.agent import Agent
from agno.knowledge.chunking.strategy import ChunkingStrategy
from agno.knowledge.document import Document
from agno.knowledge.knowledge import Knowledge
from agno.knowledge.reader.text_reader import TextReader
from agno.vectordb.pgvector import PgVector


class CustomChunking(ChunkingStrategy):
    def __init__(self, separator: str = "---"):
        self.separator = separator

    def chunk(self, document: Document) -> List[Document]:
        result = []
        for chunk_content in document.content.split(self.separator):
            chunk_content = self.clean_text(chunk_content).strip()
            if not chunk_content:
                continue

            chunk_number = len(result) + 1
            meta_data = document.meta_data.copy()
            meta_data["chunk"] = chunk_number
            meta_data["chunk_size"] = len(chunk_content)
            result.append(
                Document(
                    id=self._generate_chunk_id(
                        document,
                        chunk_number,
                        chunk_content,
                    ),
                    name=document.name,
                    meta_data=meta_data,
                    content=chunk_content,
                )
            )
        return result


db_url = "postgresql+psycopg://ai:ai@localhost:5532/ai"

knowledge = Knowledge(
    vector_db=PgVector(table_name="recipes_custom_chunking", db_url=db_url),
)

content_path = Path("tmp/recipes.txt")
content_path.parent.mkdir(parents=True, exist_ok=True)
content_path.write_text(
    "Tom kha gai uses coconut milk.\n---\nMassaman curry uses warm spices.",
    encoding="utf-8",
)

knowledge.insert(
    path=str(content_path),
    reader=TextReader(
        name="Custom Chunking Reader",
        chunking_strategy=CustomChunking(separator="---"),
    ),
)

agent = Agent(
    knowledge=knowledge,
    search_knowledge=True,
)

agent.print_response("Which recipe uses coconut milk?", markdown=True)