Skip to content

rigging.watchers

Common watcher callback makers for use with generators, chats, and completions.

write_chats_to_elastic(client: AsyncElasticsearch, index: str, *, create_index: bool = True, **kwargs: t.Any) -> WatchChatCallback #

Create a watcher to write each chat to an ElasticSearch index.

Parameters:

  • client (AsyncElasticsearch) –

    The AsyncElasticSearch client to use.

  • index (str) –

    The index to write to.

  • create_index (bool, default: True ) –

    Whether to create the index if it doesn't exist and update its mapping.

  • kwargs (Any, default: {} ) –

    Additional keyword arguments to be passed to the Elasticsearch client.

Returns:

Source code in rigging/watchers.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def write_chats_to_elastic(
    client: AsyncElasticsearch, index: str, *, create_index: bool = True, **kwargs: t.Any
) -> WatchChatCallback:
    """
    Create a watcher to write each chat to an ElasticSearch index.

    Args:
        client: The AsyncElasticSearch client to use.
        index: The index to write to.
        create_index: Whether to create the index if it doesn't exist and update its mapping.
        kwargs: Additional keyword arguments to be passed to the Elasticsearch client.

    Returns:
        A callback to use in [rigging.chat.ChatPipeline.watch][]
        or [rigging.generator.Generator.watch][].
    """

    async def _write_chats_to_elastic(chats: list[Chat]) -> None:
        await chats_to_elastic(chats, index, client, create_index=create_index, **kwargs)

    return _write_chats_to_elastic

write_chats_to_jsonl(file: str | Path, *, replace: bool = False) -> WatchChatCallback #

Create a watcher to write each chat as a single JSON line appended to a file.

Parameters:

  • file (str | Path) –

    The file to write to.

  • replace (bool, default: False ) –

    If the file should be replaced if it already exists.

Returns:

Source code in rigging/watchers.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def write_chats_to_jsonl(file: str | Path, *, replace: bool = False) -> WatchChatCallback:
    """
    Create a watcher to write each chat as a single JSON line appended to a file.

    Args:
        file: The file to write to.
        replace: If the file should be replaced if it already exists.

    Returns:
        A callback to use in [rigging.chat.ChatPipeline.watch][]
        or [rigging.generator.Generator.watch][].
    """

    # We could do the file delete externally, but I don't
    # want to produce a side effect of simply creating the
    # callback.

    file = Path(file)
    replaced: bool = False

    async def _write_chats_to_jsonl(chats: list[Chat]) -> None:
        nonlocal replaced

        if file.exists() and replace and not replaced:
            os.remove(file)
            replaced = True

        with file.open("a") as f:
            for chat in chats:
                f.write(chat.model_dump_json() + "\n")

    return _write_chats_to_jsonl

write_messages_to_jsonl(file: str | Path, *, replace: bool = False) -> WatchChatCallback #

Create a watcher to flatten chats to individual messages (like Dataframes) and append to a file.

Parameters:

  • file (str | Path) –

    The file to write to.

  • replace (bool, default: False ) –

    If the file should be replaced if it already exists.

Returns:

Source code in rigging/watchers.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def write_messages_to_jsonl(file: str | Path, *, replace: bool = False) -> WatchChatCallback:
    """
    Create a watcher to flatten chats to individual messages (like Dataframes) and append to a file.

    Args:
        file: The file to write to.
        replace: If the file should be replaced if it already exists.

    Returns:
        A callback to use in [rigging.chat.ChatPipeline.watch][]
        or [rigging.generator.Generator.watch][].
    """
    file = Path(file)
    replaced: bool = False

    async def _write_messages_to_jsonl(chats: list[Chat]) -> None:
        nonlocal replaced

        if file.exists() and replace and not replaced:
            os.remove(file)
            replaced = True

        with file.open("a") as f:
            for chat in flatten_chats(chats):
                f.write(json.dumps(chat) + "\n")

    return _write_messages_to_jsonl