fix tests

This commit is contained in:
Xi Yan 2024-12-26 18:56:23 -08:00
parent 50764d76a7
commit 562ef41ff8
3 changed files with 20 additions and 13 deletions

View file

@ -8,10 +8,9 @@ from typing import Dict, List
from pydantic import BaseModel from pydantic import BaseModel
from llama_stack.apis.inspect import Inspect from llama_stack.apis.inspect import HealthInfo, Inspect, ProviderInfo, RouteInfo
from llama_stack.distribution.datatypes import HealthInfo, ProviderInfo, RouteInfo from llama_stack.distribution.datatypes import StackRunConfig
from llama_stack.distribution.server.endpoints import get_all_api_endpoints from llama_stack.distribution.server.endpoints import get_all_api_endpoints
from llama_stack.providers.datatypes import StackRunConfig
class DistributionInspectConfig(BaseModel): class DistributionInspectConfig(BaseModel):

View file

@ -31,16 +31,9 @@ from llama_stack.distribution.client import get_client_impl
from llama_stack.distribution.datatypes import ( from llama_stack.distribution.datatypes import (
AutoRoutedProviderSpec, AutoRoutedProviderSpec,
DatasetsProtocolPrivate,
EvalTasksProtocolPrivate,
MemoryBanksProtocolPrivate,
ModelsProtocolPrivate,
Provider, Provider,
RoutingTableProviderSpec, RoutingTableProviderSpec,
ScoringFunctionsProtocolPrivate,
ShieldsProtocolPrivate,
StackRunConfig, StackRunConfig,
ToolsProtocolPrivate,
) )
from llama_stack.distribution.distribution import builtin_automatically_routed_apis from llama_stack.distribution.distribution import builtin_automatically_routed_apis
from llama_stack.distribution.store import DistributionRegistry 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 ( from llama_stack.providers.datatypes import (
Api, Api,
DatasetsProtocolPrivate,
EvalTasksProtocolPrivate,
InlineProviderSpec, InlineProviderSpec,
MemoryBanksProtocolPrivate,
ModelsProtocolPrivate,
ProviderSpec, ProviderSpec,
RemoteProviderConfig, RemoteProviderConfig,
RemoteProviderSpec, RemoteProviderSpec,
ScoringFunctionsProtocolPrivate,
ShieldsProtocolPrivate,
ToolsProtocolPrivate,
) )
log = logging.getLogger(__name__) log = logging.getLogger(__name__)

View file

@ -13,7 +13,7 @@ import secrets
import string import string
import uuid import uuid
from datetime import datetime from datetime import datetime
from typing import AsyncGenerator, Dict, List, Tuple from typing import AsyncGenerator, Dict, List, Optional, Tuple
from urllib.parse import urlparse from urllib.parse import urlparse
import httpx import httpx
@ -27,6 +27,7 @@ from llama_stack.apis.agents import (
AgentTurnResponseEvent, AgentTurnResponseEvent,
AgentTurnResponseEventType, AgentTurnResponseEventType,
AgentTurnResponseStepCompletePayload, AgentTurnResponseStepCompletePayload,
AgentTurnResponseStepProgressPayload,
AgentTurnResponseStepStartPayload, AgentTurnResponseStepStartPayload,
AgentTurnResponseStreamChunk, AgentTurnResponseStreamChunk,
AgentTurnResponseTurnCompletePayload, AgentTurnResponseTurnCompletePayload,
@ -34,11 +35,14 @@ from llama_stack.apis.agents import (
Attachment, Attachment,
CodeInterpreterToolDefinition, CodeInterpreterToolDefinition,
FunctionCallToolDefinition, FunctionCallToolDefinition,
InferenceStep,
MemoryRetrievalStep, MemoryRetrievalStep,
MemoryToolDefinition,
PhotogenToolDefinition, PhotogenToolDefinition,
SearchToolDefinition, SearchToolDefinition,
ShieldCallStep, ShieldCallStep,
StepType, StepType,
ToolExecutionStep,
Turn, Turn,
WolframAlphaToolDefinition, WolframAlphaToolDefinition,
) )
@ -49,18 +53,22 @@ from llama_stack.apis.common.content_types import (
URL, URL,
) )
from llama_stack.apis.inference import ( from llama_stack.apis.inference import (
ChatCompletionResponseEventType,
CompletionMessage, CompletionMessage,
Inference, Inference,
Message, Message,
SamplingParams, SamplingParams,
StopReason, StopReason,
SystemMessage, SystemMessage,
ToolCallDelta,
ToolCallParseStatus,
ToolDefinition, ToolDefinition,
ToolResponse,
ToolResponseMessage, ToolResponseMessage,
UserMessage, UserMessage,
) )
from llama_stack.apis.memory import Memory, MemoryBankDocument from llama_stack.apis.memory import Memory, MemoryBankDocument, QueryDocumentsResponse
from llama_stack.apis.memory_banks import MemoryBanks from llama_stack.apis.memory_banks import MemoryBanks, VectorMemoryBankParams
from llama_stack.apis.safety import Safety from llama_stack.apis.safety import Safety
from llama_stack.providers.utils.kvstore import KVStore from llama_stack.providers.utils.kvstore import KVStore