@app.post("/ask")
async def ask(data: AskModel):
    s3 = get_s3()
    chat_info = get_convo(data.chat_uuid, s3)
    access = chat_info.get('meta', {}).get('access', [])
    index_set = get_index(access, fs=s3)
    chat_history = chat_info["history"]
    chat_engine = index_set.as_chat_engine(
        similarity_top_k=3,
        streaming=True,
        chat_history=chat_history,
    )

    def stream_response():
        streaming_response = chat_engine.stream_chat(data.question)
        response_gen = streaming_response.response_gen
        first_token = next(response_gen)  # Get the first token of the response
        response = first_token
        yield f"data: {response}\n\n"  # Yield the first token separately

        for token in response_gen:
            yield f"data: {response}\n\n"
            response+=token

        chat_history.extend([
            {"role": "user", "content": data.question},
            {"role": "assistant", "content": ''.join(response)}
        ])
        chat_info['history'] = chat_history
        save_chat(data.chat_uuid, chat_info, s3)

    return StreamingResponse(stream_response(), media_type="text/event-stream")

s3 = s3fs.S3FileSystem(
    key=AWS_KEY,
    secret=AWS_SECRET,
    endpoint_url=S3_BUCKET_URL
)

def add_new_doc(name, filename, index_set):
    with open(filename, "r") as f:
        doc = f.read()
        storage_context = StorageContext.from_defaults(fs=s3)
        index_set[name] = VectorStoreIndex.from_documents(
            [Document(text=doc)],
            service_context=service_context,
            storage_context=storage_context,
        )
        storage_context.persist(persist_dir=f'bucket-name/{name}')
    return index_set