# `Lightning.AiAssistant`
[🔗](https://github.com/OpenFn/lightning/blob/main/lib/lightning/ai_assistant/ai_assistant.ex#L1)

The AI assistant context module for Lightning workflows.

This module provides a comprehensive interface for managing AI-powered chat sessions
within the Lightning platform. It handles two main types of AI assistance:

1. **Job-specific assistance** - Helps users with coding tasks, debugging, and
   adaptor-specific guidance for individual workflow jobs
2. **Workflow template generation** - Assists in creating complete workflow
   templates from natural language descriptions

# `opts`

```elixir
@type opts() :: keyword()
```

# `stream_error`

```elixir
@type stream_error() :: String.t() | {:internal, String.t()} | Ecto.Changeset.t()
```

Why a streamed query failed.

A bare string is already written for a user to read. `{:internal, text}` is
ours and is not: the caller logs it and shows something else.

# `cleanup_unsaved_job_sessions`

```elixir
@spec cleanup_unsaved_job_sessions(Lightning.Workflows.Workflow.t()) ::
  {:ok, non_neg_integer()} | {:error, term()}
```

Cleans up unsaved job data in AI sessions after workflow is saved.

When a workflow is saved and jobs get database IDs, this function:
1. Finds all sessions with matching Y.Doc job IDs in meta["unsaved_job"]
2. Updates those sessions to set job_id to the real database ID
3. Removes the meta["unsaved_job"] data

## Parameters

- `workflow` - The saved `%Workflow{}` with jobs that have database IDs

## Returns

- `{:ok, updated_count}` - Number of sessions updated
- `{:error, reason}` - If update fails

## Example

    iex> cleanup_unsaved_job_sessions(workflow)
    {:ok, 3}  # Updated 3 sessions

# `cleanup_unsaved_workflow_sessions`

```elixir
@spec cleanup_unsaved_workflow_sessions(Lightning.Workflows.Workflow.t()) ::
  {:ok, non_neg_integer()} | {:error, term()}
```

Cleans up unsaved workflow data in AI sessions after workflow is saved.

When a workflow in create mode is saved for the first time, this function:
1. Finds all sessions with matching temporary workflow IDs in meta["unsaved_workflow"]
2. Updates those sessions to set workflow_id to the real database ID
3. Removes the meta["unsaved_workflow"] data

## Parameters

- `workflow` - The saved `%Workflow{}` with a database ID

## Returns

- `{:ok, updated_count}` - Number of sessions updated
- `{:error, reason}` - If update fails

## Example

    iex> cleanup_unsaved_workflow_sessions(workflow)
    {:ok, 2}  # Updated 2 sessions

# `create_session`

```elixir
@spec create_session(
  Lightning.Workflows.Job.t(),
  Lightning.Accounts.User.t(),
  String.t(),
  opts()
) ::
  {:ok, Lightning.AiAssistant.ChatSession.t()} | {:error, Ecto.Changeset.t()}
```

Creates a new chat session for job-specific AI assistance.

Initializes a new session with:
- Association with the specified job and user
- Auto-generated title from the initial message content
- Job's expression and adaptor context (virtual fields)
- The initial user message

## Parameters

- `job` - The `%Job{}` struct this session will assist with
- `user` - The `%User{}` creating the session
- `content` - Initial message content that will become the session title
- `opts` - Keyword list of options:
  - `:meta` - Optional metadata for the session
  - `:code` - Optional workflow code for the initial message

## Returns

- `{:ok, session}` - Successfully created session with initial message
- `{:error, changeset}` - Validation errors during creation

# `create_session_for_unsaved_job`

```elixir
@spec create_session_for_unsaved_job(Lightning.Accounts.User.t(), String.t(), map()) ::
  {:ok, Lightning.AiAssistant.ChatSession.t()} | {:error, Ecto.Changeset.t()}
```

Creates a new chat session for an unsaved job (exists in Y.Doc but not DB).

This is used during collaborative editing when users want AI assistance
for jobs they've created but haven't saved to the database yet.

## Parameters

- `user` - The `%User{}` creating the session
- `content` - Initial message content that will become the session title
- `meta` - Metadata including unsaved job data with keys:
  - `"unsaved_job"` - Map with job data:
    - `"id"` - The Y.Doc job ID (for future linking when saved)
    - `"name"` - Job name
    - `"body"` - Job expression code
    - `"adaptor"` - Job adaptor
    - `"workflow_id"` - Workflow the job belongs to

## Returns

- `{:ok, session}` - Successfully created session with initial message
- `{:error, changeset}` - Validation errors during creation

# `create_workflow_session`

```elixir
@spec create_workflow_session(
  Lightning.Projects.Project.t(),
  Lightning.Workflows.Job.t() | nil,
  Lightning.Workflows.Workflow.t() | nil,
  Lightning.Accounts.User.t(),
  String.t(),
  opts()
) :: {:ok, Lightning.AiAssistant.ChatSession.t()} | {:error, Ecto.Changeset.t()}
```

Creates a new chat session for workflow template generation.

Initializes a session specifically for creating workflow templates:
- Associates with a project and optionally an existing workflow
- Sets session type to "workflow_template"
- Includes the initial user message describing the desired workflow

## Parameters

- `project` - The `%Project{}` struct where the workflow will be created
- `workflow` - The `%Workflow{}` struct to associate with the session, or `nil` for new workflows
- `user` - The `%User{}` requesting the workflow template
- `content` - Description of the desired workflow functionality
- `opts` - Keyword list of options:
  - `:meta` - Optional metadata for the session
  - `:code` - Optional workflow code for the initial message

## Returns

- `{:ok, session}` - Successfully created workflow session with initial message
- `{:error, changeset}` - Validation errors during creation

# `enabled?`

```elixir
@spec enabled?() :: boolean()
```

Checks if the AI assistant feature is enabled via application configuration.

Verifies that both the Apollo endpoint URL and API key are properly configured,
which are required for AI functionality.

## Returns

`true` if AI assistant is properly configured and enabled, `false` otherwise.

# `enrich_session_with_job_context`

```elixir
@spec enrich_session_with_job_context(Lightning.AiAssistant.ChatSession.t()) ::
  Lightning.AiAssistant.ChatSession.t()
```

Enriches a chat session with job-specific context.

Loads the associated job and enriches the session with:
- Job expression code and adaptor information
- Run logs if following a specific run (via session.meta["follow_run_id"])

For unsaved jobs (job_id is nil), uses job data from session.meta.

## Parameters
- `session` - The chat session to enrich

## Returns
The enriched session with job context loaded

# `get_session`

# `get_session!`

Retrieves a chat session by ID with all related data preloaded.

Fetches a complete chat session including:
- All non-cancelled messages ordered by creation time
- User information for each message
- Session metadata
- Project information (for workflow template sessions)

## Parameters

- `id` - UUID string of the session to retrieve

## Returns

A `ChatSession` struct with preloaded `:messages` and nested `:user` data.
For workflow template sessions, the `:project` association is also preloaded.

## Raises

`Ecto.NoResultsError` if no session exists with the given ID.

# `list_sessions`

```elixir
@spec list_sessions(
  Lightning.Projects.Project.t() | Lightning.Workflows.Job.t(),
  :asc | :desc,
  opts()
) ::
  %{
    sessions: [Lightning.AiAssistant.ChatSession.t()],
    pagination: LightningWeb.Live.AiAssistant.PaginationMeta.t()
  }
```

Returns paginated chat sessions with metadata for improved user experience.

Retrieves chat sessions associated with either a Project (workflow template sessions)
or a Job (job-specific sessions). Results are paginated and include total counts
and navigation metadata.

## Parameters

- `resource` - A `%Project{}`, `%Job{}`, or `%Snapshot.Job{}` struct to filter sessions by
- `sort_direction` - Sort order, either `:asc` or `:desc` (default: `:desc`)
- `opts` - Keyword list of options:
  - `:offset` - Number of records to skip (default: 0)
  - `:limit` - Maximum number of records to return (default: 20)

## Returns

A map containing:
- `:sessions` - List of `ChatSession` structs with preloaded data
- `:pagination` - `PaginationMeta` struct with navigation information

# `put_expression_and_adaptor`

```elixir
@spec put_expression_and_adaptor(
  Lightning.AiAssistant.ChatSession.t(),
  String.t(),
  String.t()
) ::
  Lightning.AiAssistant.ChatSession.t()
```

Adds job-specific context to a chat session for enhanced AI assistance.

Enriches a session with the job's expression code and adaptor information,
enabling the AI to provide more targeted and relevant assistance.

## Parameters

- `session` - The `%ChatSession{}` to enhance
- `expression` - The job's expression code as a string
- `adaptor` - The adaptor name/identifier for the job

## Returns

An updated `ChatSession` struct with `:expression` and `:adaptor` fields populated.
The adaptor is resolved through `Lightning.Adaptors.to_wire/1`, falling back
to the adaptor as given if that fails.

# `query_global_stream`

```elixir
@spec query_global_stream(Lightning.AiAssistant.ChatSession.t(), String.t(), opts()) ::
  {:ok, Lightning.AiAssistant.ChatSession.t()} | {:error, stream_error()}
```

Queries the Apollo global chat endpoint with SSE streaming.

The global chat unifies job and workflow assistance behind a router.
Uses the same streaming pipeline as job_chat and workflow_chat.

# `query_stream`

```elixir
@spec query_stream(Lightning.AiAssistant.ChatSession.t(), String.t(), opts()) ::
  {:ok, Lightning.AiAssistant.ChatSession.t()} | {:error, stream_error()}
```

Queries the AI service for job-specific code assistance with streaming.

Sends a user query to the Apollo AI service along with job context
(expression and adaptor) and conversation history, using SSE streaming.
Text chunks and status updates are broadcast via PubSub as they arrive.
The complete message is saved to the database when the stream finishes.

## Parameters

- `session` - The job-specific `%ChatSession{}` with expression and adaptor context
- `content` - User's question or request for assistance

## Returns

- `{:ok, session}` - AI responded successfully, session updated with response
- `{:error, reason}` - Query failed, reason is either a string error message or changeset

# `query_workflow_stream`

```elixir
@spec query_workflow_stream(Lightning.AiAssistant.ChatSession.t(), String.t(), opts()) ::
  {:ok, Lightning.AiAssistant.ChatSession.t()} | {:error, stream_error()}
```

Queries the AI service for workflow template generation with streaming.

Sends a request to generate or modify workflow templates based on user
requirements, using SSE streaming. Can include validation errors from
previous attempts to help the AI provide corrections.

## Parameters

- `session` - The workflow template `%ChatSession{}`
- `content` - User's description of desired workflow functionality or modifications
- `opts` - Keyword list of options:
  - `:code` - Current YAML to modify (default: uses latest from session)
  - `:errors` - Validation errors from previous workflow attempts

## Returns

- `{:ok, session}` - Workflow template generated successfully
- `{:error, reason}` - Generation failed, reason is either a string error message or changeset

# `retry_message`

Resets a message status to pending and enqueues it for reprocessing.
Handles both the database update and Oban job creation atomically.

# `save_message`

```elixir
@spec save_message(Lightning.AiAssistant.ChatSession.t(), map(), opts()) ::
  {:ok, Lightning.AiAssistant.ChatSession.t()} | {:error, Ecto.Changeset.t()}
```

Saves a message to an existing chat session.

Adds a new message to the session's message history and updates the session.
For assistant messages, this also triggers AI usage tracking and limit enforcement.

## Parameters

- `session` - The target `%ChatSession{}`
- `message_attrs` - Map containing message data with keys like `:role`, `:content`, `:user`
- `opts` - Keyword list of options:
  - `:usage` - Map containing AI usage metrics (default: `%{}`)
  - `:meta` - Session metadata to update (default: keeps existing)
  - `:code` - Optional workflow code to attach to the message
  - `:apollo_history` - Apollo's history to store on the session (default: keeps existing)

## Returns

- `{:ok, session}` - Successfully saved message and updated session
- `{:error, changeset}` - Validation or database errors

# `set_apply_failed`

```elixir
@spec set_apply_failed(Ecto.UUID.t(), Ecto.UUID.t(), boolean()) ::
  {:ok, Lightning.AiAssistant.ChatMessage.t()}
  | {:error, :not_found | Ecto.Changeset.t()}
```

Records whether a reply's changes reached the canvas.

The apply happens in the browser, so nothing else knows it failed. Without
this the reply reads as a success on the next page load: its diff blocks
stand as a record of changes that never landed, and it offers to undo them.

Clearing on a later success is half the point, so a retry that works leaves
nothing behind.

# `unsaved_job_session`

```elixir
@spec unsaved_job_session(Ecto.UUID.t()) ::
  Lightning.AiAssistant.ChatSession.t() | nil
```

Returns a job-code session for the job, matching either its saved `job_id` or
the id of the unsaved job in its `meta`, or `nil` when there is none.

# `update_message_status`

```elixir
@spec update_message_status(
  Lightning.AiAssistant.ChatSession.t(),
  Lightning.AiAssistant.ChatMessage.t(),
  atom()
) :: {:ok, Lightning.AiAssistant.ChatSession.t()} | {:error, Ecto.Changeset.t()}
```

Updates the status of a specific message within a chat session.

Changes the status of an individual message (e.g., from `:pending` to `:success` or `:error`)
and returns the refreshed session with updated data.

## Parameters

- `session` - The `%ChatSession{}` containing the message
- `message` - The specific `%ChatMessage{}` to update
- `status` - New status atom (e.g., `:pending`, `:success`, `:error`, `:cancelled`)

## Returns

- `{:ok, session}` - Successfully updated message, returns refreshed session
- `{:error, changeset}` - Update failed with validation errors

---

*Consult [api-reference.md](api-reference.md) for complete listing*