Site Logo

Bringing Your Python Code into LangChain: Crafting Smarter AI Apps

LangChain is a fantastic tool for creating AI applications, but sometimes you need to add a bit of your own flair—perhaps a unique calculation, a data cleanup step, or a call to a niche service. That’s where integrating Python functions shines, letting you weave your custom code into LangChain’s powerful workflows. Whether you’re building a chatbot that crunches numbers on the fly or a data pipeline that polishes raw inputs, Python functions give you the freedom to make your app do exactly what you need.

In this guide, part of the LangChain Fundamentals series, I’ll walk you through what it means to bring Python functions into LangChain, how to do it smoothly, and why it’s such a game-changer. We’ll roll up our sleeves with a hands-on example to make it concrete, keeping things clear and approachable for both newbies and seasoned developers. By the end, you’ll be ready to enhance your chatbots, document search engines, or customer support bots with your own logic. Let’s dive in!

Why Python Functions in LangChain Matter

LangChain makes it easy to work with large language models (LLMs) from providers like OpenAI or HuggingFace, but LLMs are best at generating text, not handling specialized tasks like complex math, data formatting, or external API calls. That’s where your Python functions come in, acting as custom building blocks you can plug into LangChain’s workflows to add tailored functionality.

Imagine you’re creating a chatbot that needs to calculate a discount based on a user’s input or clean up messy text before passing it to the LLM. A Python function can handle those tasks, and LangChain lets you integrate it seamlessly with prompts, chains, agents, memory, tools, and document loaders. This flexibility powers use cases like:

Adding Python functions to LangChain is like giving your app a Swiss Army knife, enabling enterprise-ready applications and workflow design patterns. Want to see the bigger picture? Check out the architecture overview or Getting Started.

Getting Your Python Functions to Play Nice with LangChain

LangChain makes it surprisingly simple to bring your Python functions into the mix, treating them as components you can call within chains, agents, or tools. Your function takes inputs, does its thing, and hands back outputs that flow into the next step of the workflow. LangChain’s LCEL (LangChain Expression Language) ties it all together, supporting both synchronous and asynchronous execution for apps that scale, as covered in performance tuning.

Here’s how it works:

  • Write Your Function: Create a Python function for your task—maybe a calculation, data cleanup, or API call.
  • Make It LangChain-Friendly: Wrap the function using utilities like RunnableLambda or tool decorators to fit it into LangChain’s ecosystem.
  • Slot It In: Add the function to a chain, agent, or tool, deciding where it runs in the workflow.
  • Let It Run: LangChain calls your function with the right inputs, processes the output, and keeps the workflow moving, often using an output parser for clean results.
  • Add Context if Needed: Pull in memory or document loaders to give your function the context it needs.

For example, in a RetrievalQA Chain, you could add a function to strip unwanted characters from retrieved text before the LLM sees it. This approach is perfect for:

  • Adding math to a chatbot for instant calculations.
  • Preprocessing data in RAG apps for cleaner inputs.
  • Extending agents with custom API calls.

The beauty of Python functions in LangChain is their ability to handle tasks LLMs struggle with, like precise calculations or complex data wrangling, while keeping your workflow smooth and integrated.

Ways to Integrate Your Python Functions

LangChain gives you several ways to bring Python functions into your app, each suited to different needs—whether it’s a quick data tweak or a full-blown agent tool. Let’s explore these approaches, what they’re good for, and how to set them up, with examples to make it click.

RunnableLambda: The Easy Way for Simple Tasks

If you’ve got a straightforward task—like formatting text or doing basic math—RunnableLambda is your go-to. It wraps your Python function so it fits neatly into a chain, taking inputs and passing outputs to the next step. Think of it as a quick way to add a custom step without much fuss.

Here’s a look at how it works:

  • Purpose: Handles simple transformations, like cleaning text or calculating values, within a chain.
  • Best For: Formatting user inputs, doing quick calculations, or preprocessing data in chatbots or data analysis agents.
  • Mechanics: Wraps your function in RunnableLambda, which LangChain calls with the chain’s input, feeding the output to the next component, like a prompt template or LLM.
  • Setup: Write your function, wrap it, and add it to the chain. Here’s an example that converts user input to uppercase before the LLM processes it:
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate

# Custom function to uppercase text
def to_uppercase(inputs):
    return {"query": inputs["query"].upper()}

# Build the chain
prompt = PromptTemplate(input_variables=["query"], template="Answer: {query}")
llm = ChatOpenAI(model="gpt-4o-mini")
chain = RunnableLambda(to_uppercase) | prompt | llm

# Test it out
result = chain.invoke({"query": "what is ai?"})
print(result.content)

Output:

AI is the development of systems that can perform tasks requiring human intelligence.
  • Example: Your chatbot needs to standardize user inputs to uppercase for consistency. The to_uppercase function ensures “what is ai?” becomes “WHAT IS AI?” before the LLM sees it, keeping things uniform.

This method is quick and clean, ideal for small, focused tweaks.

Custom Tools for Agents: Powering Dynamic Decisions

For more complex tasks—like hitting an external API or running multi-step logic—you can turn your Python function into a custom tool that an agent can use. Agents decide when to call the tool based on the user’s query, making this approach great for dynamic, decision-driven apps.

Here’s the deal:

  • Purpose: Creates a tool from your function that an agent can call when needed, handling inputs and returning results.
  • Best For: API calls for web research, calculations in e-commerce assistants, or data processing in customer support bots.
  • Mechanics: Decorates your function with @tool or defines it as a Tool object, which the agent uses, often paired with an output parser for structured results.
  • Setup: Define the function, make it a tool, and give it to an agent. Here’s an example of a discount calculator tool:
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langchain.agents import initialize_agent, AgentType

# Custom discount calculator tool
@tool
def calculate_discount(price: float, discount_rate: float) -> float:
    """Calculate the discounted price given the original price and discount rate."""
    return price * (1 - discount_rate / 100)

# Set up the agent
llm = ChatOpenAI(model="gpt-4o-mini")
agent = initialize_agent(
    tools=[calculate_discount],
    llm=llm,
    agent_type=AgentType.REACT
)

# Run the agent
result = agent.run("Calculate the price of a $100 item with a 20% discount.")
print(result)

Output:

The discounted price is $80.
  • Example: An e-commerce assistant uses the calculate_discount tool to compute sale prices, ensuring accurate math the LLM can’t do alone.

This approach is perfect for apps where agents need to make smart choices about when to use your function.

Transform Steps in Chains: Multi-Step Magic

If your workflow involves multiple steps—like cleaning data before feeding it to an LLM—you can integrate Python functions as transform steps in a chain using RunnableLambda or custom runnables. This lets you shape data as it flows through the chain.

Here’s how it goes:

  • Purpose: Adds a custom transformation step in a chain, modifying data between components like prompts and LLMs.
  • Best For: Cleaning text in RAG apps, formatting outputs in SQL query generators, or preprocessing inputs in document QA chains.
  • Mechanics: The function processes the chain’s input or intermediate data, passing the result to the next step, often with an output parser.
  • Setup: Wrap the function in RunnableLambda and slot it into the chain. Here’s an example that cleans text by removing special characters:
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate

# Function to clean text
def clean_text(inputs):
    text = inputs["text"]
    cleaned = "".join(c for c in text if c.isalnum() or c.isspace()).lower()
    return {"text": cleaned}

# Chain with cleaning step
prompt = PromptTemplate(input_variables=["text"], template="Summarize: {text}")
llm = ChatOpenAI(model="gpt-4o-mini")
chain = RunnableLambda(clean_text) | prompt | llm

# Test the chain
result = chain.invoke({"text": "AI!!! is @awesome..."})
print(result.content)

Output:

AI is awesome.
  • Example: A document QA chain cleans noisy text from a PDF before summarizing, ensuring the LLM gets clean, usable input.

This method is excellent for chains needing tailored data processing.

Callbacks: Sneaky Function Integration

You can also slip Python functions into LangChain via callbacks, which run during specific events like chain start or LLM completion. This is handy for logging, tweaking outputs, or triggering actions without changing the main workflow.

Here’s the lowdown:

  • Purpose: Executes functions at workflow events, adding logic like logging or output modification.
  • Best For: Logging metrics in customer support bots, adjusting outputs in chatbots, or sending alerts in data analysis agents.
  • Mechanics: Defines a callback handler with functions for events, called by LangChain with event data.
  • Setup: Create a custom callback handler and attach it. Here’s an example logging LLM outputs:
from langchain_core.callbacks import BaseCallbackHandler
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate

class OutputLogger(BaseCallbackHandler):
    def on_llm_end(self, response, **kwargs):
        print(f"LLM Output: {response.generations[0][0].text}")

# Simple chain with callback
prompt = PromptTemplate(input_variables=["query"], template="Answer: {query}")
llm = ChatOpenAI(model="gpt-4o-mini")
chain = prompt | llm

# Run with callback
result = chain.invoke({"query": "What is AI?"}, config={"callbacks": [OutputLogger()]})
print(result.content)

Output:

LLM Output: AI is the development of systems...
AI is the development of systems...
  • Example: A chatbot logs every LLM response to check output quality, helping you catch issues during development.

This approach is subtle but powerful for adding logic without disrupting the workflow.

Let’s Build: A Discount Calculator with a Python Function

To bring this to life, let’s create a system that calculates discounts based on user input, using a custom Python function as a tool in a ReAct Agent. We’ll include a prompt template and output parser to return structured JSON.

Get Your Environment Set Up

Follow Environment Setup to prepare your system. Install the necessary packages:

pip install langchain langchain-openai

Securely set your OpenAI API key, as outlined in security and API key management.

Write the Discount Function

Here’s a simple function to calculate discounts:

def calculate_discount(price: float, discount_rate: float) -> float:
    """Calculate the discounted price given the original price and discount rate."""
    return price * (1 - discount_rate / 100)

Turn It Into a Tool

Use the @tool decorator to make it agent-ready:

from langchain.tools import tool

@tool
def calculate_discount(price: float, discount_rate: float) -> float:
    """Calculate the discounted price given the original price and discount rate."""
    return price * (1 - discount_rate / 100)

Craft a Prompt Template

Define a Prompt Template to guide the agent:

from langchain_core.prompts import PromptTemplate

prompt = PromptTemplate(
    template="Query: {query}\nUse the discount tool if needed. Respond in JSON format with the discounted price.",
    input_variables=["query"]
)

Set Up an Output Parser

Create an Output Parser for structured JSON:

from langchain_core.output_parsers import StructuredOutputParser, ResponseSchema

schemas = [
    ResponseSchema(name="answer", description="The discounted price", type="float")
]
parser = StructuredOutputParser.from_response_schemas(schemas)

Build the ReAct Agent

Combine everything into a ReAct Agent that decides when to use the discount tool:

from langchain_openai import ChatOpenAI
from langchain.agents import initialize_agent, AgentType

prompt = PromptTemplate(
    template="Query: {query}\nUse the discount tool if needed. Respond in JSON format with the discounted price.\n{format_instructions}",
    input_variables=["query"],
    partial_variables={"format_instructions": parser.get_format_instructions()}
)

llm = ChatOpenAI(model="gpt-4o-mini")
agent = initialize_agent(
    tools=[calculate_discount],
    llm=llm,
    agent_type=AgentType.REACT,
    prompt=prompt,
    output_parser=parser
)

Test the Agent

Run a query to see the tool in action:

result = agent.run({"query": "Calculate the price of a $100 item with a 20% discount."})
print(result)

Output:

{'answer': 80.0}

Debug and Improve

If the output isn’t right—say, the discount is wrong or the JSON is off—use LangSmith for prompt debugging or visualizing evaluations. Tweak the prompt with few-shot prompting for clarity:

prompt = PromptTemplate(
    template="Query: {query}\nExamples:\nQuery: $50 item, 10% discount -> {'answer': 45.0}\nUse the discount tool if needed. Respond in JSON format with the discounted price.\n{format_instructions}",
    input_variables=["query"],
    partial_variables={"format_instructions": parser.get_format_instructions()}
)

If issues persist, check troubleshooting. To take it further, add memory for conversational flows or deploy as a Flask API for web access.

Tips to Make Your Functions Shine in LangChain

Here’s how to get the most out of integrating Python functions:

These tips help you build robust, efficient apps that align with enterprise-ready applications and workflow design patterns.

Keep Building with Python Functions

Want to take your LangChain skills further? Here are some next steps:

Wrapping It Up: Python Functions Unlock Endless Possibilities

Integrating Python functions into LangChain is like adding a turbo boost to your AI apps, letting you inject custom logic that LLMs alone can’t handle. Whether you’re using RunnableLambda for quick tweaks, custom tools for agent smarts, transform steps in chains, or callbacks for event-driven actions, functions make your app do exactly what you envision. From crunching numbers in an e-commerce assistant to cleaning data in a RAG app, they’re your ticket to building something truly unique.

Start with the discount calculator example, dive into tutorials like Build a Chatbot or Create RAG App, and share your creations with the AI Developer Community or on X with #LangChainTutorial. For more, hit up the LangChain Documentation and keep building cool stuff!