Function stream

toolbench/inference/toolbench_server.py:103–157

()

@app.route('/stream', methods=['GET', 'POST'])
@cross_origin()
def stream():
    data = json.loads(request.data)
    user_input = data["text"]
    top_k = data["top_k"]
    method = data["method"]
    print("Called stream")
    global model

    def generate(model):
        print("Called generate")
        if model.inuse:
            # send 409 error
            return Response(json.dumps({
                "method_name": "error",
                "error": "Model in use"
            }), status=409, mimetype='application/json')
            return
        model.inuse = True

        # run model.run_agent in the background
        with concurrent.futures.ThreadPoolExecutor() as executor:

            future = executor.submit(model.run_pipeline, user_input, method, top_k)
            # keep waiting for the queue to be empty
            while True:
                if model.queue.empty():
                    if future.done():
                        print("Finished with future")
                        break
                    time.sleep(0.01)
                    continue
                else:
                    obj = model.queue.get()
                if obj["method_name"] == "unknown": continue
                if obj["method_name"] == "on_request_end":
                    yield json.dumps(obj)
                    break

                try:
                    yield json.dumps(obj) + "\n"
                except Exception as e:
                    model.inuse = False
                    print(obj)
                    print(e)

            try:
                future.result()
            except Exception as e:
                model.inuse = False
                print(e)

        model.inuse = False
        return

    return Response(stream_with_context(generate(model)))

@app.route('/methods', methods=['GET'])
@cross_origin()