Add WaitResponses to PredictStreamedContext Method.

* Adds an optional method that will wait for all responses to be generated before returning for streaming responses.

PiperOrigin-RevId: 691069391

Author: kenfranko

tensorflow_serving/servables/tensorflow/mock_servable.h

@@ -37,6 +37,7 @@ class MockPredictStreamedContext : public PredictStreamedContext {
   MOCK_METHOD(absl::Status, ProcessRequest, (const PredictRequest& request),
               (final));
   MOCK_METHOD(absl::Status, Close, (), (final));
+  MOCK_METHOD(absl::Status, WaitResponses, (), (final));
 };
 
 // A mock of tensorflow::serving::Servable.

tensorflow_serving/servables/tensorflow/servable.cc

@@ -61,5 +61,9 @@ absl::Status SingleRequestPredictStreamedContext::Close() {
   return absl::OkStatus();
 }
 
+absl::Status SingleRequestPredictStreamedContext::WaitResponses() {
+  return absl::OkStatus();
+}
+
 }  // namespace serving
 }  // namespace tensorflow

tensorflow_serving/servables/tensorflow/servable.h

@@ -54,6 +54,9 @@ class PredictStreamedContext {
 
   // Closes the `PredictStreamed` session.
   virtual absl::Status Close() = 0;
+
+  // Waits for all of the responses to be generated
+  virtual absl::Status WaitResponses() = 0;
 };
 
 // A convenience wrapper for cases where the implementation allows exactly one
@@ -69,6 +72,7 @@ class SingleRequestPredictStreamedContext final
 
   absl::Status ProcessRequest(const PredictRequest& request) final;
   absl::Status Close() final;
+  absl::Status WaitResponses() final;
 
  private:
   absl::AnyInvocable<absl::Status(const PredictRequest&)> f_;