Skip to content

Add reusable pipelines #38

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hemajv wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add reusable pipelines #38

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 18 additions & 11 deletions pragmatic/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,23 @@
from pragmatic.pipelines.rag import RagPipelineWrapper
from pragmatic.pipelines.utils import produce_custom_settings


def index_path_for_rag(path, **kwargs):
def create_index_pipeline(path, **kwargs):

@ilya-kolchinsky ilya-kolchinsky Jan 29, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At the API level, I believe we still want to have a unified "one-shot" method that handles both pipeline creation and execution. It will fit cases like RamaLama where simplicity is the top priority.
This is especially true for indexing which is typically done in a single stage as opposed to querying the model. But same applies for execute_rag_query as well.

settings = produce_custom_settings(kwargs)
pipeline = LocalFileIndexingPipelineWrapper(settings, path)
pipeline.build_pipeline()
return pipeline.run()
index_pipeline = LocalFileIndexingPipelineWrapper(settings, path)
index_pipeline.build_pipeline()
return index_pipeline

def execute_rag_query(query, **kwargs):
def create_rag_pipeline(**kwargs):
settings = produce_custom_settings(kwargs)
pipeline = RagPipelineWrapper(settings, query)
pipeline.build_pipeline()
return pipeline.run()
rag_pipeline = RagPipelineWrapper(settings)
rag_pipeline.build_pipeline()
return rag_pipeline

def indexing_for_rag(index_pipeline, **kwargs):
return index_pipeline.run()

def execute_rag_query(rag_pipeline, query, **kwargs):
return rag_pipeline.run(query)

def evaluate_rag_pipeline(**kwargs):
from pragmatic.pipelines.evaluation import Evaluator
Expand All @@ -23,7 +28,9 @@ def evaluate_rag_pipeline(**kwargs):
return evaluator.evaluate_rag_pipeline()


__all__ = ["index_path_for_rag",
__all__ = ["create_index_pipeline",
"create_rag_pipeline",
"indexing_for_rag",
"execute_rag_query",
# "evaluate_rag_pipeline"
]
]
37 changes: 28 additions & 9 deletions pragmatic/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,21 +1,28 @@
import argparse

from api import index_path_for_rag, execute_rag_query, evaluate_rag_pipeline
from api import create_index_pipeline, create_rag_pipeline, indexing_for_rag, execute_rag_query, evaluate_rag_pipeline
from settings import DEFAULT_SETTINGS


def main():
"""
The tool can be executed in one of the following modes:
1) Indexing mode (-i flag) - index a collection of documents from the given path.
2) RAG query mode (-r flag) - answer a given query with RAG using the previously indexed documents.
3) Evaluation mode (-e flag) - evaluate the RAG pipeline as specified in the settings - NOT YET OFFICIALLY SUPPORTED.
1) Index pipeline creation mode (-ip flag) - create an indexing pipeline from the given path.

@ilya-kolchinsky ilya-kolchinsky Jan 29, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm really not sure these new modes should be exposed via main.py. What would a user do with a pipeline that is immediately destroyed after the program exits? IMO, this new functionality should be reserved for API only. At some future point we might introduce a more advanced main.py that interacts with the user via command line similarly to ilab chat - then these new flags would make sense.
This is of course only my opinion let's discuss if you think otherwise :)

2) RAG pipeline creatioon mode (-rp flag) - create a RAG pipeline
3) Indexing mode (-i flag) - index a collection of documents from the given path.
4) RAG query mode (-r flag) - answer a given query with RAG using the previously indexed documents.
5) Evaluation mode (-e flag) - evaluate the RAG pipeline as specified in the settings - NOT YET OFFICIALLY SUPPORTED.
"""
parser = argparse.ArgumentParser(description='RAG Pipeline PoC')

parser.add_argument('-i', '--indexing', help='Index a set of documents in the document storage', action='store_true')
parser.add_argument('-ip', '--indexpipeline', help='Initialize an indexing pipeline', action='store_true')
parser.add_argument('--path', help='The path to a directory with the documents to be indexed.')

parser.add_argument('-i', '--indexing', help='Index a set of documents in the document storage', action='store_true')
parser.add_argument('--pipeline', help='Provide the relevant pipeline for running the task, if pipeline not created initialize either index/rag pipeline depending on your task')

parser.add_argument('-rp', '--ragpipeline', help='Initialize a RAG pipeline', action='store_true')

parser.add_argument('-r', '--rag', help='Answer a given query based on the indexed documents', action='store_true')
parser.add_argument('--query', help='The query for the language model to answer.')

Expand All @@ -26,7 +33,7 @@ def main():

args = parser.parse_args()

if sum([args.indexing, args.rag, args.evaluation, args.server]) != 1:
if sum([args.indexpipeline, args.ragpipeline, args.indexing, args.rag, args.evaluation, args.server]) != 1:
print("Wrong usage: exactly one of the supported operation modes (indexing, query) must be specified.")
return

Expand All @@ -47,17 +54,29 @@ def main():
else:
print(f"Invalid format for argument: '{override}'. Expected format: key=value")

if args.indexing:
if args.indexpipeline:
if args.path is None:
print("Please specify the path containing the documents to index.")
return
index_path_for_rag(args.path, **custom_settings)
create_index_pipeline(args.path, **custom_settings)

if args.indexing:
if args.pipeline is None:
print("Please specify a relevant pipeline for running the task")
return
indexing_for_rag(args.pipeline, **custom_settings)

if args.ragpipeline:
create_rag_pipeline(**custom_settings)

if args.rag:
if args.query is None:
print("Please specify the query.")
return
print(execute_rag_query(args.query, **custom_settings))
if args.pipeline is None:
print("Please specify a relevant pipeline for running the task")
return
print(execute_rag_query(args.pipeline, args.query, **custom_settings))

if args.evaluation:
print(evaluate_rag_pipeline(**custom_settings))
Expand Down
9 changes: 6 additions & 3 deletions pragmatic/pipelines/pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,9 +50,12 @@ def _add_component(self, component_name, component_obj, component_args=None, sho
def _set_last_connect_point(self, connect_point):
self.__last_connect_point = connect_point

def run(self):
logger.debug(f"Executing the pipeline with the following arguments:\n{self._args}")
return self._pipeline.run(self._args)
def run(self, pipeline_args_dict=None):

@ilya-kolchinsky ilya-kolchinsky Jan 29, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So at pipeline creation we provide the default values for the run that are stored under self._args. Your new implementation keeps this self._args and either uses it or alternatively a dictionary provided as input. What if the caller only can/wants to provide a subset of the parameters? In this case, the two sources of parameters should be merged and the merge result should be used as input for the Haystack pipeline.
Another important point is whether we should keep the defaults provided at initialization. Wouldn't it be better to override the defaults (that will mostly be arbitrary anyway) with the new inputs? Perhaps add a Boolean parameter for run() to control that?

if pipeline_args_dict:
logger.debug(f"Executing the pipeline with the following arguments:\n{pipeline_args_dict}")
return self._pipeline.run(pipeline_args_dict)
else:
return self._pipeline.run(self._args)

def build_pipeline(self):
raise NotImplementedError()
Expand Down
57 changes: 41 additions & 16 deletions pragmatic/pipelines/rag.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,15 +32,13 @@


class RagPipelineWrapper(CommonPipelineWrapper):
def __init__(self, settings, query=None, evaluation_mode=False):
def __init__(self, settings, evaluation_mode=False):
super().__init__(settings)
self._query = query

self._evaluation_mode = evaluation_mode

def _add_embedder(self, query):
def _add_embedder(self):
embedder = SentenceTransformersTextEmbedder(model=self._settings["embedding_model_path"])
self._add_component("embedder", embedder, component_args={"text": query})
self._add_component("embedder", embedder)

def __init_sparse_retriever(self):
vector_db_type = self._settings["vector_db_type"]
Expand Down Expand Up @@ -74,16 +72,16 @@ def _add_retrievers(self):
dense_retriever = self.__init_dense_retriever()

if retriever_type == "sparse":
self._add_component("retriever", sparse_retriever, component_args={"query": self._query})
self._add_component("retriever", sparse_retriever)
elif retriever_type == "dense":
self._add_embedder(self._query)
self._add_embedder()
self._add_component("retriever", dense_retriever,
component_from_connect_point="embedder.embedding",
component_to_connect_point="retriever.query_embedding")
else: # retriever_type == "hybrid"
self._add_component("sparse_retriever", sparse_retriever, component_args={"query": self._query},
self._add_component("sparse_retriever", sparse_retriever,
should_connect=False)
self._add_embedder(self._query)
self._add_embedder()
self._add_component("dense_retriever", dense_retriever, should_connect=False)
self._add_component("document_joiner", DocumentJoiner(), should_connect=False)

Expand All @@ -97,11 +95,11 @@ def _add_ranker(self):
if not self._settings["ranker_enabled"]:
return
ranker = TransformersSimilarityRanker(model=self._settings["ranking_model"])
self._add_component("ranker", ranker, component_args={"query": self._query})
self._add_component("ranker", ranker)

def _add_prompt_builder(self):
prompt_builder = PromptBuilder(template=BASE_RAG_PROMPT)
self._add_component("prompt_builder", prompt_builder, component_args={"query": self._query},
self._add_component("prompt_builder", prompt_builder,
component_to_connect_point="prompt_builder.documents")

def _add_llm(self):
Expand Down Expand Up @@ -142,7 +140,7 @@ def _add_llm(self):
def _add_answer_builder(self):
if not self._evaluation_mode:
return
self._add_component("answer_builder", AnswerBuilder(), component_args={"query": self._query},
self._add_component("answer_builder", AnswerBuilder(),
should_connect=False)
self._pipeline.connect("llm.replies", "answer_builder.replies")
self._pipeline.connect("retriever", "answer_builder.documents")
Expand Down Expand Up @@ -170,10 +168,37 @@ def run(self, query=None):
for key in ["text", "query"]:
if key in config_dict:
config_dict[key] = query

result = super().run()


if self._evaluation_mode:
if ((self._settings["ranker_enabled"])):
if(self._settings["retriever_type"]=="dense"):
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "ranker": {"query": query}, "answer_builder": {"text": query}})
elif (self._settings["retriever_type"]=="sparse"):
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "ranker": {"query": query}, "retriever": {"query": query}, "answer_builder": {"text": query}})
else: #if retriever type is hybrid
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "ranker": {"query": query}, "sparse_retriever": {"query": query}, "answer_builder": {"text": query}})
else: #if ranker not enabled
if(self._settings["retriever_type"]=="dense"):
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "answer_builder": {"query": query}})
elif(self._settings["retriever_type"]=="sparse"):
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "retriever": {"query": query}, "answer_builder": {"text": query}})
else: #if retriever type is hybrid
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "sparse_retriever": {"query": query}, "answer_builder": {"text": query}})
return result["answer_builder"]["answers"][0]

return result["llm"]["replies"][0]
else: #not in eval mode

@ilya-kolchinsky ilya-kolchinsky Jan 29, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While this approach obviously works, it would be very challenging to maintain this code long-term going forward. It basically hard-codes the structure of the pipeline. Any attempt to introduce even a slightly different pipeline in the future will make this thing explode.
One way to resolve this would be to initialize the relevant parameters to None at pipeline creation (in all places where you removed them), then to rely on the code in the very beginning of the run() method to replace all the relevant fields with the up-to-date query. Please let me know if you'd like to handle it differently or if you have any concerns in this regard, and we will discuss.

if ((self._settings["ranker_enabled"])):
if(self._settings["retriever_type"]=="dense"):
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "ranker": {"query": query}})
elif (self._settings["retriever_type"]=="sparse"):
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "ranker": {"query": query}, "retriever": {"query": query}})
else: #if retriever type is hybrid
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "ranker": {"query": query}, "sparse_retriever": {"query": query}})
else: #if ranker not enabled
if(self._settings["retriever_type"]=="dense"):
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}})
elif(self._settings["retriever_type"]=="sparse"):
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "retriever": {"query": query}})
else: #if retriever type is hybrid
result = super().run({"embedder": {"text": query}, "prompt_builder": {"query": query}, "sparse_retriever": {"query": query}})
return result["llm"]["replies"][0]
1 change: 1 addition & 0 deletions requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,3 +2,4 @@ sentence-transformers>=3.0.0
docling>=2.9.0
milvus_haystack==0.0.11
docling_haystack==0.1.1
pymilvus>=2.4.2
36 changes: 17 additions & 19 deletions test/sanity_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@

from docling.document_converter import DocumentConverter

from pragmatic import index_path_for_rag, execute_rag_query
from pragmatic import create_index_pipeline, create_rag_pipeline, indexing_for_rag, execute_rag_query

SOURCE_PDF_URLS = [
"https://docs.redhat.com/en/documentation/red_hat_build_of_microshift/4.12/pdf/cli_tools/Red_Hat_build_of_MicroShift-4.12-CLI_tools-en-US.pdf",
Expand Down Expand Up @@ -52,39 +52,37 @@ def main():
docling_convert(SOURCE_PDF_URLS)

# index the JSONs under a Milvus Lite instance
print("Step 2: Embedding the JSONs and indexing them into Milvus vector database \n")
index_path_for_rag(DOCS_LOCAL_DIR_NAME,
milvus_deployment_type="lite",
print("Step 2: Initialize an indexing pipeline \n")
index_pipeline = create_index_pipeline(DOCS_LOCAL_DIR_NAME,milvus_deployment_type="lite",
milvus_file_path="./milvus.db",
embedding_model_path="sentence-transformers/all-MiniLM-L12-v2",
input_document_formats=['pdf'] if TEST_PDF_TO_JSON_CONVERSION else ['json'])

print("Step 3: Indexing documents into Milvus vector database using the index pipeline created \n")
indexing_for_rag(index_pipeline)

# execute a simple RAG query
print("Step 3: Executing simple RAG queries \n")
print("Question: How to install OpenShift CLI on macOS?")
result1 = execute_rag_query("How to install OpenShift CLI on macOS?",
milvus_file_path="./milvus.db",
print("Step 3: Initialize a RAG pipeline \n")
rag_pipeline = create_rag_pipeline(milvus_file_path="./milvus.db",
embedding_model_path="sentence-transformers/all-MiniLM-L12-v2",
llm_base_url="http://vllm-service:8000/v1",
llm_base_url="https://vllm-inference-predictor-mixture-of-experts-bots--runtime-int.apps.stc-ai-e1-pp.imap.p1.openshiftapps.com/v1",
llm="/mnt/models/",
top_k=3)
# To run the pipeline in eval mode
#rag_pipeline.set_evaluation_mode(True)
print("Step 4: Executing a query using the RAG pipeline \n")
print("Question: How to install OpenShift CLI on macOS?")
result1 = execute_rag_query(rag_pipeline, "How to install OpenShift CLI on macOS?")
print("Response generated:")
print(f"\n{result1}")
print("\n")
print("Question: What are the two deployment options in OpenShift AI?")
result2 = execute_rag_query("What are the two deployment options in OpenShift AI?",
milvus_file_path="./milvus.db",
embedding_model_path="sentence-transformers/all-MiniLM-L12-v2",
llm_base_url="http://vllm-service:8000/v1",
top_k=3)
result2 = execute_rag_query(rag_pipeline, "What are the two deployment options in OpenShift AI?")
print("Response generated:")
print(f"\n{result2}")
print("\n")
print("Question: What is OpenShift AI?")
result3 = execute_rag_query("What is OpenShift AI?",
milvus_file_path="./milvus.db",
embedding_model_path="sentence-transformers/all-MiniLM-L12-v2",
llm_base_url="http://vllm-service:8000/v1",
top_k=3)
result3 = execute_rag_query(rag_pipeline, "What is OpenShift AI?")
print("Response generated:")
print(f"\n{result3}")

Expand Down