Skip to content

wandb.weave

Integration with Weave from Weights & Biases

with_weave(cls)

Wraps base classes to automatically use weave.

Supported base classes: BaseCall, BaseExtractor, BaseVectorStore, BaseChunker, BaseEmbedder

Example:

import weave

from mirascope.openai import OpenAICall
from mirascope.wandb import with_weave

weave.init("my-project")


@with_weave
class BookRecommender(OpenAICall):
    prompt_template = "Please recommend some {genre} books"

    genre: str


recommender = BookRecommender(genre="fantasy")
response = recommender.call()  # this will automatically get logged with weave
print(response.content)
Source code in mirascope/wandb/weave.py 
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 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
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def with_weave(cls):
    """Wraps base classes to automatically use weave.

    Supported base classes: `BaseCall`, `BaseExtractor`, `BaseVectorStore`,
    `BaseChunker`, `BaseEmbedder`

    Example:

    ```python
    import weave

    from mirascope.openai import OpenAICall
    from mirascope.wandb import with_weave

    weave.init("my-project")


    @with_weave
    class BookRecommender(OpenAICall):
        prompt_template = "Please recommend some {genre} books"

        genre: str


    recommender = BookRecommender(genre="fantasy")
    response = recommender.call()  # this will automatically get logged with weave
    print(response.content)
    ```
    """
    if hasattr(cls, "call"):
        setattr(cls, "call", weave.op()(cls.call))
    if hasattr(cls, "call_async"):
        setattr(cls, "call_async", weave.op()(cls.call_async))

    # VectorStore
    if hasattr(cls, "retrieve"):
        setattr(cls, "retrieve", weave.op()(cls.retrieve))
    if hasattr(cls, "add"):
        setattr(cls, "add", weave.op()(cls.add))

    # Chunker
    if hasattr(cls, "chunk"):
        setattr(cls, "chunk", weave.op()(cls.chunk))

    # Embedder
    if hasattr(cls, "embed"):
        setattr(cls, "embed", weave.op()(cls.embed))
    if hasattr(cls, "embed_async"):
        setattr(cls, "embed_async", weave.op()(cls.embed_async))

    # It appears Weave does not yet support streaming or does it in a different way? :(
    # Our calls will be tracked, but the sub-calls don't since the streaming happens
    # when iterating through the generator after the call.
    if hasattr(cls, "stream"):
        setattr(cls, "stream", weave.op()(cls.stream))
    if hasattr(cls, "stream_async"):
        setattr(cls, "stream_async", weave.op()(cls.stream_async))

    if hasattr(cls, "extract"):
        setattr(cls, "extract", weave.op()(cls.extract))
    if hasattr(cls, "extract_async"):
        setattr(cls, "extract_async", weave.op()(cls.extract_async))

    if hasattr(cls, "call_params"):
        cls.call_params.weave = weave.op()
    if hasattr(cls, "vectorstore_params"):
        cls.vectorstore_params.weave = weave.op()
    return cls