How to Work with Tables

The Table class provides a structured to store and manage tabular data. It is particularly useful to store intermediate results from your LLM-powered agents. When building complex agents, you often need to track various outputs like agent responses across multiple steps. The Table class offers a simple interface to store this structured data with features like unique indexing, dynamic column addition, and multiple export formats (CSV, Markdown, etc.).

Table Basics

Creating Tables

Every table requires an index column that acts as a unique identifier for rows:

from zav.agents_sdk import Table

# Create an empty table
table = Table(index_column="id")

# Create a table with initial data
data = [
    {"id": "1", "name": "First Row", "value": "100"},
    {"id": "2", "name": "Second Row", "value": "200"}
]
table = Table(data=data, index_column="id")

Core Operations

Adding Columns

# Add a new column (all existing rows will get the default value)
table.add_column("status", default_value="pending")

Adding Rows

# Add a new row (must include the index column, other columns will get default values if not provided)
table.add_row({
    "id": "3",
    "name": "Third Row",
    "value": "300",
    "status": "active"
})

Updating Values

# Update a specific cell
table.update_cell(row_index="1", column_name="status", value="completed")

Retrieving Data

# Get all column names
columns = table.get_columns()

# Get a specific row
row = table.get_row("1")
# {"id": "1", "name": "First Row", "value": "100", "status": "completed"}

# Get all rows
rows = table.to_rows()
# [
#     {"id": "1", "name": "First Row", "value": "100", "status": "completed"},
#     {"id": "2", "name": "Second Row", "value": "200", "status": "pending"},
#     {"id": "3", "name": "Third Row", "value": "300", "status": "active"}
# ]

# Get data in columnar format
columns = table.to_columnar()
# {
#    "id": ["1", "2", "3"],
#    "name": ["First Row", "Second Row", "Third Row"],
#    "value": ["100", "200", "300"],
#    "status": ["completed", "pending", "active"]
# }

Export Formats

You can export tables in various formats. content_part is particularly useful to include to a ChatMessage.

# Export as CSV
csv_string = table.to_csv()

# Export as markdown table
markdown_string = table.to_markdown()

# Export as content part (for agent responses)
content_part = table.to_content_part(format="row")
# ContentPart(
#     table=ContentPartTable(
#         rows=[
#             {"id": "1", "name": "First Row", "value": "100", "status": "completed"},
#             {"id": "2", "name": "Second Row", "value": "200", "status": "pending"},
#             {"id": "3", "name": "Third Row", "value": "300", "status": "active"},
#         ],
#         headers=["id", "name", "value", "status"],
#     ),
#     content_type="table",
# )

Practical Example: Multi-Document Question-Answering

Prerequisites

Before you begin, make sure you have completed the Getting Started with the Agents SDK and Building and Running Your First RAG Agent tutorials.

Example

Here's a complete example of using tables to track RAG agent outputs:

import json
from typing import Dict, List, Optional

from zav.agents_sdk import (
    ChatAgent,
    ChatAgentClassRegistry,
    ChatMessage,
    ChatMessageSender,
    Table,
)
from zav.agents_sdk.adapters import ZAVChatCompletionClient, ZAVRetriever
from zav.agents_sdk.domain.table import escape_markdown


@ChatAgentClassRegistry.register()
class MultiDocQAAgent(ChatAgent):
    agent_name = "multi_doc_qa_agent"

    def __init__(
        self,
        retriever: ZAVRetriever,
        client: ZAVChatCompletionClient,
    ):
        self.retriever = retriever
        self.client = client

    async def answer_question(self, query: str, document: Dict) -> str:
        prompt = f"""Write a concise answer to the question based on these documents:

        Question: {query}

        Context: {json.dumps(document, indent=2)}
        """

        response = await self.client.complete(
            messages=[ChatMessage(sender=ChatMessageSender.USER, content=prompt)],
            max_tokens=2048,
        )
        return response.chat_completion.content if response.chat_completion else ""

    async def execute(self, conversation: List[ChatMessage]) -> Optional[ChatMessage]:
        query = conversation[-1].content

        # Initialize tracking table
        table = Table(index_column="title")
        table.add_column("date")
        table.add_column("document_score")
        table.add_column(query)

        # Retrieve relevant documents
        search_result = await self.retriever.search(
            query_string=query,
            retrieval_unit="document",
            page_size=5,
            document_types=["document"],
            collapse=None,
        )
        documents = search_result.get("hits", [])

        # Answer the question for each document and update the table
        for doc in documents:
            title = (
                doc.get("custom_metadata", {})
                .get("metadata", {})
                .get("title", "Unknown")
            )
            table.add_row(
                {
                    "title": title,
                    "date": doc.get("custom_metadata", {})
                    .get("metadata", {})
                    .get("date", "Unknown"),
                    "document_score": doc.get("score"),
                }
            )
            answer = await self.answer_question(query, doc)
            table.update_cell(title, query, escape_markdown(answer))

        # Return the table as a markdown message and the raw table as a content part
        return ChatMessage(
            sender=ChatMessageSender.BOT,
            content=table.to_markdown(),
            content_parts=[table.to_content_part()],
        )

This is the table it generates for the question what are some alternatives to the transformer architecture?:

title	date	document_score	what are some alternatives to the transformer architecture?
Swin Transformer: Hierarchical Vision Transformer using Shifted Windows	2021	1.89319	Some alternatives to the transformer architecture include the Swin Transformer, which is designed as a hierarchical vision transformer that utilizes shifted windows for efficient computation. This architecture allows for flexibility in modeling at various scales and maintains linear computational complexity relative to image size, making it suitable for a range of vision tasks. Other alternatives may include various convolutional neural networks (CNNs) and hybrid models that combine CNNs with transformer elements.
Swin Transformer: Hierarchical Vision Transformer Using Shifted Windows	2021	1.85141	Some alternatives to the transformer architecture include the Swin Transformer, which is designed for computer vision tasks. It features a hierarchical structure that utilizes shifted windows for efficient self-attention computation, allowing it to model various scales with linear complexity relative to image size. This architecture has shown superior performance in tasks like image classification and object detection compared to traditional transformers.
RoFormer: Enhanced Transformer with Rotary Position Embedding	2021	1.83838	Some alternatives to the transformer architecture include the RoFormer, which utilizes Rotary Position Embedding (RoPE) to enhance positional information in self-attention mechanisms. Other alternatives may involve variations of attention mechanisms, such as linear attention models like Performer, or different architectures that incorporate relative position encodings.
Exploring Plain Vision Transformer Backbones for Object Detection	2022	1.82201	Some alternatives to the transformer architecture include plain, non-hierarchical Vision Transformers (ViT) and other backbone networks designed for object detection. The ViT architecture can be fine-tuned for specific tasks without needing a hierarchical backbone, achieving competitive results with minimal adaptations. Additionally, methods like Masked Autoencoders (MAE) can be used for pre-training these architectures, allowing them to compete with traditional hierarchical methods.
SwinIR: Image Restoration Using Swin Transformer	2021	1.81191	Some alternatives to the transformer architecture include convolutional neural networks (CNNs), which have traditionally been the primary models for tasks like image restoration. Additionally, variations of transformer models, such as the Swin Transformer, are being explored for specific applications like image restoration, super-resolution, and denoising.