This commit is contained in:
Xi Yan 2024-12-26 18:34:00 -08:00
parent 27da763af9
commit 74de9bebd1
2 changed files with 25 additions and 6 deletions

View file

@ -8,11 +8,14 @@ import os
import pytest import pytest
import pytest_asyncio import pytest_asyncio
from llama_stack.distribution.store import * # noqa F403
from llama_stack.apis.inference import Model from llama_stack.apis.inference import Model
from llama_stack.apis.memory_banks import VectorMemoryBank from llama_stack.apis.memory_banks import VectorMemoryBank
from llama_stack.distribution.store.registry import (
CachedDiskDistributionRegistry,
DiskDistributionRegistry,
)
from llama_stack.providers.utils.kvstore import kvstore_impl, SqliteKVStoreConfig from llama_stack.providers.utils.kvstore import kvstore_impl, SqliteKVStoreConfig
from llama_stack.distribution.datatypes import * # noqa F403
@pytest.fixture @pytest.fixture

View file

@ -8,10 +8,26 @@ from typing import AsyncIterator, List, Optional, Union
import pytest import pytest
from llama_stack.apis.inference import * # noqa: F403 from llama_stack.apis.agents import (
from llama_stack.apis.memory import * # noqa: F403 AgentConfig,
from llama_stack.apis.safety import * # noqa: F403 AgentTurnCreateRequest,
from llama_stack.apis.agents import * # noqa: F403 AgentTurnResponseTurnCompletePayload,
)
from llama_stack.apis.inference import (
ChatCompletionResponse,
ChatCompletionResponseEvent,
ChatCompletionResponseStreamChunk,
CompletionMessage,
Message,
ResponseFormat,
SamplingParams,
ToolChoice,
ToolDefinition,
UserMessage,
)
from llama_stack.apis.memory import MemoryBank
from llama_stack.apis.safety import RunShieldResponse
from ..agents import ( from ..agents import (
AGENT_INSTANCES_BY_ID, AGENT_INSTANCES_BY_ID,