Vector Store and Embedding Model Integration for Flink Agents

[alnzng]

Overview

Vector stores and embedding models together enable agents to perform semantic search and knowledge retrieval over large document collections. This powerful combination unlocks new agent patterns including RAG (Retrieval-Augmented Generation) where agents search relevant documents before generating responses, knowledge base agents that answer questions from company documentation, and context-aware conversations that maintain relevant context across long interactions. These capabilities will also serve as a foundation for future long-term memory features.

We propose adding Vector Store and Embedding Model as new resource types following the same architectural pattern as ChatModel, with multiple built-in implementations to support different use cases and deployment preferences.

Design Overview

User Input → RAG Agent → Query Embedding → Vector Search → Context Retrieval → Enhanced Generation → Response

Detailed Event Flow:
┌─────────────┐    ┌──────────────────┐    ┌─────────────────────┐    ┌─────────────────┐    ┌──────────────────┐    ┌─────────────┐
│ InputEvent  │───▶│ VectorSearchEvent│───▶│ VectorSearchResult  │───▶│ ChatRequestEvent│───▶│ ChatResponseEvent│───▶│ OutputEvent │
│             │    │                  │    │       Event         │    │                 │    │                  │    │             │
└─────────────┘    └──────────────────┘    └─────────────────────┘    └─────────────────┘    └──────────────────┘    └─────────────┘
                            │                           ▲                        │                         ▲
                            ▼                           │                        ▼                         │
                   ┌─────────────────┐                  │               ┌─────────────────┐                │
                   │ Vector Store    │──────────────────┘               │ Chat Model      │────────────────┘
                   │ Backend         │                                  │ (OpenAI/etc)    │
                   └─────────────────┘                                  └─────────────────┘
                            │                    
                            ▼                    
                   ┌─────────────────┐           
                   │ Embedding Model │           
                   │ (text→vectors)  │           
                   └─────────────────┘           

APIs

Base Classes

class BaseEmbeddingModelConnection(Resource, ABC):
    """Base class for embedding model connections.
    
    Similar to BaseChatModelConnection, manages connection configuration
    for embedding service providers (OpenAI, Azure, local models, etc.).
    """
    
    @abstractmethod
    def embed_query(self, text: str) -> List[float]:
        """Generate embedding for a single query text."""
        pass

class BaseEmbeddingModelSetup(Resource, ABC):
    """Base class for embedding model setup.
    
    Similar to BaseChatModelSetup, configures the embedding model with
    specific model names, dimensions, and processing parameters.
    
    Attributes:
    ----------
    connection_name : str
        Name of the connection resource to use
    model_name : str
        Name of the specific embedding model (e.g., "text-embedding-3-small")
    """
    
    connection_name: str
    model_name: str

class BaseVectorStoreConnection(Resource, ABC):
    """Base class for vector store connections.
    
    Manages connection configuration such as endpoints, authentication, and connection parameters.
  
    """

    @abstractmethod
    def search_with_embedding(self, embedding: List[float], k: int, **kwargs) -> List[VectorSearchResult]:
        """Perform search using pre-computed embedding vector."""
        pass

class BaseVectorStoreSetup(Resource, ABC):
    """Base class for vector store setup.
    
    Similar to BaseChatModelSetup, configures the vector store with
    specific collections, embedding models, and search parameters.
    
    Attributes:
    ----------
    connection_name : str
        Name of the connection resource to use
    embedding_model : str
        Name of the embedding model resource to use
    """

    connection_name: str
    embedding_model: str

Event Types

New event types for vector operations:

class VectorSearchEvent(Event):
    """Event to request semantic search from vector store.
    
    Attributes:
    ----------
    vector_store : str
        Name of the vector store resource to use
    query : str  
        Search query text
    k : int
        Number of results to return (default: 5)
    """

    vector_store: str
    query: str
    k: int = 5

class VectorSearchResultEvent(Event):
    """Event containing vector search results.
    
    Attributes:
    ----------
    request : VectorSearchEvent
        The original search request
    results : List[VectorSearchResult] 
        Search results with documents and scores
    """

    request: VectorSearchEvent
    results: List[VectorSearchResult]

RAG Agent Example

class RAGAgent(Agent):
    @embedding_model_connection
    @staticmethod
    def openai_embedding_connection() -> Tuple[Type[BaseEmbeddingModelConnection], Dict[str, Any]]:
        """OpenAI connection for embedding models."""
        return OpenAIEmbeddingConnection, {
            "name": "openai_embed_conn",
            "api_key": os.getenv("OPENAI_API_KEY")
        }
    
    @embedding_model
    @staticmethod
    def text_embeddings() -> Tuple[Type[BaseEmbeddingModelSetup], Dict[str, Any]]:
        """Text embedding model for semantic search."""
        return OpenAIEmbeddingModel, {
            "name": "text_embeddings",
            "connection": "openai_embed_conn",
            "model_name": "text-embedding-3-small"
        }

    @vector_store_connection
    @staticmethod
    def elasticsearch_connection() -> Tuple[Type[BaseVectorStoreConnection], Dict[str, Any]]:
        """Elasticsearch connection for vector store."""
        return ElasticsearchConnection, {
            "name": "es_conn",
            "hosts": ["http://localhost:9200"]
        }

    @vector_store
    @staticmethod
    def knowledge_base() -> Tuple[Type[BaseVectorStoreSetup], Dict[str, Any]]:
        """Knowledge base vector store using Elasticsearch."""
        return ElasticsearchVectorStore, {
            "name": "knowledge_base",
            "connection": "es_conn",
            "index": "documents",
            "embedding_model": "text_embeddings"  # Reference to embedding model resource
        }

    @chat_model_server
    @staticmethod
    def openai_chat_connection() -> Tuple[Type[BaseChatModelConnection], Dict[str, Any]]:
        """OpenAI connection for chat model."""
        return OpenAIChatModelConnection, {
            "name": "openai_chat_conn",
            "api_key": os.getenv("OPENAI_API_KEY")
        }

    @chat_model
    @staticmethod
    def rag_chat_model() -> Tuple[Type[BaseChatModelSetup], Dict[str, Any]]:
        """Chat model for RAG responses."""
        return OpenAIChatModelSetup, {
            "name": "rag_model",
            "connection": "openai_chat_conn",
            "model": "gpt-4o-mini"
        }

    @action(InputEvent)
    @staticmethod
    def search_knowledge(event: InputEvent, ctx: RunnerContext):
        # Search for relevant context first
        ctx.send_event(VectorSearchEvent(
            vector_store="knowledge_base",
            query=event.input,
            k=3
        ))

    @action(VectorSearchResultEvent)
    @staticmethod
    def generate_response(event: VectorSearchResultEvent, ctx: RunnerContext):
        # Use search results as context for chat model
        context = "\\n".join([r.document.content for r in event.results])
        enhanced_prompt = f"Context: {context}\\nQuestion: {event.request.query}"

        ctx.send_event(ChatRequestEvent(
            model="rag_model",
            messages=[ChatMessage(content=enhanced_prompt)]
        ))

[alnzng]

Hi @xintongsong! This is a high-level design proposal to add vector store support to Flink Agents, enabling RAG capabilities. The design follows our existing ChatModel resource provider pattern and introduces event-driven vector search integration. Can you take a look and share your thoughts?

Hi @coderplay! The proposal focuses on semantic search capabilities that might be related for future long-term memory features. Please also take a look. thanks.

[xintongsong]

@alnzng
Thanks for preparing this design doc. In general, the design looks really good to me. There might be some details worth discussion (attribute names, parameters to include, etc.), which don't really affect the overall design and can be discussed during the code review process.

For the high level design, I just have a few comments.

1. Do you think we need an abstraction of EmbeddingModel, similar to the ChatModel?
2. The RAG Agent design looks really nice. I think we might consider making it another built-in agent of the framework, just like ReAct Agent. The user only need to specify the prompt, chat model, embedding model, and vector store, and everything else are ready to use.

[alnzng]

Thank you for reviewing this design proposal @xintongsong! Yes, let's discuss implementation details (e.g., attribute/parameter naming) during the code review process.

Regarding your comments:

• EmbeddingModel abstraction: Excellent point! I agree we should have this abstraction, as there are multiple embedding service providers (e.g., OpenAI, Voyage, etc.) that offer different APIs for interaction. This would maintain consistency with our ChatModel approach and provide better extensibility.

• Built-in RAG Agent: I think this is a great idea. We can provide a built-in RAG agent that offers a high-level abstraction for ease of use, where users only need to specify the core components (prompt, chat model, embedding model, and vector store). At the same time, we should maintain the flexibility for users to build custom RAG agents when they need more specialized behavior or want to implement their own logic.

Let me update the doc by adding this EmbeddingModel abstraction. And let me create a issue for offerring a Built-in RAG Agent.

[alnzng]

Updated the doc and created a issue for built-in RAG agent: #144

[xintongsong]

@alnzng

Thanks for addressing my comments. The updates look good to me. There's no more questions from my side.

[wenjin272]

@alnzng
Thanks for proposing this design. The overall design looks good to me, it's really a grate work.
I agree to provide EmbeddingModel abstraction, and I see the doc has been updated, maybe we could also update the discussion name to show the scope of this discussion.

[alnzng]

Good suggestion, thanks @wenjin272 . Let me update the titile to Vector Store and Embedding Model Integration for Flink Agents, which should be more clear. I will also update the content accordingly.