LangChain Agents

Overview

LangChain agents that call tools, browse the web, or execute code benefit from running in isolated Chalk Compute workloads with dedicated compute and network access.

This tutorial deploys a function-shaped LangChain ReAct agent that uses tool-calling to answer research questions, with an optional volume-backed function for persistent state.

Write the agent

Create agent.py — a self-contained LangChain agent implementation used by the Function:

# agent.py
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor, create_react_agent
from langchain_core.prompts import PromptTemplate
from langchain_community.tools import DuckDuckGoSearchRun

llm = ChatOpenAI(model="gpt-4o", temperature=0)
tools = [DuckDuckGoSearchRun()]
prompt = PromptTemplate.from_template(
    "Answer the following question using the tools available to you.\n\n"
    "Tools: {tools}\nTool names: {tool_names}\n\n"
    "Question: {input}\n{agent_scratchpad}"
)
agent = create_react_agent(llm, tools, prompt)
executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

Define the function

Create answer.py:

# answer.py
import chalkcompute
from chalkcompute import Image, Secret

agent_image = (
    Image.base("python:3.12-slim")
    .pip_install([
        "langchain",
        "langchain-openai",
        "langchain-community",
        "duckduckgo-search",
    ])
    .add_local_file("agent.py", "/app/agent.py")
    .workdir("/app")
)

@chalkcompute.function(
    name="langchain-answer",
    image=agent_image,
    secrets=[Secret.from_env("OPENAI_API_KEY")],
    min_instances=0,
    max_instances=1,
)
def answer(question: str) -> str:
    from agent import executor

    return executor.invoke({"input": question})["output"]

Deploy it

# deploy_langchain.py
from answer import answer

answer.deploy()
answer.wait_ready()
print(answer.remote("What is Chalk?"))

python deploy_langchain.py
# Chalk is a feature platform for building and serving machine learning systems.

Once the function is ready, invoke it remotely:

from answer import answer

print(answer.remote("What is the capital of France?"))
# The capital of France is Paris.

Add persistent memory with a volume

LangChain agents can persist conversation history or vector store data across restarts using a Volume. Create agent_memory.py:

import uuid

import chalkcompute
from chalkcompute import Image, Secret

memory_image = (
    Image.base("python:3.12-slim")
    .pip_install([
        "langchain-openai",
        "chromadb",
    ])
)


@chalkcompute.function(
    name="langchain-memory-answer",
    image=memory_image,
    secrets=[Secret.from_env("OPENAI_API_KEY")],
    volumes=[("agent-memory", "/app/memory")],
    min_instances=0,
    max_instances=1,
)
def answer(question: str, session_id: str) -> str:
    import chromadb
    from langchain_openai import ChatOpenAI

    client = chromadb.PersistentClient(path="/app/memory")
    collection = client.get_or_create_collection("conversation-history")
    history = collection.get(
        where={"session_id": session_id},
        include=["documents"],
    )["documents"]

    prompt = "Answer the question using this conversation history:\n"
    prompt += "\n".join(history)
    prompt += f"\nUser: {question}"
    response = ChatOpenAI(model="gpt-4o").invoke(prompt)

    collection.add(
        ids=[str(uuid.uuid4())],
        documents=[f"User: {question}\nAssistant: {response.content}"],
        metadatas=[{"session_id": session_id}],
    )
    return response.content

Deploy and invoke the function with an explicit session ID:

from agent_memory import answer

answer.deploy()
answer.wait_ready()
print(answer.remote("What did we discuss?", session_id="session-123"))

The volume at /app/memory survives function instance restarts. The explicit session_id keeps separate conversations from sharing history.

Example: transaction risk scoring with Chalk features

This example builds a LangChain agent tool that receives a financial transaction, enriches it with features from Chalk, runs a PyTorch risk model, and escalates high-risk transactions to a Kinesis review queue.

Define the feature namespace

Assume your Chalk project defines features like these:

from chalk.features import features, Features

@features
class Transaction:
    id: str
    merchant_id: str
    amount: float
    merchant_category: str
    merchant_risk_tier: int
    customer_avg_spend_30d: float
    customer_transaction_count_7d: int
    country_code: str

Write the agent tool

Create risk_tool.py — a LangChain tool that queries Chalk, scores the transaction, and posts flagged results to Kinesis:

# risk_tool.py
import json
import boto3
import torch
import torch.nn as nn
from chalkpy import ChalkClient
from langchain_core.tools import tool

RISK_THRESHOLD = 0.85
KINESIS_STREAM = "transaction-review-queue"

chalk = ChalkClient()
kinesis = boto3.client("kinesis", region_name="us-east-1")


class RiskModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(4, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


# Load pre-trained weights
model = RiskModel()
model.load_state_dict(torch.load("/app/models/risk_model.pt", weights_only=True))
model.eval()


@tool
def score_transaction(transaction_id: str) -> str:
    """Score a financial transaction for fraud risk.

    Retrieves enriched features from Chalk, runs a risk model,
    and escalates to a review queue if the score exceeds the threshold.
    """
    # 1. Query Chalk for enriched transaction features
    result = chalk.query(
        input={"transaction.id": transaction_id},
        output=[
            "transaction.amount",
            "transaction.merchant_risk_tier",
            "transaction.customer_avg_spend_30d",
            "transaction.customer_transaction_count_7d",
            "transaction.merchant_category",
            "transaction.country_code",
        ],
    )

    amount = result.get_feature_value("transaction.amount")
    merchant_risk_tier = result.get_feature_value("transaction.merchant_risk_tier")
    avg_spend = result.get_feature_value("transaction.customer_avg_spend_30d")
    txn_count = result.get_feature_value("transaction.customer_transaction_count_7d")
    merchant_category = result.get_feature_value("transaction.merchant_category")
    country = result.get_feature_value("transaction.country_code")

    # 2. Run the risk model
    features = torch.tensor([[
        amount / max(avg_spend, 1.0),  # spend ratio
        float(merchant_risk_tier),
        float(txn_count),
        amount,
    ]])

    with torch.no_grad():
        risk_score = model(features).item()

    # 3. Escalate if above threshold
    if risk_score > RISK_THRESHOLD:
        kinesis.put_record(
            StreamName=KINESIS_STREAM,
            Data=json.dumps({
                "transaction_id": transaction_id,
                "risk_score": round(risk_score, 4),
                "amount": amount,
                "merchant_category": merchant_category,
                "country": country,
                "reason": "automated_risk_score_exceeded",
            }),
            PartitionKey=transaction_id,
        )
        return (
            f"Transaction {transaction_id}: risk score {risk_score:.2%} "
            f"EXCEEDS threshold. Escalated to review queue."
        )

    return (
        f"Transaction {transaction_id}: risk score {risk_score:.2%}. "
        f"Below threshold — no action required."
    )

Wire it into a function

# fraud_review.py
import chalkcompute
from chalkcompute import Image, Secret

review_image = (
    Image.debian_slim("3.12")
    .pip_install([
        "langchain",
        "langchain-openai",
        "chalkpy",
        "torch",
        "boto3",
    ])
    .add_local_file("risk_tool.py", "/app/risk_tool.py")
    .workdir("/app")
)


@chalkcompute.function(
    name="fraud-review",
    image=review_image,
    secrets=[
        Secret.from_env("OPENAI_API_KEY"),
        Secret.from_env("CHALK_CLIENT_ID"),
        Secret.from_env("CHALK_CLIENT_SECRET"),
    ],
    volumes=[("risk-models", "/app/models")],
    min_instances=0,
    max_instances=1,
)
def review(transaction_id: str) -> str:
    from langchain_openai import ChatOpenAI
    from langchain.agents import AgentExecutor, create_react_agent
    from langchain_core.prompts import PromptTemplate
    from risk_tool import score_transaction

    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    prompt = PromptTemplate.from_template(
        "You are a fraud analyst assistant. Use the score_transaction tool "
        "to evaluate transactions when asked.\n\n"
        "Tools: {tools}\nTool names: {tool_names}\n\n"
        "Question: {input}\n{agent_scratchpad}"
    )
    agent = create_react_agent(llm, [score_transaction], prompt)
    executor = AgentExecutor(agent=agent, tools=[score_transaction], verbose=True)
    result = executor.invoke({
        "input": f"Score transaction {transaction_id} for fraud risk."
    })
    return result["output"]

Deploy the function

from chalkcompute import Volume

vol = Volume("risk-models")
vol.put_file("risk_model.pt", open("risk_model.pt", "rb").read())

from fraud_review import review

review.deploy()
review.wait_ready()
print(review.remote("txn_8a3f2c"))

Query the agent:

from fraud_review import review

print(review.remote("txn_8a3f2c"))
# Transaction txn_8a3f2c: risk score 92.31% EXCEEDS threshold. Escalated to review queue.

​Overview

​Write the agent

​Define the function

​Deploy it

​Add persistent memory with a volume

​Example: transaction risk scoring with Chalk features

​Define the feature namespace

​Write the agent tool

​Wire it into a function

​Deploy the function

On this page