feat: Expose document store Python API in instructlab/instructlab rag submodule

[dmartinol]

References

This PR relates to the task "[Dev][RAG] Expose document store Python API in instructlab/instructlab rag submodule" in the RAG project.
See #2872

Proposed interfaces

We propose two interfaces to model the ingestion and retrieval interactions with a document store, in the context of RAG pipelines:

class DocumentStoreRetriever(ABC):
    @abstractmethod
    def augmented_context(self, user_query: str) -> str:...

class DocumentStoreIngestor(ABC):
    @abstractmethod
    def ingest_documents(self, input_dir) -> tuple[bool, int]:...

Design notes

The actual instances are generated by factory functions which delegate the actual implementation to a factory function identified with some hardcoded conventions.

Factory functions:

def create_document_retriever(
    document_store_config: DocumentStoreConfig,
    retriever_config: RetrieverConfig,
) -> DocumentStoreRetriever:
...

def create_document_store_ingestor(
    document_store_config: DocumentStoreConfig,
    embedding_config: EmbeddingModelConfig,
) -> DocumentStoreIngestor:
...

Initial implementation uses a Milvus db, and there is no option to configure it differently.
The implementations are based on Haystack pipelines whose details are abstracted by the interface defined in the document_store module.

Test functions

Test functions in tests/test_document_store_factory.py can help to clarify how to consume these APIs.

Updated dependencies

Requirements file has been updated to reflect the required dependencies.

Minimize the PR

The changes include the implementation of the Haystack document stores for Milvus mainly to prove the effectiveness of the proposal. If you prefer, we can start with a minimal commit to only include the interfaces and the minimal dependencies (configuration and unit test).

@anastasds @jwm4

Checklist:

• Commit Message Formatting: Commit titles and messages follow guidelines in the
  conventional commits.
• Changelog updated with breaking and/or notable changes for the next minor release.
• Documentation has been updated, if necessary.
• Unit tests have been added, if necessary.
• Functional tests have been added, if necessary.
• E2E Workflow tests have been added, if necessary.

[anastasds]

Suggested some simplification and had a few questions.

[anastasds]

Are we sure we want to apply Haystack's document cleaning?

[dmartinol]

I'm not sure, but if we assume to feed the pipelines with un-processed user documentation (pdf), are we also assuming that they're all ready-to-use in the RAG ingestion pipeline?
BTW: I might be wrong, but I don't think Docling applies any cleanup to the original documents.

[dmartinol]

@anastasds Please note that the document_splitter might be replaced by the one available in the docling-haystack, whenever it will be available.
ATM we cannot integrate it because of a dependency conflict with the instructlab-sdg package:

    instructlab-sdg 0.6.2 depends on docling<=2.8.3 and >=2.4.2
    docling-haystack 0.1.1 depends on docling<3.0.0 and >=2.9.0

I've added one comment here to track the request of upgrading the instructlab-sdg dependencies.

[anastasds]

Pending the discussion regarding the new dependencies, looks good to me!

[dmartinol]

@anastasds

• can we make the PR ready for review?
• what about the issue with the test functions that require the download of the default embedding model?
  • make it a real unit test with mocked embedder? (not easy but probably doable)
  • make it an integration test instead?

[github-actions]

E2E (NVIDIA L40S x4) workflow launched on this PR: View run

[courtneypacheco]

Looks good! Had a few comments (mostly about test coverage)

[github-actions]

e2e workflow succeeded on this PR: View run, congrats!

[dmartinol]

@courtneypacheco added more unit tests to increase coverage

src/instructlab/rag/__init__.py                                    0      0   100%
src/instructlab/rag/document_store.py                              7      0   100%
src/instructlab/rag/document_store_factory.py                      9      0   100%
src/instructlab/rag/haystack/component_factory.py                 22      0   100%
src/instructlab/rag/haystack/components/document_splitter.py      52      7    87%
src/instructlab/rag/haystack/document_store_factory.py             9      0   100%
src/instructlab/rag/haystack/document_store_ingestor.py           42      3    93%
src/instructlab/rag/haystack/document_store_retriever.py          27      0   100%

[cdoern]

@dmartinol this is looking great! can you squash the commits before merging?

[dmartinol]

@dmartinol this is looking great! can you squash the commits before merging?
✅

[cdoern]

thanks for the hard work on this!