llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/inline/batches/reference/config.py

@@ -0,0 +1,40 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from pydantic import BaseModel, Field
+
+from llama_stack.core.storage.datatypes import KVStoreReference
+
+
+class ReferenceBatchesImplConfig(BaseModel):
+    """Configuration for the Reference Batches implementation."""
+
+    kvstore: KVStoreReference = Field(
+        description="Configuration for the key-value store backend.",
+    )
+
+    max_concurrent_batches: int = Field(
+        default=1,
+        description="Maximum number of concurrent batches to process simultaneously.",
+        ge=1,
+    )
+
+    max_concurrent_requests_per_batch: int = Field(
+        default=10,
+        description="Maximum number of concurrent requests to process per batch.",
+        ge=1,
+    )
+
+    # TODO: add a max requests per second rate limiter
+
+    @classmethod
+    def sample_run_config(cls, __distro_dir__: str) -> dict:
+        return {
+            "kvstore": KVStoreReference(
+                backend="kv_default",
+                namespace="batches",
+            ).model_dump(exclude_none=True),
+        }

llama_stack/providers/inline/datasetio/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_stack/providers/inline/datasetio/localfs/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any
+
+from .config import LocalFSDatasetIOConfig
+
+
+async def get_provider_impl(
+    config: LocalFSDatasetIOConfig,
+    _deps: dict[str, Any],
+):
+    from .datasetio import LocalFSDatasetIOImpl
+
+    impl = LocalFSDatasetIOImpl(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/inline/datasetio/localfs/config.py

@@ -0,0 +1,23 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+from typing import Any
+
+from pydantic import BaseModel
+
+from llama_stack.core.storage.datatypes import KVStoreReference
+
+
+class LocalFSDatasetIOConfig(BaseModel):
+    kvstore: KVStoreReference
+
+    @classmethod
+    def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
+        return {
+            "kvstore": KVStoreReference(
+                backend="kv_default",
+                namespace="datasetio::localfs",
+            ).model_dump(exclude_none=True)
+        }

llama_stack/providers/inline/datasetio/localfs/datasetio.py

@@ -0,0 +1,113 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+from typing import Any
+
+from llama_stack.apis.common.responses import PaginatedResponse
+from llama_stack.apis.datasetio import DatasetIO
+from llama_stack.apis.datasets import Dataset
+from llama_stack.providers.datatypes import DatasetsProtocolPrivate
+from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_uri
+from llama_stack.providers.utils.kvstore import kvstore_impl
+from llama_stack.providers.utils.pagination import paginate_records
+
+from .config import LocalFSDatasetIOConfig
+
+DATASETS_PREFIX = "localfs_datasets:"
+
+
+class PandasDataframeDataset:
+    def __init__(self, dataset_def: Dataset, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.dataset_def = dataset_def
+        self.df = None
+
+    def __len__(self) -> int:
+        assert self.df is not None, "Dataset not loaded. Please call .load() first"
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        assert self.df is not None, "Dataset not loaded. Please call .load() first"
+        if isinstance(idx, slice):
+            return self.df.iloc[idx].to_dict(orient="records")
+        else:
+            return self.df.iloc[idx].to_dict()
+
+    async def load(self) -> None:
+        if self.df is not None:
+            return
+
+        if self.dataset_def.source.type == "uri":
+            self.df = await get_dataframe_from_uri(self.dataset_def.source.uri)
+        elif self.dataset_def.source.type == "rows":
+            import pandas
+
+            self.df = pandas.DataFrame(self.dataset_def.source.rows)
+        else:
+            raise ValueError(f"Unsupported dataset source type: {self.dataset_def.source.type}")
+
+        if self.df is None:
+            raise ValueError(f"Failed to load dataset from {self.dataset_def.url}")
+
+
+class LocalFSDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
+    def __init__(self, config: LocalFSDatasetIOConfig) -> None:
+        self.config = config
+        # local registry for keeping track of datasets within the provider
+        self.dataset_infos = {}
+        self.kvstore = None
+
+    async def initialize(self) -> None:
+        self.kvstore = await kvstore_impl(self.config.kvstore)
+        # Load existing datasets from kvstore
+        start_key = DATASETS_PREFIX
+        end_key = f"{DATASETS_PREFIX}\xff"
+        stored_datasets = await self.kvstore.values_in_range(start_key, end_key)
+
+        for dataset in stored_datasets:
+            dataset = Dataset.model_validate_json(dataset)
+            self.dataset_infos[dataset.identifier] = dataset
+
+    async def shutdown(self) -> None: ...
+
+    async def register_dataset(
+        self,
+        dataset_def: Dataset,
+    ) -> None:
+        # Store in kvstore
+        key = f"{DATASETS_PREFIX}{dataset_def.identifier}"
+        await self.kvstore.set(
+            key=key,
+            value=dataset_def.model_dump_json(),
+        )
+        self.dataset_infos[dataset_def.identifier] = dataset_def
+
+    async def unregister_dataset(self, dataset_id: str) -> None:
+        key = f"{DATASETS_PREFIX}{dataset_id}"
+        await self.kvstore.delete(key=key)
+        del self.dataset_infos[dataset_id]
+
+    async def iterrows(
+        self,
+        dataset_id: str,
+        start_index: int | None = None,
+        limit: int | None = None,
+    ) -> PaginatedResponse:
+        dataset_def = self.dataset_infos[dataset_id]
+        dataset_impl = PandasDataframeDataset(dataset_def)
+        await dataset_impl.load()
+
+        records = dataset_impl.df.to_dict("records")
+        return paginate_records(records, start_index, limit)
+
+    async def append_rows(self, dataset_id: str, rows: list[dict[str, Any]]) -> None:
+        import pandas
+
+        dataset_def = self.dataset_infos[dataset_id]
+        dataset_impl = PandasDataframeDataset(dataset_def)
+        await dataset_impl.load()
+
+        new_rows_df = pandas.DataFrame(rows)
+        dataset_impl.df = pandas.concat([dataset_impl.df, new_rows_df], ignore_index=True)

llama_stack/providers/inline/eval/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_stack/providers/inline/eval/meta_reference/__init__.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+from typing import Any
+
+from llama_stack.core.datatypes import Api
+
+from .config import MetaReferenceEvalConfig
+
+
+async def get_provider_impl(
+    config: MetaReferenceEvalConfig,
+    deps: dict[Api, Any],
+):
+    from .eval import MetaReferenceEvalImpl
+
+    impl = MetaReferenceEvalImpl(
+        config,
+        deps[Api.datasetio],
+        deps[Api.datasets],
+        deps[Api.scoring],
+        deps[Api.inference],
+        deps[Api.agents],
+    )
+    await impl.initialize()
+    return impl

llama_stack/providers/inline/eval/meta_reference/config.py

@@ -0,0 +1,23 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+from typing import Any
+
+from pydantic import BaseModel
+
+from llama_stack.core.storage.datatypes import KVStoreReference
+
+
+class MetaReferenceEvalConfig(BaseModel):
+    kvstore: KVStoreReference
+
+    @classmethod
+    def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
+        return {
+            "kvstore": KVStoreReference(
+                backend="kv_default",
+                namespace="eval",
+            ).model_dump(exclude_none=True)
+        }

llama_stack/providers/inline/eval/meta_reference/eval.py

@@ -0,0 +1,259 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+import json
+from typing import Any
+
+from tqdm import tqdm
+
+from llama_stack.apis.agents import Agents, StepType
+from llama_stack.apis.benchmarks import Benchmark
+from llama_stack.apis.datasetio import DatasetIO
+from llama_stack.apis.datasets import Datasets
+from llama_stack.apis.inference import (
+    Inference,
+    OpenAIChatCompletionRequestWithExtraBody,
+    OpenAICompletionRequestWithExtraBody,
+    OpenAISystemMessageParam,
+    OpenAIUserMessageParam,
+    UserMessage,
+)
+from llama_stack.apis.scoring import Scoring
+from llama_stack.providers.datatypes import BenchmarksProtocolPrivate
+from llama_stack.providers.inline.agents.meta_reference.agent_instance import (
+    MEMORY_QUERY_TOOL,
+)
+from llama_stack.providers.utils.common.data_schema_validator import ColumnName
+from llama_stack.providers.utils.kvstore import kvstore_impl
+
+from .....apis.common.job_types import Job, JobStatus
+from .....apis.eval.eval import BenchmarkConfig, Eval, EvaluateResponse
+from .config import MetaReferenceEvalConfig
+
+EVAL_TASKS_PREFIX = "benchmarks:"
+
+
+class MetaReferenceEvalImpl(
+    Eval,
+    BenchmarksProtocolPrivate,
+):
+    def __init__(
+        self,
+        config: MetaReferenceEvalConfig,
+        datasetio_api: DatasetIO,
+        datasets_api: Datasets,
+        scoring_api: Scoring,
+        inference_api: Inference,
+        agents_api: Agents,
+    ) -> None:
+        self.config = config
+        self.datasetio_api = datasetio_api
+        self.datasets_api = datasets_api
+        self.scoring_api = scoring_api
+        self.inference_api = inference_api
+        self.agents_api = agents_api
+
+        # TODO: assume sync job, will need jobs API for async scheduling
+        self.jobs = {}
+
+        self.benchmarks = {}
+
+    async def initialize(self) -> None:
+        self.kvstore = await kvstore_impl(self.config.kvstore)
+        # Load existing benchmarks from kvstore
+        start_key = EVAL_TASKS_PREFIX
+        end_key = f"{EVAL_TASKS_PREFIX}\xff"
+        stored_benchmarks = await self.kvstore.values_in_range(start_key, end_key)
+
+        for benchmark in stored_benchmarks:
+            benchmark = Benchmark.model_validate_json(benchmark)
+            self.benchmarks[benchmark.identifier] = benchmark
+
+    async def shutdown(self) -> None: ...
+
+    async def register_benchmark(self, task_def: Benchmark) -> None:
+        # Store in kvstore
+        key = f"{EVAL_TASKS_PREFIX}{task_def.identifier}"
+        await self.kvstore.set(
+            key=key,
+            value=task_def.model_dump_json(),
+        )
+        self.benchmarks[task_def.identifier] = task_def
+
+    async def unregister_benchmark(self, benchmark_id: str) -> None:
+        if benchmark_id in self.benchmarks:
+            del self.benchmarks[benchmark_id]
+
+        key = f"{EVAL_TASKS_PREFIX}{benchmark_id}"
+        await self.kvstore.delete(key)
+
+    async def run_eval(
+        self,
+        benchmark_id: str,
+        benchmark_config: BenchmarkConfig,
+    ) -> Job:
+        task_def = self.benchmarks[benchmark_id]
+        dataset_id = task_def.dataset_id
+        scoring_functions = task_def.scoring_functions
+
+        # TODO (xiyan): validate dataset schema
+        # dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
+
+        all_rows = await self.datasetio_api.iterrows(
+            dataset_id=dataset_id,
+            limit=(-1 if benchmark_config.num_examples is None else benchmark_config.num_examples),
+        )
+        res = await self.evaluate_rows(
+            benchmark_id=benchmark_id,
+            input_rows=all_rows.data,
+            scoring_functions=scoring_functions,
+            benchmark_config=benchmark_config,
+        )
+
+        # TODO: currently needs to wait for generation before returning
+        # need job scheduler queue (ray/celery) w/ jobs api
+        job_id = str(len(self.jobs))
+        self.jobs[job_id] = res
+        return Job(job_id=job_id, status=JobStatus.completed)
+
+    async def _run_agent_generation(
+        self, input_rows: list[dict[str, Any]], benchmark_config: BenchmarkConfig
+    ) -> list[dict[str, Any]]:
+        candidate = benchmark_config.eval_candidate
+        create_response = await self.agents_api.create_agent(candidate.config)
+        agent_id = create_response.agent_id
+
+        generations = []
+        for i, x in tqdm(enumerate(input_rows)):
+            assert ColumnName.chat_completion_input.value in x, "Invalid input row"
+            input_messages = json.loads(x[ColumnName.chat_completion_input.value])
+            input_messages = [UserMessage(**x) for x in input_messages if x["role"] == "user"]
+
+            # NOTE: only single-turn agent generation is supported. Create a new session for each input row
+            session_create_response = await self.agents_api.create_agent_session(agent_id, f"session-{i}")
+            session_id = session_create_response.session_id
+
+            turn_request = dict(
+                agent_id=agent_id,
+                session_id=session_id,
+                messages=input_messages,
+                stream=True,
+            )
+            turn_response = [chunk async for chunk in await self.agents_api.create_agent_turn(**turn_request)]
+            final_event = turn_response[-1].event.payload
+
+            # check if there's a memory retrieval step and extract the context
+            memory_rag_context = None
+            for step in final_event.turn.steps:
+                if step.step_type == StepType.tool_execution.value:
+                    for tool_response in step.tool_responses:
+                        if tool_response.tool_name == MEMORY_QUERY_TOOL:
+                            memory_rag_context = " ".join(x.text for x in tool_response.content)
+
+            agent_generation = {}
+            agent_generation[ColumnName.generated_answer.value] = final_event.turn.output_message.content
+            if memory_rag_context:
+                agent_generation[ColumnName.context.value] = memory_rag_context
+
+            generations.append(agent_generation)
+
+        return generations
+
+    async def _run_model_generation(
+        self, input_rows: list[dict[str, Any]], benchmark_config: BenchmarkConfig
+    ) -> list[dict[str, Any]]:
+        candidate = benchmark_config.eval_candidate
+        assert candidate.sampling_params.max_tokens is not None, "SamplingParams.max_tokens must be provided"
+        sampling_params = {"max_tokens": candidate.sampling_params.max_tokens}
+
+        generations = []
+        for x in tqdm(input_rows):
+            if ColumnName.completion_input.value in x:
+                if candidate.sampling_params.stop:
+                    sampling_params["stop"] = candidate.sampling_params.stop
+
+                input_content = json.loads(x[ColumnName.completion_input.value])
+                params = OpenAICompletionRequestWithExtraBody(
+                    model=candidate.model,
+                    prompt=input_content,
+                    **sampling_params,
+                )
+                response = await self.inference_api.openai_completion(params)
+                generations.append({ColumnName.generated_answer.value: response.choices[0].text})
+            elif ColumnName.chat_completion_input.value in x:
+                chat_completion_input_json = json.loads(x[ColumnName.chat_completion_input.value])
+                input_messages = [
+                    OpenAIUserMessageParam(**x) for x in chat_completion_input_json if x["role"] == "user"
+                ]
+
+                messages = []
+                if candidate.system_message:
+                    messages.append(candidate.system_message)
+
+                messages += [OpenAISystemMessageParam(**x) for x in chat_completion_input_json if x["role"] == "system"]
+
+                messages += input_messages
+                params = OpenAIChatCompletionRequestWithExtraBody(
+                    model=candidate.model,
+                    messages=messages,
+                    **sampling_params,
+                )
+                response = await self.inference_api.openai_chat_completion(params)
+                generations.append({ColumnName.generated_answer.value: response.choices[0].message.content})
+            else:
+                raise ValueError("Invalid input row")
+
+        return generations
+
+    async def evaluate_rows(
+        self,
+        benchmark_id: str,
+        input_rows: list[dict[str, Any]],
+        scoring_functions: list[str],
+        benchmark_config: BenchmarkConfig,
+    ) -> EvaluateResponse:
+        candidate = benchmark_config.eval_candidate
+        if candidate.type == "agent":
+            generations = await self._run_agent_generation(input_rows, benchmark_config)
+        elif candidate.type == "model":
+            generations = await self._run_model_generation(input_rows, benchmark_config)
+        else:
+            raise ValueError(f"Invalid candidate type: {candidate.type}")
+
+        # scoring with generated_answer
+        score_input_rows = [
+            input_r | generated_r for input_r, generated_r in zip(input_rows, generations, strict=False)
+        ]
+
+        if benchmark_config.scoring_params is not None:
+            scoring_functions_dict = {
+                scoring_fn_id: benchmark_config.scoring_params.get(scoring_fn_id, None)
+                for scoring_fn_id in scoring_functions
+            }
+        else:
+            scoring_functions_dict = dict.fromkeys(scoring_functions)
+
+        score_response = await self.scoring_api.score(
+            input_rows=score_input_rows, scoring_functions=scoring_functions_dict
+        )
+
+        return EvaluateResponse(generations=generations, scores=score_response.results)
+
+    async def job_status(self, benchmark_id: str, job_id: str) -> Job:
+        if job_id in self.jobs:
+            return Job(job_id=job_id, status=JobStatus.completed)
+
+        raise ValueError(f"Job {job_id} not found")
+
+    async def job_cancel(self, benchmark_id: str, job_id: str) -> None:
+        raise NotImplementedError("Job cancel is not implemented yet")
+
+    async def job_result(self, benchmark_id: str, job_id: str) -> EvaluateResponse:
+        job = await self.job_status(benchmark_id, job_id)
+        status = job.status
+        if not status or status != JobStatus.completed:
+            raise ValueError(f"Job is not completed, Status: {status.value}")
+
+        return self.jobs[job_id]

llama_stack/providers/inline/files/localfs/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any
+
+from llama_stack.core.datatypes import AccessRule, Api
+
+from .config import LocalfsFilesImplConfig
+from .files import LocalfsFilesImpl
+
+__all__ = ["LocalfsFilesImpl", "LocalfsFilesImplConfig"]
+
+
+async def get_provider_impl(config: LocalfsFilesImplConfig, deps: dict[Api, Any], policy: list[AccessRule]):
+    impl = LocalfsFilesImpl(config, policy)
+    await impl.initialize()
+    return impl

llama_stack/providers/inline/files/localfs/config.py

@@ -0,0 +1,31 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from llama_stack.core.storage.datatypes import SqlStoreReference
+
+
+class LocalfsFilesImplConfig(BaseModel):
+    storage_dir: str = Field(
+        description="Directory to store uploaded files",
+    )
+    metadata_store: SqlStoreReference = Field(
+        description="SQL store configuration for file metadata",
+    )
+    ttl_secs: int = 365 * 24 * 60 * 60  # 1 year
+
+    @classmethod
+    def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]:
+        return {
+            "storage_dir": "${env.FILES_STORAGE_DIR:=" + __distro_dir__ + "/files}",
+            "metadata_store": SqlStoreReference(
+                backend="sql_default",
+                table_name="files_metadata",
+            ).model_dump(exclude_none=True),
+        }