Merge branch 'main' into chore/standard-unsupported-model-err-msg-2517

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

@@ -14,6 +14,8 @@ from opentelemetry.sdk.trace import SpanProcessor
 from opentelemetry.trace import Span
 from opentelemetry.trace.span import format_span_id, format_trace_id
 
+from llama_stack.providers.utils.telemetry.tracing import LOCAL_ROOT_SPAN_MARKER
+
 
 class SQLiteSpanProcessor(SpanProcessor):
     def __init__(self, conn_string):
@@ -124,7 +126,7 @@ class SQLiteSpanProcessor(SpanProcessor):
                 (
                     trace_id,
                     service_name,
-                    (span_id if span.attributes.get("__root_span__") == "true" else None),
+                    (span_id if span.attributes.get(LOCAL_ROOT_SPAN_MARKER) else None),
                     datetime.fromtimestamp(span.start_time / 1e9, UTC).isoformat(),
                     datetime.fromtimestamp(span.end_time / 1e9, UTC).isoformat(),
                 ),

llama_stack/providers/remote/vector_io/chroma/chroma.py

@@ -137,6 +137,9 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
 
     async def initialize(self) -> None:
         if isinstance(self.config, RemoteChromaVectorIOConfig):
+            if not self.config.url:
+                raise ValueError("URL is a required parameter for the remote Chroma provider's config")
+
             log.info(f"Connecting to Chroma server at: {self.config.url}")
             url = self.config.url.rstrip("/")
             parsed = urlparse(url)

llama_stack/providers/remote/vector_io/chroma/config.py

@@ -10,7 +10,7 @@ from pydantic import BaseModel
 
 
 class ChromaVectorIOConfig(BaseModel):
-    url: str
+    url: str | None
 
     @classmethod
     def sample_run_config(cls, url: str = "${env.CHROMADB_URL}", **kwargs: Any) -> dict[str, Any]:

llama_stack/providers/remote/vector_io/pgvector/config.py

@@ -13,11 +13,11 @@ from llama_stack.schema_utils import json_schema_type
 
 @json_schema_type
 class PGVectorVectorIOConfig(BaseModel):
-    host: str = Field(default="localhost")
-    port: int = Field(default=5432)
-    db: str = Field(default="postgres")
-    user: str = Field(default="postgres")
-    password: str = Field(default="mysecretpassword")
+    host: str | None = Field(default="localhost")
+    port: int | None = Field(default=5432)
+    db: str | None = Field(default="postgres")
+    user: str | None = Field(default="postgres")
+    password: str | None = Field(default="mysecretpassword")
 
     @classmethod
     def sample_run_config(

llama_stack/providers/remote/vector_io/pgvector/pgvector.py

@@ -15,7 +15,21 @@ from pydantic import BaseModel, TypeAdapter
 
 from llama_stack.apis.inference import InterleavedContent
 from llama_stack.apis.vector_dbs import VectorDB
-from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO
+from llama_stack.apis.vector_io import (
+    Chunk,
+    QueryChunksResponse,
+    SearchRankingOptions,
+    VectorIO,
+    VectorStoreChunkingStrategy,
+    VectorStoreDeleteResponse,
+    VectorStoreFileContentsResponse,
+    VectorStoreFileObject,
+    VectorStoreFileStatus,
+    VectorStoreListFilesResponse,
+    VectorStoreListResponse,
+    VectorStoreObject,
+    VectorStoreSearchResponsePage,
+)
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.utils.memory.vector_store import (
     EmbeddingIndex,
@@ -222,3 +236,108 @@ class PGVectorVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
         index = PGVectorIndex(vector_db, vector_db.embedding_dimension, self.conn)
         self.cache[vector_db_id] = VectorDBWithIndex(vector_db, index, self.inference_api)
         return self.cache[vector_db_id]
+
+    async def openai_create_vector_store(
+        self,
+        name: str,
+        file_ids: list[str] | None = None,
+        expires_after: dict[str, Any] | None = None,
+        chunking_strategy: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | None = None,
+        embedding_model: str | None = None,
+        embedding_dimension: int | None = 384,
+        provider_id: str | None = None,
+        provider_vector_db_id: str | None = None,
+    ) -> VectorStoreObject:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_list_vector_stores(
+        self,
+        limit: int | None = 20,
+        order: str | None = "desc",
+        after: str | None = None,
+        before: str | None = None,
+    ) -> VectorStoreListResponse:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_retrieve_vector_store(
+        self,
+        vector_store_id: str,
+    ) -> VectorStoreObject:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_update_vector_store(
+        self,
+        vector_store_id: str,
+        name: str | None = None,
+        expires_after: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> VectorStoreObject:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_delete_vector_store(
+        self,
+        vector_store_id: str,
+    ) -> VectorStoreDeleteResponse:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_search_vector_store(
+        self,
+        vector_store_id: str,
+        query: str | list[str],
+        filters: dict[str, Any] | None = None,
+        max_num_results: int | None = 10,
+        ranking_options: SearchRankingOptions | None = None,
+        rewrite_query: bool | None = False,
+        search_mode: str | None = "vector",
+    ) -> VectorStoreSearchResponsePage:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_attach_file_to_vector_store(
+        self,
+        vector_store_id: str,
+        file_id: str,
+        attributes: dict[str, Any] | None = None,
+        chunking_strategy: VectorStoreChunkingStrategy | None = None,
+    ) -> VectorStoreFileObject:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_list_files_in_vector_store(
+        self,
+        vector_store_id: str,
+        limit: int | None = 20,
+        order: str | None = "desc",
+        after: str | None = None,
+        before: str | None = None,
+        filter: VectorStoreFileStatus | None = None,
+    ) -> VectorStoreListFilesResponse:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_retrieve_vector_store_file(
+        self,
+        vector_store_id: str,
+        file_id: str,
+    ) -> VectorStoreFileObject:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_retrieve_vector_store_file_contents(
+        self,
+        vector_store_id: str,
+        file_id: str,
+    ) -> VectorStoreFileContentsResponse:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_update_vector_store_file(
+        self,
+        vector_store_id: str,
+        file_id: str,
+        attributes: dict[str, Any] | None = None,
+    ) -> VectorStoreFileObject:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")
+
+    async def openai_delete_vector_store_file(
+        self,
+        vector_store_id: str,
+        file_id: str,
+    ) -> VectorStoreFileObject:
+        raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")

llama_stack/providers/utils/telemetry/tracing.py

@@ -35,6 +35,9 @@ INVALID_SPAN_ID = 0x0000000000000000
 INVALID_TRACE_ID = 0x00000000000000000000000000000000
 
 ROOT_SPAN_MARKERS = ["__root__", "__root_span__"]
+# The logical root span may not be visible to this process if a parent context
+# is passed in. The local root span is the first local span in a trace.
+LOCAL_ROOT_SPAN_MARKER = "__local_root_span__"
 
 
 def trace_id_to_str(trace_id: int) -> str:
@@ -180,7 +183,13 @@ async def start_trace(name: str, attributes: dict[str, Any] = None) -> TraceCont
 
     trace_id = generate_trace_id()
     context = TraceContext(BACKGROUND_LOGGER, trace_id)
-    attributes = dict.fromkeys(ROOT_SPAN_MARKERS, True) | (attributes or {})
+    # Mark this span as the root for the trace for now. The processing of
+    # traceparent context if supplied comes later and will result in the
+    # ROOT_SPAN_MARKERS being removed. Also mark this is the 'local' root,
+    # i.e. the root of the spans originating in this process as this is
+    # needed to ensure that we insert this 'local' root span's id into
+    # the trace record in sqlite store.
+    attributes = dict.fromkeys(ROOT_SPAN_MARKERS, True) | {LOCAL_ROOT_SPAN_MARKER: True} | (attributes or {})
     context.push_span(name, attributes)
 
     CURRENT_TRACE_CONTEXT.set(context)