FastAPI异步创建ChromaDB嵌入失败，求状态跟踪功能修复

我希望在FastAPI中创建一个/status API端点，用于跟踪ChromaDB嵌入的生成进度，同时以异步模式生成这些嵌入。当前使用llama-index 0.10.12版本，编写的代码执行时出现RuntimeError：asyncio.run()无法在运行的事件循环中调用，还有协程未等待的警告。期望能成功执行asyncio.create_task(define_rag())，并让/status端点在任务执行时返回正确的状态值。

尝试的代码 #

import os
import asyncio
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import VectorStoreIndex, StorageContext, SimpleDirectoryReader
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.query_engine import RetrieverQueryEngine

app = FastAPI()

# 补全全局状态变量定义
class GlobalVariable:
    def __init__(self):
        self.upload_in_progress = 0
        self.embed_model = None  # 需自行初始化嵌入模型
        self.llm = None  # 需自行初始化LLM
        self.retriever_chunk = None
        self.query_engine_chunk = None
        self.text_qa_template = None  # 需自行初始化问答模板

global_variable = GlobalVariable()

@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    try:
        # 确保docs目录存在
        if not os.path.exists("docs"):
           os.makedirs("docs")
    
        # 将文件写入服务器
        file_path = os.path.join("docs", file.filename)
        with open(file_path, "wb") as f:
           f.write(await file.read())
        
        asyncio.create_task(define_rag())
        return JSONResponse(content={"message": "文件上传成功"})

    except Exception as e:
       return JSONResponse(content={"error": str(e)}, status_code=500)

@app.get("/status")
async def status_checker():
   return {"upload_in_progress": global_variable.upload_in_progress}

async def define_rag():
    global_variable.upload_in_progress = 1  # 开始处理时标记状态
    try:
        documents = SimpleDirectoryReader(input_dir="./docs", required_exts=[".docx",".doc",".pdf",".txt"]).load_data()
        if os.path.exists('./chroma_db'):
            print("******************************************使用chromadb文件夹中已生成的嵌入")
            chroma_client = chromadb.PersistentClient(path="./chroma_db")
            chroma_collection = chroma_client.get_or_create_collection("quickstart")
            vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
            vector_index_chunk = VectorStoreIndex.from_vector_store(
                vector_store, 
                embed_model=global_variable.embed_model,
                use_async=True, 
                show_progress=True
            )
        else:
            chroma_client = chromadb.PersistentClient(path="./chroma_db")
            chroma_collection = chroma_client.get_or_create_collection("quickstart")
            vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
            storage_context = StorageContext.from_defaults(vector_store=vector_store)
            # 使用异步类方法替代直接实例化
            vector_index_chunk = await VectorStoreIndex.from_documents(
                documents,
                embed_model=global_variable.embed_model,
                storage_context=storage_context,
                use_async=True, 
                show_progress=True
            )
        
        # 补全节点相关定义
        all_nodes = list(vector_index_chunk.index_struct.nodes.values())
        all_nodes_dict = {node.node_id: node for node in all_nodes}
        
        vector_retriever_chunk = vector_index_chunk.as_retriever(similarity_top_k=5)
        global_variable.retriever_chunk = RecursiveRetriever(
            "vector",
            retriever_dict={"vector": vector_retriever_chunk},
            node_dict=all_nodes_dict,
            verbose=True,
        )
        print("向量存储创建完成")
        global_variable.query_engine_chunk = RetrieverQueryEngine.from_args(
            global_variable.retriever_chunk,
            llm=global_variable.llm, 
            text_qa_template=global_variable.text_qa_template
        )
    finally:
        global_variable.upload_in_progress = 0  # 处理完成/失败后重置状态

报错信息 #

vector_index_chunk = await VectorStoreIndex(all_nodes, embed_model=global_variable.embed_model, storage_context=storage_context, use_async=True, show_progress=True)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lib/python3.11/site-packages/llama_index/core/indices/vector_store/base.py", line 74, in __init__
    super().__init__(
  File "/home/lib/python3.11/site-packages/llama_index/core/indices/base.py", line 91, in __init__
    index_struct = self.build_index_from_nodes(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lib/python3.11/site-packages/llama_index/core/indices/vector_store/base.py", line 307, in build_index_from_nodes
    return self._build_index_from_nodes(nodes, **insert_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lib/python3.11/site-packages/llama_index/core/indices/vector_store/base.py", line 277, in _build_index_from_nodes
    run_async_tasks(tasks)
  File "/homelib/python3.11/site-packages/llama_index/core/async_utils.py", line 49, in run_async_tasks
    outputs: List[Any] = asyncio.run(_gather())
                         ^^^^^^^^^^^^^^^^^^^^^^
  File "/homelib/python3.11/asyncio/runners.py", line 186, in run
    raise RuntimeError(
RuntimeError: asyncio.run() cannot be called from a running event loop
/home/miniconda3/envs/lib/python3.11/site-packages/uvicorn/protocols/http/httptools_impl.py:-1: RuntimeWarning: coroutine 'run_async_tasks.<locals>._gather' was never awaited
RuntimeWarning: Enable tracemalloc to get the object allocation traceback
/home/lib/python3.11/site-packages/uvicorn/protocols/http/httptools_impl.py:-1: RuntimeWarning: coroutine 'VectorStoreIndex._async_add_nodes_to_index' was never awaited
RuntimeWarning: Enable tracemalloc to get the object allocation traceback

解决方案 #

核心问题分析 #

报错根源是llama-index 0.10.12版本的run_async_tasks函数内部嵌套调用了asyncio.run()，而FastAPI已经在运行事件循环，不允许嵌套启动新循环。此外代码还存在以下问题：

• 错误地直接实例化VectorStoreIndex并await，未使用官方异步类方法
• all_nodes和all_nodes_dict未定义
• 状态变量更新时机不合理，未覆盖任务异常场景
• /status使用POST方法不符合REST查询规范

具体修复步骤 #

1. 修改llama-index异步工具函数：
   找到llama_index/core/async_utils.py中的run_async_tasks函数，修改为异步版本并替换asyncio.run为直接await：

   async def run_async_tasks(tasks: List[Awaitable]) -> List[Any]:
       """Run async tasks in parallel."""
       async def _gather() -> List[Any]:
           return await asyncio.gather(*tasks)
   
       outputs: List[Any] = await _gather()
       return outputs
   

   同时在调用该函数的_build_index_from_nodes方法中，添加await关键字：

   outputs = await run_async_tasks(tasks)
   

2. 修正define_rag逻辑：

   • 使用await VectorStoreIndex.from_documents()替代直接实例化
   • 补全all_nodes和all_nodes_dict的定义
   • 用finally块确保状态变量在任务结束（包括异常）时重置

3. 调整接口规范：将/status改为GET方法，返回结构化JSON

额外说明 #

• 需自行初始化global_variable中的embed_model、llm和text_qa_template
• 生产环境建议用Redis等分布式存储替代全局变量管理状态
• 优先考虑升级llama-index到更高版本，新版本已修复该异步循环冲突问题

内容的提问来源于stack exchange，提问作者PADALA LIKHITH RISHI