From cfc97df6d5041f2e9363de77bca28fa376722467 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Wed, 16 Oct 2024 15:37:05 -0700
Subject: [PATCH] dell tgi fix

---
 .../docker/llamastack-local-cpu/build.yaml    |  4 +-
 .../providers/adapters/inference/tgi/tgi.py   |  4 ++
 tests/examples/tgi-run.yaml                   | 48 +++++++++++++++++++
 3 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 tests/examples/tgi-run.yaml

diff --git a/llama_stack/distribution/templates/docker/llamastack-local-cpu/build.yaml b/llama_stack/distribution/templates/docker/llamastack-local-cpu/build.yaml
index 9db019454..c0fabbf4d 100644
--- a/llama_stack/distribution/templates/docker/llamastack-local-cpu/build.yaml
+++ b/llama_stack/distribution/templates/docker/llamastack-local-cpu/build.yaml
@@ -10,6 +10,8 @@ distribution_spec:
     - remote::fireworks
     safety: meta-reference
     agents: meta-reference
-    memory: meta-reference
+    memory:
+    - remote::chromadb
+    - meta-reference
     telemetry: meta-reference
 image_type: docker
diff --git a/llama_stack/providers/adapters/inference/tgi/tgi.py b/llama_stack/providers/adapters/inference/tgi/tgi.py
index cd0afad0c..92cb9ba6a 100644
--- a/llama_stack/providers/adapters/inference/tgi/tgi.py
+++ b/llama_stack/providers/adapters/inference/tgi/tgi.py
@@ -52,6 +52,10 @@ class _HfAdapter(Inference, ModelsProtocolPrivate):
 
     async def list_models(self) -> List[ModelDef]:
         repo = self.model_id
+        # tmp hack to support Dell
+        if repo not in self.huggingface_repo_to_llama_model_id:
+            repo = "meta-llama/Llama-3.1-8B-Instruct"
+
         identifier = self.huggingface_repo_to_llama_model_id[repo]
         return [
             ModelDef(
diff --git a/tests/examples/tgi-run.yaml b/tests/examples/tgi-run.yaml
new file mode 100644
index 000000000..a398b20a3
--- /dev/null
+++ b/tests/examples/tgi-run.yaml
@@ -0,0 +1,48 @@
+version: '2'
+built_at: '2024-10-08T17:40:45.325529'
+image_name: local
+docker_image: null
+conda_env: local
+apis:
+- shields
+- agents
+- models
+- memory
+- memory_banks
+- inference
+- safety
+providers:
+  inference:
+  - provider_id: remote::tgi
+    provider_type: remote::tgi
+    config:
+      url: http://127.0.0.1:5009
+  safety:
+  - provider_id: meta-reference
+    provider_type: meta-reference
+    config:
+      llama_guard_shield:
+        model: Llama-Guard-3-1B
+        excluded_categories: []
+        disable_input_check: false
+        disable_output_check: false
+      prompt_guard_shield:
+        model: Prompt-Guard-86M
+  memory:
+  - provider_id: remote::chromadb
+    provider_type: remote::chromadb
+    config:
+      host: localhost
+      port: 6000
+  agents:
+  - provider_id: meta-reference
+    provider_type: meta-reference
+    config:
+      persistence_store:
+        namespace: null
+        type: sqlite
+        db_path: /home/xiyan/.llama/runtime/kvstore.db
+  telemetry:
+  - provider_id: meta-reference
+    provider_type: meta-reference
+    config: {}