CockroachDB integrations

pip install langchain-cockroachdb

docker run -d --name cockroachdb -p 26257:26257 \
  cockroachdb/cockroach:latest start-single-node --insecure

from langchain_cockroachdb import AsyncCockroachDBVectorStore, CockroachDBEngine
from langchain_openai import OpenAIEmbeddings

# Initialize
engine = CockroachDBEngine.from_connection_string(
    "cockroachdb://user:pass@host:26257/db"
)

await engine.ainit_vectorstore_table(
    table_name="documents",
    vector_dimension=1536,
)

vectorstore = AsyncCockroachDBVectorStore(
    engine=engine,
    embeddings=OpenAIEmbeddings(),
    collection_name="documents",
)

# Use it
ids = await vectorstore.aadd_texts(["Hello world"])
results = await vectorstore.asimilarity_search("Hi", k=1)

from langchain_cockroachdb import HybridSearchConfig

# create_tsvector=True adds the TSVECTOR column and GIN index
await engine.ainit_vectorstore_table(
    table_name="documents",
    vector_dimension=1536,
    create_tsvector=True,
)

vectorstore = AsyncCockroachDBVectorStore(
    engine=engine,
    embeddings=OpenAIEmbeddings(),
    collection_name="documents",
    hybrid_search_config=HybridSearchConfig(),
)

# The regular search methods now do hybrid search transparently
results = await vectorstore.asimilarity_search("distributed database performance", k=5)

from langchain_cockroachdb import CockroachDBChatMessageHistory, CockroachDBEngine
import uuid

engine = CockroachDBEngine.from_connection_string(CONNECTION_STRING)

chat_history = CockroachDBChatMessageHistory(
    session_id=str(uuid.uuid4()),
    engine=engine,
    table_name="chat_history",
)

from langchain.messages import HumanMessage, AIMessage

await chat_history.aadd_message(HumanMessage(content="Hello!"))
await chat_history.aadd_message(AIMessage(content="Hi there!"))

messages = await chat_history.aget_messages()

with CockroachDBSaver.from_conn_string(DB_URI) as checkpointer:
    checkpointer.setup()

    # Expire checkpoints older than 30 days, clean up daily
    checkpointer.enable_ttl(ttl_interval="30 days", cron="@daily")

    # Use the checkpointer normally -- old data is cleaned up automatically
    graph = builder.compile(checkpointer=checkpointer)

    # To disable TTL later:
    # checkpointer.disable_ttl()

import os

from langchain_cockroachdb import CockroachDBStore

DB_URI = os.environ["COCKROACHDB_URI"]

with CockroachDBStore.from_conn_string(DB_URI, pool_config={"max_size": 10}) as store:
    store.setup()

    namespace = ("user", "alice", "memories")
    store.put(namespace, "seat", {"text": "Prefers window seats"})
    item = store.get(namespace, "seat")
    results = store.search(("user", "alice"))

from langchain_openai import OpenAIEmbeddings

with CockroachDBStore.from_conn_string(
    DB_URI,
    index={
        "dims": 1536,
        "embed": OpenAIEmbeddings(model="text-embedding-3-small"),
        "fields": ["text"],
    },
) as store:
    store.setup()
    store.put(("user", "alice", "memories"), "seat", {"text": "Prefers window seats"})
    results = store.search(("user", "alice", "memories"), query="Where does Alice sit?")

with CockroachDBStore.from_conn_string(DB_URI, ttl={"default_ttl": 60 * 24 * 30}) as store:
    store.setup()
    store.put(("user", "alice", "memories"), "status", {"text": "In a meeting"}, ttl=10)

deleted = store.delete_namespace(("user", "alice"))

from dataclasses import dataclass

from langchain.agents import create_agent
from langchain.tools import ToolRuntime, tool
from langchain_cockroachdb import CockroachDBStore
from langchain_cockroachdb.store import first_of, skip_duplicates, supersede_same_slot

POLICY = first_of(skip_duplicates(), supersede_same_slot())


@dataclass
class Context:
    user_id: str


@tool
def remember(text: str, topic: str, runtime: ToolRuntime[Context, dict]) -> str:
    """Save a fact about the user. topic names what it is about, such as "seat"."""
    assert isinstance(runtime.store, CockroachDBStore)
    namespace = ("user", runtime.context.user_id, "memories")
    result = runtime.store.reconcile(
        namespace,
        None,  # the store generates the key
        {"topic": topic, "text": text},
        slot="topic",
        policy=POLICY,
        provenance={"source": "chat"},
    )
    if result.new != "written":
        return f"Already known: {text}"
    replaced = f", replacing {len(result.superseded)}" if result.superseded else ""
    return f"Saved: {text}{replaced}"


@tool
def past_memories(topic: str, runtime: ToolRuntime[Context, dict]) -> str:
    """List what the user used to say about a topic, newest first."""
    assert isinstance(runtime.store, CockroachDBStore)
    namespace = ("user", runtime.context.user_id, "memories")
    lines = [
        f"{version.ended_at:%Y-%m-%d}: {version.value['text']} ({version.reason})"
        for version in runtime.store.history(namespace, limit=100)
        if version.value.get("topic") == topic
    ]
    return "\n".join(lines) or f"No earlier memories about {topic}."


with CockroachDBStore.from_conn_string(DB_URI, history=True) as store:
    store.setup()
    agent = create_agent(
        "claude-sonnet-4-6",
        tools=[remember, past_memories],
        context_schema=Context,
        store=store,
    )

from langchain_cockroachdb import AsyncCockroachDBVectorStore, CockroachDBEngine
from langchain_cockroachdb.indexes import CSPANNIndex
from langchain_openai import OpenAIEmbeddings

engine = CockroachDBEngine.from_connection_string(CONNECTION_STRING)

# Create the table with a namespace column
await engine.ainit_vectorstore_table(
    table_name="documents",
    vector_dimension=1536,
    namespace_column="namespace",
)

# Create a vectorstore scoped to a specific tenant
vectorstore = AsyncCockroachDBVectorStore(
    engine=engine,
    embeddings=OpenAIEmbeddings(),
    collection_name="documents",
    namespace="tenant_a",
)

# Index the vectors with the namespace as the prefix column, before writes start
await vectorstore.aapply_vector_index(CSPANNIndex(), prefix_columns=["namespace"])

# All operations are scoped to tenant_a
ids = await vectorstore.aadd_texts(
    ["Tenant A document"], metadatas=[{"category": "news"}]
)
results = await vectorstore.asimilarity_search("query", k=5)
results = await vectorstore.asimilarity_search(
    "query", k=5, filter={"category": "news"}
)

from langchain_cockroachdb import RETRYABLE_SQLSTATES, CockroachDBEngine, RetryPolicy

engine = CockroachDBEngine.from_connection_string(
    CONNECTION_STRING,
    retry_policy=RetryPolicy(
        max_attempts=8,
        # Optional for multi-region apps that set statement_timeout:
        retryable_sqlstates=RETRYABLE_SQLSTATES | {"57014"},
        attempt_timeout=10,
    ),
)