'commit'

dsLightRag/Start.py

@@ -45,7 +45,10 @@ async def lifespan(app: FastAPI):
     warnings.filterwarnings('ignore', message='Unverified HTTPS request is being made to host')
     yield
 
-
+async def print_stream(stream):
+    async for chunk in stream:
+        if chunk:
+            print(chunk, end="", flush=True)
 app = FastAPI(lifespan=lifespan)
 
 # 挂载静态文件目录
@@ -57,34 +60,28 @@ async def rag(request: fastapi.Request):
     query = data.get('query', '')
     working_dir = "./Test/Math"  # 使用与T2_Query.py相同的目录
 
-    async def generate_response_stream():
-        rag = None
+    async def generate_response_stream(query: str):
         try:
-            logger.info(f"Initializing RAG with working_dir: {working_dir}")
-            rag = await initialize_rag(working_dir)
-            logger.info("RAG initialized successfully")
-            
-            resp = await rag.aquery(
-                query,
-                param=QueryParam(mode="hybrid", stream=True),
-            )
-            logger.info("Query execution started")
-            
+            rag = await initialize_rag(working_dir="./Test/Math")
+            resp = await rag.aquery(query=query, stream=True)
+
             async for chunk in resp:
-                if not chunk:  # Skip empty chunks
+                if not chunk:
                     continue
-                try:
-                    yield f"data: {json.dumps({'reply': chunk}, ensure_ascii=False)}\n\n"
-                except Exception as e:
-                    logger.error(f"Error processing chunk: {e}")
-                    yield f"data: {json.dumps({'error': str(e)})}\n\n"
+
+                # 正确的SSE格式应该是单层data:前缀
+                yield f"data: {json.dumps({'reply': chunk})}\n\n"
+
+                # 打印到控制台用于调试
+                print(chunk, end='', flush=True)
+
         except Exception as e:
             yield f"data: {json.dumps({'error': str(e)})}\n\n"
         finally:
-            if rag:
-                await rag.finalize_storages()
+            if 'rag' in locals():
+                await rag.finalize()
 
-    return EventSourceResponse(generate_response_stream())
+    return EventSourceResponse(generate_response_stream(query=query))
 
 
 if __name__ == "__main__":