llama-stack-mirror/llama_stack/providers/registry/datasetio.py
Josh Salomon 8e3b579df2 Initial commit for postgresql dataset provider
Core implementation of postgresql dataset for llama stack

Signed-off-by: Josh Salomon <jsalomon@redhat.com>
2025-04-14 14:22:54 +03:00

53 lines
1.7 KiB
Python

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import List
from llama_stack.providers.datatypes import (
AdapterSpec,
Api,
InlineProviderSpec,
ProviderSpec,
remote_provider_spec,
)
def available_providers() -> List[ProviderSpec]:
return [
InlineProviderSpec(
api=Api.datasetio,
provider_type="inline::localfs",
pip_packages=["pandas"],
module="llama_stack.providers.inline.datasetio.localfs",
config_class="llama_stack.providers.inline.datasetio.localfs.LocalFSDatasetIOConfig",
api_dependencies=[],
),
remote_provider_spec(
api=Api.datasetio,
adapter=AdapterSpec(
adapter_type="huggingface",
pip_packages=[
"datasets",
],
module="llama_stack.providers.remote.datasetio.huggingface",
config_class="llama_stack.providers.remote.datasetio.huggingface.HuggingfaceDatasetIOConfig",
),
),
remote_provider_spec(
api=Api.datasetio,
adapter=AdapterSpec(
adapter_type="postgresql",
pip_packages=[
# "asyncpg",
# "datasets",
"psycopg",
"psycopg[pool]",
],
module="llama_stack.providers.remote.datasetio.postgresql",
config_class="llama_stack.providers.remote.datasetio.postgresql.PostgreSQLDatasetIOConfig",
),
),
]