|
9 | 9 | from typing import Any, Protocol |
10 | 10 | from unittest.mock import AsyncMock, MagicMock |
11 | 11 |
|
12 | | -from pydantic import BaseModel, Field |
13 | | - |
14 | | -from llama_stack.apis.inference import Inference |
15 | | -from llama_stack.core.datatypes import ( |
16 | | - Api, |
17 | | - Provider, |
18 | | - StackRunConfig, |
19 | | -) |
| 12 | +from llama_stack.apis.inference import Inference, SamplingParams |
| 13 | +from llama_stack.core.datatypes import Api, Provider, StackRunConfig |
20 | 14 | from llama_stack.core.resolver import resolve_impls |
21 | 15 | from llama_stack.core.routers.inference import InferenceRouter |
22 | 16 | from llama_stack.core.routing_tables.models import ModelsRoutingTable |
23 | 17 | from llama_stack.providers.datatypes import InlineProviderSpec, ProviderSpec |
24 | 18 |
|
| 19 | +from pydantic import BaseModel, Field |
| 20 | + |
25 | 21 |
|
26 | 22 | def add_protocol_methods(cls: type, protocol: type[Protocol]) -> None: |
27 | 23 | """Dynamically add protocol methods to a class by inspecting the protocol.""" |
@@ -54,7 +50,12 @@ def sample_run_config(cls, **kwargs: Any) -> dict[str, Any]: |
54 | 50 |
|
55 | 51 |
|
56 | 52 | class SampleImpl: |
57 | | - def __init__(self, config: SampleConfig, deps: dict[Api, Any], provider_spec: ProviderSpec = None): |
| 53 | + def __init__( |
| 54 | + self, |
| 55 | + config: SampleConfig, |
| 56 | + deps: dict[Api, Any], |
| 57 | + provider_spec: ProviderSpec = None, |
| 58 | + ): |
58 | 59 | self.__provider_id__ = "test_provider" |
59 | 60 | self.__provider_spec__ = provider_spec |
60 | 61 | self.__provider_config__ = config |
|
0 commit comments