Function chat_stream

server/server.py:111–156

Streaming version of the chat endpoint using Server-Sent Events (SSE). Streams pipeline steps and LLM tokens as they are generated.

()

        "history": new_history,
        "documents": documents,
        "rewritten": rewritten,
        "question": prompt,
        "fetched_new_documents": fetched_new_documents
    }
    if provenance_scores is not None:
        for i, doc in enumerate(response_dict["documents"]):
            response_dict["documents"][i]["provenance"] = provenance_scores[i]["score"]
    return jsonify(response_dict)

@app.route("/chat_stream", methods=['POST'])
def chat_stream():
    """
    Streaming version of the chat endpoint using Server-Sent Events (SSE).
    Streams pipeline steps and LLM tokens as they are generated.
    """
    json_data = request.get_json()
    prompt = json_data.get('prompt')
    history = json_data.get('history', [])
    original_docs = json_data.get('docs', [])
    datasets = json_data.get('datasets', [])

    def generate():
        try:
            for event_type, event_data in raghelper.handle_user_interaction_stream(prompt, history, datasets):
                if event_type == "step":
                    yield f"event: step\ndata: {safe_json_dumps({'step': event_data})}\n\n"
                elif event_type == "token":
                    yield f"event: token\ndata: {safe_json_dumps({'token': event_data})}\n\n"
                elif event_type == "documents":
                    yield f"event: documents\ndata: {safe_json_dumps({'documents': event_data})}\n\n"
                elif event_type == "done":
                    metadata = event_data
                    documents = metadata.get("documents") or original_docs
                    provenance_scores = metadata.get("provenance_scores")
                    if provenance_scores is not None and documents:
                        for i, doc in enumerate(documents):
                            if i < len(provenance_scores):
                                documents[i]["provenance"] = provenance_scores[i]["score"]
                    done_data = {
                        "reply": metadata["reply"],
                        "history": metadata["history"],
                        "documents": documents,
                        "rewritten": metadata["rewritten"],
                        "question": prompt,
                        "fetched_new_documents": metadata["fetched_new_documents"],
                    }
                    yield f"event: done\ndata: {safe_json_dumps(done_data)}\n\n"
        except Exception as e:
            logger.error(f"Streaming error: {e}", exc_info=True)