Skip to content

Commit 8c8b6b1

Browse files
mjnoviceclaude
andcommitted
feat: add Agent Episodic Memory support (recall node, tracing)
- Add MemoryConfig with field_weights, threshold, result_count, folder_path, memory_space_name - Add memory_recall node that queries UiPath Memory service before INIT - Emit OTel trace spans (Find previous memories, Apply dynamic few shot) - Extract OTel helpers to _utils/_otel.py for reuse across tools - Prefix key_path with field type (agent-input) matching backend - Add definition_system_prompt to memory search request - Filter empty string values in _build_search_fields Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent b40328d commit 8c8b6b1

10 files changed

Lines changed: 500 additions & 9 deletions

File tree

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ requires-python = ">=3.11"
77
dependencies = [
88
"uipath>=2.10.53, <2.11.0",
99
"uipath-core>=0.5.2, <0.6.0",
10-
"uipath-platform>=0.1.35, <0.2.0",
10+
"uipath-platform>=0.1.36, <0.2.0",
1111
"uipath-runtime>=0.10.0, <0.11.0",
1212
"langgraph>=1.1.8, <2.0.0",
1313
"langchain-core>=1.2.11, <2.0.0",
Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
11
from ._environment import get_execution_folder_path
2+
from ._otel import set_span_attribute
23
from ._request_mixin import UiPathRequestMixin
34

4-
__all__ = ["UiPathRequestMixin", "get_execution_folder_path"]
5+
__all__ = [
6+
"UiPathRequestMixin",
7+
"get_execution_folder_path",
8+
"set_span_attribute",
9+
]
Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
"""OpenTelemetry utilities for span and trace context."""
2+
3+
from typing import Any
4+
5+
6+
def set_span_attribute(name: str, value: Any) -> None:
7+
"""Set an attribute on the current OTel span (no-op if unavailable)."""
8+
try:
9+
from opentelemetry import trace
10+
11+
span = trace.get_current_span()
12+
if span.is_recording():
13+
span.set_attribute(name, value)
14+
except ImportError:
15+
pass

src/uipath_langchain/agent/react/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
"""UiPath ReAct Agent implementation"""
22

33
from .agent import create_agent
4-
from .types import AgentGraphConfig, AgentGraphNode, AgentGraphState
4+
from .types import AgentGraphConfig, AgentGraphNode, AgentGraphState, MemoryConfig
55
from .utils import resolve_input_model, resolve_output_model
66

77
__all__ = [
@@ -11,4 +11,5 @@
1111
"AgentGraphNode",
1212
"AgentGraphState",
1313
"AgentGraphConfig",
14+
"MemoryConfig",
1415
]

src/uipath_langchain/agent/react/agent.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from .llm_node import (
2525
create_llm_node,
2626
)
27+
from .memory_node import create_memory_recall_node
2728
from .router import (
2829
create_route_agent,
2930
)
@@ -36,6 +37,7 @@
3637
AgentGraphConfig,
3738
AgentGraphNode,
3839
AgentGraphState,
40+
MemoryConfig,
3941
)
4042
from .utils import create_state_with_input
4143

@@ -53,6 +55,7 @@ def create_agent(
5355
output_schema: Type[OutputT] | None = None,
5456
config: AgentGraphConfig | None = None,
5557
guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None = None,
58+
memory: MemoryConfig | None = None,
5659
) -> StateGraph[AgentGraphState, None, InputT, OutputT]:
5760
"""Build agent graph with INIT -> AGENT (subgraph) <-> TOOLS loop, terminated by control flow tools.
5861
@@ -122,7 +125,13 @@ def create_agent(
122125
)
123126
builder.add_node(AgentGraphNode.TERMINATE, terminate_with_guardrails_subgraph)
124127

125-
builder.add_edge(START, AgentGraphNode.INIT)
128+
if memory:
129+
memory_recall = create_memory_recall_node(memory, input_schema=input_schema)
130+
builder.add_node(AgentGraphNode.MEMORY_RECALL, memory_recall)
131+
builder.add_edge(START, AgentGraphNode.MEMORY_RECALL)
132+
builder.add_edge(AgentGraphNode.MEMORY_RECALL, AgentGraphNode.INIT)
133+
else:
134+
builder.add_edge(START, AgentGraphNode.INIT)
126135

127136
llm_node = create_llm_node(
128137
model,

src/uipath_langchain/agent/react/init_node.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,19 @@ def graph_state_init(state: Any) -> Any:
2525
resolved_messages = list(messages(state))
2626
else:
2727
resolved_messages = list(messages)
28+
29+
# Append memory injection from the MEMORY_RECALL node (if present)
30+
memory_injection = ""
31+
if hasattr(state, "inner_state") and hasattr(
32+
state.inner_state, "memory_injection"
33+
):
34+
memory_injection = state.inner_state.memory_injection or ""
35+
if memory_injection and resolved_messages:
36+
first = resolved_messages[0]
37+
if isinstance(first, SystemMessage):
38+
resolved_messages[0] = SystemMessage(
39+
content=str(first.content) + memory_injection
40+
)
2841
if is_conversational:
2942
# For conversational agents we need to reorder the messages so that the system message is first, followed by
3043
# the initial user message. When resuming the conversation, the state will have the entire message history,
Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
"""Memory recall node for Agent Episodic Memory.
2+
3+
Queries the UiPath Memory service (via LLMOps) for similar past episodes
4+
and stores the server-generated systemPromptInjection in graph state so
5+
the INIT node can append it to the system prompt.
6+
"""
7+
8+
import logging
9+
from contextlib import contextmanager
10+
from typing import Any
11+
12+
from pydantic import BaseModel
13+
from uipath.platform import UiPath
14+
from uipath.platform.memory import (
15+
FieldSettings,
16+
MemorySearchRequest,
17+
SearchField,
18+
SearchMode,
19+
SearchSettings,
20+
)
21+
22+
from .types import MemoryConfig
23+
from .utils import extract_input_data_from_state
24+
25+
logger = logging.getLogger(__name__)
26+
27+
28+
@contextmanager
29+
def _noop_context():
30+
"""No-op context manager when OTel is unavailable."""
31+
yield None
32+
33+
34+
def create_memory_recall_node(
35+
memory_config: MemoryConfig,
36+
input_schema: type[BaseModel] | None = None,
37+
):
38+
"""Create an async graph node that retrieves memory injection.
39+
40+
The node queries ``sdk.memory.search_async()`` and writes the
41+
``systemPromptInjection`` string into ``inner_state.memory_injection``.
42+
On failure it logs a warning and continues with an empty injection.
43+
44+
Args:
45+
memory_config: Memory configuration with space ID and search settings.
46+
47+
Returns:
48+
An async callable suitable for ``builder.add_node()``.
49+
"""
50+
51+
async def memory_recall_node(state: Any) -> dict[str, Any]:
52+
input_model = input_schema if input_schema is not None else BaseModel
53+
input_arguments = extract_input_data_from_state(state, input_model)
54+
if not input_arguments:
55+
logger.debug("Memory recall: no user inputs found in state")
56+
return {}
57+
58+
fields = _build_search_fields(
59+
input_arguments, field_weights=memory_config.field_weights or None
60+
)
61+
if not fields:
62+
logger.debug(
63+
"Memory recall: no search fields after filtering (inputs=%s, weights=%s)",
64+
list(input_arguments.keys()),
65+
memory_config.field_weights,
66+
)
67+
return {}
68+
69+
request = MemorySearchRequest(
70+
fields=fields,
71+
settings=SearchSettings(
72+
threshold=memory_config.threshold,
73+
result_count=memory_config.result_count,
74+
search_mode=SearchMode.Hybrid,
75+
),
76+
definition_system_prompt="",
77+
)
78+
79+
results_count = 0
80+
# Wrap the search in OTel spans so "Find previous memories" and
81+
# "Apply dynamic few shot" appear in the Execution Trace with
82+
# correct timing. The LlmOpsHttpExporter picks these up.
83+
injection = ""
84+
try:
85+
from opentelemetry import trace as otel_trace
86+
87+
tracer = otel_trace.get_tracer("uipath_langchain.memory")
88+
except ImportError:
89+
tracer = None
90+
otel_trace = None # type: ignore[assignment]
91+
92+
# Span attribute keys matching what the LlmOpsHttpExporter and
93+
# Studio UI expect. "openinference.span.kind" sets SpanType.
94+
lookup_span_ctx = (
95+
tracer.start_as_current_span(
96+
"Find previous memories",
97+
attributes={
98+
"openinference.span.kind": "agentMemoryLookup",
99+
"type": "agentMemoryLookup",
100+
"span_type": "agentMemoryLookup",
101+
"uipath.custom_instrumentation": True,
102+
"memorySpaceName": memory_config.memory_space_name or "",
103+
"memorySpaceId": memory_config.memory_space_id,
104+
"strategy": "DynamicFewShotPrompt",
105+
},
106+
)
107+
if tracer
108+
else _noop_context()
109+
)
110+
111+
with lookup_span_ctx as lookup_span:
112+
fewshot_span_ctx = (
113+
tracer.start_as_current_span(
114+
"Apply dynamic few shot",
115+
attributes={
116+
"openinference.span.kind": "applyDynamicFewShot",
117+
"type": "applyDynamicFewShot",
118+
"span_type": "applyDynamicFewShot",
119+
"uipath.custom_instrumentation": True,
120+
"memorySpaceName": memory_config.memory_space_name or "",
121+
"memorySpaceId": memory_config.memory_space_id,
122+
},
123+
)
124+
if tracer
125+
else _noop_context()
126+
)
127+
128+
with fewshot_span_ctx as fewshot_span:
129+
try:
130+
sdk = UiPath()
131+
folder_key = memory_config.folder_key
132+
if not folder_key and memory_config.folder_path:
133+
folder_key = await sdk.folders.retrieve_folder_key_async(
134+
memory_config.folder_path
135+
)
136+
response = await sdk.memory.search_async(
137+
memory_space_id=memory_config.memory_space_id,
138+
request=request,
139+
folder_key=folder_key,
140+
)
141+
injection = response.system_prompt_injection
142+
results_count = len(response.results)
143+
logger.info(
144+
"Memory recall returned %d results for space '%s'",
145+
results_count,
146+
memory_config.memory_space_id,
147+
)
148+
# Set request/response on fewshot span as JSON strings.
149+
# The exporter parses JSON strings back to objects.
150+
# The UI reads "response" to display matched memory items.
151+
if fewshot_span and hasattr(fewshot_span, "set_attribute"):
152+
import json
153+
154+
fewshot_span.set_attribute(
155+
"request",
156+
json.dumps(
157+
request.model_dump(by_alias=True, exclude_none=True)
158+
),
159+
)
160+
fewshot_span.set_attribute(
161+
"response",
162+
json.dumps(
163+
response.model_dump(by_alias=True, exclude_none=True)
164+
),
165+
)
166+
except Exception as e:
167+
from uipath.platform.errors import EnrichedException
168+
169+
if isinstance(e, EnrichedException):
170+
error_detail = (
171+
f"{e} | status={e.status_code} body={e.response_content}"
172+
)
173+
else:
174+
error_detail = repr(e)
175+
logger.warning(
176+
"Memory recall failed for space '%s': %s",
177+
memory_config.memory_space_id,
178+
error_detail,
179+
)
180+
if lookup_span and hasattr(lookup_span, "set_status"):
181+
lookup_span.set_status(
182+
otel_trace.StatusCode.ERROR, error_detail
183+
)
184+
185+
# Set result attributes after search completes
186+
if lookup_span and hasattr(lookup_span, "set_attribute"):
187+
lookup_span.set_attribute("memoryItemsMatched", results_count)
188+
if injection:
189+
lookup_span.set_attribute("result", injection)
190+
191+
if not injection:
192+
return {}
193+
194+
return {"inner_state": {"memory_injection": injection}}
195+
196+
return memory_recall_node
197+
198+
199+
def _build_search_fields(
200+
input_arguments: dict[str, Any],
201+
field_weights: dict[str, float] | None = None,
202+
field_type: str = "agent-input",
203+
) -> list[SearchField]:
204+
"""Convert agent input arguments to SearchField objects.
205+
206+
The key_path must be prefixed with the field type:
207+
keyPath = [fieldType, fieldName]
208+
e.g. ["agent-input", "a"] for episodic memory.
209+
"""
210+
fields: list[SearchField] = []
211+
for name, value in input_arguments.items():
212+
value_str = str(value) if value is not None else ""
213+
if not value_str or name.startswith("uipath__"):
214+
continue
215+
# When field_weights is specified, only include fields with configured weights
216+
if field_weights and name not in field_weights:
217+
continue
218+
settings = FieldSettings()
219+
if field_weights and name in field_weights:
220+
settings = FieldSettings(weight=field_weights[name])
221+
fields.append(
222+
SearchField(key_path=[field_type, name], value=value_str, settings=settings)
223+
)
224+
return fields

src/uipath_langchain/agent/react/types.py

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
from langchain_core.messages import AnyMessage
55
from langgraph.graph.message import add_messages
6-
from pydantic import BaseModel, Field
6+
from pydantic import BaseModel, Field, model_validator
77
from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL
88
from uipath.platform.attachments import Attachment
99

@@ -19,6 +19,7 @@ class InnerAgentGraphState(BaseModel):
1919
job_attachments: Annotated[dict[str, Attachment], merge_dicts] = {}
2020
initial_message_count: int | None = None
2121
tools_storage: Annotated[dict[Hashable, Any], merge_dicts] = {}
22+
memory_injection: str = ""
2223

2324

2425
class InnerAgentGuardrailsGraphState(InnerAgentGraphState):
@@ -55,6 +56,43 @@ class AgentGraphNode(StrEnum):
5556
TOOLS = "tools"
5657
TERMINATE = "terminate"
5758
GUARDED_TERMINATE = "guarded-terminate"
59+
MEMORY_RECALL = "memory_recall"
60+
61+
62+
class MemoryConfig(BaseModel):
63+
"""Configuration for Agent Episodic Memory.
64+
65+
When passed to ``create_agent()``, a MEMORY_RECALL node is added before
66+
INIT that queries the memory service and stores the server-generated
67+
systemPromptInjection in ``inner_state.memory_injection``.
68+
"""
69+
70+
memory_space_id: str = Field(description="GUID of the memory space to query.")
71+
memory_space_name: str = Field(
72+
default="", description="Name of the memory space (for tracing)."
73+
)
74+
folder_key: str | None = Field(
75+
default=None, description="Folder key for the memory resource."
76+
)
77+
folder_path: str | None = Field(
78+
default=None,
79+
description="Folder path for the memory resource. Resolved to folder_key at runtime if folder_key is not set.",
80+
)
81+
# Defaults match FE episodic memory settings (agentEditor.ts:324-328)
82+
result_count: int = Field(default=3, ge=1, le=10)
83+
threshold: float = Field(default=0.0, ge=0.0, le=1.0)
84+
field_weights: dict[str, float] = Field(
85+
description=(
86+
"Per-field search weights. Keys are input field names, values are "
87+
"weights between 0.0 and 1.0. At least one field must be specified."
88+
),
89+
)
90+
91+
@model_validator(mode="after")
92+
def _validate_field_weights(self) -> "MemoryConfig":
93+
if not self.field_weights:
94+
raise ValueError("field_weights must contain at least one field")
95+
return self
5896

5997

6098
class AgentGraphConfig(BaseModel):

0 commit comments

Comments
 (0)