rigging

repository·main·Indexed 19 days ago

https://github.com/dreadnode/rigging

A lightweight, production-ready LLM interaction framework for Python. It simplifies using language models via LiteLLM, supporting structured Pydantic outputs, async integration, and type hints. Key features include the @rg.prompt decorator for defining prompts as functions, ChatPipeline for orchestrating conversations, and a robust Chat object for managing message history, metadata, and exports to formats like Pandas DataFrames and Elasticsearch.

Tokens
98.4K
Snippets
366
Records
419
Agent score
64%

What's inside rigging

  1. Tokenize chats with rigging.tokenizer

    main
    The rigging.tokenizer module provides tools to encode chat conversations into token sequences suitable for LLM training or inference. It supports structured tokenization where a Chat object is transformed into a TokenizedChat, which includes specific TokenSlice objects mapping tokens back to original messages, tool calls, or other parts of the chat.
  2. Get started with Rigging

    main

    Rigging is a lightweight LLM framework designed for production code. It allows you to use language models with structured Pydantic models, LiteLLM-powered generators, and Pythonic prompt functions.

    Key features include:

    • Structured Output: Use Pydantic models for type-safe responses.
    • LiteLLM Integration: Instant access to many models via connection strings.
    • Prompt Functions: Define prompts as Python functions with type hints and docstrings.
    • Tool Use: Simple tool integration even for models without native API support.
    • Async Support: Fully asynchronous execution for high-performance generation.
  3. What is Message Slicing and how does it work?

    main

    Message slicing allows you to mark specific text ranges within messages and attach structured data, types, and metadata to them. These slices act as "smart bookmarks" that maintain their positions and context even as message content evolves through processing pipelines.

    Unlike simple text processing, slices establish a bidirectional relationship between Python objects and their exact character locations in the message. This enables:

    • Tracking tool calls through transforms.
    • Preserving training signals (like rewards) during tokenization.
    • Managing overlapping annotations.
    • Performing content manipulation while maintaining data integrity.

    Every slice contains:

    • Text range: start and stop positions.
    • Type: e.g., "model", "tool_call", or "other".
    • Associated object: An optional parsed rg.Model or ToolCall.
    • Metadata: A dictionary for additional context (IDs, confidence scores, etc.).
  4. Define the types of content in a Message

    main

    A Message can contain various types of content, which are represented by the Content type. The supported content types are:

    • ContentText: Plain text content.
    • ContentImageUrl: Image content provided via a URL (supports base64-encoded data URLs).
    • ContentAudioInput: Audio content provided as input.
    Content = ContentText | ContentImageUrl | ContentAudioInput
  5. Configure XML structure using attributes and elements

    main

    When a model has multiple fields, you must explicitly define whether each field should be an XML attribute or a child element.

    Best Practice: LLMs generally perform better when data is structured as elements rather than attributes.

    Use the following helpers to define the structure:

    • rg.attr(): Defines the field as an XML attribute.
    • rg.element(): Defines the field as a child XML element.
    • rg.wrapped(tag_name, element_definition): Wraps a list of elements inside a parent tag.
    import rigging as rg
    
    class Person(rg.Model):
        # name and age will be child elements
        name: str = rg.element()
        age: int = rg.element()
    
    # Resulting XML:
    # <person>
    #     <name>Will</name>
    #     <age>30</age>
    # </person>
  6. Define structured data models with `rg.Model`

    main

    Rigging uses XML to weave unstructured text with structured data. To define a schema for the LLM to follow, inherit from rigging.Model. This class is a wrapper around pydantic-xml's BaseXMLModel.

    By default, tag names are auto-generated from the class name. You can override the XML tag by passing the tag argument to the class definition.

    To guide the LLM, use the .xml_example() class method to generate the expected XML tag structure for your prompt.

    import rigging as rg
    
    # Standard model
    class FunFact(rg.Model):
        fact: str
    
    # Model with a custom XML tag
    class LongNameForThing(rg.Model, tag="short"):
        data: str
    
    # Usage in a prompt
    # Generates: "Provide a fun fact between <fun-fact></fun-fact> tags."
    prompt_text = f"Provide a fun fact between {FunFact.xml_example()} tags."
  7. Understand generation completion reasons with StopReason

    main

    The StopReason type indicates why a model stopped generating text. Possible values include:

    • stop: Normal completion.
    • length: Reached the max_tokens limit.
    • content_filter: Stopped due to safety/content filters.
    • tool_calls: Stopped to perform a tool call.
    • unknown: The reason could not be determined.
    StopReason = Literal[
        "stop",
        "length",
        "content_filter",
        "tool_calls",
        "unknown",
    ]
  8. Manage state and hooks in HTTPGenerator

    main

    Stateful Generators

    Use the .state dictionary to store dynamic information (like session IDs or temporary credentials) that can be accessed within your request templates via {{ state.key_name }}.

    Using Hooks for dynamic updates

    You can provide an async hook function to an HTTPGenerator. The hook is called after every HTTP request. It can inspect the httpx.Response and return a HookAction (e.g., "retry" or "continue"). This is useful for handling expiring credentials: if a 401 error occurs, the hook can update generator.state and return "retry" to transparently attempt the request again with the new state.

    import rigging as rg
    import httpx
    
    async def refresh_token_hook(generator: rg.generator.HTTPGenerator, response: httpx.Response) -> rg.generator.HookAction:
        if response.status_code == 401:
            # Logic to get new token
            new_token = "new-super-secret-token"
            generator.state["access_token"] = new_token
            return "retry"
        return "continue"
    
    stateful_api = rg.HTTPGenerator.for_json_endpoint(
        url="https://api.secure.com/v1/live/chat",
        hook=refresh_token_hook,
        state={"access_token": "initial-token"},
        auth={"header": "Authorization", "format": "Bearer {{ state.access_token }}"},
        request_body={"query": "$content"},
        response={"content_path": "$.data.result"}
    )
  9. Understand TokenSlice for mapping tokens to objects

    main

    A TokenSlice represents a specific segment of the token list. It is used to map tokens back to their source components (e.g., a specific message or tool call).

    Fields:

    • start: Starting index in the token list.
    • end: Ending index in the token list.
    • type: The type of slice (e.g., message, tool_call).
    • obj: The original object (like a Message) this slice corresponds to.
    • metadata: A dictionary of additional metadata associated with this slice.
  10. Migrate from v1.x to v2.x: Renamed Pipeline Classes

    main

    To better reflect the capabilities of the system, several class names were updated in v2.x:

    • PendingChat is now ChatPipeline.
    • PendingCompletion is now CompletionPipeline.

    Most users will not be affected unless they were manually instantiating or type-checking these specific classes.

  11. Understand PipelineStep

    main

    A PipelineStep represents an intermediate state during the generation process within a pipeline. It provides context about where the current execution sits in the hierarchy and what data is available.

    Key Attributes:

    • state: The current PipelineState of the generation.
    • chats: The ChatList associated with this specific step.
    • depth: An int representing the depth of this step in the pipeline tree (useful for recursion constraints).
    • parent: The PipelineStep that is running above this one.
    • pipeline: The ChatPipeline instance associated with this step.
    • callback: The associated ThenChatCallback or MapChatCallback if the state is 'callback'.
  12. Use watch callbacks to monitor pipelines

    main

    Watch callbacks are passive listeners used for logging, monitoring, or triggering external actions without modifying the pipeline's execution. They receive Chat or Completion objects as they are finalized. You can register multiple watchers using the .watch() method on Generator, ChatPipeline, CompletionPipeline, or Prompt objects.

    Rigging provides pre-built watchers in the rigging.watchers module, such as write_chats_to_jsonl.

    import rigging as rg
    
    # Use a pre-built watcher to log chats to a file
    log_to_file = rg.watchers.write_chats_to_jsonl("chats.jsonl")
    
    # Define a custom async watch callback
    async def print_chat_ids(chats: list[rg.Chat]) -> None:
        print(f"Watched {len(chats)} chats: {[chat.uuid for chat in chats]}")
    
    pipeline = (
        rg.get_generator("openai/gpt-4o-mini")
        .chat("Explain why the sky is blue")
        .watch(log_to_file, print_chat_ids) # Register multiple watchers
    )
    
    # Watchers will be called during the run_many execution
    chats = await pipeline.run_many(3)