fix(server): add check for vector store document id

Signed-off-by: Radek Ježek <[email protected]>

Author: jezekra1

apps/agentstack-sdk-py/src/agentstack_sdk/platform/vector_store.py

@@ -5,7 +5,7 @@
 
 import typing
 import uuid
-from typing import Literal
+from typing import Literal, Self
 
 import pydantic
 
@@ -29,12 +29,24 @@ class VectorStoreDocument(pydantic.BaseModel):
 class VectorStoreItem(pydantic.BaseModel):
     id: str = pydantic.Field(default_factory=lambda: uuid.uuid4().hex)
     document_id: str
-    document_type: typing.Literal["platform_file", "external"] = "platform_file"
+    document_type: typing.Literal["platform_file", "external"] | None = "platform_file"
     model_id: str | typing.Literal["platform"] = "platform"
     text: str
     embedding: list[float]
     metadata: Metadata | None = None
 
+    @pydantic.model_validator(mode="after")
+    def validate_document_id(self) -> Self:
+        """Validate that document_id is a valid UUID when document_type is platform_file."""
+        if self.document_type == "platform_file":
+            try:
+                _ = uuid.UUID(self.document_id)
+            except ValueError as ex:
+                raise ValueError(
+                    f"document_id must be a valid UUID when document_type is platform_file, got: {self.document_id}"
+                ) from ex
+        return self
+
 
 class VectorStoreSearchResult(pydantic.BaseModel):
     item: VectorStoreItem

apps/agentstack-server/src/agentstack_server/domain/models/vector_store.py

@@ -2,10 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import StrEnum
-from typing import Literal
+from typing import Literal, Self
 from uuid import UUID, uuid4
 
-from pydantic import AwareDatetime, BaseModel, Field
+from pydantic import AwareDatetime, BaseModel, Field, model_validator
 
 from agentstack_server.domain.models.common import Metadata
 from agentstack_server.utils.utils import utc_now
@@ -67,6 +67,18 @@ class VectorStoreItem(BaseModel):
     embedding: list[float]
     metadata: Metadata | None = None
 
+    @model_validator(mode="after")
+    def validate_document_id(self) -> Self:
+        """Validate that document_id is a valid UUID when document_type is platform_file."""
+        if self.document_type == DocumentType.PLATFORM_FILE:
+            try:
+                _ = UUID(self.document_id)
+            except ValueError as ex:
+                raise ValueError(
+                    f"document_id must be a valid UUID when document_type is platform_file, got: {self.document_id}"
+                ) from ex
+        return self
+
 
 class VectorStoreSearchResult(BaseModel):
     """Result of a vector store search operation containing full item data and similarity score."""

apps/agentstack-server/src/agentstack_server/infrastructure/vector_database/vector_db.py

@@ -17,12 +17,14 @@
     String,
     Table,
     Text,
+    case,
 )
 from sqlalchemy.dialects.postgresql import JSONB, insert
 from sqlalchemy.dialects.postgresql import UUID as SQL_UUID
 from sqlalchemy.ext.asyncio import AsyncConnection
 
 from agentstack_server.domain.models.vector_store import (
+    DocumentType,
     VectorStoreDocumentInfo,
     VectorStoreItem,
     VectorStoreSearchResult,
@@ -152,6 +154,7 @@ def _to_item(self, row: Row) -> VectorStoreItem:
         return VectorStoreItem(
             id=row.id,
             document_id=row.vector_store_document_id,
+            document_type=row.document_type or DocumentType.EXTERNAL,
             embedding=row.embedding.to_list(),
             text=row.text,
             metadata=row.metadata,
@@ -177,7 +180,19 @@ async def similarity_search(
         # Select all columns plus the distance as a named column
         query = (
             table.select()
-            .add_columns(table.c.embedding.cosine_distance(query_vector).label("distance"))
+            .add_columns(
+                table.c.embedding.cosine_distance(query_vector).label("distance"),
+                case(
+                    (vector_store_documents_table.c.file_id.is_not(None), DocumentType.PLATFORM_FILE),
+                    else_=DocumentType.EXTERNAL,
+                ).label("document_type"),
+            )
+            .join(
+                vector_store_documents_table,
+                (table.c.vector_store_document_id == vector_store_documents_table.c.id)
+                & (table.c.vector_store_id == vector_store_documents_table.c.vector_store_id),
+                isouter=True,
+            )
             .where(table.c.vector_store_id == collection_id)
             .order_by(table.c.embedding.cosine_distance(query_vector))
             .limit(limit)

apps/agentstack-server/tests/integration/vector_database/test_vector_db.py

@@ -8,7 +8,7 @@
 from sqlalchemy import text
 from sqlalchemy.ext.asyncio import AsyncConnection
 
-from agentstack_server.domain.models.vector_store import VectorStoreItem, VectorStoreSearchResult
+from agentstack_server.domain.models.vector_store import DocumentType, VectorStoreItem, VectorStoreSearchResult
 from agentstack_server.infrastructure.vector_database.vector_db import VectorDatabaseRepository
 
 pytestmark = pytest.mark.integration
@@ -33,20 +33,23 @@ async def sample_vector_items(test_collection_id: UUID) -> list[VectorStoreItem]
         VectorStoreItem(
             id=uuid.uuid4(),
             document_id="doc_001",
+            document_type=DocumentType.EXTERNAL,
             embedding=[1.0] * 128,
             text="The quick brown fox jumps over the lazy dog.",
             metadata={"source": "test_doc_1.txt", "chapter": "1"},
         ),
         VectorStoreItem(
             id=uuid.uuid4(),
             document_id="doc_001",
+            document_type=DocumentType.EXTERNAL,
             embedding=[2.0] * 128,
             text="Artificial intelligence is revolutionizing technology.",
             metadata={"source": "test_doc_1.txt", "chapter": "2"},
         ),
         VectorStoreItem(
             id=uuid.uuid4(),
             document_id="doc_002",
+            document_type=DocumentType.EXTERNAL,
             embedding=[3.0] * 128,
             text="Vector databases enable efficient similarity search.",
             metadata={"source": "test_doc_2.txt", "chapter": "1"},