Skip to content

Commit b611452

Browse files
committed
chore: expose class from package level for retrieval pipe
1 parent bf38415 commit b611452

File tree

2 files changed

+8
-4
lines changed

2 files changed

+8
-4
lines changed

wren-ai-service/src/globals.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,7 @@
66
from src.config import Settings
77
from src.core.pipeline import PipelineComponent
88
from src.core.provider import EmbedderProvider, LLMProvider
9-
from src.pipelines import generation, indexing
10-
from src.pipelines.retrieval import historical_question, preprocess_sql_data, retrieval
9+
from src.pipelines import generation, indexing, retrieval
1110
from src.web.v1.services.ask import AskService
1211
from src.web.v1.services.ask_details import AskDetailsService
1312
from src.web.v1.services.question_recommendation import QuestionRecommendation
@@ -88,7 +87,7 @@ def create_service_container(
8887
table_column_retrieval_size=settings.table_column_retrieval_size,
8988
allow_using_db_schemas_without_pruning=settings.allow_using_db_schemas_without_pruning,
9089
),
91-
"historical_question": historical_question.HistoricalQuestion(
90+
"historical_question": retrieval.HistoricalQuestion(
9291
**pipe_components["historical_question_retrieval"],
9392
),
9493
"sql_generation": generation.SQLGeneration(
@@ -108,7 +107,7 @@ def create_service_container(
108107
),
109108
sql_answer_service=SqlAnswerService(
110109
pipelines={
111-
"preprocess_sql_data": preprocess_sql_data.PreprocessSqlData(
110+
"preprocess_sql_data": retrieval.PreprocessSqlData(
112111
**pipe_components["preprocess_sql_data"],
113112
),
114113
"sql_answer": generation.SQLAnswer(
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
from .historical_question import HistoricalQuestion
2+
from .preprocess_sql_data import PreprocessSqlData
3+
from .retrieval import Retrieval
4+
5+
__all__ = ["HistoricalQuestion", "PreprocessSqlData", "Retrieval"]

0 commit comments

Comments
 (0)