From 562ef41ff8f4c17c1438b1dd81614553a8c899b5 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Thu, 26 Dec 2024 18:56:23 -0800
Subject: [PATCH] fix tests

---
 llama_stack/distribution/inspect.py                |  5 ++---
 llama_stack/distribution/resolver.py               | 14 +++++++-------
 .../inline/agents/meta_reference/agent_instance.py | 14 +++++++++++---
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/llama_stack/distribution/inspect.py b/llama_stack/distribution/inspect.py
index 88d04a0ff..dbb16d8ce 100644
--- a/llama_stack/distribution/inspect.py
+++ b/llama_stack/distribution/inspect.py
@@ -8,10 +8,9 @@ from typing import Dict, List
 
 from pydantic import BaseModel
 
-from llama_stack.apis.inspect import Inspect
-from llama_stack.distribution.datatypes import HealthInfo, ProviderInfo, RouteInfo
+from llama_stack.apis.inspect import HealthInfo, Inspect, ProviderInfo, RouteInfo
+from llama_stack.distribution.datatypes import StackRunConfig
 from llama_stack.distribution.server.endpoints import get_all_api_endpoints
-from llama_stack.providers.datatypes import StackRunConfig
 
 
 class DistributionInspectConfig(BaseModel):
diff --git a/llama_stack/distribution/resolver.py b/llama_stack/distribution/resolver.py
index 76125948e..0a6eed345 100644
--- a/llama_stack/distribution/resolver.py
+++ b/llama_stack/distribution/resolver.py
@@ -31,16 +31,9 @@ from llama_stack.distribution.client import get_client_impl
 
 from llama_stack.distribution.datatypes import (
     AutoRoutedProviderSpec,
-    DatasetsProtocolPrivate,
-    EvalTasksProtocolPrivate,
-    MemoryBanksProtocolPrivate,
-    ModelsProtocolPrivate,
     Provider,
     RoutingTableProviderSpec,
-    ScoringFunctionsProtocolPrivate,
-    ShieldsProtocolPrivate,
     StackRunConfig,
-    ToolsProtocolPrivate,
 )
 from llama_stack.distribution.distribution import builtin_automatically_routed_apis
 from llama_stack.distribution.store import DistributionRegistry
@@ -48,10 +41,17 @@ from llama_stack.distribution.utils.dynamic import instantiate_class_type
 
 from llama_stack.providers.datatypes import (
     Api,
+    DatasetsProtocolPrivate,
+    EvalTasksProtocolPrivate,
     InlineProviderSpec,
+    MemoryBanksProtocolPrivate,
+    ModelsProtocolPrivate,
     ProviderSpec,
     RemoteProviderConfig,
     RemoteProviderSpec,
+    ScoringFunctionsProtocolPrivate,
+    ShieldsProtocolPrivate,
+    ToolsProtocolPrivate,
 )
 
 log = logging.getLogger(__name__)
diff --git a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
index 9e20119d1..fdf567333 100644
--- a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
+++ b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
@@ -13,7 +13,7 @@ import secrets
 import string
 import uuid
 from datetime import datetime
-from typing import AsyncGenerator, Dict, List, Tuple
+from typing import AsyncGenerator, Dict, List, Optional, Tuple
 from urllib.parse import urlparse
 
 import httpx
@@ -27,6 +27,7 @@ from llama_stack.apis.agents import (
     AgentTurnResponseEvent,
     AgentTurnResponseEventType,
     AgentTurnResponseStepCompletePayload,
+    AgentTurnResponseStepProgressPayload,
     AgentTurnResponseStepStartPayload,
     AgentTurnResponseStreamChunk,
     AgentTurnResponseTurnCompletePayload,
@@ -34,11 +35,14 @@ from llama_stack.apis.agents import (
     Attachment,
     CodeInterpreterToolDefinition,
     FunctionCallToolDefinition,
+    InferenceStep,
     MemoryRetrievalStep,
+    MemoryToolDefinition,
     PhotogenToolDefinition,
     SearchToolDefinition,
     ShieldCallStep,
     StepType,
+    ToolExecutionStep,
     Turn,
     WolframAlphaToolDefinition,
 )
@@ -49,18 +53,22 @@ from llama_stack.apis.common.content_types import (
     URL,
 )
 from llama_stack.apis.inference import (
+    ChatCompletionResponseEventType,
     CompletionMessage,
     Inference,
     Message,
     SamplingParams,
     StopReason,
     SystemMessage,
+    ToolCallDelta,
+    ToolCallParseStatus,
     ToolDefinition,
+    ToolResponse,
     ToolResponseMessage,
     UserMessage,
 )
-from llama_stack.apis.memory import Memory, MemoryBankDocument
-from llama_stack.apis.memory_banks import MemoryBanks
+from llama_stack.apis.memory import Memory, MemoryBankDocument, QueryDocumentsResponse
+from llama_stack.apis.memory_banks import MemoryBanks, VectorMemoryBankParams
 from llama_stack.apis.safety import Safety
 
 from llama_stack.providers.utils.kvstore import KVStore