RHAIENG-191 - Use LlamaStackDistribution from openshift wrappers (#467)

* RHAIENG-191 - Use LlamaStackDistribution from openshift wrappers

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: jiripetrlik

tests/rag/conftest.py

@@ -11,13 +11,14 @@
 from ocp_resources.namespace import Namespace
 from ocp_resources.project_project_openshift_io import Project
 from simple_logger.logger import get_logger
-from timeout_sampler import TimeoutSampler, retry
+from timeout_sampler import retry
 
 from utilities.constants import DscComponents, Timeout
 from utilities.data_science_cluster_utils import update_components_in_dsc
 from utilities.general import generate_random_name
 from utilities.infra import create_ns
-from utilities.rag_utils import LlamaStackDistribution, create_llama_stack_distribution
+from ocp_resources.llama_stack_distribution import LlamaStackDistribution
+from utilities.rag_utils import create_llama_stack_distribution
 
 LOGGER = get_logger(name=__name__)
 
@@ -38,14 +39,12 @@ def llama_stack_server() -> Dict[str, Any]:
                 {"name": "VLLM_TLS_VERIFY", "value": "false"},
                 {"name": "VLLM_API_TOKEN", "value": rag_vllm_token},
                 {"name": "VLLM_URL", "value": rag_vllm_url},
-                {"name": "MILVUS_DB_PATH", "value": "~/.llama/milvus.db"},
                 {"name": "FMS_ORCHESTRATOR_URL", "value": "http://localhost"},
             ],
             "name": "llama-stack",
             "port": 8321,
         },
-        "distribution": {"image": "quay.io/opendatahub/llama-stack:odh"},
-        "storage": {"size": "5Gi"},
+        "distribution": {"name": "rh-dev"},
     }
 
 
@@ -99,15 +98,7 @@ def llama_stack_distribution_deployment(
         name="rag-llama-stack-distribution",
     )
 
-    timeout = Timeout.TIMEOUT_15_SEC
-    sampler = TimeoutSampler(
-        wait_timeout=timeout, sleep=1, func=lambda deployment: deployment.exists is not None, deployment=deployment
-    )
-    for item in sampler:
-        if item:
-            break  # Break after first successful iteration
-
-    assert deployment.exists, f"llama stack distribution deployment doesn't exist within {timeout} seconds"
+    deployment.wait(timeout=Timeout.TIMEOUT_2MIN)
     yield deployment
 
 

utilities/rag_utils.py

@@ -1,5 +1,5 @@
 from contextlib import contextmanager
-from ocp_resources.resource import NamespacedResource
+from ocp_resources.llama_stack_distribution import LlamaStackDistribution
 from kubernetes.dynamic import DynamicClient
 from typing import Any, Dict, Generator, List, TypedDict, cast
 from llama_stack_client import Agent, AgentEventLogger
@@ -8,29 +8,6 @@
 LOGGER = get_logger(name=__name__)
 
 
-class LlamaStackDistribution(NamespacedResource):
-    api_group: str = "llamastack.io"
-
-    def __init__(self, replicas: int, server: Dict[str, Any], **kwargs: Any):
-        """
-        Args:
-            kwargs: Keyword arguments to pass to the LlamaStackDistribution constructor
-        """
-        super().__init__(
-            **kwargs,
-        )
-        self.replicas = replicas
-        self.server = server
-
-    def to_dict(self) -> None:
-        super().to_dict()
-        if not self.kind_dict and not self.yaml_file:
-            self.res["spec"] = {}
-            _spec = self.res["spec"]
-            _spec["replicas"] = self.replicas
-            _spec["server"] = self.server
-
-
 @contextmanager
 def create_llama_stack_distribution(
     client: DynamicClient,