RAGWrangler - A simple RAG (Retrieval-Augumented Generation) Task Manager

RAGWrangler is a Python project designed to streamline the management of RAG tasks using generative language models and the Weaviate database.

With this tool, you can automatically create, retrieve, and store the outputs of RAG tasks in a structured manner in Weaviate, allowing for easy management and tracking of outputs, ultimately saving time and resources.

Features

• Automated Task Handling: Simplify the creation and management of tasks with a straightforward Python class interface.
• Weaviate Integration: Seamlessly store and retrieve task outputs using Weaviate database integration.
• Support for Multiple Language Models: You can easily swap out language models as desired.
• Logging: Utilize integrated logging for effortless debugging and tracking of task statuses.

Installation

The tool should be compatible with Python 3.8 and higher, although development primarily utilized Python 3.9.

To get started, install the necessary Python packages using the command below:

pip install ragwrangler

Quickstart

See example_usage.py to see a brief example of how to use the tool.

Usage

You can connect to any Weaviate instance to save your data. Refer to the Weaviate documentation for more information on how to set up & connect to a Weaviate instance.

The primary class you will interact with is RAGTask.

RAGTask

Instantiate a RAGTask for each generative task to be handled.

Provide it with a Weaviate instance, and a builder function. You will also need to set the OpenAI API key.

import weaviate, os
from ragwrangler import RAGTask, set_openai_api_key

client = weaviate.Client(
    url=os.environ['WCS_URL'],
    auth_client_secret=weaviate.AuthApiKey(os.environ['WCS_ADMIN_KEY']),
    additional_headers={"X-OpenAI-Api-Key": os.environ["OPENAI_APIKEY"]}
)

set_openai_api_key(os.environ["OPENAI_APIKEY"])
quiz_rag = RAGTask(client=client, task_prompt_builder=revision_quiz_json_builder)

Where revision_quiz_json_builder is a function that takes in a source text and returns a task prompt for the LLM.

Extend RAGTask by defining custom task prompt builder functions that dictate how to generate prompts from the source text, which are then used to derive outputs.

Working with Tasks

To create a new task, instantiate an object of RAGTask (or its extension) with a source text:

quiz_rag = RAGTask(client=client, task_prompt_builder=revision_quiz_json_builder)

To obtain the task output, utilize the get_output method, specifying a model name if desired:

output = task.get_output(source_text="Your source text here", model_name="gpt-3.5-turbo")

License

This project is licensed under the MIT License.

Copyright

© 2023 JP Hwang