From 4813d478c15cfefd95f500efee76111e1eab2234 Mon Sep 17 00:00:00 2001 From: Cristian Pufu Date: Thu, 19 Feb 2026 13:28:31 +0200 Subject: [PATCH] feat: add SQLite session persistence for agent-framework Add SqliteSessionStore to persist AgentSession state between turns, keyed by runtime_id. Each runtime gets isolated conversation history via Agent Framework's auto-injected InMemoryHistoryProvider. The factory manages a shared store singleton and the runtime loads/saves sessions around each execute/stream call. Co-Authored-By: Claude Opus 4.6 --- .../uipath-agent-framework/pyproject.toml | 11 +- .../samples/quickstart-agent/pyproject.toml | 4 +- .../uipath_agent_framework/chat/anthropic.py | 2 +- .../uipath_agent_framework/runtime/factory.py | 45 ++++- .../uipath_agent_framework/runtime/runtime.py | 38 +++- .../uipath_agent_framework/runtime/storage.py | 116 ++++++++++++ .../tests/test_storage.py | 170 ++++++++++++++++++ packages/uipath-agent-framework/uv.lock | 28 ++- 8 files changed, 405 insertions(+), 9 deletions(-) create mode 100644 packages/uipath-agent-framework/src/uipath_agent_framework/runtime/storage.py create mode 100644 packages/uipath-agent-framework/tests/test_storage.py diff --git a/packages/uipath-agent-framework/pyproject.toml b/packages/uipath-agent-framework/pyproject.toml index 021bf37..9c73f40 100644 --- a/packages/uipath-agent-framework/pyproject.toml +++ b/packages/uipath-agent-framework/pyproject.toml @@ -1,11 +1,12 @@ [project] name = "uipath-agent-framework" -version = "0.0.1" +version = "0.0.2" description = "Python SDK that enables developers to build and deploy Microsoft Agent Framework agents to the UiPath Cloud Platform" readme = "README.md" requires-python = ">=3.11" dependencies = [ "agent-framework-core>=1.0.0b260212", + "aiosqlite>=0.20.0", "openinference-instrumentation-agent-framework>=0.1.0", "uipath>=2.8.41, <2.9.0", "uipath-runtime>=0.9.0, <0.10.0", @@ -91,6 +92,14 @@ module = "openinference.*" ignore_missing_imports = true ignore_errors = true +[[tool.mypy.overrides]] +module = "aiosqlite.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = ["anthropic", "anthropic.*"] +ignore_missing_imports = true + [[tool.mypy.overrides]] module = "agent_framework_anthropic.*" ignore_missing_imports = true diff --git a/packages/uipath-agent-framework/samples/quickstart-agent/pyproject.toml b/packages/uipath-agent-framework/samples/quickstart-agent/pyproject.toml index a755970..52dbf51 100644 --- a/packages/uipath-agent-framework/samples/quickstart-agent/pyproject.toml +++ b/packages/uipath-agent-framework/samples/quickstart-agent/pyproject.toml @@ -19,6 +19,4 @@ dev = [ [tool.uv] prerelease = "allow" -[tool.uv.sources] -uipath-dev = { path = "../../../../../uipath-dev-python", editable = true } -uipath-agent-framework = { path = "../../", editable = true } + diff --git a/packages/uipath-agent-framework/src/uipath_agent_framework/chat/anthropic.py b/packages/uipath-agent-framework/src/uipath_agent_framework/chat/anthropic.py index fe605b1..322c6eb 100644 --- a/packages/uipath-agent-framework/src/uipath_agent_framework/chat/anthropic.py +++ b/packages/uipath-agent-framework/src/uipath_agent_framework/chat/anthropic.py @@ -103,7 +103,7 @@ def __new__( _check_anthropic_dependency() from agent_framework_anthropic import AnthropicClient - from anthropic import AsyncAnthropic # type: ignore[import-not-found] + from anthropic import AsyncAnthropic # type: ignore[import-untyped] uipath_url, token = get_uipath_config() gateway_url = build_gateway_url("awsbedrock", model, uipath_url) diff --git a/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/factory.py b/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/factory.py index 3627d70..8741b6b 100644 --- a/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/factory.py +++ b/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/factory.py @@ -1,6 +1,7 @@ """Factory for creating Agent Framework runtimes from agent_framework.json configuration.""" import asyncio +import os from typing import Any from agent_framework import BaseAgent @@ -25,6 +26,7 @@ ) from uipath_agent_framework.runtime.loader import AgentFrameworkAgentLoader from uipath_agent_framework.runtime.runtime import UiPathAgentFrameworkRuntime +from uipath_agent_framework.runtime.storage import SqliteSessionStore class UiPathAgentFrameworkRuntimeFactory: @@ -47,6 +49,9 @@ def __init__( self._agent_loaders: dict[str, AgentFrameworkAgentLoader] = {} self._agent_lock = asyncio.Lock() + self._session_store: SqliteSessionStore | None = None + self._session_store_lock = asyncio.Lock() + self._setup_instrumentation() def _setup_instrumentation(self) -> None: @@ -64,6 +69,32 @@ def _load_config(self) -> AgentFrameworkConfig: self._config = AgentFrameworkConfig() return self._config + def _get_db_path(self) -> str: + """Get the database path for session persistence. + + Uses UiPathRuntimeContext to resolve the state file path. + Cleans up stale state files when not resuming. + """ + path = self.context.resolved_state_file_path + # Delete previous state file if not resuming + if ( + not self.context.resume + and self.context.job_id is None + and not self.context.keep_state_file + ): + if os.path.exists(path): + os.remove(path) + return path + + async def _get_session_store(self) -> SqliteSessionStore: + """Get or create the shared session store instance.""" + async with self._session_store_lock: + if self._session_store is None: + db_path = self._get_db_path() + self._session_store = SqliteSessionStore(db_path) + await self._session_store.setup() + return self._session_store + async def _load_agent(self, entrypoint: str) -> BaseAgent: """ Load an agent for the given entrypoint. @@ -182,11 +213,19 @@ async def _create_runtime_instance( runtime_id: str, entrypoint: str, ) -> UiPathRuntimeProtocol: - """Create a runtime instance from an agent.""" + """Create a runtime instance from an agent. + + Creates the runtime with a shared SqliteSessionStore for persistent + conversation history. Sessions are isolated by runtime_id — each + runtime instance gets its own conversation state. + """ + session_store = await self._get_session_store() + return UiPathAgentFrameworkRuntime( agent=agent, runtime_id=runtime_id, entrypoint=entrypoint, + session_store=session_store, ) async def new_runtime( @@ -218,3 +257,7 @@ async def dispose(self) -> None: self._agent_loaders.clear() self._agent_cache.clear() + + if self._session_store: + await self._session_store.dispose() + self._session_store = None diff --git a/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/runtime.py b/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/runtime.py index 2949da5..570f0aa 100644 --- a/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/runtime.py +++ b/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/runtime.py @@ -8,6 +8,7 @@ from agent_framework import ( AgentResponse, AgentResponseUpdate, + AgentSession, BaseAgent, Content, FunctionTool, @@ -36,6 +37,7 @@ get_agent_tools, get_entrypoints_schema, ) +from .storage import SqliteSessionStore logger = logging.getLogger(__name__) @@ -48,11 +50,13 @@ def __init__( agent: BaseAgent, runtime_id: str | None = None, entrypoint: str | None = None, + session_store: SqliteSessionStore | None = None, ): self.agent: BaseAgent = agent self.runtime_id: str = runtime_id or "default" self.entrypoint: str | None = entrypoint self.chat = AgentFrameworkChatMessagesMapper() + self._session_store = session_store @staticmethod def _build_agent_tool_names(agent: BaseAgent) -> set[str]: @@ -84,6 +88,30 @@ def _build_tool_name_to_agent(agent: BaseAgent) -> dict[str, str]: mapping[tool.name] = agent_name return mapping + async def _load_session(self) -> AgentSession: + """Load or create an AgentSession for this runtime_id. + + If a session store is configured, loads the persisted session state. + Otherwise creates a fresh session each time. + """ + if self._session_store: + session_data = await self._session_store.load_session(self.runtime_id) + if session_data is not None: + logger.debug( + "Restoring session from store for runtime_id=%s", + self.runtime_id, + ) + return AgentSession.from_dict(session_data) # type: ignore[attr-defined] + + return self.agent.create_session(session_id=self.runtime_id) # type: ignore[attr-defined] + + async def _save_session(self, session: AgentSession) -> None: + """Persist the session state after execution.""" + if self._session_store: + session_data = session.to_dict() # type: ignore[attr-defined] + await self._session_store.save_session(self.runtime_id, session_data) + logger.debug("Saved session to store for runtime_id=%s", self.runtime_id) + async def execute( self, input: dict[str, Any] | None = None, @@ -92,7 +120,9 @@ async def execute( """Execute the agent with the provided input and return the result.""" try: user_input = self._prepare_input(input) - response = await self.agent.run(user_input) # type: ignore[attr-defined] + session = await self._load_session() + response = await self.agent.run(user_input, session=session) # type: ignore[attr-defined] + await self._save_session(session) output = self._extract_output(response) return self._create_success_result(output) except Exception as e: @@ -115,6 +145,7 @@ async def stream( """ try: user_input = self._prepare_input(input) + session = await self._load_session() agent_name = self.agent.name or "agent" # Pre-compute which tool names correspond to sub-agents @@ -132,7 +163,7 @@ async def stream( active_tools: str | None = None final_text = "" - response_stream = self.agent.run(user_input, stream=True) # type: ignore[attr-defined] + response_stream = self.agent.run(user_input, stream=True, session=session) # type: ignore[attr-defined] async for update in response_stream: if not isinstance(update, AgentResponseUpdate): continue @@ -290,6 +321,9 @@ async def stream( for msg_event in self.chat.close_message(): yield UiPathRuntimeMessageEvent(payload=msg_event) + # Persist session state after streaming completes + await self._save_session(session) + # Get final response final_response = await response_stream.get_final_response() output = self._extract_output(final_response) diff --git a/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/storage.py b/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/storage.py new file mode 100644 index 0000000..f56bbb8 --- /dev/null +++ b/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/storage.py @@ -0,0 +1,116 @@ +"""SQLite session store for Agent Framework agents. + +Persists AgentSession state between turns using SQLite, keyed by runtime_id. +Each runtime_id maps to an isolated session — conversation history accumulates +across calls via the InMemoryHistoryProvider that Agent Framework auto-injects. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +from typing import Any + +import aiosqlite + +logger = logging.getLogger(__name__) + + +class SqliteSessionStore: + """SQLite-backed store for Agent Framework session state. + + Stores serialized AgentSession dicts (via to_dict/from_dict) in a single + table, keyed by runtime_id. Thread-safe via asyncio lock. + """ + + def __init__(self, db_path: str) -> None: + self.db_path = db_path + self._conn: aiosqlite.Connection | None = None + self._lock = asyncio.Lock() + self._initialized = False + + async def setup(self) -> None: + """Ensure storage directory and database table exist.""" + dir_name = os.path.dirname(self.db_path) + if dir_name: + os.makedirs(dir_name, exist_ok=True) + + conn = await self._get_conn() + async with self._lock: + await conn.execute( + """ + CREATE TABLE IF NOT EXISTS sessions ( + runtime_id TEXT PRIMARY KEY, + session_data TEXT NOT NULL + ) + """ + ) + await conn.commit() + self._initialized = True + logger.debug("Session store initialized at %s", self.db_path) + + async def _get_conn(self) -> aiosqlite.Connection: + """Get or create the database connection.""" + if self._conn is None: + self._conn = await aiosqlite.connect(self.db_path, timeout=30.0) + await self._conn.execute("PRAGMA journal_mode=WAL") + await self._conn.execute("PRAGMA busy_timeout=30000") + await self._conn.execute("PRAGMA synchronous=NORMAL") + await self._conn.commit() + return self._conn + + async def load_session(self, runtime_id: str) -> dict[str, Any] | None: + """Load a serialized session dict for the given runtime_id. + + Returns None if no session exists for this runtime_id. + """ + if not self._initialized: + await self.setup() + + conn = await self._get_conn() + async with self._lock: + cursor = await conn.execute( + "SELECT session_data FROM sessions WHERE runtime_id = ?", + (runtime_id,), + ) + row = await cursor.fetchone() + + if not row: + logger.debug("No session found for runtime_id=%s", runtime_id) + return None + + logger.debug("Loaded session for runtime_id=%s", runtime_id) + return json.loads(row[0]) + + async def save_session(self, runtime_id: str, session_data: dict[str, Any]) -> None: + """Save a serialized session dict for the given runtime_id.""" + if not self._initialized: + await self.setup() + + data_json = json.dumps(session_data) + conn = await self._get_conn() + async with self._lock: + await conn.execute( + """ + INSERT INTO sessions (runtime_id, session_data) + VALUES (?, ?) + ON CONFLICT(runtime_id) DO UPDATE SET + session_data = excluded.session_data + """, + (runtime_id, data_json), + ) + await conn.commit() + + logger.debug("Saved session for runtime_id=%s", runtime_id) + + async def dispose(self) -> None: + """Close the database connection.""" + if self._conn: + await self._conn.close() + self._conn = None + self._initialized = False + + +__all__ = ["SqliteSessionStore"] diff --git a/packages/uipath-agent-framework/tests/test_storage.py b/packages/uipath-agent-framework/tests/test_storage.py new file mode 100644 index 0000000..72a4659 --- /dev/null +++ b/packages/uipath-agent-framework/tests/test_storage.py @@ -0,0 +1,170 @@ +"""Tests for SQLite session store.""" + +import os +import tempfile + +from uipath_agent_framework.runtime.storage import SqliteSessionStore + + +class TestSqliteSessionStore: + """Tests for SqliteSessionStore.""" + + async def test_setup_creates_db_file(self): + """Setup creates the SQLite database file.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + store = SqliteSessionStore(db_path) + await store.setup() + + assert os.path.exists(db_path) + await store.dispose() + + async def test_setup_creates_nested_directories(self): + """Setup creates parent directories if they don't exist.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "nested", "dir", "test.db") + store = SqliteSessionStore(db_path) + await store.setup() + + assert os.path.exists(db_path) + await store.dispose() + + async def test_load_returns_none_for_missing_session(self): + """Loading a non-existent session returns None.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + store = SqliteSessionStore(db_path) + await store.setup() + + result = await store.load_session("nonexistent") + assert result is None + await store.dispose() + + async def test_save_and_load_session(self): + """Saved session data can be loaded back.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + store = SqliteSessionStore(db_path) + await store.setup() + + session_data = { + "session_id": "runtime-123", + "state": {"memory": {"messages": [{"role": "user", "content": "hi"}]}}, + } + await store.save_session("runtime-123", session_data) + loaded = await store.load_session("runtime-123") + + assert loaded == session_data + await store.dispose() + + async def test_save_overwrites_existing_session(self): + """Saving with the same runtime_id overwrites the previous data.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + store = SqliteSessionStore(db_path) + await store.setup() + + await store.save_session("rt-1", {"version": 1}) + await store.save_session("rt-1", {"version": 2}) + loaded = await store.load_session("rt-1") + + assert loaded == {"version": 2} + await store.dispose() + + async def test_sessions_isolated_by_runtime_id(self): + """Different runtime_ids have independent sessions.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + store = SqliteSessionStore(db_path) + await store.setup() + + await store.save_session("rt-a", {"agent": "alpha"}) + await store.save_session("rt-b", {"agent": "beta"}) + + assert await store.load_session("rt-a") == {"agent": "alpha"} + assert await store.load_session("rt-b") == {"agent": "beta"} + assert await store.load_session("rt-c") is None + await store.dispose() + + async def test_dispose_allows_reconnect(self): + """After dispose, the store can be set up again and data persists.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + + store = SqliteSessionStore(db_path) + await store.setup() + await store.save_session("rt-1", {"data": "persisted"}) + await store.dispose() + + # Reconnect to the same DB + store2 = SqliteSessionStore(db_path) + await store2.setup() + loaded = await store2.load_session("rt-1") + + assert loaded == {"data": "persisted"} + await store2.dispose() + + async def test_auto_setup_on_load(self): + """Loading without explicit setup triggers auto-setup.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + store = SqliteSessionStore(db_path) + + # No explicit setup() call + result = await store.load_session("any-id") + assert result is None + await store.dispose() + + async def test_auto_setup_on_save(self): + """Saving without explicit setup triggers auto-setup.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + store = SqliteSessionStore(db_path) + + # No explicit setup() call + await store.save_session("rt-1", {"key": "value"}) + loaded = await store.load_session("rt-1") + + assert loaded == {"key": "value"} + await store.dispose() + + async def test_complex_session_data_roundtrip(self): + """Complex nested session data survives serialization roundtrip.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + store = SqliteSessionStore(db_path) + await store.setup() + + session_data = { + "session_id": "abc-123", + "state": { + "memory": { + "messages": [ + { + "role": "user", + "content": "What is the weather?", + "metadata": {"timestamp": "2025-01-01T00:00:00Z"}, + }, + { + "role": "assistant", + "content": "It's sunny!", + "tool_calls": [ + { + "id": "call_1", + "name": "get_weather", + "arguments": {"city": "SF"}, + } + ], + }, + ] + }, + "custom_key": [1, 2, 3], + "nested": {"a": {"b": {"c": True}}}, + }, + } + + await store.save_session("rt-complex", session_data) + loaded = await store.load_session("rt-complex") + + assert loaded == session_data + await store.dispose() diff --git a/packages/uipath-agent-framework/uv.lock b/packages/uipath-agent-framework/uv.lock index b6d5c97..50ccef3 100644 --- a/packages/uipath-agent-framework/uv.lock +++ b/packages/uipath-agent-framework/uv.lock @@ -5,6 +5,19 @@ requires-python = ">=3.11" [options] prerelease-mode = "allow" +[[package]] +name = "agent-framework-anthropic" +version = "1.0.0b260212" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "agent-framework-core" }, + { name = "anthropic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ff/fb/db7157cd342ed079a467a4080f3d5aadd403b798e25b6c52b638c317f44c/agent_framework_anthropic-1.0.0b260212.tar.gz", hash = "sha256:46adac3ff7cfedbf97d76fafa9c6a461b4ad1b53c075336d300800a283289f21", size = 12977, upload-time = "2026-02-13T00:27:13.58Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f0/1f/0c016c683b922919c5aabb8591fd04cffec7432a224169b029049e90967c/agent_framework_anthropic-1.0.0b260212-py3-none-any.whl", hash = "sha256:8c2a8bb5474b7984994b7e6cf0318f06338765f406c50a3d7336f38ed44c9444", size = 13025, upload-time = "2026-02-13T00:38:15.841Z" }, +] + [[package]] name = "agent-framework-core" version = "1.0.0b260212" @@ -151,6 +164,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, ] +[[package]] +name = "aiosqlite" +version = "0.22.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/8a/64761f4005f17809769d23e518d915db74e6310474e733e3593cfc854ef1/aiosqlite-0.22.1.tar.gz", hash = "sha256:043e0bd78d32888c0a9ca90fc788b38796843360c855a7262a532813133a0650", size = 14821, upload-time = "2025-12-23T19:25:43.997Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/00/b7/e3bf5133d697a08128598c8d0abc5e16377b51465a33756de24fa7dee953/aiosqlite-0.22.1-py3-none-any.whl", hash = "sha256:21c002eb13823fad740196c5a2e9d8e62f6243bd9e7e4a1f87fb5e44ecb4fceb", size = 17405, upload-time = "2025-12-23T19:25:42.139Z" }, +] + [[package]] name = "annotated-types" version = "0.7.0" @@ -2426,10 +2448,11 @@ wheels = [ [[package]] name = "uipath-agent-framework" -version = "0.0.1" +version = "0.0.2" source = { editable = "." } dependencies = [ { name = "agent-framework-core" }, + { name = "aiosqlite" }, { name = "openinference-instrumentation-agent-framework" }, { name = "uipath" }, { name = "uipath-runtime" }, @@ -2437,6 +2460,7 @@ dependencies = [ [package.optional-dependencies] anthropic = [ + { name = "agent-framework-anthropic" }, { name = "anthropic" }, ] @@ -2453,7 +2477,9 @@ dev = [ [package.metadata] requires-dist = [ + { name = "agent-framework-anthropic", marker = "extra == 'anthropic'", specifier = ">=1.0.0b260212" }, { name = "agent-framework-core", specifier = ">=1.0.0b260212" }, + { name = "aiosqlite", specifier = ">=0.20.0" }, { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.43.0" }, { name = "openinference-instrumentation-agent-framework", specifier = ">=0.1.0" }, { name = "uipath", specifier = ">=2.8.41,<2.9.0" },