[API Updates] Model / shield / memory-bank routing + agent persistence + support for private headers (#92)

This is yet another of those large PRs (hopefully we will have less and less of them as things mature fast). This one introduces substantial improvements and some simplifications to the stack.

Most important bits:

* Agents reference implementation now has support for session / turn persistence. The default implementation uses sqlite but there's also support for using Redis.

* We have re-architected the structure of the Stack APIs to allow for more flexible routing. The motivating use cases are:
  - routing model A to ollama and model B to a remote provider like Together
  - routing shield A to local impl while shield B to a remote provider like Bedrock
  - routing a vector memory bank to Weaviate while routing a keyvalue memory bank to Redis

* Support for provider specific parameters to be passed from the clients. A client can pass data using `x_llamastack_provider_data` parameter which can be type-checked and provided to the Adapter implementations.

llama_stack/providers/adapters/agents/__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/adapters/agents/sample/__init__.py

@@ -0,0 +1,17 @@
+# 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 SampleConfig
+
+
+async def get_adapter_impl(config: SampleConfig, _deps) -> Any:
+    from .sample import SampleAgentsImpl
+
+    impl = SampleAgentsImpl(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/adapters/agents/sample/config.py

@@ -0,0 +1,12 @@
+# 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
+
+
+class SampleConfig(BaseModel):
+    host: str = "localhost"
+    port: int = 9999

llama_stack/providers/adapters/agents/sample/sample.py

@@ -0,0 +1,18 @@
+# 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 .config import SampleConfig
+
+
+from llama_stack.apis.agents import *  # noqa: F403
+
+
+class SampleAgentsImpl(Agents):
+    def __init__(self, config: SampleConfig):
+        self.config = config
+
+    async def initialize(self):
+        pass

llama_stack/providers/adapters/inference/fireworks/fireworks.py

@@ -6,14 +6,14 @@
 
 from typing import AsyncGenerator
 
+from fireworks.client import Fireworks
+
 from llama_models.llama3.api.chat_format import ChatFormat
 
 from llama_models.llama3.api.datatypes import Message, StopReason
 from llama_models.llama3.api.tokenizer import Tokenizer
 from llama_models.sku_list import resolve_model
 
-from fireworks.client import Fireworks
-
 from llama_stack.apis.inference import *  # noqa: F403
 from llama_stack.providers.utils.inference.prepare_messages import prepare_messages
 
@@ -42,7 +42,14 @@ class FireworksInferenceAdapter(Inference):
     async def shutdown(self) -> None:
         pass
 
-    async def completion(self, request: CompletionRequest) -> AsyncGenerator:
+    async def completion(
+        self,
+        model: str,
+        content: InterleavedTextMedia,
+        sampling_params: Optional[SamplingParams] = SamplingParams(),
+        stream: Optional[bool] = False,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> AsyncGenerator:
         raise NotImplementedError()
 
     def _messages_to_fireworks_messages(self, messages: list[Message]) -> list:

llama_stack/providers/adapters/inference/ollama/ollama.py

@@ -38,6 +38,7 @@ class OllamaInferenceAdapter(Inference):
         return AsyncClient(host=self.url)
 
     async def initialize(self) -> None:
+        print("Initializing Ollama, checking connectivity to server...")
         try:
             await self.client.ps()
         except httpx.ConnectError as e:
@@ -48,7 +49,14 @@ class OllamaInferenceAdapter(Inference):
     async def shutdown(self) -> None:
         pass
 
-    async def completion(self, request: CompletionRequest) -> AsyncGenerator:
+    async def completion(
+        self,
+        model: str,
+        content: InterleavedTextMedia,
+        sampling_params: Optional[SamplingParams] = SamplingParams(),
+        stream: Optional[bool] = False,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> AsyncGenerator:
         raise NotImplementedError()
 
     def _messages_to_ollama_messages(self, messages: list[Message]) -> list:

llama_stack/providers/adapters/inference/sample/__init__.py

@@ -0,0 +1,17 @@
+# 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 SampleConfig
+
+
+async def get_adapter_impl(config: SampleConfig, _deps) -> Any:
+    from .sample import SampleInferenceImpl
+
+    impl = SampleInferenceImpl(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/adapters/inference/sample/config.py

@@ -0,0 +1,12 @@
+# 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
+
+
+class SampleConfig(BaseModel):
+    host: str = "localhost"
+    port: int = 9999

llama_stack/providers/adapters/inference/sample/sample.py

@@ -0,0 +1,18 @@
+# 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 .config import SampleConfig
+
+
+from llama_stack.apis.inference import *  # noqa: F403
+
+
+class SampleInferenceImpl(Inference):
+    def __init__(self, config: SampleConfig):
+        self.config = config
+
+    async def initialize(self):
+        pass

llama_stack/providers/adapters/inference/tgi/tgi.py

@@ -54,7 +54,14 @@ class TGIAdapter(Inference):
     async def shutdown(self) -> None:
         pass
 
-    async def completion(self, request: CompletionRequest) -> AsyncGenerator:
+    async def completion(
+        self,
+        model: str,
+        content: InterleavedTextMedia,
+        sampling_params: Optional[SamplingParams] = SamplingParams(),
+        stream: Optional[bool] = False,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> AsyncGenerator:
         raise NotImplementedError()
 
     def get_chat_options(self, request: ChatCompletionRequest) -> dict:

llama_stack/providers/adapters/inference/together/__init__.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from .config import TogetherImplConfig
+from .config import TogetherImplConfig, TogetherHeaderExtractor
 
 
 async def get_adapter_impl(config: TogetherImplConfig, _deps):

llama_stack/providers/adapters/inference/together/config.py

@@ -4,9 +4,18 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
+from llama_models.schema_utils import json_schema_type
+
+from llama_stack.distribution.request_headers import annotate_header
+
+
+class TogetherHeaderExtractor(BaseModel):
+    api_key: annotate_header(
+        "X-LlamaStack-Together-ApiKey", str, "The API Key for the request"
+    )
+
 
 @json_schema_type
 class TogetherImplConfig(BaseModel):

llama_stack/providers/adapters/inference/together/together.py

@@ -42,7 +42,14 @@ class TogetherInferenceAdapter(Inference):
     async def shutdown(self) -> None:
         pass
 
-    async def completion(self, request: CompletionRequest) -> AsyncGenerator:
+    async def completion(
+        self,
+        model: str,
+        content: InterleavedTextMedia,
+        sampling_params: Optional[SamplingParams] = SamplingParams(),
+        stream: Optional[bool] = False,
+        logprobs: Optional[LogProbConfig] = None,
+    ) -> AsyncGenerator:
         raise NotImplementedError()
 
     def _messages_to_together_messages(self, messages: list[Message]) -> list:

llama_stack/providers/adapters/memory/chroma/chroma.py

@@ -31,9 +31,6 @@ class ChromaIndex(EmbeddingIndex):
             embeddings
         ), f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
 
-        for i, chunk in enumerate(chunks):
-            print(f"Adding chunk #{i} tokens={chunk.token_count}")
-
         await self.collection.add(
             documents=[chunk.json() for chunk in chunks],
             embeddings=embeddings,

llama_stack/providers/adapters/memory/pgvector/pgvector.py

@@ -80,7 +80,6 @@ class PGVectorIndex(EmbeddingIndex):
 
         values = []
         for i, chunk in enumerate(chunks):
-            print(f"Adding chunk #{i} tokens={chunk.token_count}")
             values.append(
                 (
                     f"{chunk.document_id}:chunk-{i}",

llama_stack/providers/adapters/memory/sample/__init__.py

@@ -0,0 +1,17 @@
+# 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 SampleConfig
+
+
+async def get_adapter_impl(config: SampleConfig, _deps) -> Any:
+    from .sample import SampleMemoryImpl
+
+    impl = SampleMemoryImpl(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/adapters/memory/sample/config.py

@@ -0,0 +1,12 @@
+# 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
+
+
+class SampleConfig(BaseModel):
+    host: str = "localhost"
+    port: int = 9999

llama_stack/providers/adapters/memory/sample/sample.py

@@ -0,0 +1,18 @@
+# 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 .config import SampleConfig
+
+
+from llama_stack.apis.memory import *  # noqa: F403
+
+
+class SampleMemoryImpl(Memory):
+    def __init__(self, config: SampleConfig):
+        self.config = config
+
+    async def initialize(self):
+        pass

llama_stack/providers/adapters/safety/__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/adapters/safety/sample/__init__.py

@@ -0,0 +1,17 @@
+# 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 SampleConfig
+
+
+async def get_adapter_impl(config: SampleConfig, _deps) -> Any:
+    from .sample import SampleSafetyImpl
+
+    impl = SampleSafetyImpl(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/adapters/safety/sample/config.py

@@ -0,0 +1,12 @@
+# 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
+
+
+class SampleConfig(BaseModel):
+    host: str = "localhost"
+    port: int = 9999

llama_stack/providers/adapters/safety/sample/sample.py

@@ -0,0 +1,18 @@
+# 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 .config import SampleConfig
+
+
+from llama_stack.apis.safety import *  # noqa: F403
+
+
+class SampleSafetyImpl(Safety):
+    def __init__(self, config: SampleConfig):
+        self.config = config
+
+    async def initialize(self):
+        pass

llama_stack/providers/adapters/telemetry/__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/adapters/telemetry/opentelemetry/__init__.py

@@ -0,0 +1,15 @@
+# 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 .config import OpenTelemetryConfig
+
+
+async def get_adapter_impl(config: OpenTelemetryConfig, _deps):
+    from .opentelemetry import OpenTelemetryAdapter
+
+    impl = OpenTelemetryAdapter(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/adapters/telemetry/opentelemetry/config.py

@@ -0,0 +1,12 @@
+# 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
+
+
+class OpenTelemetryConfig(BaseModel):
+    jaeger_host: str = "localhost"
+    jaeger_port: int = 6831

llama_stack/providers/adapters/telemetry/opentelemetry/opentelemetry.py

@@ -0,0 +1,201 @@
+# 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 datetime import datetime
+
+from opentelemetry import metrics, trace
+from opentelemetry.exporter.jaeger.thrift import JaegerExporter
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import (
+    ConsoleMetricExporter,
+    PeriodicExportingMetricReader,
+)
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.semconv.resource import ResourceAttributes
+
+from llama_stack.apis.telemetry import *  # noqa: F403
+
+from .config import OpenTelemetryConfig
+
+
+def string_to_trace_id(s: str) -> int:
+    # Convert the string to bytes and then to an integer
+    return int.from_bytes(s.encode(), byteorder="big", signed=False)
+
+
+def string_to_span_id(s: str) -> int:
+    # Use only the first 8 bytes (64 bits) for span ID
+    return int.from_bytes(s.encode()[:8], byteorder="big", signed=False)
+
+
+def is_tracing_enabled(tracer):
+    with tracer.start_as_current_span("check_tracing") as span:
+        return span.is_recording()
+
+
+class OpenTelemetryAdapter(Telemetry):
+    def __init__(self, config: OpenTelemetryConfig):
+        self.config = config
+
+        self.resource = Resource.create(
+            {ResourceAttributes.SERVICE_NAME: "foobar-service"}
+        )
+
+        # Set up tracing with Jaeger exporter
+        jaeger_exporter = JaegerExporter(
+            agent_host_name=self.config.jaeger_host,
+            agent_port=self.config.jaeger_port,
+        )
+        trace_provider = TracerProvider(resource=self.resource)
+        trace_processor = BatchSpanProcessor(jaeger_exporter)
+        trace_provider.add_span_processor(trace_processor)
+        trace.set_tracer_provider(trace_provider)
+        self.tracer = trace.get_tracer(__name__)
+
+        # Set up metrics
+        metric_reader = PeriodicExportingMetricReader(ConsoleMetricExporter())
+        metric_provider = MeterProvider(
+            resource=self.resource, metric_readers=[metric_reader]
+        )
+        metrics.set_meter_provider(metric_provider)
+        self.meter = metrics.get_meter(__name__)
+
+    async def initialize(self) -> None:
+        pass
+
+    async def shutdown(self) -> None:
+        trace.get_tracer_provider().shutdown()
+        metrics.get_meter_provider().shutdown()
+
+    async def log_event(self, event: Event) -> None:
+        if isinstance(event, UnstructuredLogEvent):
+            self._log_unstructured(event)
+        elif isinstance(event, MetricEvent):
+            self._log_metric(event)
+        elif isinstance(event, StructuredLogEvent):
+            self._log_structured(event)
+
+    def _log_unstructured(self, event: UnstructuredLogEvent) -> None:
+        span = trace.get_current_span()
+        span.add_event(
+            name=event.message,
+            attributes={"severity": event.severity.value, **event.attributes},
+            timestamp=event.timestamp,
+        )
+
+    def _log_metric(self, event: MetricEvent) -> None:
+        if isinstance(event.value, int):
+            self.meter.create_counter(
+                name=event.metric,
+                unit=event.unit,
+                description=f"Counter for {event.metric}",
+            ).add(event.value, attributes=event.attributes)
+        elif isinstance(event.value, float):
+            self.meter.create_gauge(
+                name=event.metric,
+                unit=event.unit,
+                description=f"Gauge for {event.metric}",
+            ).set(event.value, attributes=event.attributes)
+
+    def _log_structured(self, event: StructuredLogEvent) -> None:
+        if isinstance(event.payload, SpanStartPayload):
+            context = trace.set_span_in_context(
+                trace.NonRecordingSpan(
+                    trace.SpanContext(
+                        trace_id=string_to_trace_id(event.trace_id),
+                        span_id=string_to_span_id(event.span_id),
+                        is_remote=True,
+                    )
+                )
+            )
+            span = self.tracer.start_span(
+                name=event.payload.name,
+                kind=trace.SpanKind.INTERNAL,
+                context=context,
+                attributes=event.attributes,
+            )
+
+            if event.payload.parent_span_id:
+                span.set_parent(
+                    trace.SpanContext(
+                        trace_id=string_to_trace_id(event.trace_id),
+                        span_id=string_to_span_id(event.payload.parent_span_id),
+                        is_remote=True,
+                    )
+                )
+        elif isinstance(event.payload, SpanEndPayload):
+            span = trace.get_current_span()
+            span.set_status(
+                trace.Status(
+                    trace.StatusCode.OK
+                    if event.payload.status == SpanStatus.OK
+                    else trace.StatusCode.ERROR
+                )
+            )
+            span.end(end_time=event.timestamp)
+
+    async def get_trace(self, trace_id: str) -> Trace:
+        # we need to look up the root span id
+        raise NotImplementedError("not yet no")
+
+
+# Usage example
+async def main():
+    telemetry = OpenTelemetryTelemetry("my-service")
+    await telemetry.initialize()
+
+    # Log an unstructured event
+    await telemetry.log_event(
+        UnstructuredLogEvent(
+            trace_id="trace123",
+            span_id="span456",
+            timestamp=datetime.now(),
+            message="This is a log message",
+            severity=LogSeverity.INFO,
+        )
+    )
+
+    # Log a metric event
+    await telemetry.log_event(
+        MetricEvent(
+            trace_id="trace123",
+            span_id="span456",
+            timestamp=datetime.now(),
+            metric="my_metric",
+            value=42,
+            unit="count",
+        )
+    )
+
+    # Log a structured event (span start)
+    await telemetry.log_event(
+        StructuredLogEvent(
+            trace_id="trace123",
+            span_id="span789",
+            timestamp=datetime.now(),
+            payload=SpanStartPayload(name="my_operation"),
+        )
+    )
+
+    # Log a structured event (span end)
+    await telemetry.log_event(
+        StructuredLogEvent(
+            trace_id="trace123",
+            span_id="span789",
+            timestamp=datetime.now(),
+            payload=SpanEndPayload(status=SpanStatus.OK),
+        )
+    )
+
+    await telemetry.shutdown()
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())

llama_stack/providers/adapters/telemetry/sample/__init__.py

@@ -0,0 +1,17 @@
+# 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 SampleConfig
+
+
+async def get_adapter_impl(config: SampleConfig, _deps) -> Any:
+    from .sample import SampleTelemetryImpl
+
+    impl = SampleTelemetryImpl(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/adapters/telemetry/sample/config.py

@@ -0,0 +1,12 @@
+# 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
+
+
+class SampleConfig(BaseModel):
+    host: str = "localhost"
+    port: int = 9999

llama_stack/providers/adapters/telemetry/sample/sample.py

@@ -0,0 +1,18 @@
+# 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 .config import SampleConfig
+
+
+from llama_stack.apis.telemetry import *  # noqa: F403
+
+
+class SampleTelemetryImpl(Telemetry):
+    def __init__(self, config: SampleConfig):
+        self.config = config
+
+    async def initialize(self):
+        pass