LangChain Custom Tools

LangChain Custom tools are classes that can be used to interact with the world.

Imagine SDK can leverage the tools available on LangChain Community which already contains many third-party integrations.

Perform an Internet search and summarize the results

The following example performs an internet search with the DuckDuckGo search engine and then summarizes the results:

from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.globals import set_debug
from langchain_core.prompts.prompt import PromptTemplate

from imagine.langchain import ImagineChat


set_debug(True)

search = DuckDuckGoSearchRun()

title_input = "What are the olympic medals tally for china and USA in 2024"


title_template = PromptTemplate(
    input_variables=["subject"],
    template="I want to know about following - {subject}. Give me a question i can search on google",
)


outline_template = PromptTemplate(
    input_variables=["DuckDuckGo_Search"],
    template="Can you give me a one line short summary for this search result - {DuckDuckGo_Search}",
)

llm = ImagineChat(model="Llama-3-8B", max_tokens=500, temperature=0.0)


title_chain = title_template | llm
outline_chain = outline_template | llm
title = title_chain.invoke(title_input)
search_result = search.run(title_input)
outline = outline_chain.invoke(input={"DuckDuckGo_Search": search_result})

print(f"Duck Duck Go search summary: {outline}")

Count the rows of a CSV file with a custom tool

The following example implements a custom tool to count the number of rows in a CSV file:

import re

import pandas as pd

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts.chat import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough

from imagine.langchain import ImagineChat


# Define the LLM and prompt template
llm = ImagineChat(max_tokens=1000)
prompt = ChatPromptTemplate.from_template("""You are a data analysis expert.
Generate Python code in one single chunk in ```python``` syntax such that it saves the
output as a string variable `result`. assume that dataframe is already loaded as df
variable. Write this code to perform the following analysis on a Pandas DataFrame:

{task_description}
The dataframe looks like -
```{data_head}```
""")


def extract_code(text):
    # Regular expression pattern to match code blocks
    pattern = r"\`\`\`[\w\s]+?\n([\s\S]*?)\n\`\`\`"

    # Use re.DOTALL flag to allow . to match newline characters
    match = re.search(pattern, text, re.DOTALL)

    if match:
        # Remove the language specification and triple backticks from the matched code
        code = match.group(1).strip()

        return code  # Return the cleaned-up code
    else:
        return None  # No code block found


# Define the tool's functionality
def execute_code(context: dict) -> str:
    code = extract_code(context.get("code"))

    dataframe = context.get("dataframe")
    try:
        exec_globals = {"df": dataframe["dataframe"]}
        exec(code, exec_globals)
        return str(
            exec_globals.get("result", "Execution completed without returning a result")
        )
    except Exception as e:
        return f"Error during execution: {str(e)}"


# Modify the runnable to handle the dataframe
code_generation_chain = (
    {"task_description": RunnablePassthrough(), "data_head": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)
execution_chain = RunnableLambda(execute_code)
# Combine chains to pass both code and dataframe to execute_code
final_chain = {
    "code": code_generation_chain,
    "dataframe": RunnablePassthrough(),
} | execution_chain


# Example usage
task_description = "How many rows are there?"
your_dataframe = pd.read_csv("../assets/file.csv")
result = final_chain.invoke(
    {
        "task_description": task_description,
        "dataframe": your_dataframe,
        "data_head": your_dataframe.head(2),
    }
)
print(result)

Using GitHub with a custom tool

The following example implements a custom tool to interact with GitHub. The tools used are defined here

import os
import re

from github_api_request import fetch_issue_details, fetch_issue_labels
from langchain_core.prompts.prompt import PromptTemplate

from imagine.langchain import ImagineChat
from langchain.tools import BaseTool


# Initialize the language model
llm = ImagineChat(
    model="Llama-3-8B",
    max_tokens=500,
    temperature=0.0,
)

os.environ["GITHUB_TOKEN"] = "<your_github_token>"
os.environ["REPO"] = "PyGithub/PyGithub"  # for example
issue_number = 2892  # Replace with the actual issue number


def extract_query(text):
    """
    Extracts the query from the given text using a regular expression pattern.
    """
    pattern = r"```([\s\S]*?)```"
    match = re.search(pattern, text, re.DOTALL)

    if match:
        query = match.group(1).strip()
        return query
    else:
        return None  # No query found


# Define tool classes
class GitHubIssueDetailsTool(BaseTool):
    name = "github_issue_details"
    description = "Fetches details of a GitHub issue"

    def _run(self, issue_number: int) -> str:
        return fetch_issue_details(issue_number, print_output=False)


class GitHubLabelsTool(BaseTool):
    name = "github_labels"
    description = "Fetches labels from a GitHub repository"

    def _run(self) -> str:
        labels = fetch_issue_labels()
        return ", ".join(labels)


# Instantiate the tools
github_issue_details_tool = GitHubIssueDetailsTool()
github_labels_tool = GitHubLabelsTool()

# Define prompt templates
title_template = PromptTemplate(
    input_variables=["issue"],
    template="Please summarize the GitHub issue \n{issue}. Do a small paragraph to explain the goal of the issue then discuss of the changes made in the commits",
)

outline_template = PromptTemplate(
    input_variables=["issue_summary", "labels"],
    template="Here is the summary of an issue - \n{issue_summary} suggest between 1 and 3 appropriate labels from this list - \n{labels}",
)

# Create chains for title and outline
title_chain = title_template | llm
outline_chain = outline_template | llm

# search for the issue
issue = github_issue_details_tool.run({"issue_number": issue_number})
print(" issue ".center(100, "="))
print(issue)

# Generate the title
title = title_chain.invoke(input={"issue": issue})
print(" title ".center(100, "="))
print(title)

# Fetch issue labels
labels = github_labels_tool.run({})
print(" labels ".center(100, "="))
print(labels)

# Generate the outline
outline = outline_chain.invoke(input={"issue_summary": title, "labels": labels})
print(" outline ".center(100, "="))
print(outline.content)

Using Jira with a custom tool

The following example implements a custom tool to interact with JIRA using this Jira lib

import getpass
import os
import re

import urllib3

from jira import JIRA
from langchain_core.prompts.prompt import PromptTemplate

from imagine.langchain import ImagineChat
from langchain.tools import BaseTool


urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

# Set the variables below to be able to use Jira
username = getpass.getuser()
jira_server = "jira_server_endpoint"
pwd = "your_password"

# Initialize the language model
llm = ImagineChat(
    model="Llama-3-8B",
    max_tokens=500,
    temperature=0.0,
)

# Set environment variables for Jira authentication
os.environ["JIRA_USER"] = username
os.environ["JIRA_PASSWORD"] = pwd


# Basic class that extends langchain BaseTool
class JiraSearchTool(BaseTool):
    name = "jira_search"
    description = "Searches Jira issues"

    def _run(self, query: str, server: str, verify: str = True) -> str:
        user = os.environ["JIRA_USER"]
        pwd = os.environ["JIRA_PASSWORD"]
        jira = JIRA(server=server, basic_auth=(user, pwd), options={"verify": verify})
        q = query + " ORDER BY createdDate DESC"
        return jira.search_issues(q, maxResults=0)


# Instantiate the JiraSearchTool
jira_search_tool = JiraSearchTool(server=jira_server, user=username, pwd=pwd)


def extract_query(text):
    """
    Extracts the query from the given text using a regular expression pattern.
    """
    pattern = r"```([\s\S]*?)```"
    match = re.search(pattern, text, re.DOTALL)

    if match:
        query = match.group(1).strip()
        return query
    else:
        return None  # No query found


# Define prompt templates
title_template = PromptTemplate(
    input_variables=["username"],
    template="Please give a jira query for the jira issues assigned to {username} created in the last 180 days. put the query as in answer between ``` like this: ```query```",
)

outline_template = PromptTemplate(
    input_variables=["Jira_Search"],
    template="Can you give me a short summary for this search result - {Jira_Search}",
)

# Create chains for title and outline
title_chain = title_template | llm
outline_chain = outline_template | llm

# Generate the title
title = title_chain.invoke(input={"username": username})
print(" title ".center(100, "="))
print(title)

# Extract the query from the title
query = extract_query(title.content)
print(" query ".center(100, "="))
print(query)

# Run the Jira search tool with the extracted query
search_result = jira_search_tool.run({"query": query, "server": jira_server})


def create_issue_dict(issue):
    """
    Creates a dictionary representation of a Jira issue.
    """
    return {
        "issue": issue,
        "jira_id": issue.key,
        "summary": issue.fields.summary,
        "description": issue.fields.description,
        "status": issue.fields.status,
        "resolution": issue.fields.resolution
        if hasattr(issue.fields, "resolution")
        else None,
        "reporter": issue.fields.reporter.displayName,
        "assigned_to": issue.fields.assignee.displayName
        if issue.fields.assignee
        else None,
        "reporter_username": issue.fields.reporter.name
        if hasattr(issue.fields.reporter, "name")
        else None,
        "assignee_username": (
            issue.fields.assignee.name
            if issue.fields.assignee and hasattr(issue.fields.assignee, "name")
            else None
        ),
        "created_date": issue.fields.created,
        "updated_date": issue.fields.updated,
        "components": [component.name for component in issue.fields.components],
    }


# Convert search results to a list of dictionaries
dict_result = [create_issue_dict(issue) for issue in search_result]

# Generate the outline
outline = outline_chain.invoke(input={"Jira_Search": dict_result})
print(" summary ".center(100, "="))
print(outline.content)

Using tools with Agents

Check out how to use tools with crewAI agents.