fix bug

Canner · Sep 18, 2024 · 1c03e97 · 1c03e97
1 parent a9bee1f
commit 1c03e97
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 8 deletions.
diff --git a/wren-ai-service/src/__main__.py b/wren-ai-service/src/__main__.py
@@ -39,7 +39,6 @@ async def lifespan(app: FastAPI):
         engine_config=EngineConfig(provider=os.getenv("ENGINE", "wren_ui"))
     )
     app.state.service_container = create_service_container(
-        app,
         *providers,
         should_force_deploy=bool(os.getenv("SHOULD_FORCE_DEPLOY", "")),
         column_indexing_batch_size=(
@@ -64,7 +63,7 @@ async def lifespan(app: FastAPI):
             "ttl": int(os.getenv("QUERY_CACHE_TTL") or 120),
         },
     )
-    app.state.service_metadata = create_service_metadata(app, *providers)
+    app.state.service_metadata = create_service_metadata(*providers)
     init_langfuse()
 
     yield

diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py
@@ -3,7 +3,6 @@
 from typing import Optional
 
 import toml
-from fastapi import FastAPI
 
 from src.core.engine import Engine
 from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider
@@ -49,7 +48,6 @@ class ServiceMetadata:
 
 
 def create_service_container(
-    app: FastAPI,
     llm_provider: LLMProvider,
     embedder_provider: EmbedderProvider,
     document_store_provider: DocumentStoreProvider,
@@ -180,7 +178,6 @@ def get_service_container():
 
 
 def create_service_metadata(
-    app: FastAPI,
     llm_provider: LLMProvider,
     embedder_provider: EmbedderProvider,
     *_,

diff --git a/wren-ai-service/src/pipelines/generation/sql_expansion.py b/wren-ai-service/src/pipelines/generation/sql_expansion.py
@@ -38,7 +38,7 @@
 """
 
 sql_expansion_user_prompt_template = """
-SQL: {{history.sql}}
+SQL: {{sql}}
 
 User's input: {{query}}
 
@@ -61,7 +61,7 @@ def prompt(
     logger.debug(f"query: {query}")
     logger.debug(f"documents: {documents}")
     logger.debug(f"history: {history}")
-    return prompt_builder.run(query=query, documents=documents, history=history)
+    return prompt_builder.run(query=query, documents=documents, sql=history.sql)
 
 
 @async_timer

diff --git a/wren-ai-service/src/web/v1/routers.py b/wren-ai-service/src/web/v1/routers.py
@@ -177,7 +177,7 @@ async def sql_expansion(
     service_metadata: ServiceMetadata = Depends(get_service_metadata),
 ) -> SqlExpansionResponse:
     query_id = str(uuid.uuid4())
-    sql_expansion_request.query = query_id
+    sql_expansion_request.query_id = query_id
     service_container.sql_expansion_service._sql_expansion_results[
         query_id
     ] = SqlExpansionResultResponse(