# Source: https://dev.writer.com/framework/ai-module.md

# Writer AI module

This module leverages the [Writer Python SDK](https://pypi.org/project/writer-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](/home/quickstart) 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:

<CodeGroup>
  ```bash For macOS and Linux theme={null}
  export WRITER_API_KEY=your_api_key_here
  ```

  ```bash For Windows theme={null}
  set WRITER_API_KEY=your_api_key_here
  ```
</CodeGroup>

You can manage your environment variables using methods that best suit your setup, such as employing tools like [python-dotenv](https://pypi.org/project/python-dotenv/).

## Chat completion with the Conversation class

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

```python  theme={null}
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.

```python  theme={null}
# 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.

```python  theme={null}
# 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.

<CodeGroup>
  ```python complete theme={null}
  # Using `complete` to get a single response
  response = conversation.complete()
  print("LLM Response:", response)
  ```

  ```python stream_complete theme={null}
  # Using `stream_complete` to get streamed responses
  for chunk in conversation.stream_complete():
      print("Streamed Message:", chunk)
      # Manually adding to the conversation
      conversation += chunk
  ```
</CodeGroup>

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

```python  theme={null}
# 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:

```python  theme={null}
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.

```python  theme={null}
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.

```python  theme={null}
# 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:

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

response = conversation.complete(tools=tool)
```

### Using Function Calls with Conversations

<Warning>
  Function tools are only available with `palmyra-x4` and `palmyra-x5` models.
</Warning>

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.

```python  theme={null}
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:

```python  theme={null}
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:

```python  theme={null}
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.

```python  theme={null}
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.

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

  ```python stream_complete theme={null}
  # Using `stream_complete` for streamed text completions
  for text_chunk in stream_complete("Explore the benefits of AI.", config={'temperature': 0.3}):
      print("Streamed Text:", text_chunk)
  ```
</CodeGroup>

### 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.

<Note>
  •	Use graph-level methods when working with a single graph instance.
  •	Use module-level methods when querying multiple graphs or when graph IDs are your primary input.
</Note>

#### 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.

<CodeGroup>
  ```python ask theme={null}
  # 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)
  ```

  ```python stream_ask theme={null}
  # Retrieve a specific graph
  graph = retrieve_graph("f47ac10b-58cc-4372-a567-0e02b2c3d479")

  # Pose a question and stream the response in chunks
  for chunk in graph.stream_ask("Explain the history of solar energy."):
      print(chunk)
  ```
</CodeGroup>

#### 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.

<CodeGroup>
  ```python ask theme={null}
  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)
  ```

  ```python stream_ask theme={null}
  from writer.ai import stream_ask

  # Stream responses from multiple graphs
  for chunk in stream_ask(
      question="Describe the key features of renewable energy sources.",
      graphs_or_graph_ids=[
          "550e8400-e29b-41d4-a716-446655440000",
          "123e4567-e89b-12d3-a456-426614174000"
          ]
  ):
      print(chunk)
  ```
</CodeGroup>

## Using the `Tools` class

<Note>
  This document outlines the use of `Tools` in the Writer Framework; for more thorough documentation, check out (this guide)\[[https://dev.writer.com/tools](https://dev.writer.com/tools)].
</Note>

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.

```python  theme={null}
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.

```python  theme={null}
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.

<CodeGroup>
  ```python file_id theme={null}
  from writer.ai import tools

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

  ```python upload_file theme={null}
  from writer.ai import tools

  file_object = upload_file(
      data=pdf_content,
      type="application/pdf",
      name="uploaded_file.pdf"
  )
  parsed_content = tools.parse_pdf(file_object, format="text")
  print(parsed_content)
  ```
</CodeGroup>

**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.