Merge branch 'dev'

README.md

@@ -13,7 +13,7 @@ Rigging is a lightweight LLM interaction framework built on Pydantic XML. The go
 import rigging as rg
 from rigging.model import CommaDelimitedAnswer as Answer
 
-answer = rg.get_generator('gpt-4') \
+chat = rg.get_generator('gpt-4') \
     .chat(f"Give me 3 famous authors between {Answer.xml_tags()} tags.") \
     .until_parsed_as(Answer) \
     .run()

docs/index.md

@@ -11,7 +11,7 @@ Rigging is a lightweight LLM interaction framework built on Pydantic XML. The go
 import rigging as rg
 from rigging.model import CommaDelimitedAnswer as Answer
 
-answer = rg.get_generator('gpt-4') \
+chat = rg.get_generator('gpt-4') \
     .chat(f"Give me 3 famous authors between {Answer.xml_tags()} tags.") \
     .until_parsed_as(Answer) \
     .run()
@@ -310,10 +310,13 @@ and we have a few options:
 
 === "Option 2 - Until"
 
-    ```py hl_lines="3"
-    chat = rg.get_generator('gpt-3.5-turbo').chat(
-        f"Provide a fun fact between {FunFact.xml_example()} tags."
-    ).until_parsed_as(FunFact).run()
+    ```py hl_lines="4"
+    chat = (
+        rg.get_generator('gpt-3.5-turbo')
+        .chat(f"Provide a fun fact between {FunFact.xml_example()} tags.")
+        .until_parsed_as(FunFact)
+        .run()
+    )
 
     fun_fact = chat.last.parse(FunFact) # This call should never fail
 
@@ -357,4 +360,8 @@ Assuming we wanted to extend our example to produce a set of interesting facts,
 
     for fun_fact in chat.last.parse_set(FunFact):
         print(fun_fact.fact)
-    ```
+    ```
+
+### Keep Going
+
+Check out the **[topics section](topics/workflow.md)** for more in-depth explanations and examples.

docs/topics/callbacks-and-mapping.md

@@ -35,6 +35,9 @@ print(chat.conversation)
 # [assistant]: <joke>Why did the duck go to the doctor? Because he was feeling a little down!</joke>
 # [user]: Please include a cat in your joke
 # [assistant]: <joke>Why was the cat sitting on the computer? Because it wanted to keep an eye on the mouse!</joke>
+
+print(chat.last.parse(Joke))
+# Joke(content='Why was the cat sitting on the computer? Because it wanted to keep an eye on the mouse!')
 ```
 
 1. Returning `True` from this callback tells Rigging to go back to the generator with the supplied

docs/topics/chats-and-messages.md

@@ -133,7 +133,7 @@ meaning = rg.get_generator("claude-2.1").chat([
     },
 ]).run()
 
-# Gracefully handle mising models
+# Gracefully handle missing models
 reasoning = meaning.last.try_parse(Reasoning)
 if reasoning:
     print("Reasoning:", reasoning.content)
@@ -143,7 +143,7 @@ if reasoning:
 without_reasons = meaning.strip(Reasoning)
 print("Meaning of life:", without_reasons.last.content)
 
-# follow_up = without_thoughts.continue_(...)
+follow_up = without_reasons.continue_(...).run()
 ```
 
 ## Metadata
@@ -155,4 +155,17 @@ store things like tags, metrics, and supporting data for storage, sorting, and f
 - [`Chat.meta()`][rigging.chat.Chat.meta] adds to [`Chat.metadata`][rigging.chat.Chat.metadata]
 
 Metadata will carry forward from a PendingChat to a Chat object when generation completes. This
-metadata is also maintained in the [serialization process](serialization.md).
+metadata is also maintained in the [serialization process](serialization.md).
+
+```py
+import rigging as rg
+
+pending = rg.get_generator("claude-2.1").chat("Hello!").meta(prompt_version=1)
+chat = pending.run().meta(user="Will")
+
+print(chat.metadata)
+# {
+#   'prompt_version': 1, 
+#   'user': 'Will'
+# }
+```

docs/topics/completions.md

@@ -19,10 +19,10 @@ While we try to maintain parity between the "Chat" and "Completions" interfaces
 find some deviations here and there. Completions should be a simple transition if you are familiar
 with the other code in rigging. Here are the highlights:
 
-- [`chat`][rigging.generator.Generator.chat] ~= [`complete`][rigging.generator.Generator.complete]
-- [`Chat`][rigging.chat.Chat] ~= [`Completion`][rigging.completion.Completion]
-- [`PendingChat`][rigging.chat.PendingChat] ~= [`PendingCompletion`][rigging.completion.PendingCompletion]
-- [`generate_messages`][rigging.generator.Generator.generate_messages] ~= [`generate_texts`][rigging.generator.Generator.generate_texts]
+- [`chat`][rigging.generator.Generator.chat] -> [`complete`][rigging.generator.Generator.complete]
+- [`Chat`][rigging.chat.Chat] -> [`Completion`][rigging.completion.Completion]
+- [`PendingChat`][rigging.chat.PendingChat] -> [`PendingCompletion`][rigging.completion.PendingCompletion]
+- [`generate_messages`][rigging.generator.Generator.generate_messages] -> [`generate_texts`][rigging.generator.Generator.generate_texts]
 
 On all of these interfaces, you'll note that sequences of [`Message`][rigging.message.Message] objects have been
 replaced with basic `str` objects for both inputs and ouputs.
@@ -45,9 +45,11 @@ Output: [translated text]
 Input: $input
 Output: """
 
-translator = rg.get_generator('gpt-3.5-turbo') \
-    .complete(PROMPT) \
-    .with_(stop=["---", "Input:", "\n\n"])
+translator = (
+    rg.get_generator('gpt-3.5-turbo') # (1)!
+    .complete(PROMPT)
+    .with_(stop=["---", "Input:", "\n\n"]) # (2)!
+)
 
 text = "Could you please tell me where the nearest train station is?"
 
@@ -63,6 +65,14 @@ for language in ["spanish", "french", "german"]:
 # [german]:  Könnten Sie mir bitte sagen, wo sich der nächste Bahnhof befindet?
 ```
 
+1. OpenAPI supports the same model IDs for both completions and chats, but other
+   providers might require you to specify a specific model ID used for text completions.
+2. We use [`.with_()`][rigging.completion.PendingCompletion.with_] to set stop tokens
+   and prevent the generation from simply continuing until our max tokens are reached. This
+   is a very common and often required pattern when doing completions over chats. Here, we 
+   aren't totally sure what the model might generate after our translation, so
+   we use a few different token sequences to be safe.
+
 !!! tip "Using .apply()"
 
     Text completion is a great place to use the [`.apply`][rigging.completion.PendingCompletion.apply]

docs/topics/generators.md

@@ -120,7 +120,16 @@ Generators operate in a batch context by default, taking in groups of message li
 your implementation takes advantage of this batching is up to you, but where possible you
 should be optimizing as much as possible.
 
+!!! tip "Generators are Flexible"
+
+    Generators don't make any assumptions about the underlying mechanism that completes text.
+    You might use a local model, API endpoint, or static code, etc. The base class is designed
+    to be flexible and support a wide variety of use cases. You'll obviously find that the inclusion
+    of `api_key`, `model`, and generation params are common enough that they are included in the base class.
+
 ```py
+from rigging import Generator, GenerateParams, Message
+
 class Custom(Generator):
     # model: str
     # api_key: str

docs/topics/models.md

@@ -84,6 +84,9 @@ that the content between tags conforms to any constraints we need. Take this exa
 Rigging model for instance:
 
 ```py
+from pydantic import field_validator
+import typing as t
+
 class YesNoAnswer(Model):
     "Yes/No answer answer with coercion"
 
@@ -102,7 +105,8 @@ class YesNoAnswer(Model):
 
 You can see the interior field of the model is now a `bool` type, which means pydantic will accept standard
 values which could be reasonably interpreted as a boolean. We also add a custom field validator to
-check for instances of `yes/no` as text strings. All of these XML values will parse correctly:
+check for instances of `yes/no` as text strings. All of these XML values will parse correctly
+into the `YesNoAnswer` model, so you can handle cases where the LLM outputs a variety of different
 
 ```xml
 <yes-no-answer>true</yes-no-answer>
@@ -112,6 +116,13 @@ check for instances of `yes/no` as text strings. All of these XML values will pa
 <yes-no-answer>1</yes-no-answer>
 ```
 
+```py
+YesNoAnswer(boolean="true")
+YesNoAnswer(boolean=" NO ")
+YesNoAnswer(boolean="1")
+# ...
+```
+
 The choice to build on Pydantic offers an incredible amount of flexibility for controlling exactly
 how data is validated in your models. This kind of parsing work is exactly what these libraries were designed
 to do. The sky is the limit, and **everything you find in Pydantic and Pydantic XML are compatible
@@ -184,7 +195,35 @@ and use [`.to_prety_xml()`][rigging.model.Model.to_pretty_xml]
 
 
 ```py
+import rigging as rg
+from typing import Annotated
+from pydantic import PlainSerializer
+
+newline_str = Annotated[str, PlainSerializer(lambda x: f'\n{x}\n', return_type=str)] # (1)!
+
+class SaveMemory(rg.Model):
+    key: str = rg.attr()
+    content: newline_str
+
+    @classmethod
+    def xml_example(cls) -> str:
+        return SaveMemory(
+            key="my-note",
+            content="Lots of custom data\nKeep this for later."
+        ).to_pretty_xml()
+
+print(f"Use the following format:\n{SaveMemory.xml_example()}")
+
+# Use the following format:
+# <save-memory key="my-note">
+# Lots of custom data
+# Keep this for later.
+# </save-memory>
+```
 
+1. Using pydantic serializer annotations is an easy way to introduce subtle content changes
+   before they are formed into their XML form. Here we're injecting newlines to make the
+   XML more readable.
 
 ## Complex Models
 

docs/topics/serialization.md

@@ -10,7 +10,7 @@ The following objects in Rigging have great serialization support for storage an
 Most of this stems from our use of Pydantic for core models, and we've included some helpful
 fields for reconstructing Chats and Completions.
 
-## Serializing Chats
+## JSON Serialization
 
 Let's build a joke pipeline and serialize the final chat into JSON.
 
@@ -95,7 +95,7 @@ You'll notice that every Chat gets a unique `id` field to help track them in a d
 assign a `timestamp` to understand when the generation took place. We are also taking advantage of the
 [`.meta()`][rigging.chat.PendingChat.meta] to add a tracking tag for filtering later.
 
-## Deserializing Chats
+## JSON Deserialization
 
 The JSON has everything required to reconstruct a Chat including a `generator_id` dynamically
 constructed to perserve the parameters used to create the generated message(s). We can now
@@ -118,4 +118,55 @@ print(continued.last)
 # The math book is described as being sad because it has "too many problems," which could be
 # interpreted as having both mathematical problems (equations to solve) and emotional difficulties.
 # This play on words adds humor to the joke.
+```
+
+## Pandas DataFrames
+
+Rigging also has helpers in the [`rigging.data`][] module for performing conversions
+between Chat objects and other storage formats like Pandas. In [`chats_to_df`][rigging.data.chats_to_df]
+the messages are flattened and stored with a `chat_id` column for grouping.
+[`df_to_chats`][rigging.data.df_to_chats] allows you to reconstruct a list of Chat objects back from a DataFrame.
+
+```py
+import rigging as rg
+
+chats = (
+    rg.get_generator("claude-3-haiku-20240307")
+    .chat("Write me a haiku.")
+    .run_many(3)
+)
+
+df = rg.chats_to_df(chats)
+
+print(df.info())
+
+# RangeIndex: 6 entries, 0 to 5
+# Data columns (total 9 columns):
+#  #   Column             Non-Null Count  Dtype         
+# ---  ------             --------------  -----         
+#  0   chat_id            6 non-null      string        
+#  1   chat_metadata      6 non-null      string        
+#  2   chat_generator_id  6 non-null      string        
+#  3   chat_timestamp     6 non-null      datetime64[ms]
+#  4   generated          6 non-null      bool          
+#  5   role               6 non-null      category      
+#  6   parts              6 non-null      string        
+#  7   content            6 non-null      string        
+#  8   message_id         6 non-null      string        
+# dtypes: bool(1), category(1), datetime64[ms](1), string(6)
+
+df.content.apply(lambda x: len(x)).mean()
+
+# 60.166666666666664
+
+back = rg.df_to_chats(df)
+print(back[0].conversation)
+
+# [user]: Write me a haiku.
+# 
+# [assistant]: Here's a haiku for you:
+# 
+# Gentle breeze whispers,
+# Flowers bloom in vibrant hues,
+# Nature's simple bliss.
 ```

docs/topics/workflow.md

@@ -17,11 +17,21 @@ of the many [`.run()`][rigging.chat.PendingChat.run] functions, the generator is
 message (or many messages) based on the prior context and any constraints you have in place. Once you have a 
 [`Chat`][rigging.chat.Chat] object, the interation is "done" and you can inspect and operate on the messages.
 
+??? tip "Chats vs Completions"
+
+    Rigging supports both Chat objects (messages with roles in a "conversation" format), as well
+    as raw text completions. While we use Chat objects in most of our examples, you can check
+    out the [Completions](completions.md) section to learn more about their feature parity.
+
 You'll often see us use functional styling chaining as most of our
 utility functions return the object back to you.
 
-```py
-chat = generator.chat(...) \
-    .using(...).until(...).with_(...) \
+```go
+chat = (
+    generator.chat(...)
+    .using(...)
+    .until(...)
+    .with_(...)
     .run()
+)
 ```

examples/bandit.py

@@ -8,7 +8,7 @@
 
 import asyncssh
 import click
-import requests  # type: ignore
+import requests
 from loguru import logger
 from pydantic import StringConstraints
 

mkdocs.yml

@@ -80,7 +80,6 @@ markdown_extensions:
   - admonition
   - pymdownx.highlight:
       anchor_linenums: true
-      line_spans: __span
       pygments_lang_class: true
   - pymdownx.inlinehilite
   - pymdownx.snippets

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "rigging"
-version = "1.0.0"
+version = "1.0.1"
 description = "LLM Interaction Framework"
 authors = ["Nick Landers <monoxgas@gmail.com>"]
 license = "MIT"

rigging/__init__.py

@@ -1,5 +1,6 @@
 from rigging.chat import Chat, PendingChat
 from rigging.completion import Completion, PendingCompletion
+from rigging.data import chats_to_df, df_to_chats
 from rigging.generator import GenerateParams, Generator, chat, complete, get_generator, register_generator
 from rigging.message import Message, MessageDict, Messages
 from rigging.model import Model, attr, element, wrapped
@@ -24,6 +25,8 @@
     "Completion",
     "PendingCompletion",
     "register_generator",
+    "chats_to_df",
+    "df_to_chats",
 ]
 
 from loguru import logger

rigging/chat.py

@@ -213,7 +213,10 @@ def apply(self, **kwargs: str) -> "Chat":
         Returns:
             The modified Chat object.
         """
-        self.last.apply(**kwargs)
+        if self.generated:
+            self.generated[-1] = self.generated[-1].apply(**kwargs)
+        else:
+            self.messages[-1] = self.messages[-1].apply(**kwargs)
         return self
 
     def apply_to_all(self, **kwargs: str) -> "Chat":
@@ -226,9 +229,8 @@ def apply_to_all(self, **kwargs: str) -> "Chat":
         Returns:
             The modified chat object.
         """
-        Message.apply_to_list(self.all, **kwargs)
-        for message in self.all:
-            message.apply(**kwargs)
+        self.messages = Message.apply_to_list(self.messages, **kwargs)
+        self.generated = Message.apply_to_list(self.generated, **kwargs)
         return self
 
     def strip(self, model_type: type[Model], fail_on_missing: bool = False) -> "Chat":