Lab 04 · Remembering customers

langchain==1.4.2
langchain-core==1.6.4
langchain-openai==1.6.2
python-dotenv==1.2.3
langgraph==1.2.11
langgraph-checkpoint==4.2.0
langgraph-checkpoint-sqlite==3.1.1

# Azure OpenAI settings. Your instructor gives you these three values.
AZURE_OPENAI_ENDPOINT=https://<your-resource>.openai.azure.com/openai/v1
AZURE_OPENAI_API_KEY=<paste-your-key-here>
AZURE_OPENAI_DEPLOYMENT=<your-chat-deployment-name>

# Only used when your endpoint does NOT end in /openai/v1.
OPENAI_API_VERSION=2025-04-01-preview

"""Azure OpenAI model setup shared by every lab.

Import it in lab code as:  from model_setup import get_model, get_embeddings
"""

import asyncio
import os
import sys

from dotenv import load_dotenv
from langchain_core.embeddings import Embeddings
from langchain_core.language_models import BaseChatModel

# Read AZURE_OPENAI_* values from the .env file next to your lab code.
load_dotenv()

# Windows consoles default to cp1252 and crash on the curly quotes models love to print.
for _stream in (sys.stdout, sys.stderr):
    if hasattr(_stream, "reconfigure"):
        _stream.reconfigure(encoding="utf-8", errors="replace")

def _setting(name: str) -> str:
    """Return a required setting from .env, or stop with a readable message."""
    value = os.getenv(name, "").strip()
    if not value:
        sys.exit(f"Missing {name} in your .env file. Open .env and fill it in.")
    return value


def _uses_v1_api(endpoint: str) -> bool:
    """True for the OpenAI-compatible Azure endpoint that ends in /openai/v1."""
    return endpoint.rstrip("/").endswith("/openai/v1")

def get_model(deployment: str | None = None, **kwargs) -> BaseChatModel:
    """Return a chat model for your Azure OpenAI deployment."""
    endpoint = _setting("AZURE_OPENAI_ENDPOINT")
    api_key = _setting("AZURE_OPENAI_API_KEY")
    name = deployment or _setting("AZURE_OPENAI_DEPLOYMENT")

    if _uses_v1_api(endpoint):
        from langchain_openai import ChatOpenAI

        return ChatOpenAI(model=name, base_url=endpoint, api_key=api_key, **kwargs)

    from langchain_openai import AzureChatOpenAI

    return AzureChatOpenAI(
        azure_deployment=name,
        azure_endpoint=endpoint,
        api_key=api_key,
        api_version=os.getenv("OPENAI_API_VERSION", "2025-04-01-preview"),
        **kwargs,
    )

def get_embeddings(deployment: str | None = None) -> Embeddings:
    """Return an embeddings model for your Azure OpenAI embeddings deployment."""
    endpoint = _setting("AZURE_OPENAI_ENDPOINT")
    api_key = _setting("AZURE_OPENAI_API_KEY")
    name = deployment or _setting("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")

    if _uses_v1_api(endpoint):
        from langchain_openai import OpenAIEmbeddings

        return OpenAIEmbeddings(model=name, base_url=endpoint, api_key=api_key)

    from langchain_openai import AzureOpenAIEmbeddings

    return AzureOpenAIEmbeddings(
        azure_deployment=name,
        azure_endpoint=endpoint,
        api_key=api_key,
        api_version=os.getenv("OPENAI_API_VERSION", "2025-04-01-preview"),
    )

def use_stable_event_loop() -> None:
    """On Windows, use the Selector event loop so async MCP code exits cleanly."""
    if sys.platform == "win32":
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

"""Lab 04: remembering customers with SQLite threads and a store-backed profile.

Two different thread_ids, one shared long-term memory.

Run:  python remembering_customers.py
"""

import sqlite3
from dataclasses import dataclass
from typing import Any

from langchain.agents import create_agent
from langchain.tools import ToolRuntime, tool
from langchain_core.language_models import BaseChatModel
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.store.base import BaseStore
from langgraph.store.memory import InMemoryStore

from model_setup import get_model

ORDERS = {"A-100": "shipped via DHL, ETA 2026-09-15", "A-101": "processing, ETA 2026-09-18"}


@dataclass
class SupportContext:
    """Who this run serves. Per-run, and deliberately never persisted in the thread."""

    org_id: str
    customer_id: str

def preference_namespace(org_id: str, customer_id: str) -> tuple[str, str, str]:
    """(org, user, kind): tenant first, so a prefix search can never cross organisations."""
    return (org_id, customer_id, "preferences")

@tool
def get_order_status(order_id: str) -> str:
    """Look up the shipping status of an order by id, such as A-100."""
    return f"Order {order_id.upper()}: {ORDERS.get(order_id.upper(), 'no such order')}."

@tool
def save_preference(kind: str, value: str, runtime: ToolRuntime[SupportContext]) -> str:
    """Store a lasting preference so future conversations can use it. `kind` is a short slug
    such as "contact_channel"; restating a preference updates it. Never store secrets."""
    namespace = preference_namespace(runtime.context.org_id, runtime.context.customer_id)
    runtime.store.put(namespace, kind, {"text": f"{kind}: {value}", "value": value})
    return f"Saved preference {kind}: {value}"

@tool
def recall_preferences(query: str, runtime: ToolRuntime[SupportContext]) -> str:
    """Look up what we already know about this customer. Call this at the start of every
    conversation, before asking them to repeat themselves."""
    namespace = preference_namespace(runtime.context.org_id, runtime.context.customer_id)
    hits = runtime.store.search(namespace, query=query, limit=5)
    return "\n".join(f"- {hit.value['text']}" for hit in hits) or "No stored preferences yet."

SYSTEM_PROMPT = (
    "You are the Northwind support assistant. Call `recall_preferences` at the start of a "
    "conversation, `save_preference` when the customer states a lasting preference, and "
    "`get_order_status` for shipping. Never store secrets. Reply in one sentence."
)

def build_agent(
    model: BaseChatModel | None = None,
    *,
    checkpointer: BaseCheckpointSaver | None = None,
    store: BaseStore | None = None,
) -> Any:
    """Build the agent with short-term (checkpointer) and long-term (store) memory."""
    return create_agent(
        model=model or get_model(),
        tools=[get_order_status, save_preference, recall_preferences],
        system_prompt=SYSTEM_PROMPT,
        context_schema=SupportContext,
        checkpointer=checkpointer,
        store=store,
    )

CONVERSATIONS = [
    ("mon-1", "Hi, I'm Ada. Always email me, never phone."),
    ("thu-1", "Hi again - where is order A-100, and how will you update me?"),
]

if __name__ == "__main__":
    connection = sqlite3.connect("lab04_threads.db", check_same_thread=False)
    store = InMemoryStore()
    agent = build_agent(checkpointer=SqliteSaver(connection), store=store)
    ada = SupportContext(org_id="acme-corp", customer_id="ada")

    for thread_id, text in CONVERSATIONS:
        config = {"configurable": {"thread_id": thread_id}}
        before = len(agent.get_state(config).values.get("messages", []))
        result = agent.invoke({"messages": [{"role": "user", "content": text}]},
                              config=config, context=ada)
        new_messages = result["messages"][before:]
        tools = [c["name"] for m in new_messages for c in getattr(m, "tool_calls", [])]
        print(f"[{thread_id}] you   : {text}")
        print(f"[{thread_id}] tools : {', '.join(tools) or 'none'}")
        print(f"[{thread_id}] agent : {result['messages'][-1].text}\n")

    # Short-term memory: each thread has its own message history in SQLite.
    for thread_id, _ in CONVERSATIONS:
        state = agent.get_state({"configurable": {"thread_id": thread_id}})
        print(f"thread {thread_id}: {len(state.values.get('messages', []))} messages saved")

    # Long-term memory: one profile per (org, customer), shared by every thread.
    print("acme-corp/ada profile:",
          [i.value["text"] for i in store.search(preference_namespace("acme-corp", "ada"))])
    print("globex/ada profile   :",
          [i.value["text"] for i in store.search(preference_namespace("globex", "ada"))])
    connection.close()

[mon-1] you   : Hi, I'm Ada. Always email me, never phone.
[mon-1] tools : recall_preferences, save_preference
[mon-1] agent : Nice to meet you, Ada — I’ll email you and won’t call.

[thu-1] you   : Hi again - where is order A-100, and how will you update me?
[thu-1] tools : recall_preferences, get_order_status
[thu-1] agent : Order A-100 has shipped via DHL with an ETA of 2026-09-15, and I will send updates to your preferred contact channel (email).

thread mon-1: 6 messages saved
thread thu-1: 5 messages saved
acme-corp/ada profile: ['contact_channel: email']
globex/ada profile   : []