-
Notifications
You must be signed in to change notification settings - Fork 41
feat(templates): add conversation history persistence to HTTP agent templates #794
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -113,21 +113,39 @@ def ensure_credentials_loaded(): | |
| ) | ||
|
|
||
|
|
||
| # Session and Runner | ||
| async def setup_session_and_runner(user_id, session_id): | ||
| ensure_credentials_loaded() | ||
| session_service = InMemorySessionService() | ||
| session = await session_service.create_session( | ||
| # Module-level session service and runner (preserves history across invocations) | ||
| _session_service = InMemorySessionService() | ||
| _runner = None | ||
|
|
||
|
|
||
| def get_or_create_runner(): | ||
| global _runner | ||
| if _runner is None: | ||
| ensure_credentials_loaded() | ||
| _runner = Runner( | ||
| agent=agent, | ||
| app_name=APP_NAME, | ||
| session_service=_session_service, | ||
| ) | ||
| return _runner | ||
|
|
||
|
|
||
| async def get_or_create_session(user_id, session_id): | ||
| session = await _session_service.get_session( | ||
| app_name=APP_NAME, user_id=user_id, session_id=session_id | ||
| ) | ||
| runner = Runner(agent=agent, app_name=APP_NAME, session_service=session_service) | ||
| return session, runner | ||
| if session is None: | ||
| session = await _session_service.create_session( | ||
| app_name=APP_NAME, user_id=user_id, session_id=session_id | ||
| ) | ||
| return session | ||
|
|
||
|
|
||
| # Agent Interaction | ||
| async def call_agent_async(query, user_id, session_id): | ||
| content = types.Content(role="user", parts=[types.Part(text=query)]) | ||
| session, runner = await setup_session_and_runner(user_id, session_id) | ||
| runner = get_or_create_runner() | ||
| session = await get_or_create_session(user_id, session_id) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Cross-user conversation access via client-supplied
That means any caller can now set |
||
| events = runner.run_async( | ||
| user_id=user_id, session_id=session.id, new_message=content | ||
| ) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2,6 +2,7 @@ | |
| from typing import Any | ||
|
|
||
| from langchain_core.messages import HumanMessage{{#if hasConfigBundle}}, SystemMessage{{/if}} | ||
| from langgraph.checkpoint.memory import InMemorySaver | ||
| from langgraph.prebuilt import create_react_agent | ||
| from langchain.tools import tool | ||
| {{#if hasConfigBundle}} | ||
|
|
@@ -49,6 +50,9 @@ def add_numbers(a: int, b: int) -> int: | |
| # Define a collection of tools used by the model | ||
| tools = [add_numbers] | ||
|
|
||
| # Module-level checkpointer preserves conversation history across invocations | ||
| _checkpointer = InMemorySaver() | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Unbounded in-memory checkpoint growth (DoS / OOM in long-lived containers).
The sibling Strands and OpenAIAgents templates in this same PR explicitly cap their session caches with |
||
|
|
||
| {{#if sessionStorageMountPath}} | ||
| SESSION_STORAGE_PATH = "{{sessionStorageMountPath}}" | ||
|
|
||
|
|
@@ -141,29 +145,44 @@ async def invoke(payload, context): | |
| if mcp_client: | ||
| mcp_tools = await mcp_client.get_tools() | ||
|
|
||
| # Define the agent using create_react_agent | ||
| # Define the agent using create_react_agent (checkpointer is shared across invocations) | ||
| {{#if hasConfigBundle}} | ||
| graph = create_react_agent(get_or_create_model(), tools=mcp_tools + tools, prompt=DEFAULT_SYSTEM_PROMPT) | ||
| graph = create_react_agent( | ||
| get_or_create_model(), | ||
| tools=mcp_tools + tools, | ||
| prompt=DEFAULT_SYSTEM_PROMPT, | ||
| checkpointer=_checkpointer, | ||
| ) | ||
| callback = ConfigBundleCallback() | ||
|
|
||
| # Process the user prompt | ||
| prompt = payload.get("prompt", "What can you help me with?") | ||
| session_id = getattr(context, "session_id", "default-session") | ||
| log.info(f"Agent input: {prompt}") | ||
|
|
||
| # Run the agent with config bundle callback | ||
| # Run the agent with config bundle callback (checkpointer auto-loads/saves history per session) | ||
| result = await graph.ainvoke( | ||
| {"messages": [HumanMessage(content=prompt)]}, | ||
| config={"callbacks": [callback]}, | ||
| config={"callbacks": [callback], "configurable": {"thread_id": session_id}}, | ||
| ) | ||
| {{else}} | ||
| graph = create_react_agent(get_or_create_model(), tools=mcp_tools + tools, prompt=DEFAULT_SYSTEM_PROMPT) | ||
| graph = create_react_agent( | ||
| get_or_create_model(), | ||
| tools=mcp_tools + tools, | ||
| prompt=DEFAULT_SYSTEM_PROMPT, | ||
| checkpointer=_checkpointer, | ||
| ) | ||
|
|
||
| # Process the user prompt | ||
| prompt = payload.get("prompt", "What can you help me with?") | ||
| session_id = getattr(context, "session_id", "default-session") | ||
| log.info(f"Agent input: {prompt}") | ||
|
|
||
| # Run the agent | ||
| result = await graph.ainvoke({"messages": [HumanMessage(content=prompt)]}) | ||
| # Run the agent (checkpointer auto-loads/saves history per session) | ||
| result = await graph.ainvoke( | ||
| {"messages": [HumanMessage(content=prompt)]}, | ||
| config={"configurable": {"thread_id": session_id}}, | ||
| ) | ||
| {{/if}} | ||
|
|
||
| # Return result | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Unbounded in-memory session growth (DoS / OOM in long-lived containers).
InMemorySessionServiceis a plain in-memory store with no built-in eviction. With this PR it is now process-lifetime scoped, and every distinct(user_id, session_id)invocation creates a new entry that is never freed. An attacker (or a noisy client) can rotatesession_ids and grow the process's heap until it OOMs — the runtime container typically has a fixed memory budget, so this is a denial-of-service vector.The sibling Strands and OpenAIAgents templates in this same PR explicitly cap their session caches with
@lru_cache(maxsize=128)and call out LRU eviction in a comment. The googleadk template (and the langchain template, withInMemorySaver()) should apply the same mitigation, or call out durable-store guidance just as prominently. As written, these two templates are noticeably less safe to ship into production than the other two.