diff --git a/llama_stack/templates/cerebras/cerebras.py b/llama_stack/templates/cerebras/cerebras.py index 17fc26632..df3b55ddd 100644 --- a/llama_stack/templates/cerebras/cerebras.py +++ b/llama_stack/templates/cerebras/cerebras.py @@ -9,15 +9,11 @@ from pathlib import Path from llama_models.sku_list import all_registered_models from llama_stack.apis.models.models import ModelType -from llama_stack.distribution.datatypes import ( - ModelInput, - Provider, - ShieldInput, - ToolGroupInput, -) +from llama_stack.distribution.datatypes import ModelInput, Provider, ToolGroupInput from llama_stack.providers.inline.inference.sentence_transformers import ( SentenceTransformersInferenceConfig, ) +from llama_stack.providers.inline.memory.faiss.config import FaissImplConfig from llama_stack.providers.remote.inference.cerebras import CerebrasImplConfig from llama_stack.providers.remote.inference.cerebras.cerebras import model_aliases from llama_stack.templates.template import DistributionTemplate, RunConfigSettings @@ -41,6 +37,7 @@ def get_distribution_template() -> DistributionTemplate: ], } + name = "cerebras" inference_provider = Provider( provider_id="cerebras", provider_type="remote::cerebras", @@ -71,6 +68,11 @@ def get_distribution_template() -> DistributionTemplate: "embedding_dimension": 384, }, ) + memory_provider = Provider( + provider_id="faiss", + provider_type="inline::faiss", + config=FaissImplConfig.sample_run_config(f"distributions/{name}"), + ) default_tool_groups = [ ToolGroupInput( toolgroup_id="builtin::websearch", @@ -98,9 +100,10 @@ def get_distribution_template() -> DistributionTemplate: "run.yaml": RunConfigSettings( provider_overrides={ "inference": [inference_provider, embedding_provider], + "memory": [memory_provider], }, default_models=default_models + [embedding_model], - default_shields=[ShieldInput(shield_id="meta-llama/Llama-Guard-3-8B")], + default_shields=[], default_tool_groups=default_tool_groups, ), }, diff --git a/llama_stack/templates/cerebras/run.yaml b/llama_stack/templates/cerebras/run.yaml index e0beab9cc..bfc492bda 100644 --- a/llama_stack/templates/cerebras/run.yaml +++ b/llama_stack/templates/cerebras/run.yaml @@ -32,12 +32,6 @@ providers: type: sqlite namespace: null db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/cerebras}/faiss_store.db - - provider_id: chromadb - provider_type: remote::chromadb - config: {} - - provider_id: pgvector - provider_type: remote::pgvector - config: {} agents: - provider_id: meta-reference provider_type: inline::meta-reference @@ -111,8 +105,7 @@ models: model_id: all-MiniLM-L6-v2 provider_id: sentence-transformers model_type: embedding -shields: -- shield_id: meta-llama/Llama-Guard-3-8B +shields: [] memory_banks: [] datasets: [] scoring_fns: []