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

llama_stack/providers/utils/responses/responses_store.py

@@ -0,0 +1,292 @@
+# 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 llama_stack.apis.agents import (
+    Order,
+)
+from llama_stack.apis.agents.openai_responses import (
+    ListOpenAIResponseInputItem,
+    ListOpenAIResponseObject,
+    OpenAIDeleteResponseObject,
+    OpenAIResponseInput,
+    OpenAIResponseObject,
+    OpenAIResponseObjectWithInput,
+)
+from llama_stack.apis.inference import OpenAIMessageParam
+from llama_stack.core.datatypes import AccessRule
+from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqlStoreReference
+from llama_stack.log import get_logger
+
+from ..sqlstore.api import ColumnDefinition, ColumnType
+from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from ..sqlstore.sqlstore import sqlstore_impl
+
+logger = get_logger(name=__name__, category="openai_responses")
+
+
+class _OpenAIResponseObjectWithInputAndMessages(OpenAIResponseObjectWithInput):
+    """Internal class for storing responses with chat completion messages.
+
+    This extends the public OpenAIResponseObjectWithInput with messages field
+    for internal storage. The messages field is not exposed in the public API.
+
+    The messages field is optional for backward compatibility with responses
+    stored before this feature was added.
+    """
+
+    messages: list[OpenAIMessageParam] | None = None
+
+
+class ResponsesStore:
+    def __init__(
+        self,
+        reference: ResponsesStoreReference | SqlStoreReference,
+        policy: list[AccessRule],
+    ):
+        if isinstance(reference, ResponsesStoreReference):
+            self.reference = reference
+        else:
+            self.reference = ResponsesStoreReference(**reference.model_dump())
+
+        self.policy = policy
+        self.sql_store = None
+
+    async def initialize(self):
+        """Create the necessary tables if they don't exist."""
+        base_store = sqlstore_impl(self.reference)
+        self.sql_store = AuthorizedSqlStore(base_store, self.policy)
+
+        await self.sql_store.create_table(
+            "openai_responses",
+            {
+                "id": ColumnDefinition(type=ColumnType.STRING, primary_key=True),
+                "created_at": ColumnType.INTEGER,
+                "response_object": ColumnType.JSON,
+                "model": ColumnType.STRING,
+            },
+        )
+
+        await self.sql_store.create_table(
+            "conversation_messages",
+            {
+                "conversation_id": ColumnDefinition(type=ColumnType.STRING, primary_key=True),
+                "messages": ColumnType.JSON,
+            },
+        )
+
+    async def shutdown(self) -> None:
+        return
+
+    async def flush(self) -> None:
+        """Maintained for compatibility; no-op now that writes are synchronous."""
+        return
+
+    async def store_response_object(
+        self,
+        response_object: OpenAIResponseObject,
+        input: list[OpenAIResponseInput],
+        messages: list[OpenAIMessageParam],
+    ) -> None:
+        await self._write_response_object(response_object, input, messages)
+
+    async def _write_response_object(
+        self,
+        response_object: OpenAIResponseObject,
+        input: list[OpenAIResponseInput],
+        messages: list[OpenAIMessageParam],
+    ) -> None:
+        if self.sql_store is None:
+            raise ValueError("Responses store is not initialized")
+
+        data = response_object.model_dump()
+        data["input"] = [input_item.model_dump() for input_item in input]
+        data["messages"] = [msg.model_dump() for msg in messages]
+
+        await self.sql_store.insert(
+            "openai_responses",
+            {
+                "id": data["id"],
+                "created_at": data["created_at"],
+                "model": data["model"],
+                "response_object": data,
+            },
+        )
+
+    async def list_responses(
+        self,
+        after: str | None = None,
+        limit: int | None = 50,
+        model: str | None = None,
+        order: Order | None = Order.desc,
+    ) -> ListOpenAIResponseObject:
+        """
+        List responses from the database.
+
+        :param after: The ID of the last response to return.
+        :param limit: The maximum number of responses to return.
+        :param model: The model to filter by.
+        :param order: The order to sort the responses by.
+        """
+        if not self.sql_store:
+            raise ValueError("Responses store is not initialized")
+
+        if not order:
+            order = Order.desc
+
+        where_conditions = {}
+        if model:
+            where_conditions["model"] = model
+
+        paginated_result = await self.sql_store.fetch_all(
+            table="openai_responses",
+            where=where_conditions if where_conditions else None,
+            order_by=[("created_at", order.value)],
+            cursor=("id", after) if after else None,
+            limit=limit,
+        )
+
+        data = [OpenAIResponseObjectWithInput(**row["response_object"]) for row in paginated_result.data]
+        return ListOpenAIResponseObject(
+            data=data,
+            has_more=paginated_result.has_more,
+            first_id=data[0].id if data else "",
+            last_id=data[-1].id if data else "",
+        )
+
+    async def get_response_object(self, response_id: str) -> _OpenAIResponseObjectWithInputAndMessages:
+        """
+        Get a response object with automatic access control checking.
+        """
+        if not self.sql_store:
+            raise ValueError("Responses store is not initialized")
+
+        row = await self.sql_store.fetch_one(
+            "openai_responses",
+            where={"id": response_id},
+        )
+
+        if not row:
+            # SecureSqlStore will return None if record doesn't exist OR access is denied
+            # This provides security by not revealing whether the record exists
+            raise ValueError(f"Response with id {response_id} not found") from None
+
+        return _OpenAIResponseObjectWithInputAndMessages(**row["response_object"])
+
+    async def delete_response_object(self, response_id: str) -> OpenAIDeleteResponseObject:
+        if not self.sql_store:
+            raise ValueError("Responses store is not initialized")
+
+        row = await self.sql_store.fetch_one("openai_responses", where={"id": response_id})
+        if not row:
+            raise ValueError(f"Response with id {response_id} not found")
+        await self.sql_store.delete("openai_responses", where={"id": response_id})
+        return OpenAIDeleteResponseObject(id=response_id)
+
+    async def list_response_input_items(
+        self,
+        response_id: str,
+        after: str | None = None,
+        before: str | None = None,
+        include: list[str] | None = None,
+        limit: int | None = 20,
+        order: Order | None = Order.desc,
+    ) -> ListOpenAIResponseInputItem:
+        """
+        List input items for a given response.
+
+        :param response_id: The ID of the response to retrieve input items for.
+        :param after: An item ID to list items after, used for pagination.
+        :param before: An item ID to list items before, used for pagination.
+        :param include: Additional fields to include in the response.
+        :param limit: A limit on the number of objects to be returned.
+        :param order: The order to return the input items in.
+        """
+        if include:
+            raise NotImplementedError("Include is not supported yet")
+        if before and after:
+            raise ValueError("Cannot specify both 'before' and 'after' parameters")
+
+        response_with_input_and_messages = await self.get_response_object(response_id)
+        items = response_with_input_and_messages.input
+
+        if order == Order.desc:
+            items = list(reversed(items))
+
+        start_index = 0
+        end_index = len(items)
+
+        if after or before:
+            for i, item in enumerate(items):
+                item_id = getattr(item, "id", None)
+                if after and item_id == after:
+                    start_index = i + 1
+                if before and item_id == before:
+                    end_index = i
+                    break
+
+            if after and start_index == 0:
+                raise ValueError(f"Input item with id '{after}' not found for response '{response_id}'")
+            if before and end_index == len(items):
+                raise ValueError(f"Input item with id '{before}' not found for response '{response_id}'")
+
+        items = items[start_index:end_index]
+
+        # Apply limit
+        if limit is not None:
+            items = items[:limit]
+
+        return ListOpenAIResponseInputItem(data=items)
+
+    async def store_conversation_messages(self, conversation_id: str, messages: list[OpenAIMessageParam]) -> None:
+        """Store messages for a conversation.
+
+        :param conversation_id: The conversation identifier.
+        :param messages: List of OpenAI message parameters to store.
+        """
+        if not self.sql_store:
+            raise ValueError("Responses store is not initialized")
+
+        # Serialize messages to dict format for JSON storage
+        messages_data = [msg.model_dump() for msg in messages]
+
+        # Upsert: try insert first, update if exists
+        try:
+            await self.sql_store.insert(
+                table="conversation_messages",
+                data={"conversation_id": conversation_id, "messages": messages_data},
+            )
+        except Exception:
+            # If insert fails due to ID conflict, update existing record
+            await self.sql_store.update(
+                table="conversation_messages",
+                data={"messages": messages_data},
+                where={"conversation_id": conversation_id},
+            )
+
+        logger.debug(f"Stored {len(messages)} messages for conversation {conversation_id}")
+
+    async def get_conversation_messages(self, conversation_id: str) -> list[OpenAIMessageParam] | None:
+        """Get stored messages for a conversation.
+
+        :param conversation_id: The conversation identifier.
+        :returns: List of OpenAI message parameters, or None if no messages stored.
+        """
+        if not self.sql_store:
+            raise ValueError("Responses store is not initialized")
+
+        record = await self.sql_store.fetch_one(
+            table="conversation_messages",
+            where={"conversation_id": conversation_id},
+        )
+
+        if record is None:
+            return None
+
+        # Deserialize messages from JSON storage
+        from pydantic import TypeAdapter
+
+        adapter = TypeAdapter(list[OpenAIMessageParam])
+        return adapter.validate_python(record["messages"])

llama_stack/providers/utils/scheduler.py

@@ -0,0 +1,270 @@
+# 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 abc
+import asyncio
+import functools
+import threading
+from collections.abc import Callable, Coroutine, Iterable
+from datetime import UTC, datetime
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel
+
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="providers::utils")
+
+
+# TODO: revisit the list of possible statuses when defining a more coherent
+# Jobs API for all API flows; e.g. do we need new vs scheduled?
+class JobStatus(Enum):
+    new = "new"
+    scheduled = "scheduled"
+    running = "running"
+    failed = "failed"
+    completed = "completed"
+
+
+type JobID = str
+type JobType = str
+
+
+class JobArtifact(BaseModel):
+    type: JobType
+    name: str
+    # TODO: uri should be a reference to /files API; revisit when /files is implemented
+    uri: str | None = None
+    metadata: dict[str, Any]
+
+
+JobHandler = Callable[
+    [Callable[[str], None], Callable[[JobStatus], None], Callable[[JobArtifact], None]], Coroutine[Any, Any, None]
+]
+
+
+type LogMessage = tuple[datetime, str]
+
+
+_COMPLETED_STATUSES = {JobStatus.completed, JobStatus.failed}
+
+
+class Job:
+    def __init__(self, job_type: JobType, job_id: JobID, handler: JobHandler):
+        super().__init__()
+        self.id = job_id
+        self._type = job_type
+        self._handler = handler
+        self._artifacts: list[JobArtifact] = []
+        self._logs: list[LogMessage] = []
+        self._state_transitions: list[tuple[datetime, JobStatus]] = [(datetime.now(UTC), JobStatus.new)]
+
+    @property
+    def handler(self) -> JobHandler:
+        return self._handler
+
+    @property
+    def status(self) -> JobStatus:
+        return self._state_transitions[-1][1]
+
+    @status.setter
+    def status(self, status: JobStatus):
+        if status in _COMPLETED_STATUSES and self.status in _COMPLETED_STATUSES:
+            raise ValueError(f"Job is already in a completed state ({self.status})")
+        if self.status == status:
+            return
+        self._state_transitions.append((datetime.now(UTC), status))
+
+    @property
+    def artifacts(self) -> list[JobArtifact]:
+        return self._artifacts
+
+    def register_artifact(self, artifact: JobArtifact) -> None:
+        self._artifacts.append(artifact)
+
+    def _find_state_transition_date(self, status: Iterable[JobStatus]) -> datetime | None:
+        for date, s in reversed(self._state_transitions):
+            if s in status:
+                return date
+        return None
+
+    @property
+    def scheduled_at(self) -> datetime | None:
+        return self._find_state_transition_date([JobStatus.scheduled])
+
+    @property
+    def started_at(self) -> datetime | None:
+        return self._find_state_transition_date([JobStatus.running])
+
+    @property
+    def completed_at(self) -> datetime | None:
+        return self._find_state_transition_date(_COMPLETED_STATUSES)
+
+    @property
+    def logs(self) -> list[LogMessage]:
+        return self._logs[:]
+
+    def append_log(self, message: LogMessage) -> None:
+        self._logs.append(message)
+
+    # TODO: implement
+    def cancel(self) -> None:
+        raise NotImplementedError
+
+
+class _SchedulerBackend(abc.ABC):
+    @abc.abstractmethod
+    def on_log_message_cb(self, job: Job, message: LogMessage) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def on_status_change_cb(self, job: Job, status: JobStatus) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def on_artifact_collected_cb(self, job: Job, artifact: JobArtifact) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def shutdown(self) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def schedule(
+        self,
+        job: Job,
+        on_log_message_cb: Callable[[str], None],
+        on_status_change_cb: Callable[[JobStatus], None],
+        on_artifact_collected_cb: Callable[[JobArtifact], None],
+    ) -> None:
+        raise NotImplementedError
+
+
+class _NaiveSchedulerBackend(_SchedulerBackend):
+    def __init__(self, timeout: int = 5):
+        self._timeout = timeout
+        self._loop = asyncio.new_event_loop()
+        # There may be performance implications of using threads due to Python
+        # GIL; may need to measure if it's a real problem though
+        self._thread = threading.Thread(target=self._run_loop, daemon=True)
+        self._thread.start()
+
+    def _run_loop(self) -> None:
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+        # TODO: When stopping the loop, give tasks a chance to finish
+        # TODO: should we explicitly inform jobs of pending stoppage?
+
+        # cancel all tasks
+        for task in asyncio.all_tasks(self._loop):
+            if not task.done():
+                task.cancel()
+
+        self._loop.close()
+
+    async def shutdown(self) -> None:
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        self._thread.join()
+
+    # TODO: decouple scheduling and running the job
+    def schedule(
+        self,
+        job: Job,
+        on_log_message_cb: Callable[[str], None],
+        on_status_change_cb: Callable[[JobStatus], None],
+        on_artifact_collected_cb: Callable[[JobArtifact], None],
+    ) -> None:
+        async def do():
+            try:
+                job.status = JobStatus.running
+                await job.handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb)
+            except Exception as e:
+                on_log_message_cb(str(e))
+                job.status = JobStatus.failed
+                logger.exception(f"Job {job.id} failed.")
+
+        asyncio.run_coroutine_threadsafe(do(), self._loop)
+
+    def on_log_message_cb(self, job: Job, message: LogMessage) -> None:
+        pass
+
+    def on_status_change_cb(self, job: Job, status: JobStatus) -> None:
+        pass
+
+    def on_artifact_collected_cb(self, job: Job, artifact: JobArtifact) -> None:
+        pass
+
+
+_BACKENDS = {
+    "naive": _NaiveSchedulerBackend,
+}
+
+
+def _get_backend_impl(backend: str) -> _SchedulerBackend:
+    try:
+        return _BACKENDS[backend]()
+    except KeyError as e:
+        raise ValueError(f"Unknown backend {backend}") from e
+
+
+class Scheduler:
+    def __init__(self, backend: str = "naive"):
+        # TODO: if server crashes, job states are lost; we need to persist jobs on disc
+        self._jobs: dict[JobID, Job] = {}
+        self._backend = _get_backend_impl(backend)
+
+    def _on_log_message_cb(self, job: Job, message: str) -> None:
+        msg = (datetime.now(UTC), message)
+        # At least for the time being, until there's a better way to expose
+        # logs to users, log messages on console
+        logger.info(f"Job {job.id}: {message}")
+        job.append_log(msg)
+        self._backend.on_log_message_cb(job, msg)
+
+    def _on_status_change_cb(self, job: Job, status: JobStatus) -> None:
+        job.status = status
+        self._backend.on_status_change_cb(job, status)
+
+    def _on_artifact_collected_cb(self, job: Job, artifact: JobArtifact) -> None:
+        job.register_artifact(artifact)
+        self._backend.on_artifact_collected_cb(job, artifact)
+
+    def schedule(self, type_: JobType, job_id: JobID, handler: JobHandler) -> JobID:
+        job = Job(type_, job_id, handler)
+        if job.id in self._jobs:
+            raise ValueError(f"Job {job.id} already exists")
+
+        self._jobs[job.id] = job
+        job.status = JobStatus.scheduled
+        self._backend.schedule(
+            job,
+            functools.partial(self._on_log_message_cb, job),
+            functools.partial(self._on_status_change_cb, job),
+            functools.partial(self._on_artifact_collected_cb, job),
+        )
+
+        return job.id
+
+    def cancel(self, job_id: JobID) -> None:
+        self.get_job(job_id).cancel()
+
+    def get_job(self, job_id: JobID) -> Job:
+        try:
+            return self._jobs[job_id]
+        except KeyError as e:
+            raise ValueError(f"Job {job_id} not found") from e
+
+    def get_jobs(self, type_: JobType | None = None) -> list[Job]:
+        jobs = list(self._jobs.values())
+        if type_:
+            jobs = [job for job in jobs if job._type == type_]
+        return jobs
+
+    async def shutdown(self):
+        # TODO: also cancel jobs once implemented
+        await self._backend.shutdown()

llama_stack/providers/utils/scoring/__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/utils/scoring/aggregation_utils.py

@@ -0,0 +1,75 @@
+# 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 statistics
+from typing import Any
+
+from llama_stack.apis.scoring import ScoringResultRow
+from llama_stack.apis.scoring_functions import AggregationFunctionType
+
+
+def aggregate_accuracy(scoring_results: list[ScoringResultRow]) -> dict[str, Any]:
+    num_correct = sum(result["score"] for result in scoring_results)
+    avg_score = num_correct / len(scoring_results)
+
+    return {
+        "accuracy": avg_score,
+        "num_correct": num_correct,
+        "num_total": len(scoring_results),
+    }
+
+
+def aggregate_average(scoring_results: list[ScoringResultRow]) -> dict[str, Any]:
+    return {
+        "average": sum(result["score"] for result in scoring_results if result["score"] is not None)
+        / len([_ for _ in scoring_results if _["score"] is not None]),
+    }
+
+
+def aggregate_weighted_average(scoring_results: list[ScoringResultRow]) -> dict[str, Any]:
+    return {
+        "weighted_average": sum(
+            result["score"] * result["weight"]
+            for result in scoring_results
+            if result["score"] is not None and result["weight"] is not None
+        )
+        / sum(result["weight"] for result in scoring_results if result["weight"] is not None),
+    }
+
+
+def aggregate_categorical_count(
+    scoring_results: list[ScoringResultRow],
+) -> dict[str, Any]:
+    scores = [str(r["score"]) for r in scoring_results]
+    unique_scores = sorted(set(scores))
+    return {"categorical_count": {s: scores.count(s) for s in unique_scores}}
+
+
+def aggregate_median(scoring_results: list[ScoringResultRow]) -> dict[str, Any]:
+    scores = [r["score"] for r in scoring_results if r["score"] is not None]
+    median = statistics.median(scores) if scores else None
+    return {"median": median}
+
+
+# TODO: decide whether we want to make aggregation functions as a registerable resource
+AGGREGATION_FUNCTIONS = {
+    AggregationFunctionType.accuracy: aggregate_accuracy,
+    AggregationFunctionType.average: aggregate_average,
+    AggregationFunctionType.weighted_average: aggregate_weighted_average,
+    AggregationFunctionType.categorical_count: aggregate_categorical_count,
+    AggregationFunctionType.median: aggregate_median,
+}
+
+
+def aggregate_metrics(
+    scoring_results: list[ScoringResultRow], metrics: list[AggregationFunctionType]
+) -> dict[str, Any]:
+    agg_results = {}
+    for metric in metrics:
+        if metric not in AGGREGATION_FUNCTIONS:
+            raise ValueError(f"Aggregation function {metric} not found")
+        agg_fn = AGGREGATION_FUNCTIONS[metric]
+        agg_results[metric] = agg_fn(scoring_results)
+    return agg_results