From 7954741206e2b587e241e60a772e9a3cb4fd847d Mon Sep 17 00:00:00 2001 From: Anuj Tambwekar Date: Wed, 24 Dec 2025 01:56:40 +0530 Subject: [PATCH 1/3] Add streaming example --- Streaming_and_Tools/.env.example | 3 + Streaming_and_Tools/.gitignore | 5 + Streaming_and_Tools/.gradient/agent.yml | 3 + Streaming_and_Tools/README.md | 130 +++++++++++++++++++++ Streaming_and_Tools/main.py | 145 ++++++++++++++++++++++++ Streaming_and_Tools/requirements.txt | 8 ++ 6 files changed, 294 insertions(+) create mode 100644 Streaming_and_Tools/.env.example create mode 100644 Streaming_and_Tools/.gitignore create mode 100644 Streaming_and_Tools/.gradient/agent.yml create mode 100644 Streaming_and_Tools/README.md create mode 100644 Streaming_and_Tools/main.py create mode 100644 Streaming_and_Tools/requirements.txt diff --git a/Streaming_and_Tools/.env.example b/Streaming_and_Tools/.env.example new file mode 100644 index 0000000..16a21e9 --- /dev/null +++ b/Streaming_and_Tools/.env.example @@ -0,0 +1,3 @@ +DIGITALOCEAN_INFERENCE_KEY= +TAVILY_API_KEY= +DIGITALOCEAN_API_TOKEN= diff --git a/Streaming_and_Tools/.gitignore b/Streaming_and_Tools/.gitignore new file mode 100644 index 0000000..8ae1eb9 --- /dev/null +++ b/Streaming_and_Tools/.gitignore @@ -0,0 +1,5 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ + +# Environments +.env diff --git a/Streaming_and_Tools/.gradient/agent.yml b/Streaming_and_Tools/.gradient/agent.yml new file mode 100644 index 0000000..f6af294 --- /dev/null +++ b/Streaming_and_Tools/.gradient/agent.yml @@ -0,0 +1,3 @@ +agent_environment: main +agent_name: streaming-agent-with-tool +entrypoint_file: main.py diff --git a/Streaming_and_Tools/README.md b/Streaming_and_Tools/README.md new file mode 100644 index 0000000..21fa7a3 --- /dev/null +++ b/Streaming_and_Tools/README.md @@ -0,0 +1,130 @@ +# Streaming Agents and using Python Functions as Tools + +This example demonstrates an ADK agent that streams responses and uses a tool that is a Python function. Using Langgraph's `tool` decorator, we can construct an agent that can execute python functions as a tool. In this example, we use the Tavily search libary to enable our agent to search the web. +The runtime builds a LangGraph `StateGraph` where the first node decides if tools need to be used. If so, the agent queries the tools and then passes along the results from the tool to a final answer node. If a tool is not required, the answer node responds directly. +``` + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------------+ + | decision_model_call | + +---------------------+ + .. .. + .. . + . .. ++-------+ . +| tools | .. ++-------+ . + ** .. + ** .. + * . + +-------------------+ + | answer_model_call | + +-------------------+ + * + * + * + +---------+ + | __end__ | + +---------+ +``` + +Additionally, this agent yields streaming outputs instead of providing an entire response. Enabling streaming responses from your agent can help improve the user experience by allowing users of the agent to start getting a response quickly without waiting for the entire response to complete. + + +## Quickstart + +1. Create and activate a virtual environment. +2. Install dependencies: + + pip install -r Templates/ToolCalling/requirements.txt + +3. Set the required enviornment variables in the .env file (`TAVILY_API_KEY` and `DIGITALOCEAN_INFERENCE_KEY`). You can obtain a Tavily API key for free at [this link](https://app.tavily.com/home). + +4. Set your DIGITALOCEAN_API_TOKEN via + + ``` + export DIGITALOCEAN_API_TOKEN= # On MacOS/Linux + + set DIGITALOCEAN_API_TOKEN= # On Windows + ``` + +5. Run your agent locally via + + `gradient agent run` + + Since this agent produces streaming responses, you can invoke the agent via code to see the streaming + + ```python + import requests + import json + + def stream_endpoint(url: str, body : dict, headers: dict = {}, chunk_size: int = 1024): + payload = json.dumps(body) + with requests.post(url, data = payload, headers=headers, stream=True) as resp: + resp.raise_for_status() + for chunk in resp.iter_content(chunk_size=chunk_size): + if chunk: # filter keep-alive chunks + yield chunk + + if __name__ == "__main__": + body = { + "prompt" : { + "messages" : "Who were winners and runners up of the women's 2025 cricket world cup?" + } + } + + buffer = "" + for chunk in stream_endpoint("http://localhost:8080/run", body=body): + buffer += chunk.decode("utf-8") + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + if not line.strip(): + continue + response = json.loads(line) + print(response["response"], end="", flush=True) + + ``` + +6. Change the name of the agent if you need to in `gradient/agent.yaml` and then deploy with + + ``` + gradient agent deploy + ``` + + You can the invoke the agent via the same script, just using your deployed agent's URL instead. Make sure you include your API key for authorization + + ```python + import requests + import json + + def stream_endpoint(url: str, body : dict, headers: dict = {}, chunk_size: int = 1024): + payload = json.dumps(body) + with requests.post(url, data = payload, headers=headers, stream=True) as resp: + resp.raise_for_status() + for chunk in resp.iter_content(chunk_size=chunk_size): + if chunk: # filter keep-alive chunks + yield chunk + + if __name__ == "__main__": + body = { + "prompt" : { + "messages" : "Who were winners and runners up of the women's 2025 cricket world cup?" + } + } + + headers = {'Authorization': 'Bearer '} + + buffer = "" + for chunk in stream_endpoint('https://agents.do-ai.run//main/run', body=body, headers=headers): + buffer += chunk.decode("utf-8") + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + if not line.strip(): + continue + response = json.loads(line) + print(response["response"], end="", flush=True) + ``` diff --git a/Streaming_and_Tools/main.py b/Streaming_and_Tools/main.py new file mode 100644 index 0000000..d845e6b --- /dev/null +++ b/Streaming_and_Tools/main.py @@ -0,0 +1,145 @@ +import os +import json +from langgraph.graph import StateGraph, MessagesState, START, END +from langgraph.prebuilt import ToolNode +from langchain.messages import SystemMessage +from langchain_core.tools import tool +from langchain_openai import ChatOpenAI +from dotenv import load_dotenv +from typing import Dict, Literal +from gradient_adk import entrypoint +from tavily import TavilyClient +from langgraph.config import get_stream_writer + +load_dotenv() + +TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") +tavily_client = TavilyClient(api_key=TAVILY_API_KEY) + + +@tool +def search_tool(query: str) -> dict: + """ + A tool to search a query on the web using Tavily for detailed and up to date information. + If you need to look up recent information beyond your knowledge cutoff date, use this tool. + """ + response = tavily_client.search(query) + return response + + +# This function determines whether to branch to the tools node or the answer model node +def tools_branch_condition( + state, + messages_key: str = "messages", +) -> Literal["tools", "answer_model_call"]: + if isinstance(state, list): + ai_message = state[-1] + elif isinstance(state, dict) and (messages := state.get(messages_key, [])): + ai_message = messages[-1] + elif messages := getattr(state, messages_key, []): + ai_message = messages[-1] + else: + raise ValueError(f"No messages found in input state to tool_edge: {state}") + if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0: + return "tools" + return "answer_model_call" + + +def build_graph(): + + # We have two models - one for the initial decision making (non-streaming), and one for the final answer (streaming) + decision_model = ChatOpenAI( + model="openai-gpt-4.1", + base_url="https://inference.do-ai.run/v1", + api_key=os.getenv("DIGITALOCEAN_INFERENCE_KEY"), + ) + answer_model = ChatOpenAI( + model="openai-gpt-4.1", + base_url="https://inference.do-ai.run/v1", + api_key=os.getenv("DIGITALOCEAN_INFERENCE_KEY"), + streaming=True, + ) + + tools = [search_tool] + decision_model = decision_model.bind_tools(tools) + + async def decision_model_call(state: MessagesState): + messages = state["messages"] + system_prompt = SystemMessage( + "Determine if the user query is out of your scope of knowledge and requires a web search. If so, invoke the tool. Otherwise, simply repeat the original query so that the next model can answer it." + ) + response = await decision_model.ainvoke([system_prompt] + messages) + return {"messages": response} + + # In order to stream from the answer model, we use a stream writer to yield chunks as they arrive + # We store the streamed chunks in a custom key in the state dict, which we can then access in the entrypoint. + # For more information, see https://docs.langchain.com/oss/python/langgraph/streaming#llm-tokens + async def answer_model_call(state: MessagesState): + messages = state["messages"] + writer = get_stream_writer() + async for response in answer_model.astream(messages): + writer({"custom_llm_chunk": response}) + return {"result": "completed"} + + builder = StateGraph(MessagesState) + builder.add_node("decision_model_call", decision_model_call) + builder.add_node("tools", ToolNode(tools)) + builder.add_node("answer_model_call", answer_model_call) + builder.add_edge(START, "decision_model_call") + builder.add_conditional_edges( + "decision_model_call", + tools_branch_condition, + ) + builder.add_edge("tools", "answer_model_call") + builder.add_edge("answer_model_call", END) + graph = builder.compile() + return graph + + +AGENT_GRAPH = build_graph() + + +@entrypoint +async def main(input: Dict, context: Dict): + """Entrypoint""" + + input_request = input.get("prompt") + + # As mentioned earlier, we stream the response from the answer model by accessing the custom key in the state dict + async for _, chunk in AGENT_GRAPH.astream(input_request, stream_mode=["custom"]): + chunk_data = chunk.get("custom_llm_chunk", None) + if chunk_data and chunk_data.content: + response_text = chunk_data.content + yield json.dumps({"response": response_text}) + "\n" + + +# # You can stream responses from this agent by sending a streaming request to the endpoint. +# import requests +# import json + +# def stream_endpoint(url: str, body : dict, headers : dict = {}, chunk_size: int = 1024): +# payload = json.dumps(body) +# with requests.post(url, data = payload, headers=headers, stream=True) as resp: +# resp.raise_for_status() +# for chunk in resp.iter_content(chunk_size=chunk_size): +# if chunk: # filter keep-alive chunks +# yield chunk + +# if __name__ == "__main__": +# body = { +# "prompt" : { +# "messages" : "Who were winners and runners up of the women's 2025 cricket world cup?" +# } +# } +# +# headers = {"Authorization": f"Bearer {os.getenv('DIGITALOCEAN_API_TOKEN')}"} + +# buffer = "" +# for chunk in stream_endpoint(url, body=body, headers=headers): +# buffer += chunk.decode("utf-8") +# while "\n" in buffer: +# line, buffer = buffer.split("\n", 1) +# if not line.strip(): +# continue +# response = json.loads(line) +# print(response["response"], end="", flush=True) diff --git a/Streaming_and_Tools/requirements.txt b/Streaming_and_Tools/requirements.txt new file mode 100644 index 0000000..62cd8ac --- /dev/null +++ b/Streaming_and_Tools/requirements.txt @@ -0,0 +1,8 @@ +gradient-adk +langgraph +langchain +langchain-core +gradient +langchain-community +langchain_openai +tavily From 242e36ac5a133a9531e4d4c42db81a585c550683 Mon Sep 17 00:00:00 2001 From: Anuj Tambwekar Date: Tue, 23 Dec 2025 15:50:36 -0500 Subject: [PATCH 2/3] Readme update --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index fba966e..40ff868 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ Templates included: - `Crew/` — A Crew AI example that creates a small crew of agents to research news on a topic+date and produce interesting trivia snippets. - `WebSearch/` - A simple LangGraph agent that uses a web-search tool. - `KnowledgeBaseRAG/` - A simple LangGraph agent that queries your DigitalOcean built Knowledge Base. +- `Streaming_and_Tools/` - A LangGraph agent that uses a web-search tool and streams back responses Each template directory includes a `main.py` (the ADK entrypoint) and a `requirements.txt`. See the README in each template for details and quickstart instructions: @@ -17,5 +18,5 @@ Each template directory includes a `main.py` (the ADK entrypoint) and a `require - `Crew/README.md` — details for the Crew AI trivia generator and its search tool. - `WebSearch/README.md` — details for the WebSearch agent. - `KnowledgeBaseRAG/README.md` — details for the Agent that queries your DigitalOcean Knowledge Base. - +- `Streaming_and_Tools/README.md` - details for the agent that uses a search tool and streams responses. From 371b109a898c81e46e4df2c76b69364a2ad63db1 Mon Sep 17 00:00:00 2001 From: Anuj Tambwekar Date: Wed, 24 Dec 2025 12:45:36 -0500 Subject: [PATCH 3/3] Fix typo --- Streaming_and_Tools/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Streaming_and_Tools/README.md b/Streaming_and_Tools/README.md index 21fa7a3..17ba90c 100644 --- a/Streaming_and_Tools/README.md +++ b/Streaming_and_Tools/README.md @@ -42,7 +42,7 @@ Additionally, this agent yields streaming outputs instead of providing an entire pip install -r Templates/ToolCalling/requirements.txt -3. Set the required enviornment variables in the .env file (`TAVILY_API_KEY` and `DIGITALOCEAN_INFERENCE_KEY`). You can obtain a Tavily API key for free at [this link](https://app.tavily.com/home). +3. Set the required environment variables in the .env file (`TAVILY_API_KEY` and `DIGITALOCEAN_INFERENCE_KEY`). You can obtain a Tavily API key for free at [this link](https://app.tavily.com/home). 4. Set your DIGITALOCEAN_API_TOKEN via