Using Writer with LangChain - Writer AI Studio

Writer LangChain Integration

The Writer LangChain integration allows you to leverage Writer’s capabilities within the LangChain ecosystem, making it easy to build sophisticated AI applications. In this tutorial, you’ll explore each component of the integration and understand how they work.

Prerequisites

Before you begin, make sure you have:

Installation and setup

First, install the necessary packages:

pip install langchain-writer python-dotenv

Next, create a .env file with WRITER_API_KEY set to your Writer API key:

WRITER_API_KEY=<your-api-key>

Components of the Writer LangChain integration

The langchain-writer package provides several key components:

  1. ChatWriter for text generation
  2. Tool calling capabilities, including:
    • GraphTool for Knowledge Graph integration
    • NoCodeAppTool for no-code applications
    • LLMTool for model delegation
  3. Additional tools like PDFParser for parsing PDFs

ChatWriter

ChatWriter is a LangChain chat model that provides access to Writer’s AI capabilities for text generation. It supports streaming, non-streaming, batching, and asynchronous operations. You can use any of the Palmyra chat models available in AI Studio. See the full documentation for ChatWriter to learn more about the available parameters.

Usage

This example uses ChatWriter to ask the Palmyra X5 model to explain what LangChain is in plain terms.

from langchain_writer import ChatWriter
from dotenv import load_dotenv

load_dotenv()

# Initialize the chat model
# These are optional parameters with default values listed here
chat = ChatWriter(
    model="palmyra-x5",  # default model
    temperature=0.7,        # controls randomness (0-1)
    max_tokens=None,        # maximum number of tokens to generate
    timeout=None,           # request timeout
    max_retries=2          # number of retries on failure
)

# Generate a response
response = chat.invoke("Explain what LangChain is in simple terms.")
print(response.content)

Streaming

Streaming allows you to receive the generated text in chunks as it’s being produced. This example shows how to stream a response from the Palmyra X5 model using synchronous streaming:

from langchain_writer import ChatWriter
import asyncio
from dotenv import load_dotenv

load_dotenv()

chat = ChatWriter()

# Streaming (synchronous)
for chunk in chat.stream("Write a short poem about artificial intelligence."):
    print(chunk.content, end="")

You can also use asynchronous streaming with the async for loop and the astream method:

from langchain_writer import ChatWriter
import asyncio
from dotenv import load_dotenv

load_dotenv()

chat = ChatWriter()

# Streaming (asynchronous)
async def async_stream():
    async for chunk in chat.astream("Write a short poem about artificial intelligence."):
        print(chunk.content, end="")

asyncio.run(async_stream())

Batch processing

You can batch process multiple prompts for efficient processing. The following example batches three individual LLM invocations and runs them in parallel:

from langchain_writer import ChatWriter
from dotenv import load_dotenv

load_dotenv()

chat = ChatWriter()

questions = [
    "What is a three-sentence definition of retrieval-augmented generation?",
    "When did transformer architecture first appear in the literature?",
    "What are 3 benefits of using AI to assist in software development?"
]

# Process multiple prompts in parallel
responses = chat.batch(questions)

for question, response in zip(questions, responses):
    print(f"Q: {question}")
    print(f"A: {response.content}\n")

Note that batch returns results in the same order as the inputs. You can use batch_as_completed to return results as they complete. Results may arrive out of order, but each includes the input index for matching. You can also optionally set the max_concurrency parameter to control the number of concurrent requests.

responses = chat.batch(questions,
    config={"max_concurrency": 2}
)

See the LangChain documentation on parallel execution for more information.

Tool calling

ChatWriter supports tool calling, which allows the model to use external functions to enhance its capabilities. Tool calling is available with Palmyra X4 and later.

Tool calling basics

To use tool calling, follow these steps:

  1. Define a function that will be called by the model and decorate it with the @tool decorator.
  2. Bind the tool to the chat model using the bind_tools method.
  3. Use the tool in a chat and append the response to the messages list.
  4. Execute the tool call with the arguments given by the model and append the response to the messages list.
  5. Invoke the chat model with the updated messages list to receive the final response.

Here’s an example of how to use tool calling:

from langchain_writer import ChatWriter
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage
from dotenv import load_dotenv

load_dotenv()

@tool
def get_weather(location: str) -> str:
    """Get the current weather for a location."""
    return f"The weather in {location} is sunny and 75°F"

chat = ChatWriter()
chat_with_tools = chat.bind_tools([get_weather])

messages = [
    HumanMessage(
        "What's the weather like in San Francisco?"
    )
]

response = chat_with_tools.invoke(messages)
messages.append(response)

for tool_call in response.tool_calls:
    selected_tool = {
        "get_weather": get_weather,
    }[tool_call["name"].lower()]
    tool_msg = selected_tool.invoke(tool_call)
    messages.append(tool_msg)

response = chat_with_tools.invoke(messages)
print(response.content)

GraphTool

GraphTool is a LangChain tool that allows you to retrieve information from a Knowledge Graph to enhance its responses. For more details on the built-in Knowledge Graph chat tool in Writer, see the Knowledge Graph chat support guide.

Usage

from langchain_writer import ChatWriter
from langchain_writer.tools import GraphTool
from dotenv import load_dotenv

load_dotenv()

# Initialize the chat model
chat = ChatWriter()

# Create a graph tool with your knowledge graph ID
graph_tool = GraphTool(graph_ids=["your-knowledge-graph-id"])

# Bind the tool to the chat model
chat_with_tools = chat.bind_tools([graph_tool])

# Ask a question that can be answered using the knowledge graph
response = chat_with_tools.invoke("What information do you have about product X?")
print(response.content)

NoCodeAppTool

NoCodeAppTool is a specialized tool that enables access to Writer’s no-code applications as LLM tools. Here’s an example of how to use it:

from langchain_writer import ChatWriter
from langchain_writer.tools import NoCodeAppTool
from langchain_core.messages import HumanMessage
import os
from dotenv import load_dotenv

load_dotenv()

chat = ChatWriter()

# Create a NoCodeAppTool
app_tool = NoCodeAppTool(
    app_id=os.getenv("APP_ID"),
    name="Social post generator",
    description="No-code app that generates social posts from product descriptions"
)

# Bind the tool ChatWriter
chat_with_tools = chat.bind_tools([app_tool])

product_description = "The Terra running shoe is a high-performance, lightweight shoe that offers a comfortable and durable experience."

# Create a conversation
messages = [
    HumanMessage("Can you help me generate a social post about a new running shoe? Here is the product description:\n\n" + product_description)
]

# Get the model's response
response = chat_with_tools.invoke(messages)
messages.append(response)

if response.tool_calls:
    for tool_call in response.tool_calls:
        inputs = {}
        for arg in tool_call['args']:
            inputs[arg] = tool_call['args'][arg]

try:
            tool_response = app_tool.run(tool_input={"inputs": inputs})
            messages.append(tool_response.suggestion)
        except Exception as e:
            print(f"Error running tool: {e}")
    try:
        final_response = chat_with_tools.invoke(messages)
        print("\nFinal response:", final_response.content)
    except Exception as e:
        print(f"Error getting final response: {e}")
else:
    print("No tool calls were requested.")

LLMTool

LLMTool is a specialized tool that enables delegation to another model. Here’s an example of how to use the LLMTool:

from langchain_writer import ChatWriter
from langchain_writer.tools import LLMTool
from dotenv import load_dotenv

load_dotenv()

chat = ChatWriter()

# Create an LLMTool that delegates coding tasks to a model specialized in code generation
llm_tool = LLMTool(
    model_name="anthropic.claude-3-sonnet-20240229-v1:0",
    description="A function that invokes a model specialized in writing and debugging code."
)

# Bind the tool to the ChatWriter
chat_with_tools = chat.bind_tools([llm_tool])

# The primary model delegates the coding request to the specialized model
response = chat_with_tools.invoke([
    ("system", "You are a helpful assistant. Delegate coding tasks to the code generation tool."),
    ("human", "Write a Python function that returns the nth Fibonacci number using memoization.")
])

print(response.content)

Additional tools

PDFParser

PDFParser is a document loader that uses Writer’s PDF parsing capabilities to extract text from PDF documents. Here’s an example of how to use the PDFParser:

from langchain_writer import PDFParser
from langchain_core.documents.base import Blob
from dotenv import load_dotenv

load_dotenv()

# Initialize the PDF parser with the desired output format
parser = PDFParser(output_format="markdown")  # Options: "text", "markdown"

# Load a PDF file
file = Blob.from_path("path/to/your/document.pdf")

parsed_pages = parser.parse(blob=file)
print(parsed_pages)

Output formats

The PDFParser supports different output formats:

Conclusion

In this tutorial, you’ve explored the LangChain integration with Writer, covering each of its components:

This integration provides a foundation for building AI applications with Writer and LangChain. Check out the package README and documentation, as well as the LangChain Documentation for more information on how to use LangChain with Writer.