Address Ragu's comments

Author: gyliu513

examples/agents/simple_chat.py

@@ -5,13 +5,11 @@
 # the root directory of this source tree.
 import os
 
-import inspect
-
 import fire
 from llama_stack_client import LlamaStackClient, Agent, AgentEventLogger
 from termcolor import colored
 
-from .utils import check_model_is_available, get_any_available_model
+from .utils import check_model_is_available, get_any_available_chat_model
 
 
 def main(host: str, port: int, model_id: str | None = None):
@@ -29,14 +27,8 @@ def main(host: str, port: int, model_id: str | None = None):
         provider_data={"tavily_search_api_key": os.getenv("TAVILY_SEARCH_API_KEY")},
     )
 
-    available_shields = [shield.identifier for shield in client.shields.list()]
-    if not available_shields:
-        print(colored("No available shields. Disabling safety.", "yellow"))
-    else:
-        print(f"Available shields found: {available_shields}")
-
     if model_id is None:
-        model_id = get_any_available_model(client)
+        model_id = get_any_available_chat_model(client)
         if model_id is None:
             return
     else:
@@ -45,25 +37,13 @@ def main(host: str, port: int, model_id: str | None = None):
 
     print(f"Using model: {model_id}")
 
-    agent_kwargs = {
-        "model": model_id,
-        "instructions": "",
+    agent = Agent(
+        client,
+        model=model_id,
+        instructions="",
         # OpenAI Responses tool schema requires a type discriminator.
-        "tools": [{"type": "web_search"}],
-        "input_shields": available_shields,
-        "output_shields": available_shields,
-        "enable_session_persistence": False,
-    }
-    allowed_params = set(inspect.signature(Agent.__init__).parameters)
-    filtered_kwargs = {k: v for k, v in agent_kwargs.items() if k in allowed_params}
-    try:
-        agent = Agent(client, **filtered_kwargs)
-    except TypeError as exc:
-        # Fallback for older clients that only accept string tool names.
-        if "Unsupported tool type" not in str(exc):
-            raise
-        filtered_kwargs["tools"] = ["builtin::websearch"]
-        agent = Agent(client, **filtered_kwargs)
+        tools=[{"type": "web_search"}],
+    )
     user_prompts = [
         "Hello",
         "Search web for which players played in the winning team of the NBA western conference semifinals of 2024",

examples/agents/utils.py

@@ -3,10 +3,12 @@
 
 
 def _get_model_type(model) -> str | None:
-    for attr in ("model_type", "type", "model_kind", "kind", "model_family"):
-        value = getattr(model, attr, None)
-        if isinstance(value, str):
-            return value
+    for metadata_attr in ("custom_metadata", "metadata"):
+        metadata = getattr(model, metadata_attr, None)
+        if isinstance(metadata, dict):
+            value = metadata.get("model_type") or metadata.get("type")
+            if isinstance(value, str):
+                return value
     return None
 
 
@@ -56,3 +58,34 @@ def get_any_available_model(client: LlamaStackClient):
         return None
 
     return available_models[0]
+
+
+def can_model_chat(client: LlamaStackClient, model_id: str) -> bool:
+    # Lightweight probe to ensure the model supports chat completions.
+    try:
+        client.chat.completions.create(
+            model=model_id,
+            messages=[{"role": "user", "content": "ping"}],
+            max_tokens=1,
+        )
+    except Exception:
+        return False
+    return True
+
+def get_any_available_chat_model(client: LlamaStackClient):
+    available_models = [
+        model_id
+        for model in client.models.list()
+        for model_id in [_get_model_id(model)]
+        if model_id and _is_llm_model(model) and "guard" not in model_id
+    ]
+    if not available_models:
+        print(colored("No available models.", "red"))
+        return None
+
+    for model_id in available_models:
+        if can_model_chat(client, model_id):
+            return model_id
+
+    print(colored("No available chat-capable models.", "red"))
+    return None