phoenix-oss/llama-stack
1
0
Fork 0
forked from phoenix-oss/llama-stack-mirror
Code Pull requests Releases Packages 2 Activity Actions

Add Chroma and PGVector adapters (#56)

Browse source 
Co-authored-by: Ashwin Bharambe <ashwin@meta.com>
This commit is contained in:
Ashwin Bharambe 2024-09-06 18:53:17 -07:00 • committed by GitHub
parent 5de6ed946e
commit 3f090d1975
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 628 additions and 119 deletions
Download patch file Download diff file Expand all files Collapse all files

30
llama_toolchain/memory/providers.py
View file

@ -6,7 +6,12 @@
from typing import List
from llama_toolchain.core.datatypes import Api, InlineProviderSpec, ProviderSpec
from llama_toolchain.core.datatypes import * # noqa: F403
EMBEDDING_DEPS = [
"blobfile",
"sentence-transformers",
]
def available_memory_providers() -> List[ProviderSpec]:
@ -14,12 +19,25 @@ def available_memory_providers() -> List[ProviderSpec]:
InlineProviderSpec(
api=Api.memory,
provider_id="meta-reference-faiss",
pip_packages=[
"blobfile",
"faiss-cpu",
"sentence-transformers",
],
pip_packages=EMBEDDING_DEPS + ["faiss-cpu"],
module="llama_toolchain.memory.meta_reference.faiss",
config_class="llama_toolchain.memory.meta_reference.faiss.FaissImplConfig",
),
remote_provider_spec(
api=Api.memory,
adapter=AdapterSpec(
adapter_id="chromadb",
pip_packages=EMBEDDING_DEPS + ["chromadb-client"],
module="llama_toolchain.memory.adapters.chroma",
),
),
remote_provider_spec(
api=Api.memory,
adapter=AdapterSpec(
adapter_id="pgvector",
pip_packages=EMBEDDING_DEPS + ["psycopg2-binary"],
module="llama_toolchain.memory.adapters.pgvector",
config_class="llama_toolchain.memory.adapters.pgvector.PGVectorConfig",
),
),
]