using a property for Chunk.chunk_id

Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>

llama_stack/apis/vector_io/vector_io.py

@@ -8,6 +8,7 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+import uuid
 from typing import Annotated, Any, Literal, Protocol, runtime_checkable
 
 from pydantic import BaseModel, Field
@@ -15,6 +16,7 @@ from pydantic import BaseModel, Field
 from llama_stack.apis.inference import InterleavedContent
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+from llama_stack.providers.utils.vector_io.chunk_utils import generate_chunk_id
 from llama_stack.schema_utils import json_schema_type, webmethod
 from llama_stack.strong_typing.schema import register_schema
 
@@ -23,10 +25,12 @@ from llama_stack.strong_typing.schema import register_schema
 class ChunkMetadata(BaseModel):
     """
     `ChunkMetadata` is backend metadata for a `Chunk` that is used to store additional information about the chunk that
-        will NOT be inserted into the context during inference, but is required for backend functionality.
-        Use `metadata` in `Chunk` for metadata that will be used during inference.
+        will not be used in the context during inference, but is required for backend functionality. The `ChunkMetadata`
+        is set during chunk creation in `MemoryToolRuntimeImpl().insert()`and is not expected to change after.
+        Use `Chunk.metadata` for metadata that will be used in the context during inference.
+    :param chunk_id: The ID of the chunk. If not set, it will be generated based on the document ID and content.
     :param document_id: The ID of the document this chunk belongs to.
-    :param source: The source of the content, such as a URL or file path.
+    :param source: The source of the content, such as a URL, file path, or other identifier.
     :param created_timestamp: An optional timestamp indicating when the chunk was created.
     :param updated_timestamp: An optional timestamp indicating when the chunk was last updated.
     :param chunk_window: The window of the chunk, which can be used to group related chunks together.
@@ -37,8 +41,8 @@ class ChunkMetadata(BaseModel):
     :param metadata_token_count: The number of tokens in the metadata of the chunk.
     """
 
+    chunk_id: str = None
     document_id: str | None = None
-    chunk_id: str | None = None
     source: str | None = None
     created_timestamp: int | None = None
     updated_timestamp: int | None = None
@@ -56,16 +60,37 @@ class Chunk(BaseModel):
     A chunk of content that can be inserted into a vector database.
     :param content: The content of the chunk, which can be interleaved text, images, or other types.
     :param embedding: Optional embedding for the chunk. If not provided, it will be computed later.
-    :param metadata: Metadata associated with the chunk that will be used during inference.
-    :param chunk_metadata: Metadata for the chunk that will NOT be inserted into the context during inference
-        that is required backend functionality.
+    :param metadata: Metadata associated with the chunk that will be used in the model context during inference.
+    :param stored_chunk_id: The chunk ID that is stored in the vector database. Used for backend functionality.
+    :param chunk_metadata: Metadata for the chunk that will NOT be used in the context during inference.
+        The `chunk_metadata` is required backend functionality.
     """
 
     content: InterleavedContent
     metadata: dict[str, Any] = Field(default_factory=dict)
     embedding: list[float] | None = None
+    # The alias parameter serializes the field as "chunk_id" in JSON but keeps the internal name as "stored_chunk_id"
+    stored_chunk_id: str | None = Field(default=None, alias="chunk_id")
     chunk_metadata: ChunkMetadata | None = None
 
+    model_config = {"populate_by_name": True}
+
+    def model_post_init(self, __context):
+        # Extract chunk_id from metadata if present
+        if self.metadata and "chunk_id" in self.metadata:
+            self.stored_chunk_id = self.metadata.pop("chunk_id")
+
+    @property
+    def chunk_id(self) -> str:
+        """Returns the chunk ID, which is either an input `chunk_id` or a generated one if not set."""
+        if self.stored_chunk_id:
+            return self.stored_chunk_id
+
+        if "document_id" in self.metadata:
+            return generate_chunk_id(self.metadata["document_id"], str(self.content))
+
+        return generate_chunk_id(str(uuid.uuid4()), str(self.content))
+
 
 @json_schema_type
 class QueryChunksResponse(BaseModel):