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 langchain_cockroachdb import AsyncCockroachDBVectorStore, CockroachDBEngine
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",
)

# All operations are scoped to tenant_a
ids = await vectorstore.aadd_texts(["Tenant A document"])
results = await vectorstore.asimilarity_search("query", k=5)

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,
    ),
)