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

llama_stack/providers/inline/agents/meta_reference/persistence.py

@@ -1,228 +0,0 @@
-# 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
-import uuid
-from datetime import UTC, datetime
-
-from llama_stack.apis.agents import AgentConfig, Session, ToolExecutionStep, Turn
-from llama_stack.apis.common.errors import SessionNotFoundError
-from llama_stack.core.access_control.access_control import AccessDeniedError, is_action_allowed
-from llama_stack.core.access_control.datatypes import AccessRule
-from llama_stack.core.datatypes import User
-from llama_stack.core.request_headers import get_authenticated_user
-from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore import KVStore
-
-log = get_logger(name=__name__, category="agents::meta_reference")
-
-
-class AgentSessionInfo(Session):
-    # TODO: is this used anywhere?
-    vector_db_id: str | None = None
-    started_at: datetime
-    owner: User | None = None
-    identifier: str | None = None
-    type: str = "session"
-
-
-class AgentInfo(AgentConfig):
-    created_at: datetime
-
-
-class AgentPersistence:
-    def __init__(self, agent_id: str, kvstore: KVStore, policy: list[AccessRule]):
-        self.agent_id = agent_id
-        self.kvstore = kvstore
-        self.policy = policy
-
-    async def create_session(self, name: str) -> str:
-        session_id = str(uuid.uuid4())
-
-        # Get current user's auth attributes for new sessions
-        user = get_authenticated_user()
-
-        session_info = AgentSessionInfo(
-            session_id=session_id,
-            session_name=name,
-            started_at=datetime.now(UTC),
-            owner=user,
-            turns=[],
-            identifier=name,  # should this be qualified in any way?
-        )
-        if not is_action_allowed(self.policy, "create", session_info, user):
-            raise AccessDeniedError("create", session_info, user)
-
-        await self.kvstore.set(
-            key=f"session:{self.agent_id}:{session_id}",
-            value=session_info.model_dump_json(),
-        )
-        return session_id
-
-    async def get_session_info(self, session_id: str) -> AgentSessionInfo:
-        value = await self.kvstore.get(
-            key=f"session:{self.agent_id}:{session_id}",
-        )
-        if not value:
-            raise SessionNotFoundError(session_id)
-
-        session_info = AgentSessionInfo(**json.loads(value))
-
-        # Check access to session
-        if not self._check_session_access(session_info):
-            return None
-
-        return session_info
-
-    def _check_session_access(self, session_info: AgentSessionInfo) -> bool:
-        """Check if current user has access to the session."""
-        # Handle backward compatibility for old sessions without access control
-        if not hasattr(session_info, "access_attributes") and not hasattr(session_info, "owner"):
-            return True
-
-        return is_action_allowed(self.policy, "read", session_info, get_authenticated_user())
-
-    async def get_session_if_accessible(self, session_id: str) -> AgentSessionInfo | None:
-        """Get session info if the user has access to it. For internal use by sub-session methods."""
-        session_info = await self.get_session_info(session_id)
-        if not session_info:
-            return None
-
-        return session_info
-
-    async def add_vector_db_to_session(self, session_id: str, vector_db_id: str):
-        session_info = await self.get_session_if_accessible(session_id)
-        if session_info is None:
-            raise SessionNotFoundError(session_id)
-
-        session_info.vector_db_id = vector_db_id
-        await self.kvstore.set(
-            key=f"session:{self.agent_id}:{session_id}",
-            value=session_info.model_dump_json(),
-        )
-
-    async def add_turn_to_session(self, session_id: str, turn: Turn):
-        if not await self.get_session_if_accessible(session_id):
-            raise SessionNotFoundError(session_id)
-
-        await self.kvstore.set(
-            key=f"session:{self.agent_id}:{session_id}:{turn.turn_id}",
-            value=turn.model_dump_json(),
-        )
-
-    async def get_session_turns(self, session_id: str) -> list[Turn]:
-        if not await self.get_session_if_accessible(session_id):
-            raise SessionNotFoundError(session_id)
-
-        values = await self.kvstore.values_in_range(
-            start_key=f"session:{self.agent_id}:{session_id}:",
-            end_key=f"session:{self.agent_id}:{session_id}:\xff\xff\xff\xff",
-        )
-        turns = []
-        for value in values:
-            try:
-                turn = Turn(**json.loads(value))
-                turns.append(turn)
-            except Exception as e:
-                log.error(f"Error parsing turn: {e}")
-                continue
-
-        # The kvstore does not guarantee order, so we sort by started_at
-        # to ensure consistent ordering of turns.
-        turns.sort(key=lambda t: t.started_at)
-
-        return turns
-
-    async def get_session_turn(self, session_id: str, turn_id: str) -> Turn | None:
-        if not await self.get_session_if_accessible(session_id):
-            raise SessionNotFoundError(session_id)
-
-        value = await self.kvstore.get(
-            key=f"session:{self.agent_id}:{session_id}:{turn_id}",
-        )
-        if not value:
-            return None
-        return Turn(**json.loads(value))
-
-    async def set_in_progress_tool_call_step(self, session_id: str, turn_id: str, step: ToolExecutionStep):
-        if not await self.get_session_if_accessible(session_id):
-            raise SessionNotFoundError(session_id)
-
-        await self.kvstore.set(
-            key=f"in_progress_tool_call_step:{self.agent_id}:{session_id}:{turn_id}",
-            value=step.model_dump_json(),
-        )
-
-    async def get_in_progress_tool_call_step(self, session_id: str, turn_id: str) -> ToolExecutionStep | None:
-        if not await self.get_session_if_accessible(session_id):
-            return None
-
-        value = await self.kvstore.get(
-            key=f"in_progress_tool_call_step:{self.agent_id}:{session_id}:{turn_id}",
-        )
-        return ToolExecutionStep(**json.loads(value)) if value else None
-
-    async def set_num_infer_iters_in_turn(self, session_id: str, turn_id: str, num_infer_iters: int):
-        if not await self.get_session_if_accessible(session_id):
-            raise SessionNotFoundError(session_id)
-
-        await self.kvstore.set(
-            key=f"num_infer_iters_in_turn:{self.agent_id}:{session_id}:{turn_id}",
-            value=str(num_infer_iters),
-        )
-
-    async def get_num_infer_iters_in_turn(self, session_id: str, turn_id: str) -> int | None:
-        if not await self.get_session_if_accessible(session_id):
-            return None
-
-        value = await self.kvstore.get(
-            key=f"num_infer_iters_in_turn:{self.agent_id}:{session_id}:{turn_id}",
-        )
-        return int(value) if value else None
-
-    async def list_sessions(self) -> list[Session]:
-        values = await self.kvstore.values_in_range(
-            start_key=f"session:{self.agent_id}:",
-            end_key=f"session:{self.agent_id}:\xff\xff\xff\xff",
-        )
-        sessions = []
-        for value in values:
-            try:
-                data = json.loads(value)
-                if "turn_id" in data:
-                    continue
-
-                session_info = Session(**data)
-                sessions.append(session_info)
-            except Exception as e:
-                log.error(f"Error parsing session info: {e}")
-                continue
-        return sessions
-
-    async def delete_session_turns(self, session_id: str) -> None:
-        """Delete all turns and their associated data for a session.
-
-        Args:
-            session_id: The ID of the session whose turns should be deleted.
-        """
-        turns = await self.get_session_turns(session_id)
-        for turn in turns:
-            await self.kvstore.delete(key=f"session:{self.agent_id}:{session_id}:{turn.turn_id}")
-
-    async def delete_session(self, session_id: str) -> None:
-        """Delete a session and all its associated turns.
-
-        Args:
-            session_id: The ID of the session to delete.
-
-        Raises:
-            ValueError: If the session does not exist.
-        """
-        session_info = await self.get_session_info(session_id)
-        if session_info is None:
-            raise SessionNotFoundError(session_id)
-
-        await self.kvstore.delete(key=f"session:{self.agent_id}:{session_id}")

llama_stack/providers/inline/telemetry/__init__.py

@@ -1,5 +0,0 @@
-# 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/telemetry/meta_reference/__init__.py

@@ -1,21 +0,0 @@
-# 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 TelemetryConfig, TelemetrySink
-
-__all__ = ["TelemetryConfig", "TelemetrySink"]
-
-
-async def get_provider_impl(config: TelemetryConfig, deps: dict[Api, Any]):
-    from .telemetry import TelemetryAdapter
-
-    impl = TelemetryAdapter(config, deps)
-    await impl.initialize()
-    return impl

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

@@ -1,47 +0,0 @@
-# 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 enum import StrEnum
-from typing import Any
-
-from pydantic import BaseModel, Field, field_validator
-
-
-class TelemetrySink(StrEnum):
-    OTEL_TRACE = "otel_trace"
-    OTEL_METRIC = "otel_metric"
-    CONSOLE = "console"
-
-
-class TelemetryConfig(BaseModel):
-    otel_exporter_otlp_endpoint: str | None = Field(
-        default=None,
-        description="The OpenTelemetry collector endpoint URL (base URL for traces, metrics, and logs). If not set, the SDK will use OTEL_EXPORTER_OTLP_ENDPOINT environment variable.",
-    )
-    service_name: str = Field(
-        # service name is always the same, use zero-width space to avoid clutter
-        default="\u200b",
-        description="The service name to use for telemetry",
-    )
-    sinks: list[TelemetrySink] = Field(
-        default_factory=list,
-        description="List of telemetry sinks to enable (possible values: otel_trace, otel_metric, console)",
-    )
-
-    @field_validator("sinks", mode="before")
-    @classmethod
-    def validate_sinks(cls, v):
-        if isinstance(v, str):
-            return [TelemetrySink(sink.strip()) for sink in v.split(",")]
-        return v or []
-
-    @classmethod
-    def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]:
-        return {
-            "service_name": "${env.OTEL_SERVICE_NAME:=\u200b}",
-            "sinks": "${env.TELEMETRY_SINKS:=}",
-            "otel_exporter_otlp_endpoint": "${env.OTEL_EXPORTER_OTLP_ENDPOINT:=}",
-        }

llama_stack/providers/inline/telemetry/meta_reference/telemetry.py

@@ -1,252 +0,0 @@
-# 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 os
-import threading
-from typing import Any
-
-from opentelemetry import metrics, trace
-from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
-from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
-from opentelemetry.sdk.metrics import MeterProvider
-from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
-from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import BatchSpanProcessor
-from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
-
-from llama_stack.apis.telemetry import (
-    Event,
-    MetricEvent,
-    SpanEndPayload,
-    SpanStartPayload,
-    SpanStatus,
-    StructuredLogEvent,
-    Telemetry,
-    UnstructuredLogEvent,
-)
-from llama_stack.core.datatypes import Api
-from llama_stack.log import get_logger
-from llama_stack.providers.utils.telemetry.tracing import ROOT_SPAN_MARKERS
-
-from .config import TelemetryConfig
-
-_GLOBAL_STORAGE: dict[str, dict[str | int, Any]] = {
-    "active_spans": {},
-    "counters": {},
-    "gauges": {},
-    "up_down_counters": {},
-}
-_global_lock = threading.Lock()
-_TRACER_PROVIDER = None
-
-logger = get_logger(name=__name__, category="telemetry")
-
-
-def is_tracing_enabled(tracer):
-    with tracer.start_as_current_span("check_tracing") as span:
-        return span.is_recording()
-
-
-class TelemetryAdapter(Telemetry):
-    def __init__(self, _config: TelemetryConfig, deps: dict[Api, Any]) -> None:
-        self.datasetio_api = deps.get(Api.datasetio)
-        self.meter = None
-
-        global _TRACER_PROVIDER
-        # Initialize the correct span processor based on the provider state.
-        # This is needed since once the span processor is set, it cannot be unset.
-        # Recreating the telemetry adapter multiple times will result in duplicate span processors.
-        # Since the library client can be recreated multiple times in a notebook,
-        # the kernel will hold on to the span processor and cause duplicate spans to be written.
-        if os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"):
-            if _TRACER_PROVIDER is None:
-                provider = TracerProvider()
-                trace.set_tracer_provider(provider)
-                _TRACER_PROVIDER = provider
-
-                # Use single OTLP endpoint for all telemetry signals
-
-                # Let OpenTelemetry SDK handle endpoint construction automatically
-                # The SDK will read OTEL_EXPORTER_OTLP_ENDPOINT and construct appropriate URLs
-                # https://opentelemetry.io/docs/languages/sdk-configuration/otlp-exporter
-                span_exporter = OTLPSpanExporter()
-                span_processor = BatchSpanProcessor(span_exporter)
-                trace.get_tracer_provider().add_span_processor(span_processor)
-
-                metric_reader = PeriodicExportingMetricReader(OTLPMetricExporter())
-                metric_provider = MeterProvider(metric_readers=[metric_reader])
-                metrics.set_meter_provider(metric_provider)
-            self.is_otel_endpoint_set = True
-        else:
-            logger.warning("OTEL_EXPORTER_OTLP_ENDPOINT is not set, skipping telemetry")
-            self.is_otel_endpoint_set = False
-
-        self.meter = metrics.get_meter(__name__)
-        self._lock = _global_lock
-
-    async def initialize(self) -> None:
-        pass
-
-    async def shutdown(self) -> None:
-        if self.is_otel_endpoint_set:
-            trace.get_tracer_provider().force_flush()
-
-    async def log_event(self, event: Event, ttl_seconds: int = 604800) -> None:
-        if isinstance(event, UnstructuredLogEvent):
-            self._log_unstructured(event, ttl_seconds)
-        elif isinstance(event, MetricEvent):
-            self._log_metric(event)
-        elif isinstance(event, StructuredLogEvent):
-            self._log_structured(event, ttl_seconds)
-        else:
-            raise ValueError(f"Unknown event type: {event}")
-
-    def _log_unstructured(self, event: UnstructuredLogEvent, ttl_seconds: int) -> None:
-        with self._lock:
-            # Use global storage instead of instance storage
-            span_id = int(event.span_id, 16)
-            span = _GLOBAL_STORAGE["active_spans"].get(span_id)
-
-            if span:
-                timestamp_ns = int(event.timestamp.timestamp() * 1e9)
-                span.add_event(
-                    name=event.type.value,
-                    attributes={
-                        "message": event.message,
-                        "severity": event.severity.value,
-                        "__ttl__": ttl_seconds,
-                        **(event.attributes or {}),
-                    },
-                    timestamp=timestamp_ns,
-                )
-            else:
-                print(f"Warning: No active span found for span_id {span_id}. Dropping event: {event}")
-
-    def _get_or_create_counter(self, name: str, unit: str) -> metrics.Counter:
-        assert self.meter is not None
-        if name not in _GLOBAL_STORAGE["counters"]:
-            _GLOBAL_STORAGE["counters"][name] = self.meter.create_counter(
-                name=name,
-                unit=unit,
-                description=f"Counter for {name}",
-            )
-        return _GLOBAL_STORAGE["counters"][name]
-
-    def _get_or_create_gauge(self, name: str, unit: str) -> metrics.ObservableGauge:
-        assert self.meter is not None
-        if name not in _GLOBAL_STORAGE["gauges"]:
-            _GLOBAL_STORAGE["gauges"][name] = self.meter.create_gauge(
-                name=name,
-                unit=unit,
-                description=f"Gauge for {name}",
-            )
-        return _GLOBAL_STORAGE["gauges"][name]
-
-    def _log_metric(self, event: MetricEvent) -> None:
-        # Add metric as an event to the current span
-        try:
-            with self._lock:
-                # Only try to add to span if we have a valid span_id
-                if event.span_id:
-                    try:
-                        span_id = int(event.span_id, 16)
-                        span = _GLOBAL_STORAGE["active_spans"].get(span_id)
-
-                        if span:
-                            timestamp_ns = int(event.timestamp.timestamp() * 1e9)
-                            span.add_event(
-                                name=f"metric.{event.metric}",
-                                attributes={
-                                    "value": event.value,
-                                    "unit": event.unit,
-                                    **(event.attributes or {}),
-                                },
-                                timestamp=timestamp_ns,
-                            )
-                    except (ValueError, KeyError):
-                        # Invalid span_id or span not found, but we already logged to console above
-                        pass
-        except Exception:
-            # Lock acquisition failed
-            logger.debug("Failed to acquire lock to add metric to span")
-
-        # Log to OpenTelemetry meter if available
-        if self.meter is None:
-            return
-        if isinstance(event.value, int):
-            counter = self._get_or_create_counter(event.metric, event.unit)
-            counter.add(event.value, attributes=event.attributes)
-        elif isinstance(event.value, float):
-            up_down_counter = self._get_or_create_up_down_counter(event.metric, event.unit)
-            up_down_counter.add(event.value, attributes=event.attributes)
-
-    def _get_or_create_up_down_counter(self, name: str, unit: str) -> metrics.UpDownCounter:
-        assert self.meter is not None
-        if name not in _GLOBAL_STORAGE["up_down_counters"]:
-            _GLOBAL_STORAGE["up_down_counters"][name] = self.meter.create_up_down_counter(
-                name=name,
-                unit=unit,
-                description=f"UpDownCounter for {name}",
-            )
-        return _GLOBAL_STORAGE["up_down_counters"][name]
-
-    def _log_structured(self, event: StructuredLogEvent, ttl_seconds: int) -> None:
-        with self._lock:
-            span_id = int(event.span_id, 16)
-            tracer = trace.get_tracer(__name__)
-            if event.attributes is None:
-                event.attributes = {}
-            event.attributes["__ttl__"] = ttl_seconds
-
-            # Extract these W3C trace context attributes so they are not written to
-            # underlying storage, as we just need them to propagate the trace context.
-            traceparent = event.attributes.pop("traceparent", None)
-            tracestate = event.attributes.pop("tracestate", None)
-            if traceparent:
-                # If we have a traceparent header value, we're not the root span.
-                for root_attribute in ROOT_SPAN_MARKERS:
-                    event.attributes.pop(root_attribute, None)
-
-            if isinstance(event.payload, SpanStartPayload):
-                # Check if span already exists to prevent duplicates
-                if span_id in _GLOBAL_STORAGE["active_spans"]:
-                    return
-
-                context = None
-                if event.payload.parent_span_id:
-                    parent_span_id = int(event.payload.parent_span_id, 16)
-                    parent_span = _GLOBAL_STORAGE["active_spans"].get(parent_span_id)
-                    context = trace.set_span_in_context(parent_span)
-                elif traceparent:
-                    carrier = {
-                        "traceparent": traceparent,
-                        "tracestate": tracestate,
-                    }
-                    context = TraceContextTextMapPropagator().extract(carrier=carrier)
-
-                span = tracer.start_span(
-                    name=event.payload.name,
-                    context=context,
-                    attributes=event.attributes or {},
-                )
-                _GLOBAL_STORAGE["active_spans"][span_id] = span
-
-            elif isinstance(event.payload, SpanEndPayload):
-                span = _GLOBAL_STORAGE["active_spans"].get(span_id)
-                if span:
-                    if event.attributes:
-                        span.set_attributes(event.attributes)
-
-                    status = (
-                        trace.Status(status_code=trace.StatusCode.OK)
-                        if event.payload.status == SpanStatus.OK
-                        else trace.Status(status_code=trace.StatusCode.ERROR)
-                    )
-                    span.set_status(status)
-                    span.end()
-                    _GLOBAL_STORAGE["active_spans"].pop(span_id, None)
-            else:
-                raise ValueError(f"Unknown structured log event: {event}")

llama_stack/providers/remote/inference/bedrock/models.py

@@ -1,29 +0,0 @@
-# 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.models.llama.sku_types import CoreModelId
-from llama_stack.providers.utils.inference.model_registry import (
-    build_hf_repo_model_entry,
-)
-
-SAFETY_MODELS_ENTRIES = []
-
-
-# https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html
-MODEL_ENTRIES = [
-    build_hf_repo_model_entry(
-        "meta.llama3-1-8b-instruct-v1:0",
-        CoreModelId.llama3_1_8b_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "meta.llama3-1-70b-instruct-v1:0",
-        CoreModelId.llama3_1_70b_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "meta.llama3-1-405b-instruct-v1:0",
-        CoreModelId.llama3_1_405b_instruct.value,
-    ),
-] + SAFETY_MODELS_ENTRIES

llama_stack/providers/utils/kvstore/sqlite/config.py

@@ -1,20 +0,0 @@
-# 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.schema_utils import json_schema_type
-
-
-@json_schema_type
-class SqliteControlPlaneConfig(BaseModel):
-    db_path: str = Field(
-        description="File path for the sqlite database",
-    )
-    table_name: str = Field(
-        default="llamastack_control_plane",
-        description="Table into which all the keys will be placed",
-    )

llama_stack/providers/utils/sqlstore/__init__.py

@@ -1,5 +0,0 @@
-# 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.