From 506fea5e7aefd7b652166977219a2aaa7ba9c6fd Mon Sep 17 00:00:00 2001
From: Sampath M <sampathm@google.com>
Date: Fri, 18 Oct 2024 23:05:04 +0200
Subject: [PATCH] feat(generative-ai): update claude_3_streaming_example.py

---
 .../claude_3_streaming_example.py             | 102 +++++++++---------
 ..._test.py => test_model_garden_examples.py} |  10 +-
 2 files changed, 56 insertions(+), 56 deletions(-)
 rename generative_ai/model_garden/{model_garden_test.py => test_model_garden_examples.py} (83%)

diff --git a/generative_ai/model_garden/claude_3_streaming_example.py b/generative_ai/model_garden/claude_3_streaming_example.py
index 2ba642e2d4c..3e817e9b584 100644
--- a/generative_ai/model_garden/claude_3_streaming_example.py
+++ b/generative_ai/model_garden/claude_3_streaming_example.py
@@ -11,54 +11,54 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-# import os
-#
-# PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-#
-#
-# def generate_text_streaming() -> str:
-#     # [START generativeaionvertexai_claude_3_streaming]
-#     # TODO(developer): Vertex AI SDK - uncomment below & run
-#     # pip3 install --upgrade --user google-cloud-aiplatform
-#     # gcloud auth application-default login
-#     # pip3 install -U 'anthropic[vertex]'
-#
-#     # TODO(developer): Update and un-comment below line
-#     # PROJECT_ID = "your-project-id"
-#
-#     from anthropic import AnthropicVertex
-#
-#     client = AnthropicVertex(project_id=PROJECT_ID, region="us-east5")
-#     result = []
-#
-#     with client.messages.stream(
-#         model="claude-3-5-sonnet@20240620",
-#         max_tokens=1024,
-#         messages=[
-#             {
-#                 "role": "user",
-#                 "content": "Send me a recipe for banana bread.",
-#             }
-#         ],
-#     ) as stream:
-#         for text in stream.text_stream:
-#             print(text, end="", flush=True)
-#             result.append(text)
-#
-#     # Example response:
-#     # Here's a simple recipe for delicious banana bread:
-#     # Ingredients:
-#     # - 2-3 ripe bananas, mashed
-#     # - 1/3 cup melted butter
-#     # ...
-#     # ...
-#     # 8. Bake for 50-60 minutes, or until a toothpick inserted into the center comes out clean.
-#     # 9. Let cool in the pan for a few minutes, then remove and cool completely on a wire rack.
-#
-#     # [END generativeaionvertexai_claude_3_streaming]
-#     return "".join(result)
-#
-#
-# if __name__ == "__main__":
-#     generate_text_streaming()
+
+import os
+
+PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+
+
+def generate_text_streaming() -> str:
+    # [START generativeaionvertexai_claude_3_streaming]
+    # TODO(developer): Vertex AI SDK - uncomment below & run
+    # pip3 install --upgrade --user google-cloud-aiplatform
+    # gcloud auth application-default login
+    # pip3 install -U 'anthropic[vertex]'
+
+    # TODO(developer): Update and un-comment below line
+    # PROJECT_ID = "your-project-id"
+
+    from anthropic import AnthropicVertex
+
+    client = AnthropicVertex(project_id=PROJECT_ID, region="us-east5")
+    result = []
+
+    with client.messages.stream(
+        model="claude-3-5-sonnet@20240620",
+        max_tokens=1024,
+        messages=[
+            {
+                "role": "user",
+                "content": "Send me a recipe for banana bread.",
+            }
+        ],
+    ) as stream:
+        for text in stream.text_stream:
+            print(text, end="", flush=True)
+            result.append(text)
+
+    # Example response:
+    # Here's a simple recipe for delicious banana bread:
+    # Ingredients:
+    # - 2-3 ripe bananas, mashed
+    # - 1/3 cup melted butter
+    # ...
+    # ...
+    # 8. Bake for 50-60 minutes, or until a toothpick inserted into the center comes out clean.
+    # 9. Let cool in the pan for a few minutes, then remove and cool completely on a wire rack.
+
+    # [END generativeaionvertexai_claude_3_streaming]
+    return "".join(result)
+
+
+if __name__ == "__main__":
+    generate_text_streaming()
diff --git a/generative_ai/model_garden/model_garden_test.py b/generative_ai/model_garden/test_model_garden_examples.py
similarity index 83%
rename from generative_ai/model_garden/model_garden_test.py
rename to generative_ai/model_garden/test_model_garden_examples.py
index f3c6c3e4df8..4f82cdf742e 100644
--- a/generative_ai/model_garden/model_garden_test.py
+++ b/generative_ai/model_garden/test_model_garden_examples.py
@@ -16,15 +16,15 @@
 
 from google.api_core.exceptions import ResourceExhausted
 
-# import claude_3_streaming_example
+import claude_3_streaming_example
 import claude_3_tool_example
 import claude_3_unary_example
 
 
-# @backoff.on_exception(backoff.expo, ResourceExhausted, max_time=10)
-# def test_generate_text_streaming() -> None:
-#     responses = claude_3_streaming_example.generate_text_streaming()
-#     assert "bread" in responses
+@backoff.on_exception(backoff.expo, ResourceExhausted, max_time=10)
+def test_generate_text_streaming() -> None:
+    responses = claude_3_streaming_example.generate_text_streaming()
+    assert "bread" in responses
 
 
 @backoff.on_exception(backoff.expo, ResourceExhausted, max_time=10)