diff --git a/agent-langgraph-basic/.env.example b/agent-langgraph-basic/.env.example new file mode 100644 index 00000000..53ddc27b --- /dev/null +++ b/agent-langgraph-basic/.env.example @@ -0,0 +1,29 @@ +# Copy to .env for local development: cp .env.example .env +# +# The only required setting for local dev is a Databricks auth profile (used to call the model). +# Everything else — Lakebase durability, MLflow tracing — is optional and off by default. + +# Databricks auth profile (from `databricks auth profiles`). Used to call the model endpoint. +DATABRICKS_CONFIG_PROFILE=DEFAULT +# Or use explicit host/token instead of a profile: +# DATABRICKS_HOST=https://.databricks.com +# DATABRICKS_TOKEN=dapi... + +# --- Optional: MLflow tracing --- +# Leave UNSET to skip tracing (local dev). To enable, set both: an experiment to log to, and the +# tracking destination (your workspace). +# MLFLOW_EXPERIMENT_ID= +# MLFLOW_TRACKING_URI="databricks" + +# --- Optional: durable conversation history (managed session store) --- +# Leave UNSET to keep the local SQLite session store. Set to a managed session store name to +# persist the transcript to its agents/v1 items API (durable, shared across replicas). +# AGENT_SESSION_STORE=my-agent-sessions + +# --- Optional: long-running background mode + crash recovery (Lakebase) --- +# Leave UNSET to serve in-request. Set the Lakebase endpoint to enable durable background mode. +# LAKEBASE_AUTOSCALING_ENDPOINT= + +# --- Optional: local SQLite session store path (when AGENT_SESSION_STORE is unset) --- +# Defaults to a file so history survives restarts. Set ":memory:" for ephemeral storage. +# LOCAL_SESSION_DB_PATH=local_agent_sessions.db diff --git a/agent-langgraph-basic/.gitignore b/agent-langgraph-basic/.gitignore new file mode 100644 index 00000000..047565f9 --- /dev/null +++ b/agent-langgraph-basic/.gitignore @@ -0,0 +1,195 @@ +# Created by https://www.toptal.com/developers/gitignore/api/python +# Edit at https://www.toptal.com/developers/gitignore?templates=python + +### Python ### +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +# VS Code +.vscode/ + +### Python Patch ### +# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration +poetry.toml + +# ruff +.ruff_cache/ + +# LSP config files +pyrightconfig.json + +# End of https://www.toptal.com/developers/gitignore/api/python + +.DS_* + +# Databricks / MLflow local artifacts +**/mlruns/ +mlflow.db +**/.databricks +.claude/ + +# Environment files +**/.env +**/.env.local + +# Local (Lakebase-free) SQLite session store +local_agent_sessions.db +local_agent_sessions.db-* +*.db-journal diff --git a/agent-langgraph-basic/AGENTS.md b/agent-langgraph-basic/AGENTS.md new file mode 100644 index 00000000..d6d81be5 --- /dev/null +++ b/agent-langgraph-basic/AGENTS.md @@ -0,0 +1,69 @@ +# Agent Development Guide + +A lean LangGraph agent backend for Databricks Apps. Local-first: runs with no database and no setup +beyond a Databricks auth profile. Lakebase durability and MLflow tracing are optional. + +See `README.md` for the full run / deploy / client-contract docs. This file is the quick map for +making changes. + +## Run it + +```bash +cp .env.example .env # set DATABRICKS_CONFIG_PROFILE= +uv run start-server # http://localhost:8000 +``` + +No database needed — conversation state uses an in-process LangGraph checkpointer by default. + +## Where things live + +| You want to… | Edit | +| --- | --- | +| Change model / instructions | `agent/agent.py` (`create_agent_graph`) | +| Add a function tool | new `*.py` in `agent/tools/` with a `@tool` function (auto-collected) | +| Add an MCP server | append a `DatabricksMCPServer` to `build_mcp_servers()` in `agent/mcps.py` | +| Change how a request maps to a run | `agent/agent.py` (`@invoke` / `@stream` handlers) | +| Change the session checkpointer | `agent/mason/session_store.py` | +| Server / durability wiring | `server/start_server.py` (rarely needed) | +| Add a test | `tests/` (hermetic; gate model calls on a workspace profile — see `test_agent.py`) | + +`agent/mason/` holds plumbing (session checkpointer, tracing, MCP tool loading, wire translation) +slated to move into Databricks SDKs — grouped so that migration is localized. You rarely edit it; +build the agent in `agent/agent.py`, `agent/tools/`, and `agent/mcps.py`. + +## How tools register + +`agent/tools/all_tools()` auto-imports every module in the package and collects every +`@tool`-decorated `BaseTool` it finds. So a tool registers just by existing in a file there — +`create_agent_graph()` calls `all_tools()`. **Do not** edit `agent/agent.py` to add a tool — just +add a file to `agent/tools/`. + +## Sessions & durability + +- Default: `agent/mason/session_store.py`'s `checkpointer()` returns an in-process `InMemorySaver`, + keyed per request by `thread_config(session_id)` — no database, multi-turn works in-process. +- **Two independent durable stores:** + - Conversation history → swap the checkpointer for a `PostgresSaver` over Lakebase (durable, + shared across replicas). + - `LAKEBASE_AUTOSCALING_ENDPOINT` → `start_server.py` passes it into `LongRunningAgentServer` for + its durable server store (background mode + crash recovery). + - Enable either/both/neither. + +## MLflow tracing + +Optional. Set both `MLFLOW_EXPERIMENT_ID` and `MLFLOW_TRACKING_URI` to enable (`mlflow.langchain.autolog()`); +leave either unset to skip (the server boots with tracing disabled). + +## Quick commands + +| Task | Command | +| --- | --- | +| Run locally | `uv run start-server` | +| Run via CLI local App runner | `databricks apps run-local --prepare-environment -p ` | +| Test | `uv run pytest` (hermetic; live model test runs only with a profile) | +| Deploy | `databricks apps deploy agent-langgraph-basic --source-code-path ` | + +## Notes for maintainers + +- `agent/mason/wire/` is LangGraph-specific (inbound request→session id; outbound `astream` + `updates`/`messages` events→Responses wire events). diff --git a/agent-langgraph-basic/CLAUDE.md b/agent-langgraph-basic/CLAUDE.md new file mode 100644 index 00000000..43c994c2 --- /dev/null +++ b/agent-langgraph-basic/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/agent-langgraph-basic/README.md b/agent-langgraph-basic/README.md new file mode 100644 index 00000000..f13f9d73 --- /dev/null +++ b/agent-langgraph-basic/README.md @@ -0,0 +1,182 @@ +# Agent — LangGraph (Basic) + +A lean [LangGraph](https://langchain-ai.github.io/langgraph/) agent **backend** for Databricks Apps. +It runs locally with **no database and no setup** — just an auth profile — and exposes the +[OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) (`POST /responses`, +`POST /invocations`). Durable long-running background execution is **opt-in** via +[Databricks Lakebase](https://docs.databricks.com/aws/en/lakebase/); without it the agent keeps +conversation state in an in-process checkpointer and runs in-request. + +This template is API-first (no bundled UI). Call it with the OpenAI SDK, `curl`, or from your own +frontend / model-serving client. + +## Project layout + +``` +agent/ # the agent (reasoning plane) — this is what you edit + agent.py # @invoke / @stream handlers + create_agent_graph() + tools/ # function tools — drop a *.py file here to add one (auto-collected) + sample_tool.py # get_current_time — a working example (@tool) + mcps.py # MCP servers: none by default; add to build_mcp_servers() to offer some + mason/ # plumbing that will move into Databricks SDKs later — rarely edited + session_store.py # LangGraph checkpointer: in-memory by default; swap for a durable one + memory.py # remember / recall — memory_tools() returns them when AGENT_MEMORY_STORE is set + tracing.py # MLflow tracing setup (on only when both MLFLOW_* vars are set) + mcp_runtime.py # loads tools from the servers in mcps.build_mcp_servers() + wire/ # Responses <-> LangGraph translation + inbound.py # request -> session id (LangGraph thread_id) + outbound.py # LangGraph astream events -> Responses wire events +server/ # the durable plane (LongRunningAgentServer wiring) — rarely edited + start_server.py # builds the server; passes LAKEBASE_AUTOSCALING_ENDPOINT for durability if set +tests/ + test_agent.py # hermetic smoke tests + one gated live model call +``` + +You edit `agent/agent.py`, `agent/tools/`, and `agent/mcps.py`; everything in `agent/mason/` is +plumbing (session checkpointer, tracing, MCP tool loading, wire translation) that's slated to move +into Databricks SDKs, grouped so that migration is a localized change. `tools/` is a drop-in +package: add a `*.py` with a `@tool` function and it's auto-collected (no edits to existing code). +`mcps.py` exposes `build_mcp_servers()` (empty by default — add servers to offer them). + +## Run locally + +No database required. Conversation state is kept in an in-process LangGraph checkpointer. + +```bash +# 1. Configure a Databricks auth profile (used only to call the model) +cp .env.example .env +# edit .env: set DATABRICKS_CONFIG_PROFILE= + +# 2. Start the server (installs deps via uv on first run) +uv run start-server # serves at http://localhost:8000 + +# 3. Send a request +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{"input": [{"role": "user", "content": "What time is it? Use your tool."}]}' +``` + +You can also launch it through the Databricks CLI's local App runner: + +```bash +databricks apps run-local --prepare-environment +``` + +The model call goes to your Databricks workspace (via the profile). Everything else — durable +background mode, tracing — is off by default and requires no setup. + +## Client contract + +`POST /responses` (and its alias `POST /invocations`) implement the OpenAI Responses API. Replace +`` with `http://localhost:8000` locally, or `https://.databricksapps.com` (with an +`Authorization: Bearer ` header) when deployed. + +**Non-streaming:** + +```bash +curl -X POST /responses \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }] }' +``` + +**Streaming** (add `"stream": true`) returns an SSE stream ending with `[DONE]`. + +**Multi-turn** — pass the `session_id` returned in `custom_outputs` back on the next request: + +```bash +# First turn returns: "custom_outputs": { "session_id": "..." } +curl -X POST /responses -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "My name is Alice" }] }' + +# Second turn — agent remembers the first (same process; see durability note below) +curl -X POST /responses -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "What is my name?" }], + "custom_inputs": { "session_id": "" } }' +``` + +## Customize the agent + +- **Model / instructions:** `create_agent_graph()` in `agent/agent.py`. +- **Add a tool:** drop a new file in `agent/tools/` with a `@tool`-decorated function; it's + collected automatically (see `agent/tools/sample_tool.py`). No wiring to edit. +- **Add an MCP server:** append a `DatabricksMCPServer` to `build_mcp_servers()` in `agent/mcps.py`. +- **Change the session checkpointer:** `agent/mason/session_store.py` (in-memory by default; swap for + a durable `PostgresSaver` over Lakebase). +- **Add long-term memory:** set `AGENT_MEMORY_STORE` to a managed memory store name; `create_agent_graph()` + then includes the `remember`/`recall` tools from `agent/mason/memory.py` (persist/search facts across + conversations). Unset → the model isn't offered them. + +## Test + +```bash +uv run pytest # hermetic smoke tests (tools, session, wire) +``` + +The smoke tests need no auth. `tests/test_agent.py` also has one end-to-end test that calls the +model; it runs only when a workspace profile is configured (`DATABRICKS_CONFIG_PROFILE` or +`DATABRICKS_HOST`+`DATABRICKS_TOKEN`) and skips otherwise. + +## Deploy + +Deploy to Databricks Apps with the CLI, which provisions any requested resources (experiment, +Lakebase) and wires them into the app: + +```bash +databricks apps deploy agent-langgraph-basic --source-code-path +``` + +`app.yaml` carries the app's start command and env. By default the deployed app is the same lean +backend: in-process session state, in-request execution, tracing off. The features below are +independent — enable either, both, or neither. + +### Enable MLflow tracing (optional) + +Tracing turns on when MLflow has **both a destination and an experiment** — set one of each, in +whichever form you have. The app code needs no change; MLflow resolves the specific value. + +- **Destination:** `MLFLOW_TRACKING_URI` (e.g. `"databricks"`) or `MLFLOW_TRACING_DESTINATION` + (an experiment id or a `catalog.schema`). +- **Experiment:** `MLFLOW_EXPERIMENT_ID` or `MLFLOW_EXPERIMENT_NAME`. + +Set neither half → tracing stays off. Examples: + +- **Local:** `MLFLOW_TRACKING_URI="databricks"` + `MLFLOW_EXPERIMENT_ID=` (or `..._NAME=`) + in `.env`, pointing at an experiment in the workspace your profile targets. +- **Deployed:** set the same env in `app.yaml` and attach an `experiment` resource (its `valueFrom` + binding injects `MLFLOW_EXPERIMENT_ID`). + +When both halves are present the agent enables MLflow autolog (`mlflow.langchain.autolog()`) and tags +each trace with the session id. Otherwise it disables tracing outright, so the agent-server +framework's per-request span is never created and no traces are exported. + +### Enable durable background mode (optional) + +**Long-running background execution + crash recovery** — a Lakebase instance attached to the app as +a resource named `postgres` (with `CAN_CONNECT_AND_CREATE` — that permission lets the app's service +principal connect and create its `agent_server.*` tables; no manual SQL grant needed) **plus** +`LAKEBASE_AUTOSCALING_ENDPOINT` set to that instance's endpoint (the resource grants access; the env +var is the address — both are needed). `server/start_server.py` passes the endpoint into +`LongRunningAgentServer`, enabling durable background mode (survives the ~120s Apps proxy timeout, +reconnect via `GET /responses/{id}`). Unset → in-request execution. + +**Durable conversation history** is a separate concern from the background store. By default the +agent uses an in-process LangGraph checkpointer (`InMemorySaver`) — multi-turn works within a +running process but does not survive restarts or span replicas. For durable, shared history, swap +`agent/mason/session_store.py`'s checkpointer for a `PostgresSaver` over the same Lakebase. + +## Configuration + +| Variable | Default | Purpose | +| --- | --- | --- | +| `DATABRICKS_CONFIG_PROFILE` | `DEFAULT` | Auth profile used to call the model (local dev) | +| `AGENT_MEMORY_STORE` | _unset_ | Managed memory store name → registers `remember`/`recall` long-term-memory tools | +| `LAKEBASE_AUTOSCALING_ENDPOINT` | _unset_ | Lakebase endpoint → durable background mode + crash recovery (else in-request) | +| `MLFLOW_TRACKING_URI` | _unset_ | Trace destination (e.g. `databricks`). A destination + an experiment enables tracing | +| `MLFLOW_TRACING_DESTINATION` | _unset_ | Alt destination — experiment id or `catalog.schema` (either destination var works) | +| `MLFLOW_EXPERIMENT_ID` | _unset_ | Experiment to trace to (by id) | +| `MLFLOW_EXPERIMENT_NAME` | _unset_ | Experiment to trace to (by name; alternative to the id) | + +## Notes + +- **`agent/mason/wire/` is LangGraph-specific** — it converts Responses input to LangGraph messages + and maps `astream` events (completed node outputs + token chunks) back to Responses wire events. diff --git a/agent-langgraph-basic/agent/__init__.py b/agent-langgraph-basic/agent/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/agent-langgraph-basic/agent/agent.py b/agent-langgraph-basic/agent/agent.py new file mode 100644 index 00000000..bba111cf --- /dev/null +++ b/agent-langgraph-basic/agent/agent.py @@ -0,0 +1,59 @@ +from collections.abc import AsyncGenerator + +from databricks_langchain import ChatDatabricks +from langchain.agents import create_agent +from mlflow.genai.agent_server import invoke, stream +from mlflow.types.responses import ( + ResponsesAgentRequest, + ResponsesAgentResponse, + ResponsesAgentStreamEvent, + to_chat_completions_input, +) + +from agent.mason import mcp_runtime, tracing +from agent.mason.memory import memory_tools +from agent.mason.session_store import checkpointer, thread_config +from agent.mason.wire.inbound import get_session_id +from agent.mason.wire.outbound import process_agent_astream_events + +# Importing the tools package auto-registers every tool module. +from agent.tools import all_tools + +MODEL = "databricks-gpt-5-2" + + +def configure() -> None: + """Wire up global state; call once at server startup (not at import).""" + tracing.configure() + + +async def create_agent_graph(): + """Build the LangGraph agent: local tools + long-term-memory tools + any MCP tools.""" + tools = [*all_tools(), *memory_tools(), *await mcp_runtime.mcp_tools()] + return create_agent(model=ChatDatabricks(endpoint=MODEL), tools=tools, checkpointer=checkpointer()) + + +@invoke() +async def invoke_handler(request: ResponsesAgentRequest) -> ResponsesAgentResponse: + outputs = [ + event.item + async for event in stream_handler(request) + if event.type == "response.output_item.done" + ] + return ResponsesAgentResponse(output=outputs) + + +@stream() +async def stream_handler( + request: ResponsesAgentRequest, +) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + session_id = get_session_id(request) + tracing.tag_session(session_id) + + agent = await create_agent_graph() + messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} + + async for event in process_agent_astream_events( + agent.astream(input=messages, config=thread_config(session_id), stream_mode=["updates", "messages"]) + ): + yield event diff --git a/agent-langgraph-basic/agent/mason/__init__.py b/agent-langgraph-basic/agent/mason/__init__.py new file mode 100644 index 00000000..caa3337d --- /dev/null +++ b/agent-langgraph-basic/agent/mason/__init__.py @@ -0,0 +1,7 @@ +"""Plumbing that will move into Databricks SDKs later (databricks-langchain and friends). + +Nothing here is meant to be edited to build an agent — it's the session-store checkpointer, MLflow +tracing setup, MCP tool loading, and the Responses<->LangGraph wire translation. Grouped in one +place so the migration to SDK-provided equivalents is a localized change. Edit the agent in +``agent/agent.py`` and ``agent/tools/`` instead. +""" diff --git a/agent-langgraph-basic/agent/mason/mcp_runtime.py b/agent-langgraph-basic/agent/mason/mcp_runtime.py new file mode 100644 index 00000000..72802f67 --- /dev/null +++ b/agent-langgraph-basic/agent/mason/mcp_runtime.py @@ -0,0 +1,26 @@ +"""MCP tool loading — plumbing, slated to move into a Databricks SDK helper. + +``mcp_tools`` builds a ``DatabricksMultiServerMCPClient`` from the servers in ``agent/mcps.py`` and +returns their tools (LangChain tools), which the agent appends to its local tools. You configure +*which* servers to offer in ``agent/mcps.py`` — this file only fetches their tools. +""" + +import logging + +from databricks_langchain import DatabricksMultiServerMCPClient + +from agent.mcps import build_mcp_servers + +logger = logging.getLogger(__name__) + + +async def mcp_tools() -> list: + """Return the tools exposed by the configured MCP servers; empty list if none/on failure.""" + servers = build_mcp_servers() + if not servers: + return [] + try: + return await DatabricksMultiServerMCPClient(servers).get_tools() + except Exception: + logger.warning("Failed to fetch MCP tools; continuing without them.", exc_info=True) + return [] diff --git a/agent-langgraph-basic/agent/mason/memory.py b/agent-langgraph-basic/agent/mason/memory.py new file mode 100644 index 00000000..f5e67da3 --- /dev/null +++ b/agent-langgraph-basic/agent/mason/memory.py @@ -0,0 +1,62 @@ +"""Long-term memory tools — opt-in, gated on ``AGENT_MEMORY_STORE``. + +Unlike the session store (short-term transcript for one conversation), long-term memory is exposed +to the model as two tools — ``remember`` and ``recall`` — over the Databricks managed memory store's +``agents/v1`` entries API. Facts persist across conversations. + +``memory_tools()`` returns the tools when ``AGENT_MEMORY_STORE`` is set, else an empty list, so the +model never sees them when memory is unconfigured. ``create_agent_graph`` composes them into its +tool list. This is a stand-in for a future ``databricks-langchain`` memory helper — when the SDK +provides one, swap the import in ``agent.py`` and delete this file. + +Memory entries are per-actor. This uses the store name as the actor id, giving the agent one shared +long-term memory; change ``_actor_id`` to scope per user (e.g. from request context) if needed. +""" + +import os + +from databricks.sdk import WorkspaceClient +from langchain_core.tools import BaseTool, tool + +_AGENTS_V1 = "/api/agents/v1" + + +def _actor_id() -> str: + return os.getenv("AGENT_MEMORY_ACTOR_ID", "agent") + + +def _store_path() -> str: + return f"{_AGENTS_V1}/memory-stores/{os.environ['AGENT_MEMORY_STORE']}" + + +def _api(): + # Build the client lazily (needs workspace auth) so importing this module stays cheap. + return WorkspaceClient().api_client + + +@tool +def remember(fact: str, topic: str) -> str: + """Persist a durable fact about the user in long-term memory.""" + _api().do( + "POST", + f"{_store_path()}/entries", + body={"actor_id": _actor_id(), "path": f"/{topic}/{fact[:8]}.md", "content": fact}, + ) + return "stored" + + +@tool +def recall(query: str) -> str: + """Search the user's long-term memory for facts relevant to the query.""" + data = _api().do( + "POST", + f"{_store_path()}/entries:search", + body={"actor_id": _actor_id(), "query": query, "limit": 5}, + ) + entries = data.get("managed_memory_entries") or [] + return "\n".join(f"- {e.get('content')}" for e in entries) or "No relevant memories." + + +def memory_tools() -> list[BaseTool]: + """The long-term-memory tools when ``AGENT_MEMORY_STORE`` is set, else none.""" + return [remember, recall] if os.getenv("AGENT_MEMORY_STORE") else [] diff --git a/agent-langgraph-basic/agent/mason/session_store.py b/agent-langgraph-basic/agent/mason/session_store.py new file mode 100644 index 00000000..f9ed7cfa --- /dev/null +++ b/agent-langgraph-basic/agent/mason/session_store.py @@ -0,0 +1,33 @@ +"""Conversation session store for the agent. + +LangGraph persists conversation state through a **checkpointer** keyed by a ``thread_id`` (passed in +the run config), not through a session object. ``checkpointer()`` returns the checkpointer the agent +is built with, and ``thread_config(session_id)`` maps a session id onto that thread. + +Default: an in-memory checkpointer (``InMemorySaver``) — multi-turn history is preserved within a +single running process, no database. It does NOT survive restarts or span replicas; for that, swap +in a durable checkpointer (e.g. ``langgraph.checkpoint.postgres.PostgresSaver`` over the Lakebase +attached as the ``postgres`` app resource). That swap is the one edit needed here; nothing else in +the agent changes. +""" + +from functools import lru_cache + +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.checkpoint.memory import InMemorySaver + + +@lru_cache(maxsize=1) +def checkpointer() -> BaseCheckpointSaver: + """The checkpointer the agent persists conversation state to. In-memory by default. + + Cached so every request shares one saver (that's what makes multi-turn work in-process). To make + history durable + shared across replicas, return a ``PostgresSaver`` built from the Lakebase + endpoint instead. + """ + return InMemorySaver() + + +def thread_config(session_id: str) -> dict: + """Run config that anchors this request to ``session_id``'s conversation thread.""" + return {"configurable": {"thread_id": session_id}} diff --git a/agent-langgraph-basic/agent/mason/tracing.py b/agent-langgraph-basic/agent/mason/tracing.py new file mode 100644 index 00000000..4eed61c5 --- /dev/null +++ b/agent-langgraph-basic/agent/mason/tracing.py @@ -0,0 +1,44 @@ +"""MLflow tracing setup — opt-in, enabled when MLflow has both a destination and an experiment. + +Tracing turns on only when a full config is present: a destination (``MLFLOW_TRACKING_URI`` or +``MLFLOW_TRACING_DESTINATION``) AND an experiment (``MLFLOW_EXPERIMENT_ID`` or +``MLFLOW_EXPERIMENT_NAME``) — whichever pair the user or the Apps resource binding provides. MLflow +resolves the specific value itself; this only decides on/off. Requiring both halves avoids the +half-configured case where traces silently export to a local file store instead of the workspace. +When unconfigured, tracing is disabled outright so the agent-server framework's mandatory per-request +span has nothing to export to. No user decision lives here — it's all driven by env — so this whole +module is a candidate to move behind an SDK helper. +""" + +import os + +import mlflow + +# Destination and experiment can each be named more than one way; accept any combination MLflow +# understands (see mlflow.tracking.fluent._get_experiment_id_from_env for the experiment resolution). +_DESTINATION_VARS = ("MLFLOW_TRACKING_URI", "MLFLOW_TRACING_DESTINATION") +_EXPERIMENT_VARS = ("MLFLOW_EXPERIMENT_ID", "MLFLOW_EXPERIMENT_NAME") + +# Snapshotted once by configure() at startup (after .env is loaded) rather than at import, so this +# module has no import-time side effects and load order does not matter. +_enabled = False + + +def configure() -> None: + """Wire up tracing. Call once at startup.""" + global _enabled + has_destination = any(os.getenv(v) for v in _DESTINATION_VARS) + has_experiment = any(os.getenv(v) for v in _EXPERIMENT_VARS) + _enabled = has_destination and has_experiment + if _enabled: + mlflow.langchain.autolog() + else: + # The agent-server framework wraps every request in a span regardless; without an + # experiment it would try to export to a missing one (INVALID_PARAMETER_VALUE), so disable. + mlflow.tracing.disable() + + +def tag_session(session_id: str) -> None: + """Tag the active MLflow trace with the session id, when tracing is enabled.""" + if _enabled and session_id: + mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id}) diff --git a/agent-langgraph-basic/agent/mason/wire/__init__.py b/agent-langgraph-basic/agent/mason/wire/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/agent-langgraph-basic/agent/mason/wire/inbound.py b/agent-langgraph-basic/agent/mason/wire/inbound.py new file mode 100644 index 00000000..f171a145 --- /dev/null +++ b/agent-langgraph-basic/agent/mason/wire/inbound.py @@ -0,0 +1,22 @@ +"""Inbound wire translation: Responses request -> LangGraph run input. + +Extracts the session id (used as the LangGraph ``thread_id``). Message conversion itself is done in +the handler via MLflow's ``to_chat_completions_input``; this module only resolves the session id. +""" + +from mlflow.types.responses import ResponsesAgentRequest +from uuid_utils import uuid7 + + +def get_session_id(request: ResponsesAgentRequest) -> str: + """Extract session_id from the request or generate a new one.""" + # Priority: + # 1. session_id from custom_inputs + # 2. conversation_id from ChatContext + # 3. a new UUID + ci = dict(request.custom_inputs or {}) + if ci.get("session_id"): + return str(ci["session_id"]) + if request.context and getattr(request.context, "conversation_id", None): + return str(request.context.conversation_id) + return str(uuid7()) diff --git a/agent-langgraph-basic/agent/mason/wire/outbound.py b/agent-langgraph-basic/agent/mason/wire/outbound.py new file mode 100644 index 00000000..51f0abce --- /dev/null +++ b/agent-langgraph-basic/agent/mason/wire/outbound.py @@ -0,0 +1,45 @@ +"""Outbound wire translation: LangGraph astream events -> Responses wire events. + +LangGraph's ``astream(stream_mode=["updates", "messages"])`` yields two event shapes: ``updates`` +(completed node outputs — full messages, incl. tool calls/results) and ``messages`` (token-level +chunks for streaming text). We map completed messages to Responses output items and text chunks to +text deltas. +""" + +import json +import logging +from collections.abc import AsyncGenerator, AsyncIterator +from typing import Any + +from langchain.messages import AIMessageChunk, ToolMessage +from mlflow.types.responses import ( + ResponsesAgentStreamEvent, + create_text_delta, + output_to_responses_items_stream, +) + +logger = logging.getLogger(__name__) + + +async def process_agent_astream_events( + async_stream: AsyncIterator[Any], +) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + """Relay LangGraph stream events as Responses wire events.""" + async for event in async_stream: + mode, payload = event[0], event[1] + if mode == "updates": + for node_data in payload.values(): + messages = node_data.get("messages", []) if isinstance(node_data, dict) else [] + for msg in messages: + # Tool results may carry non-string content; the Responses items stream needs str. + if isinstance(msg, ToolMessage) and not isinstance(msg.content, str): + msg.content = json.dumps(msg.content) + for item in output_to_responses_items_stream(messages): + yield item + elif mode == "messages": + try: + chunk = payload[0] + if isinstance(chunk, AIMessageChunk) and (content := chunk.content): + yield ResponsesAgentStreamEvent(**create_text_delta(delta=content, item_id=chunk.id)) + except Exception: + logger.exception("Error processing agent stream chunk") diff --git a/agent-langgraph-basic/agent/mcps.py b/agent-langgraph-basic/agent/mcps.py new file mode 100644 index 00000000..868ebe69 --- /dev/null +++ b/agent-langgraph-basic/agent/mcps.py @@ -0,0 +1,26 @@ +"""MCP servers to offer the agent — this is where you configure them. + +Empty by default: the agent runs with no MCP servers. Add servers to ``build_mcp_servers`` to offer +them; ``agent/mason/mcp_runtime.py`` turns them into tools for each request. +""" + +from databricks_langchain import DatabricksMCPServer + + +def build_mcp_servers() -> list[DatabricksMCPServer]: + """Return the MCP servers to offer the agent. Empty by default — add your own. + + Example (a Databricks-managed MCP, authed as the app service principal):: + + from databricks.sdk import WorkspaceClient + + host = WorkspaceClient().config.host + return [ + DatabricksMCPServer( + name="system-ai", + url=f"{host}/api/2.0/mcp/functions/system/ai", + workspace_client=WorkspaceClient(), + ), + ] + """ + return [] diff --git a/agent-langgraph-basic/agent/tools/__init__.py b/agent-langgraph-basic/agent/tools/__init__.py new file mode 100644 index 00000000..8c8a975d --- /dev/null +++ b/agent-langgraph-basic/agent/tools/__init__.py @@ -0,0 +1,24 @@ +"""Agent tools package. + +Every module here is auto-imported, and every LangChain ``BaseTool`` it defines (via the +``@tool`` decorator) is collected by ``all_tools()``. Drop a new ``*.py`` into this folder with a +``@tool``-decorated function and it's picked up automatically — no wiring to edit. +``create_agent_graph`` uses ``all_tools()``. +""" + +import importlib +import inspect +import pkgutil + +from langchain_core.tools import BaseTool + + +def all_tools() -> list[BaseTool]: + """Every BaseTool defined across the modules in this package.""" + tools: list[BaseTool] = [] + for module in pkgutil.iter_modules(__path__): + mod = importlib.import_module(f"{__name__}.{module.name}") + for _, obj in inspect.getmembers(mod, lambda o: isinstance(o, BaseTool)): + if obj not in tools: # a tool imported into several modules is collected once + tools.append(obj) + return tools diff --git a/agent-langgraph-basic/agent/tools/sample_tool.py b/agent-langgraph-basic/agent/tools/sample_tool.py new file mode 100644 index 00000000..c0741bfa --- /dev/null +++ b/agent-langgraph-basic/agent/tools/sample_tool.py @@ -0,0 +1,15 @@ +"""Sample tool. A working example — add your own tools as new files in this package. + +Decorate a function with ``@tool`` (LangChain) and it becomes an agent tool; the package +auto-collects it via ``all_tools()``, which ``create_agent_graph`` uses. +""" + +from datetime import datetime + +from langchain_core.tools import tool + + +@tool +def get_current_time() -> str: + """Get the current date and time.""" + return datetime.now().isoformat() diff --git a/agent-langgraph-basic/app.yaml b/agent-langgraph-basic/app.yaml new file mode 100644 index 00000000..75958b44 --- /dev/null +++ b/agent-langgraph-basic/app.yaml @@ -0,0 +1,23 @@ +command: ["uv", "run", "start-server"] +# databricks apps listen by default on port 8000 +# +# No env vars are required by default. To enable tracing / a session or memory store / Lakebase +# durability, add the matching entries below (literal values): +# +# env: +# # Tracing (set both): the destination + the experiment to log to. +# - name: MLFLOW_TRACKING_URI +# value: "databricks" +# - name: MLFLOW_EXPERIMENT_ID +# value: "" +# # Durable conversation history: a managed session store name (else local SQLite). +# - name: AGENT_SESSION_STORE +# value: "" +# # Long-term memory tools: a managed memory store name (registers remember/recall). +# - name: AGENT_MEMORY_STORE +# value: "" +# # Long-running background mode + crash recovery: the Lakebase autoscaling endpoint. Also attach +# # a Lakebase resource to the app (CAN_CONNECT_AND_CREATE) so the app SP can connect — the env +# # var is the address, the resource is the authorization. +# - name: LAKEBASE_AUTOSCALING_ENDPOINT +# value: "" diff --git a/agent-langgraph-basic/manifest.yaml b/agent-langgraph-basic/manifest.yaml new file mode 100644 index 00000000..d1d4a54e --- /dev/null +++ b/agent-langgraph-basic/manifest.yaml @@ -0,0 +1,13 @@ +version: 1 +name: "Agent - LangGraph Basic" +description: "A lean LangGraph agent backend that runs locally with no database. Exposes the Responses API; optionally adds Lakebase-backed long-running background execution and durable session history when configured." + +resource_specs: + - name: "experiment" + description: "Optional. The destination experiment for MLflow traces from the agent's execution." + experiment_spec: + permission: "CAN_EDIT" + - name: "postgres" + description: "Optional. A Lakebase Autoscaling database backing LongRunningAgentServer's durable store — long-running background execution and crash recovery. Omit to run in-request. (Conversation history is separate: in-memory by default, or a durable LangGraph checkpointer.)" + postgres_spec: + permission: "CAN_CONNECT_AND_CREATE" diff --git a/agent-langgraph-basic/pyproject.toml b/agent-langgraph-basic/pyproject.toml new file mode 100644 index 00000000..9552b79e --- /dev/null +++ b/agent-langgraph-basic/pyproject.toml @@ -0,0 +1,49 @@ +[project] +name = "agent-langgraph-basic" +version = "0.1.0" +description = "Lean LangGraph agent backend for Databricks Apps — local-first, optional Lakebase durability" +readme = "README.md" +authors = [ + { name = "Agent Developer", email = "developer@example.com" } +] +requires-python = ">=3.11" +dependencies = [ + "fastapi>=0.129.0", + "uvicorn>=0.41.0", + "mlflow>=3.10.1", + "databricks-langchain>=0.17.0", + "langgraph>=1.1.0", + "langchain>=1.0.0", + "langchain-mcp-adapters>=0.2.1", + "python-dotenv>=1.2.1", + "uuid-utils>=0.10.0", + "databricks-sdk>=0.79.0", + "databricks-agents>=1.9.3", + "databricks-ai-bridge[agent-server]>=0.19.0", + "opentelemetry-exporter-otlp-proto-grpc>=1.25.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["agent", "server"] + + +[dependency-groups] +dev = [ + "hatchling>=1.28.0", + "pytest>=9.0.2", + "pytest-asyncio>=1.3.0", +] + +[tool.uv] +default-groups = ["dev"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +asyncio_mode = "auto" + +[project.scripts] +start-server = "server.start_server:main" diff --git a/agent-langgraph-basic/server/__init__.py b/agent-langgraph-basic/server/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/agent-langgraph-basic/server/start_server.py b/agent-langgraph-basic/server/start_server.py new file mode 100644 index 00000000..67d2f5b5 --- /dev/null +++ b/agent-langgraph-basic/server/start_server.py @@ -0,0 +1,31 @@ +"""Agent server entry point.""" + +import os +from pathlib import Path + +from databricks_ai_bridge.long_running import LongRunningAgentServer +from dotenv import load_dotenv + +# Importing the agent registers its @invoke/@stream handlers; the import is side-effect-free (no env +# is read until configure()), so it can sit with the other imports. +import agent.agent + +# Load .env before the runtime steps below read env (agent client auth + tracing config). +load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env", override=True) + +agent.agent.configure() + +# Pass the Lakebase autoscaling endpoint through from the env (set by the "postgres" app resource). +# When it's set, LongRunningAgentServer enables durable background mode + crash recovery; when unset +# (local dev, no Lakebase) it's None, so the server serves in-request. +agent_server = LongRunningAgentServer( + "ResponsesAgent", + db_autoscaling_endpoint=os.getenv("LAKEBASE_AUTOSCALING_ENDPOINT"), +) + +# Module-level app so uvicorn can import it by string (and to enable multiple workers). +app = agent_server.app + + +def main(): + agent_server.run(app_import_string="server.start_server:app") diff --git a/agent-langgraph-basic/tests/__init__.py b/agent-langgraph-basic/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/agent-langgraph-basic/tests/test_agent.py b/agent-langgraph-basic/tests/test_agent.py new file mode 100644 index 00000000..ecfec295 --- /dev/null +++ b/agent-langgraph-basic/tests/test_agent.py @@ -0,0 +1,72 @@ +"""Smoke tests for the agent. + +Hermetic tests import only the leaf modules (tools, session store, wire) — no Databricks auth +needed, so they run anywhere, including in `databricks agent test`. The live test builds the full +agent and calls the model; it is skipped unless a workspace profile is configured. +""" + +import os + +import pytest +from langchain_core.tools import BaseTool +from mlflow.types.responses import ResponsesAgentRequest + +from agent.mason.session_store import checkpointer, thread_config +from agent.mason.wire.inbound import get_session_id +from agent.tools import all_tools + + +def test_tools_autoregister(): + tools = all_tools() + assert tools, "expected the sample tool to auto-register" + assert all(isinstance(t, BaseTool) for t in tools) + assert "get_current_time" in {t.name for t in tools} + + +def test_thread_config_from_session_id(): + assert thread_config("abc-123") == {"configurable": {"thread_id": "abc-123"}} + + +def test_checkpointer_is_shared(): + # Cached so multi-turn history is preserved in-process across requests. + assert checkpointer() is checkpointer() + + +def test_session_id_from_custom_inputs(): + request = ResponsesAgentRequest( + input=[{"role": "user", "content": "hi"}], + custom_inputs={"session_id": "abc-123"}, + ) + assert get_session_id(request) == "abc-123" + + +def test_session_id_generated_when_absent(): + request = ResponsesAgentRequest(input=[{"role": "user", "content": "hi"}]) + generated = get_session_id(request) + assert generated and generated != get_session_id( + ResponsesAgentRequest(input=[{"role": "user", "content": "hi"}]) + ) + + +def _has_workspace_auth() -> bool: + return bool( + os.getenv("DATABRICKS_CONFIG_PROFILE") + or (os.getenv("DATABRICKS_HOST") and os.getenv("DATABRICKS_TOKEN")) + ) + + +@pytest.mark.skipif( + not _has_workspace_auth(), + reason="no Databricks profile configured; skipping live model call", +) +@pytest.mark.asyncio +async def test_agent_responds_end_to_end(): + from agent.agent import configure, create_agent_graph + + configure() + agent = await create_agent_graph() + result = await agent.ainvoke( + {"messages": [{"role": "user", "content": "Reply with the single word: pong"}]}, + config=thread_config("test-e2e"), + ) + assert result["messages"][-1].content