modular_llm_driver.py

config = dict(
    vector_db_config=json.loads(vector_db_config),
    embedding_service=embedding_service,  # this triggers config.when() in embedding_module
    api_key=embedding_service_api_key,
    model_name=model_name,
)

dr = driver.Driver(
    config,
    data_module,
    vector_db_module,  # pinecone_module, weaviate_module or lancedb_module
    embedding_module,  # contains cohere, openai, and sentence_transformer implementations
    adapter=base.SimplePythonGraphAdapter(base.DictResult()),
)

dr.execute(
    final_vars=["initialize_vector_db_indices", "push_to_vector_db"],
    inputs=dict(
        class_name="SQuADEntry",
    ),
)