[Runtime] Simplify the vector store setup process (#231)

Author: alnzng

api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java

@@ -29,7 +29,6 @@ public enum ResourceType {
     EMBEDDING_MODEL("embedding_model"),
     EMBEDDING_MODEL_CONNECTION("embedding_model_connection"),
     VECTOR_STORE("vector_store"),
-    VECTOR_STORE_CONNECTION("vector_store_connection"),
     PROMPT("prompt"),
     TOOL("tool");
 

python/flink_agents/api/decorators.py

@@ -173,26 +173,8 @@ def mcp_server(func: Callable) -> Callable:
     return func
 
 
-def vector_store_connection(func: Callable) -> Callable:
-    """Decorator for marking a function declaring a vector store connection.
-
-    Parameters
-    ----------
-    func : Callable
-        Function to be decorated.
-
-    Returns:
-    -------
-    Callable
-        Decorator function that marks the target function declare a vector store
-        connection.
-    """
-    func._is_vector_store_connection = True
-    return func
-
-
-def vector_store_setup(func: Callable) -> Callable:
-    """Decorator for marking a function declaring a vector store setup.
+def vector_store(func: Callable) -> Callable:
+    """Decorator for marking a function declaring a vector store.
 
     Parameters
     ----------
@@ -204,5 +186,5 @@ def vector_store_setup(func: Callable) -> Callable:
     Callable
         Decorator function that marks the target function declare a vector store.
     """
-    func._is_vector_store_setup = True
+    func._is_vector_store = True
     return func

python/flink_agents/api/resource.py

@@ -35,7 +35,6 @@ class ResourceType(Enum):
     EMBEDDING_MODEL = "embedding_model"
     EMBEDDING_MODEL_CONNECTION = "embedding_model_connection"
     VECTOR_STORE = "vector_store"
-    VECTOR_STORE_CONNECTION = "vector_store_connection"
     PROMPT = "prompt"
     MCP_SERVER = "mcp_server"
 

python/flink_agents/api/vector_stores/vector_store.py

@@ -110,43 +110,14 @@ def __str__(self) -> str:
         return f"QueryResult: {len(self.documents)} documents"
 
 
-class BaseVectorStoreConnection(Resource, ABC):
-    """Base abstract class for vector store connection.
+class BaseVectorStore(Resource, ABC):
+    """Base abstract class for vector store.
 
-    Manages connection configuration and provides raw vector search operations
-    using pre-computed embeddings. One connection can be shared across multiple
-    vector store setups.
+    Provides vector store functionality that integrates embedding models
+    for text-based semantic search. Handles both connection management and
+    embedding generation internally.
     """
 
-    @classmethod
-    @override
-    def resource_type(cls) -> ResourceType:
-        """Return resource type of class."""
-        return ResourceType.VECTOR_STORE_CONNECTION
-
-    @abstractmethod
-    def query(self, embedding: List[float], limit: int = 10, **kwargs: Any) -> List[Document]:
-        """Perform vector search using pre-computed embedding.
-
-        Args:
-            embedding: Pre-computed embedding vector for semantic search
-            limit: Maximum number of results to return (default: 10)
-            **kwargs: Vector store-specific parameters (filters, distance metrics, etc.)
-
-        Returns:
-            List of documents matching the search criteria
-        """
-
-
-class BaseVectorStoreSetup(Resource, ABC):
-    """Base abstract class for vector store setup.
-
-    Coordinates between vector store connections and embedding models to provide
-    text-based semantic search. Automatically converts text queries to embeddings
-    before delegating to the connection layer.
-    """
-
-    connection: str = Field(description="Name of the referenced connection.")
     embedding_model: str = Field(description="Name of the embedding model resource to use.")
 
     @classmethod
@@ -181,19 +152,27 @@ def query(self, query: VectorStoreQuery) -> VectorStoreQueryResult:
         )
         query_embedding = embedding_model.embed(query.query_text)
 
-        # Get vector store connection resource
-        connection = self.get_resource(
-            self.connection, ResourceType.VECTOR_STORE_CONNECTION
-        )
-
         # Merge setup kwargs with query-specific args
         merged_kwargs = self.store_kwargs.copy()
         merged_kwargs.update(query.extra_args)
 
-        # Perform vector search
-        documents = connection.query(query_embedding, query.limit, **merged_kwargs)
+        # Perform vector search using the abstract method
+        documents = self.query_embedding(query_embedding, query.limit, **merged_kwargs)
 
         # Return structured result
         return VectorStoreQueryResult(
             documents=documents
         )
+
+    @abstractmethod
+    def query_embedding(self, embedding: List[float], limit: int = 10, **kwargs: Any) -> List[Document]:
+        """Perform vector search using pre-computed embedding.
+
+        Args:
+            embedding: Pre-computed embedding vector for semantic search
+            limit: Maximum number of results to return (default: 10)
+            **kwargs: Vector store-specific parameters (filters, distance metrics, etc.)
+
+        Returns:
+            List of documents matching the search criteria
+        """

python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py

@@ -24,16 +24,15 @@
 from pydantic import Field
 
 from flink_agents.api.vector_stores.vector_store import (
-    BaseVectorStoreConnection,
-    BaseVectorStoreSetup,
+    BaseVectorStore,
     Document,
 )
 
 DEFAULT_COLLECTION = "flink_agents_chroma_collection"
 
 
-class ChromaVectorStoreConnection(BaseVectorStoreConnection):
-    """ChromaDB Vector Store Connection which manages connection to ChromaDB.
+class ChromaVectorStore(BaseVectorStore):
+    """ChromaDB vector store that handles connection and semantic search.
 
     Visit https://docs.trychroma.com/ for ChromaDB documentation.
 
@@ -59,8 +58,15 @@ class ChromaVectorStoreConnection(BaseVectorStoreConnection):
         ChromaDB tenant for multi-tenancy support (default: "default_tenant").
     database : str
         ChromaDB database name (default: "default_database").
+    collection : str
+        Name of the ChromaDB collection to use (default: flink_agents_collection).
+    collection_metadata : Dict[str, Any]
+        Metadata for the collection (optional).
+    create_collection_if_not_exists : bool
+        Whether to create the collection if it doesn't exist (default: True).
     """
 
+    # Connection configuration
     persist_directory: str | None = Field(
         default=None,
         description="Directory for persistent storage. If None, uses in-memory client.",
@@ -90,36 +96,60 @@ class ChromaVectorStoreConnection(BaseVectorStoreConnection):
         description="ChromaDB database name.",
     )
 
+    # Collection configuration
+    collection: str = Field(
+        default=DEFAULT_COLLECTION,
+        description="Name of the ChromaDB collection to use.",
+    )
+    collection_metadata: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Metadata for the collection.",
+    )
+    create_collection_if_not_exists: bool = Field(
+        default=True,
+        description="Whether to create the collection if it doesn't exist.",
+    )
+
     __client: ChromaClient | None = None
 
     def __init__(
         self,
+        *,
+        embedding_model: str,
         persist_directory: str | None = None,
         host: str | None = None,
         port: int | None = 8000,
         api_key: str | None = None,
         client_settings: Settings | None = None,
         tenant: str = "default_tenant",
         database: str = "default_database",
+        collection: str = DEFAULT_COLLECTION,
+        collection_metadata: Dict[str, Any] | None = None,
+        create_collection_if_not_exists: bool = True,
         **kwargs: Any,
     ) -> None:
         """Init method."""
+        if collection_metadata is None:
+            collection_metadata = {}
         super().__init__(
+            embedding_model=embedding_model,
             persist_directory=persist_directory,
             host=host,
             port=port,
             api_key=api_key,
             client_settings=client_settings,
             tenant=tenant,
             database=database,
+            collection=collection,
+            collection_metadata=collection_metadata,
+            create_collection_if_not_exists=create_collection_if_not_exists,
             **kwargs,
         )
 
     @property
     def client(self) -> ChromaClient:
         """Return ChromaDB client, creating it if necessary."""
         if self.__client is None:
-
             # Choose client type based on configuration
             if self.api_key is not None:
                 # Cloud mode
@@ -155,7 +185,16 @@ def client(self) -> ChromaClient:
 
         return self.__client
 
-    def query(self, embedding: List[float], limit: int = 10, **kwargs: Any) -> List[Document]:
+    @property
+    def store_kwargs(self) -> Dict[str, Any]:
+        """Return ChromaDB-specific setup settings."""
+        return {
+            "collection": self.collection,
+            "collection_metadata": self.collection_metadata,
+            "create_collection_if_not_exists": self.create_collection_if_not_exists,
+        }
+
+    def query_embedding(self, embedding: List[float], limit: int = 10, **kwargs: Any) -> List[Document]:
         """Perform vector search using pre-computed embedding.
 
         Args:
@@ -167,9 +206,9 @@ def query(self, embedding: List[float], limit: int = 10, **kwargs: Any) -> List[
             List of documents matching the search criteria
         """
         # Extract ChromaDB-specific parameters
-        collection_name = kwargs.get("collection", DEFAULT_COLLECTION)
-        collection_metadata = kwargs.get("collection_metadata", {})
-        create_collection_if_not_exists = kwargs.get("create_collection_if_not_exists", True)
+        collection_name = kwargs.get("collection", self.collection)
+        collection_metadata = kwargs.get("collection_metadata", self.collection_metadata)
+        create_collection_if_not_exists = kwargs.get("create_collection_if_not_exists", self.create_collection_if_not_exists)
         where = kwargs.get("where")  # Metadata filters
 
         # Get or create collection based on configuration
@@ -206,62 +245,3 @@ def query(self, embedding: List[float], limit: int = 10, **kwargs: Any) -> List[
 
         return documents
 
-
-class ChromaVectorStoreSetup(BaseVectorStoreSetup):
-    """ChromaDB vector store setup which manages collection configuration
-    and coordinates with embedding models for semantic search.
-
-    Attributes:
-    ----------
-    collection : str
-        Name of the ChromaDB collection to use (default: flink_agents_collection).
-    collection_metadata : Dict[str, Any]
-        Metadata for the collection (optional).
-    create_collection_if_not_exists : bool
-        Whether to create the collection if it doesn't exist (default: True).
-    """
-
-    collection: str = Field(
-        default=DEFAULT_COLLECTION,
-        description="Name of the ChromaDB collection to use.",
-    )
-    collection_metadata: Dict[str, Any] = Field(
-        default_factory=dict,
-        description="Metadata for the collection.",
-    )
-    create_collection_if_not_exists: bool = Field(
-        default=True,
-        description="Whether to create the collection if it doesn't exist.",
-    )
-
-    def __init__(
-        self,
-        *,
-        connection: str,
-        embedding_model: str,
-        collection: str = DEFAULT_COLLECTION,
-        collection_metadata: Dict[str, Any] | None = None,
-        create_collection_if_not_exists: bool = True,
-        **kwargs: Any,
-    ) -> None:
-        """Init method."""
-        if collection_metadata is None:
-            collection_metadata = {}
-        super().__init__(
-            connection=connection,
-            embedding_model=embedding_model,
-            collection=collection,
-            collection_metadata=collection_metadata,
-            create_collection_if_not_exists=create_collection_if_not_exists,
-            **kwargs,
-        )
-
-    @property
-    def store_kwargs(self) -> Dict[str, Any]:
-        """Return ChromaDB-specific setup settings passed to connection."""
-        return {
-            "collection": self.collection,
-            "collection_metadata": self.collection_metadata,
-            "create_collection_if_not_exists": self.create_collection_if_not_exists,
-        }
-