diff --git a/google_genai/mcp/run_worker.py b/google_genai/mcp/run_worker.py index 5098c0e8..821c3555 100644 --- a/google_genai/mcp/run_worker.py +++ b/google_genai/mcp/run_worker.py @@ -5,7 +5,6 @@ ``list_tools`` / ``call_tool`` as activities. """ -# @@@SNIPSTART python-google-genai-mcp-worker import asyncio import os import sys @@ -25,6 +24,7 @@ ECHO_SERVER = str(Path(__file__).parent / "echo_mcp_server.py") +# @@@SNIPSTART python-google-genai-mcp-worker @asynccontextmanager async def echo_session() -> AsyncIterator[ClientSession]: """Yield a connected, initialized session to the stdio echo MCP server.""" @@ -38,6 +38,7 @@ async def echo_session() -> AsyncIterator[ClientSession]: async def main() -> None: genai_client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"]) plugin = GoogleGenAIPlugin(genai_client, mcp_servers={"echo": echo_session}) + # @@@SNIPEND client = await Client.connect( os.environ.get("TEMPORAL_ADDRESS", "localhost:7233"), @@ -55,4 +56,3 @@ async def main() -> None: if __name__ == "__main__": asyncio.run(main()) -# @@@SNIPEND diff --git a/google_genai/streaming/run_workflow.py b/google_genai/streaming/run_workflow.py index 9b86910c..0cded918 100644 --- a/google_genai/streaming/run_workflow.py +++ b/google_genai/streaming/run_workflow.py @@ -1,6 +1,5 @@ """Start the streaming workflow and consume model chunks live.""" -# @@@SNIPSTART python-google-genai-streaming-run-workflow import asyncio import os from datetime import timedelta @@ -17,6 +16,7 @@ STREAM_TIMEOUT = 60.0 +# @@@SNIPSTART python-google-genai-streaming-run-workflow async def consume(client: Client, workflow_id: str) -> None: """Subscribe to the "gemini" topic and print chunks as the model produces them.""" stream = WorkflowStreamClient.create(client, workflow_id) @@ -41,6 +41,7 @@ async def main() -> None: os.environ.get("TEMPORAL_ADDRESS", "localhost:7233"), data_converter=pydantic_data_converter, ) + # @@@SNIPEND workflow_id = "google-genai-streaming" handle = await client.start_workflow( @@ -63,4 +64,3 @@ async def main() -> None: if __name__ == "__main__": asyncio.run(main()) -# @@@SNIPEND diff --git a/google_genai/tools/run_worker.py b/google_genai/tools/run_worker.py index f251542b..05951fa3 100644 --- a/google_genai/tools/run_worker.py +++ b/google_genai/tools/run_worker.py @@ -1,6 +1,5 @@ """Worker for the tools sample.""" -# @@@SNIPSTART python-google-genai-tools-worker import asyncio import os @@ -21,16 +20,17 @@ async def main() -> None: plugins=[plugin], ) + # @@@SNIPSTART python-google-genai-tools-worker worker = Worker( client, task_queue="google-genai-tools", workflows=[ToolsWorkflow], activities=[get_weather], ) + # @@@SNIPEND print("Worker started. Ctrl+C to exit.") await worker.run() if __name__ == "__main__": asyncio.run(main()) -# @@@SNIPEND diff --git a/google_genai/vertex_ai/run_worker.py b/google_genai/vertex_ai/run_worker.py index ee2f77f9..1952a1a4 100644 --- a/google_genai/vertex_ai/run_worker.py +++ b/google_genai/vertex_ai/run_worker.py @@ -5,7 +5,6 @@ ``GOOGLE_APPLICATION_CREDENTIALS`` to a service-account key file. """ -# @@@SNIPSTART python-google-genai-vertex-ai-worker import asyncio import os @@ -18,12 +17,14 @@ async def main() -> None: + # @@@SNIPSTART python-google-genai-vertex-ai-worker genai_client = genai.Client( vertexai=True, project=os.environ["GOOGLE_CLOUD_PROJECT"], location=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), ) plugin = GoogleGenAIPlugin(genai_client) + # @@@SNIPEND client = await Client.connect( os.environ.get("TEMPORAL_ADDRESS", "localhost:7233"), @@ -41,4 +42,3 @@ async def main() -> None: if __name__ == "__main__": asyncio.run(main()) -# @@@SNIPEND