Migrating raw OpenAI SDK memory management to LangChain

Moving from raw OpenAI SDK memory management to LangChain is one of those migrations that looks simple on paper but bites you in production. The openai sdk memory management langchain migration path requires translating manual message array handling into LangChain’s memory abstractions without losing context fidelity or introducing latency regressions. This guide walks through the complete migration with runnable code at each step and verification checkpoints you can automate.

Step 1: Audit your current memory implementation

Before touching LangChain, document exactly what your raw SDK code does. Most teams store messages as a list of dictionaries, append user and assistant turns, and optionally truncate or summarize when approaching token limits.

# legacy_memory.py — typical raw SDK pattern
from openai import OpenAI
import tiktoken

client = OpenAI()
encoding = tiktoken.encoding_for_model("gpt-4o")

MAX_TOKENS = 128_000  # model context window
RESERVED = 4_000      # headroom for response

class RawMemory:
    def __init__(self, system_prompt: str = ""):
        self.messages = []
        if system_prompt:
            self.messages.append({"role": "system", "content": system_prompt})

    def add_user(self, content: str):
        self.messages.append({"role": "user", "content": content})

    def add_assistant(self, content: str):
        self.messages.append({"role": "assistant", "content": content})

    def token_count(self) -> int:
        return sum(
            len(encoding.encode(m["content"])) for m in self.messages
        )

    def truncate(self):
        # Naive: drop oldest non-system messages until under budget
        while self.token_count() > MAX_TOKENS - RESERVED and len(self.messages) > 1:
            # Keep system message at index 0
            self.messages.pop(1)

    def get_messages(self) -> list[dict]:
        self.truncate()
        return self.messages

Verify: Write a test that exercises a 20-turn conversation, asserts token count stays under budget, and confirms system prompt survives truncation.

# test_legacy_memory.py
def test_truncation_preserves_system():
    mem = RawMemory("You are a helpful assistant.")
    for i in range(30):
        mem.add_user(f"User message {i} " + "x" * 500)
        mem.add_assistant(f"Assistant reply {i} " + "y" * 500)
    msgs = mem.get_messages()
    assert msgs[0]["role"] == "system"
    assert mem.token_count() <= 128_000 - 4_000

Run pytest test_legacy_memory.py -v. Green means you have a baseline.

Step 2: Map raw concepts to LangChain equivalents

LangChain separates memory (state) from chat message history (storage) and buffer management (truncation/summarization). The mapping:

Raw SDK concept	LangChain component
`messages` list	`BaseChatMessageHistory` (e.g., `InMemoryChatMessageHistory`)
Manual append	`add_user_message` / `add_ai_message`
Token counting + truncate	`ConversationBufferWindowMemory` or `ConversationTokenBufferMemory`
System prompt	`SystemMessage` at history initialization or via prompt template

Install the minimal dependencies:

pip install langchain-openai langchain-core tiktoken

Step 3: Implement drop-in replacement with ConversationBufferWindowMemory

Start with the simplest LangChain memory that mimics fixed-window behavior. This keeps the last k message pairs (user + assistant).

# langchain_memory_v1.py
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain.memory import ConversationBufferWindowMemory
from langchain_core.chat_history import InMemoryChatMessageHistory

class LangChainWindowMemory:
    def __init__(self, system_prompt: str = "", k: int = 10):
        self.history = InMemoryChatMessageHistory()
        if system_prompt:
            self.history.add_message(SystemMessage(content=system_prompt))
        self.memory = ConversationBufferWindowMemory(
            chat_memory=self.history,
            k=k,
            return_messages=True,
            memory_key="history",
            input_key="input",
            output_key="output",
        )

    def add_user(self, content: str):
        self.history.add_user_message(content)

    def add_assistant(self, content: str):
        self.history.add_ai_message(content)

    def get_messages(self) -> list:
        # Returns list of BaseMessage objects
        return self.memory.load_memory_variables({})["history"]

    def to_openai_format(self) -> list[dict]:
        """Convert for direct OpenAI SDK calls if needed."""
        mapping = {
            "system": "system",
            "human": "user",
            "ai": "assistant",
        }
        return [
            {"role": mapping[m.type], "content": m.content}
            for m in self.get_messages()
        ]

Verify: Same 20-turn test, now asserting the window size.

# test_langchain_window.py
def test_window_memory_keeps_last_k():
    mem = LangChainWindowMemory("System prompt", k=5)
    for i in range(20):
        mem.add_user(f"User {i}")
        mem.add_assistant(f"Assistant {i}")
    msgs = mem.get_messages()
    # System + 5 pairs = 11 messages
    assert len(msgs) == 11
    assert msgs[0].type == "system"
    assert msgs[-1].type == "ai"
    assert "User 19" in msgs[-2].content

Step 4: Switch to token-aware buffering with ConversationTokenBufferMemory

Fixed windows waste context or truncate aggressively. ConversationTokenBufferMemory uses a token counter to keep as much history as fits under a limit — closer to your raw implementation.

# langchain_memory_v2.py
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage
from langchain.memory import ConversationTokenBufferMemory
from langchain_core.chat_history import InMemoryChatMessageHistory
import tiktoken

class LangChainTokenMemory:
    def __init__(
        self,
        system_prompt: str = "",
        max_token_limit: int = 124_000,  # leave headroom
        model_name: str = "gpt-4o",
    ):
        self.llm = ChatOpenAI(model=model_name, temperature=0)
        self.history = InMemoryChatMessageHistory()
        if system_prompt:
            self.history.add_message(SystemMessage(content=system_prompt))

        self.memory = ConversationTokenBufferMemory(
            llm=self.llm,
            chat_memory=self.history,
            max_token_limit=max_token_limit,
            return_messages=True,
            memory_key="history",
            input_key="input",
            output_key="output",
        )

    def add_user(self, content: str):
        self.history.add_user_message(content)

    def add_assistant(self, content: str):
        self.history.add_ai_message(content)

    def get_messages(self) -> list:
        return self.memory.load_memory_variables({})["history"]

    def token_count(self) -> int:
        # LangChain uses the LLM's get_num_tokens under the hood
        return self.llm.get_num_tokens_from_messages(self.history.messages)

    def to_openai_format(self) -> list[dict]:
        mapping = {"system": "system", "human": "user", "ai": "assistant"}
        return [
            {"role": mapping[m.type], "content": m.content}
            for m in self.get_messages()
        ]

Verify: Stress test with variable-length messages.

# test_langchain_token.py
def test_token_memory_respects_limit():
    mem = LangChainTokenMemory("System prompt", max_token_limit=2000)
    # Add increasingly long messages
    for i in range(15):
        content = "x" * (200 + i * 50)
        mem.add_user(content)
        mem.add_assistant(content)
    msgs = mem.get_messages()
    assert mem.token_count() <= 2000
    assert msgs[0].type == "system"
    # Should have dropped oldest turns
    assert len(msgs) < 31  # system + 15 pairs would be 31

Step 5: Integrate with LCEL for production chains

LangChain Expression Language (LCEL) is the recommended way to compose memory with prompts and models. Replace your raw client.chat.completions.create call with a chain that injects history automatically.

# chain.py
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_memory_v2 import LangChainTokenMemory

def build_chain(memory: LangChainTokenMemory):
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful assistant."),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}"),
    ])

    llm = ChatOpenAI(model="gpt-4o", temperature=0)

    chain = (
        RunnablePassthrough.assign(
            history=lambda _: memory.get_messages()
        )
        | prompt
        | llm
        | StrOutputParser()
    )
    return chain

# Usage
memory = LangChainTokenMemory(max_token_limit=120_000)
chain = build_chain(memory)

# Turn 1
response = chain.invoke({"input": "What is the capital of France?"})
memory.add_user("What is the capital of France?")
memory.add_assistant(response)
print(response)

# Turn 2 — history automatically included
response = chain.invoke({"input": "What language do they speak there?"})
memory.add_user("What language do they speak there?")
memory.add_assistant(response)
print(response)

Verify: Run the script and confirm the second response references France without the user repeating it.

Step 6: Persist history across restarts

In-memory history evaporates on process exit. Swap InMemoryChatMessageHistory for a persistent backend — Redis, Postgres, or SQLite — without changing chain code.

# persistent_memory.py
from langchain_redis import RedisChatMessageHistory
from langchain.memory import ConversationTokenBufferMemory
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage

def build_persistent_memory(
    session_id: str,
    redis_url: str = "redis://localhost:6379/0",
    system_prompt: str = "",
    max_token_limit: int = 120_000,
):
    history = RedisChatMessageHistory(
        session_id=session_id,
        redis_url=redis_url,
    )
    # Seed system prompt only if history is empty
    if len(history.messages) == 0 and system_prompt:
        history.add_message(SystemMessage(content=system_prompt))

    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    memory = ConversationTokenBufferMemory(
        llm=llm,
        chat_memory=history,
        max_token_limit=max_token_limit,
        return_messages=True,
        memory_key="history",
        input_key="input",
        output_key="output",
    )
    return memory

Verify: Start the app, send a message, kill the process, restart, send a follow-up — the model should still have context.

Step 7: Add summarization for very long conversations

When token budgets exceed what buffering can handle, layer ConversationSummaryBufferMemory. It summarizes older turns while keeping recent messages verbatim.

# summary_memory.py
from langchain.memory import ConversationSummaryBufferMemory
from langchain_redis import RedisChatMessageHistory
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage

def build_summary_memory(
    session_id: str,
    redis_url: str = "redis://localhost:6379/0",
    system_prompt: str = "",
    max_token_limit: int = 120_000,
):
    history = RedisChatMessageHistory(session_id=session_id, redis_url=redis_url)
    if len(history.messages) == 0 and system_prompt:
        history.add_message(SystemMessage(content=system_prompt))

    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)  # cheaper for summarization
    memory = ConversationSummaryBufferMemory(
        llm=llm,
        chat_memory=history,
        max_token_limit=max_token_limit,
        return_messages=True,
        memory_key="history",
        input_key="input",
        output_key="output",
    )
    return memory

Verify: Run a 100-turn conversation, inspect Redis — you should see SystemMessage, a few recent HumanMessage/AIMessage pairs, and one AIMessage with type="summary" containing condensed earlier context.

Step 8: Handle streaming and tool calls

Raw SDK streaming with memory requires manual chunk accumulation. LangChain’s astream/astream_events handles this while preserving message objects.

# streaming_chain.py
from langchain_core.runnables import RunnableConfig
from chain import build_chain
from persistent_memory import build_persistent_memory

memory = build_persistent_memory("user-123")
chain = build_chain(memory)

async def run_stream(user_input: str):
    config = RunnableConfig(configurable={"session_id": "user-123"})
    full_response = ""
    async for chunk in chain.astream({"input": user_input}, config=config):
        print(chunk, end="", flush=True)
        full_response += chunk
    print()
    memory.add_user(user_input)
    memory.add_assistant(full_response)
    return full_response

Verify: Call run_stream in a loop, confirm tokens appear incrementally, and verify the persisted history contains complete assistant messages (not fragments).

Step 9: Observability — log token usage per turn

LangChain callbacks let you capture token counts without parsing response headers manually.

# callbacks.py
from langchain_core.callbacks import BaseCallbackHandler
from typing import Any, Dict

class TokenUsageLogger(BaseCallbackHandler):
    def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs):
        usage = response.llm_output.get("token_usage", {})
        print(f"[tokens] prompt={usage.get('prompt_tokens')} "
              f"completion={usage.get('completion_tokens')} "
              f"total={usage.get('total_tokens')}")

# Attach to chain
from langchain_core.runnables import RunnableConfig
config = RunnableConfig(callbacks=[TokenUsageLogger()])
chain.invoke({"input": "Hello"}, config=config)

Verify: Run a few turns, confirm logs show non-zero prompt/completion tokens, and cross-check against your provider dashboard.

Step 10: Migration checklist and rollback plan

Before cutting traffic, run through this checklist in staging:

Parity test: Replay 1,000 real conversation logs through both implementations; diff the assistant responses. Acceptable divergence: only whitespace or deterministic sampling differences at temperature=0.
Latency budget: p99 chain latency ≤ raw SDK p99 + 50 ms. Profile with py-spy or cProfile if exceeded.
Memory leak check: Run 10,000 turns in a loop; RSS growth < 50 MB.
Persistence failover: Kill Redis mid-conversation; verify graceful degradation (in-memory fallback) or cleartext fallback, log warning and continue).
Rollback flag: Wrap the LangChain chain behind a feature flag. If error rate spikes, flip flag to route back to raw SDK class.

# feature_flag.py
import os

USE_LANGCHAIN = os.getenv("USE_LANGCHAIN", "false").lower() == "true"

def get_chain():
    if USE_LANGCHAIN:
        from chain import build_chain
        from persistent_memory import build_persistent_memory
        memory = build_persistent_memory("default")
        return build_chain(memory)
    else:
        from legacy_chain import build_legacy_chain
        return build_legacy_chain()

Deploy with USE_LANGCHAIN=false, enable for 5% of traffic, monitor error rates and latency for 24 hours, then ramp.

The openai sdk memory management langchain migration pays off when you need persistence, summarization, or multi-turn tool use — all of which become one-line changes instead of custom code. Start with ConversationTokenBufferMemory, verify against your real traffic patterns, then layer persistence and summarization as needed. The raw SDK class stays in the repo behind a flag until you’re confident; delete it after the first incident-free month.