From 967ef9d1b6d605ce251c505dc5a9d89139bbccbc Mon Sep 17 00:00:00 2001 From: Pavel Tisnovsky Date: Wed, 21 May 2025 10:21:32 +0200 Subject: [PATCH 1/3] Use client.py --- src/app/endpoints/models.py | 7 ++++++- src/app/endpoints/query.py | 22 ++-------------------- 2 files changed, 8 insertions(+), 21 deletions(-) diff --git a/src/app/endpoints/models.py b/src/app/endpoints/models.py index f2bc390f..d6be1149 100644 --- a/src/app/endpoints/models.py +++ b/src/app/endpoints/models.py @@ -6,6 +6,8 @@ from fastapi import APIRouter, Request from llama_stack_client import LlamaStackClient +from client import get_llama_stack_client +from configuration import configuration from models.responses import ModelsResponse logger = logging.getLogger(__name__) @@ -40,7 +42,10 @@ @router.get("/models", responses=models_responses) def models_endpoint_handler(request: Request) -> ModelsResponse: - client = LlamaStackClient(base_url="http://localhost:8321") + llama_stack_config = configuration.llama_stack_configuration + logger.info("LLama stack config: %s", llama_stack_config) + + client = get_llama_stack_client(llama_stack_config) models = client.models.list() m = [dict(m) for m in models] return ModelsResponse(models=m) diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index d4c68631..15f26fdf 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -3,15 +3,14 @@ import logging from typing import Any -from llama_stack.distribution.library_client import LlamaStackAsLibraryClient from llama_stack_client.lib.agents.agent import Agent from llama_stack_client import LlamaStackClient from llama_stack_client.types import UserMessage from fastapi import APIRouter, Request +from client import get_llama_stack_client from configuration import configuration -from models.config import LLamaStackConfiguration from models.responses import QueryResponse logger = logging.getLogger("app.endpoints.handlers") @@ -27,7 +26,7 @@ @router.post("/query", responses=query_response) -def info_endpoint_handler(request: Request, query: str) -> QueryResponse: +def query_endpoint_handler(request: Request, query: str) -> QueryResponse: llama_stack_config = configuration.llama_stack_configuration logger.info("LLama stack config: %s", llama_stack_config) @@ -70,20 +69,3 @@ def retrieve_response(client: LlamaStackClient, model_id: str, prompt: str) -> s ) return str(response.output_message.content) - - -def get_llama_stack_client( - llama_stack_config: LLamaStackConfiguration, -) -> LlamaStackClient: - if llama_stack_config.use_as_library_client is True: - logger.info("Using Llama stack as library client") - client = LlamaStackAsLibraryClient( - llama_stack_config.library_client_config_path - ) - client.initialize() - return client - else: - logger.info("Using Llama stack running as a service") - return LlamaStackClient( - base_url=llama_stack_config.url, api_key=llama_stack_config.api_key - ) From 7d3df28cee704ba7c349fd93c4dba12358d5f836 Mon Sep 17 00:00:00 2001 From: Pavel Tisnovsky Date: Wed, 21 May 2025 10:21:47 +0200 Subject: [PATCH 2/3] Getting llama-stack client --- src/client.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 src/client.py diff --git a/src/client.py b/src/client.py new file mode 100644 index 00000000..e8f37cbe --- /dev/null +++ b/src/client.py @@ -0,0 +1,26 @@ +"""LLama stack client retrieval.""" + +import logging + +from llama_stack.distribution.library_client import LlamaStackAsLibraryClient +from llama_stack_client import LlamaStackClient +from models.config import LLamaStackConfiguration + +logger = logging.getLogger(__name__) + + +def get_llama_stack_client( + llama_stack_config: LLamaStackConfiguration, +) -> LlamaStackClient: + if llama_stack_config.use_as_library_client is True: + logger.info("Using Llama stack as library client") + client = LlamaStackAsLibraryClient( + llama_stack_config.library_client_config_path + ) + client.initialize() + return client + else: + logger.info("Using Llama stack running as a service") + return LlamaStackClient( + base_url=llama_stack_config.url, api_key=llama_stack_config.api_key + ) From 4bfb05a9b88ae283c105b7920fc40bcf1c76b9a2 Mon Sep 17 00:00:00 2001 From: Pavel Tisnovsky Date: Wed, 21 May 2025 10:22:18 +0200 Subject: [PATCH 3/3] Unit tests for the client.py --- tests/unit/test_client.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 tests/unit/test_client.py diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py new file mode 100644 index 00000000..888abc2e --- /dev/null +++ b/tests/unit/test_client.py @@ -0,0 +1,31 @@ +"""Unit tests for functions defined in src/client.py.""" + +import os +from unittest.mock import patch + +from client import get_llama_stack_client +from models.config import LLamaStackConfiguration + + +@patch.dict(os.environ, {"INFERENCE_MODEL": "llama3.2:3b-instruct-fp16"}) +def test_get_llama_stack_library_client(): + cfg = LLamaStackConfiguration( + url=None, + api_key=None, + use_as_library_client=True, + library_client_config_path="ollama", + ) + client = get_llama_stack_client(cfg) + assert client is not None + + +@patch.dict(os.environ, {"INFERENCE_MODEL": "llama3.2:3b-instruct-fp16"}) +def test_get_llama_stack_remote_client(): + cfg = LLamaStackConfiguration( + url="http://localhost:8321", + api_key=None, + use_as_library_client=False, + library_client_config_path="ollama", + ) + client = get_llama_stack_client(cfg) + assert client is not None