phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-04 12:07:34 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

addressing comments

Browse source
This commit is contained in:
Hardik Shah 2024-07-31 22:07:45 -07:00
parent 0e985648f5
commit fd8adc1e50
3 changed files with 9 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

8
tests/test_ollama_inference.py
View file

@ -9,6 +9,9 @@ from llama_models.llama3_1.api.datatypes import (
StopReason,
SystemMessage,
)
from llama_toolchain.inference.api_instance import (
get_inference_api_instance,
)
from llama_toolchain.inference.api.datatypes import (
ChatCompletionResponseEventType,
)
@ -16,6 +19,7 @@ from llama_toolchain.inference.api.endpoints import (
ChatCompletionRequest
)
from llama_toolchain.inference.api.config import (
InferenceConfig,
OllamaImplConfig
)
from llama_toolchain.inference.ollama import (
@ -32,7 +36,9 @@ class OllamaInferenceTests(unittest.IsolatedAsyncioTestCase):
)
# setup ollama
self.api = OllamaInference(ollama_config)
self.api = await get_inference_api_instance(
InferenceConfig(impl_config=ollama_config)
)
await self.api.initialize()
current_date = datetime.now()