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")