Writer AI module

This module leverages the Writer Python SDK to enable applications to interact with large language models (LLMs) in chat or text completion formats. It provides tools to manage conversation states and to dynamically interact with LLMs using both synchronous and asynchronous methods.

​

Getting your API key

To utilize the Writer AI module, you’ll need to configure the WRITER_API_KEY environment variable with an API key obtained from AI Studio. Here is a detailed guide to setup up this key. You will need to select an API app under Developer tools

Once you have your API key, set it as an environment variable on your system:

export WRITER_API_KEY=your_api_key_here

You can manage your environment variables using methods that best suit your setup, such as employing tools like python-dotenv.

Furthermore, when deploying an application with writer deploy, the WRITER_API_KEY environment variable is automatically configured with the API key specified during the deployment process.

​

Chat completion with the Conversation class

The Conversation class manages LLM communications within a chat framework, storing the conversation history and handling the interactions.

import writer as wf
import writer.ai

def handle_simple_message(state, payload):
    # Update the conversation state by appending the incoming user message.
    state["conversation"] += payload
    
    # Stream the complete response from the AI model in chunks.
    for chunk in state["conversation"].stream_complete():
        # Append each chunk of the model's response to the ongoing conversation state.
        state["conversation"] += chunk

# Initialize the application state with a new Conversation object.
initial_state = wf.init_state({
    "conversation": writer.ai.Conversation(),
})

​

Initializing a conversation

A Conversation can be initialized with either a system prompt or a list of previous messages. It can also accept a default configuration dictionary that sets parameters for all interactions.

# Initialize with a system prompt for a Financial Analyst specializing in balance sheets
conversation = Conversation("You assist clients with analyzing and understanding their balance sheets")

# Initialize with a history of messages related to balance sheet queries
history = [
    {"role": "user", "content": "Can you explain the liabilities section?"},
    {"role": "assistant", "content": "Certainly! Liabilities are legally binding obligations payable to another entity."}
]

conversation = Conversation(history)

# Initialize with a configuration suitable for financial analysis discussions
config = {'max_tokens': 200, 'temperature': 0.5}
conversation = Conversation("You provide detailed insights into balance sheet components", config=config)

​

Adding messages to conversation

Messages can be added to a Conversation instance using the + operator or the add method.

# Using the `+` operator to add a balance sheet-related query
conversation += {"role": "user", "content": "Can you break down the assets section of the balance sheet?"}

# Using the `add` method to add a balance sheet-related query
conversation.add(role="user", content="How should I interpret the equity section?")

​

Completing and streaming Conversations

The complete and stream_complete methods facilitate interaction with the LLM based on the accumulated messages and configuration. These methods execute calls to generate responses and return them in the form of a message object, but do not alter the conversation’s messages list, allowing you to validate or modify the output before deciding to add it to the history.

# Using `complete` to get a single response
response = conversation.complete()
print("LLM Response:", response)

Instance-wide configuration parameters can be complemented or overriden on individual call’s level, if a config dictionary is provided to the method:

# Overriding configuration for a specific call
response = conversation.complete(config={'max_tokens': 200, 'temperature': 0.5})

​

Using Graphs with Conversation

A Graph is a collection of files meant to provide their contents to the LLM during conversations. Framework allows you to create, retrieve, update, and delete graphs, as well as manage the files within them.

​

Creating and Managing Graphs

To create and manipulate graphs, use the following methods:

from writer.ai import create_graph, retrieve_graph, list_graphs, delete_graph

# Create a new graph
graph = create_graph(name="Financial Data", description="Quarterly reports")

# Retrieve an existing graph by ID
graph = retrieve_graph("d90a632b-5c1f-42b8-8748-5b7f769d9a36")

# Update a graph
graph.update(name="Updated Financial Data", description="Updated description")

# Retrieve a list of created graphs
graphs = list_graphs()
for graph in graphs:
    # Delete a graph
    delete_graph(graph)

​

Adding and Removing Files from Graphs

You can upload files, associate them with graphs, and download or remove them.

from writer.ai import upload_file

# Upload a file
file = upload_file(data=b"file content", type="application/pdf", name="Report.pdf")

# Add the file to a graph
graph.add_file(file)

# Remove the file from the graph
graph.remove_file(file)

​

Applying Graphs to Conversation completion

You can utilize graphs within conversations. For instance, you may want to provide the LLM access to a collection of files during an ongoing conversation to query or analyze the file content. When passing a graph to the conversation, the LLM can query the graph to retrieve relevant data.

# Retrieve a graph
graph = retrieve_graph("d90a632b-5c1f-42b8-8748-5b7f769d9a36")

# Pass the graph to the conversation for completion
response = conversation.complete(tools=graph)

Alternatively, you can define a graph using JSON:

tool = {
    "type": "graph",
    "graph_ids": ["d90a632b-5c1f-42b8-8748-5b7f769d9a36"]
}

response = conversation.complete(tools=tool)

​

Using Function Calls with Conversations

Framework allows you to register Python functions that can be called automatically during conversations. When the LLM determines a need for specific information or processing, it issues a request to use the local code (your function), and Framework handles that request automatically.

​

Defining Function Tools

Function tools are defined using either a Python class or a JSON configuration.

from writer.ai import create_function_tool

# Define a function tool with Python callable
def calculate_interest(principal: float, rate: float, time: float):
    return principal * rate * time

tool = create_function_tool(
    name="calculate_interest",
    callable=calculate_interest,
    parameters={
        "principal": {"type": "float", "description": "Loan principal"},
        "rate": {"type": "float", "description": "Interest rate"},
        "time": {"type": "float", "description": "Time in years"}
    }
)

response = conversation.complete(tools=tool)

Alternatively, you can define a function tool in JSON format, but the callable function must still be passed:

tool = {
    "type": "function",
    "name": "calculate_interest",
    "callable": calculate_interest,
    "parameters": {
        "principal": {"type": "float", "description": "Loan principal"},
        "rate": {"type": "float", "description": "Interest rate"},
        "time": {"type": "float", "description": "Time in years"}
    }
}

response = conversation.complete(tools=tool)

Function tools require the following properties:

• name: str: A string that defines how the function is referenced by the LLM. It should describe the function’s purpose.
• callable: Callable: A Python function that will be called automatically when needed by the LLM.
• parameters: dict: A dictionary that specifies what input the function expects. The keys should match the function’s parameter names, and each parameter should have a type, and an optional description.
  Supported types are: string, number, integer, float, boolean, array, object and null.

​

Automated Function Calling

When a conversation involves a tool (either a graph or a function), Framework automatically handles the requests from LLM to use the tools during interactions. If the tool needs multiple steps (for example, querying data and processing it), Framework will handle those steps recursively, calling functions as needed until the final result is returned.

By default, to prevent endless recursion, Framework will only handle 5 consecutive tool calls. You can expand it in case it doesn’t suit your case – both complete() and stream_complete() accept a max_tool_depth parameter, which configures the maximum allowed recursion depth:

response = conversation.complete(tools=tool, max_tool_depth=7)

​

Providing a Tool or a List of Tools

You can pass either a single tool or a list of tools to the complete() or stream_complete() methods. The tools can be a combination of FunctionTool, Graph, or JSON-defined tools.

from writer.ai import create_function_tool, retrieve_graph

# Define a function tool
tool1 = create_function_tool(
    name="get_data",
    callable=lambda x: f"Data for {x}",
    parameters={"x": {"type": "string", "description": "Input value"}}
)

# Retrieve a graph
graph = retrieve_graph("d90a632b-5c1f-42b8-8748-5b7f769d9a36")

# Provide both tools in a list
response = conversation.complete(tools=[tool1, graph])

​

Text generation without a conversation state

​

Text generation against a string prompt

complete and stream_complete methods are designed for one-off text generation without the need to manage a conversation state. They return the model’s response as a string. Each function accepts a config dictionary allowing call-specific configurations.

# Using `complete` for a single completion
text_response = complete("Explore the benefits of AI.", config={'temperature': 0.3})
print("Completion:", text_response)

​

Text generation against graphs

The ask and stream_ask methods allow you to query one or more graphs to generate responses from the information stored within them.

​

Two approaches to questioning graphs

There are two ways to query graphs, depending on your needs:

1. Graph-Level Methods (Graph.ask, Graph.stream_ask): Used when working with a single graph instance. These methods are tied directly to the Graph object, encapsulating operations within that instance.
2. Module-Level Methods (writer.ai.ask, writer.ai.stream_ask): Designed for querying multiple graphs simultaneously. These methods operate on a broader scale, allowing mixed inputs of graph objects and IDs.

​

Parameters

Both methods include: • question: str: The main query for the LLM. • Optional subqueries: bool (default: False): Allows the LLM to generate additional questions during response preparation for more detailed answers. Enabling this might increase response time.

Method-level methods require: • graphs_or_graph_ids: list[Graph | str]: A list of graphs to use for the question. You can pass Graph objects directly into the list, use graph IDs in string form, or a mix of both.

​

Graph-level methods

The graph-level methods, Graph.ask and Graph.stream_ask, are designed for interacting with a single graph. By calling these methods on a specific Graph instance, you can easily pose questions and retrieve answers tailored to that graph’s content.

# Retrieve a specific graph
graph = retrieve_graph("f47ac10b-58cc-4372-a567-0e02b2c3d479")

# Pose a question to the graph and get a complete response
response = graph.ask("What are the benefits of renewable energy?")
print(response)

​

Module-level methods

The module-level methods, writer.ai.ask and writer.ai.stream_ask, are designed for querying multiple graphs simultaneously. They are useful when you need to aggregate or compare data across multiple graphs.

from writer.ai import ask

# Pose a question to multiple graphs
response = ask(
    question="What are the latest advancements in AI?",
    graphs_or_graph_ids=[
        "550e8400-e29b-41d4-a716-446655440000",
        "123e4567-e89b-12d3-a456-426614174000"
        ]
)
print(response)

​

Using the Tools class

The writer.ai.tools instance provides access to Writer SDK tools resources, such as text splitting, medical content comprehension, and PDF parsing. Below is a guide on how to use each method.

​

Splitting Content

The split method divides text into chunks based on a selected strategy.

from writer.ai import tools

content = \
    """
    This is a long piece of text that needs to be split into smaller parts.
    Lorem ipsum dolor sit amet...
    """
chunks = tools.split(content, strategy="llm_split")
print(chunks)

Parameters:

• content (str): The text to be split.
• strategy (str): The splitting strategy (llm_split, fast_split, or hybrid_split).

Returns: A list of text chunks.

​

Medical Content Comprehension

The comprehend_medical method processes medical text and extracts relevant entities based on a specified response type.

from writer.ai import tools

medical_text = "Patient shows symptoms of hypertension and diabetes."
entities = tools.comprehend_medical(medical_text, response_type="Entities")
print(entities)

Parameters:

• content (str): The medical text to process.
• response_type (str): The type of medical response (Entities, RxNorm, ICD-10-CM, or SNOMED CT).

Returns: A list of extracted medical entities.

​

PDF Parsing

The parse_pdf method extracts text content from a PDF file. The file can be referenced by its ID, or provided as a File object.

from writer.ai import tools

file_id = "example-file-id"
parsed_content = tools.parse_pdf(file_id, format="text")
print(parsed_content)

Parameters:

• file_id_or_file (str or File): The file to parse (by ID or as an object).
• format (str): The format of the extracted content (text or markdown).

Returns: The text content of the PDF.