/ src / pipeline / steps / query_embedding_step.py
query_embedding_step.py
 1  from __future__ import annotations
 2  
 3  """Step that generates an embedding for the user query."""
 4  
 5  import logging
 6  
 7  from ...embeddings.protocol import Embeddings
 8  from ..contexts.query_context import QueryContext
 9  
10  
11  class QueryEmbeddingStep:
12      """Step that generates an embedding for the user query."""
13  
14      def __init__(self, embedding_model: Embeddings):
15          """Initialize the query embedding step.
16  
17          Parameters
18          ----------
19          embedding_model
20              Embedding model instance created by EmbeddingModelFactory.
21          """
22          self.embedding_model = embedding_model
23  
24      def run(self, context: QueryContext) -> None:
25          """Generate embedding for the user query.
26  
27          Parameters
28          ----------
29          context
30              Query context with user_query set.
31          """
32          logger = logging.getLogger(__name__)
33          logger.info(f"Generating embedding for query: {context.user_query}")
34  
35          query_vector = self.embedding_model.embed_query(context.user_query)
36          context.query_vector = query_vector
37  
38          logger.info("Query embedding generated")