feat: add generation source handling for task creation and updates

2026-03-28 15:35:49 -03:00
parent ea54858165
commit 58c3a54d4a
14 changed files with 268 additions and 14 deletions
--- a/backend/src/main/kotlin/com/condado/newsletter/dto/EntityTaskDtos.kt
+++ b/backend/src/main/kotlin/com/condado/newsletter/dto/EntityTaskDtos.kt
@@ -1,6 +1,7 @@
 package com.condado.newsletter.dto

 import com.condado.newsletter.model.EntityTask
+import com.condado.newsletter.model.TaskGenerationSource
 import jakarta.validation.constraints.NotBlank
 import jakarta.validation.constraints.NotNull
 import java.time.LocalDateTime
@@ -11,7 +12,8 @@ data class EntityTaskCreateDto(
    @field:NotBlank val name: String,
    val prompt: String,
    @field:NotBlank val scheduleCron: String,
-    @field:NotBlank val emailLookback: String
+    @field:NotBlank val emailLookback: String,
+    val generationSource: TaskGenerationSource = TaskGenerationSource.LLAMA
 )

 data class EntityTaskUpdateDto(
@@ -19,7 +21,8 @@ data class EntityTaskUpdateDto(
    @field:NotBlank val name: String,
    @field:NotBlank val prompt: String,
    @field:NotBlank val scheduleCron: String,
-    @field:NotBlank val emailLookback: String
+    @field:NotBlank val emailLookback: String,
+    val generationSource: TaskGenerationSource? = null
 )

 data class EntityTaskResponseDto(
@@ -29,6 +32,7 @@ data class EntityTaskResponseDto(
    val prompt: String,
    val scheduleCron: String,
    val emailLookback: String,
+    val generationSource: TaskGenerationSource,
    val active: Boolean,
    val createdAt: LocalDateTime?
 ) {
@@ -41,6 +45,7 @@ data class EntityTaskResponseDto(
                prompt = task.prompt,
                scheduleCron = task.scheduleCron,
                emailLookback = task.emailLookback,
+                generationSource = task.generationSource,
                active = task.active,
                createdAt = task.createdAt
            )
--- a/backend/src/main/kotlin/com/condado/newsletter/model/EntityTask.kt
+++ b/backend/src/main/kotlin/com/condado/newsletter/model/EntityTask.kt
@@ -3,6 +3,8 @@ package com.condado.newsletter.model
 import jakarta.persistence.CascadeType
 import jakarta.persistence.Column
 import jakarta.persistence.Entity
+import jakarta.persistence.EnumType
+import jakarta.persistence.Enumerated
 import jakarta.persistence.FetchType
 import jakarta.persistence.GeneratedValue
 import jakarta.persistence.GenerationType
@@ -37,6 +39,10 @@ class EntityTask(
    @Column(name = "email_lookback", nullable = false)
    val emailLookback: String,

+    @Enumerated(EnumType.STRING)
+    @Column(name = "generation_source", nullable = false)
+    val generationSource: TaskGenerationSource = TaskGenerationSource.LLAMA,
+
    @Column(nullable = false)
    val active: Boolean = true,

--- a/backend/src/main/kotlin/com/condado/newsletter/model/TaskGenerationSource.kt
+++ b/backend/src/main/kotlin/com/condado/newsletter/model/TaskGenerationSource.kt
@@ -0,0 +1,19 @@
+package com.condado.newsletter.model
+
+import com.fasterxml.jackson.annotation.JsonCreator
+import com.fasterxml.jackson.annotation.JsonValue
+
+enum class TaskGenerationSource(
+    @get:JsonValue val value: String
+) {
+    OPENAI("openai"),
+    LLAMA("llama");
+
+    companion object {
+        @JvmStatic
+        @JsonCreator
+        fun from(value: String): TaskGenerationSource =
+            entries.firstOrNull { it.value.equals(value, ignoreCase = true) }
+                ?: throw IllegalArgumentException("Invalid generationSource: $value")
+    }
+}
--- a/backend/src/main/kotlin/com/condado/newsletter/service/EntityTaskService.kt
+++ b/backend/src/main/kotlin/com/condado/newsletter/service/EntityTaskService.kt
@@ -47,6 +47,7 @@ class EntityTaskService(
            prompt = dto.prompt,
            scheduleCron = dto.scheduleCron,
            emailLookback = dto.emailLookback,
+            generationSource = dto.generationSource,
            active = true
        )

@@ -66,6 +67,7 @@ class EntityTaskService(
            prompt = dto.prompt,
            scheduleCron = dto.scheduleCron,
            emailLookback = dto.emailLookback,
+            generationSource = dto.generationSource ?: existing.generationSource,
            active = existing.active,
            createdAt = existing.createdAt
        ).apply { id = existing.id }
@@ -83,6 +85,7 @@ class EntityTaskService(
            prompt = existing.prompt,
            scheduleCron = existing.scheduleCron,
            emailLookback = existing.emailLookback,
+            generationSource = existing.generationSource,
            active = false,
            createdAt = existing.createdAt
        ).apply { id = existing.id }
@@ -100,6 +103,7 @@ class EntityTaskService(
            prompt = existing.prompt,
            scheduleCron = existing.scheduleCron,
            emailLookback = existing.emailLookback,
+            generationSource = existing.generationSource,
            active = true,
            createdAt = existing.createdAt
        ).apply { id = existing.id }
--- a/backend/src/main/kotlin/com/condado/newsletter/service/TaskGeneratedMessageService.kt
+++ b/backend/src/main/kotlin/com/condado/newsletter/service/TaskGeneratedMessageService.kt
@@ -3,6 +3,7 @@ package com.condado.newsletter.service
 import com.condado.newsletter.dto.GeneratedMessageHistoryResponseDto
 import com.condado.newsletter.dto.TaskPreviewGenerateRequestDto
 import com.condado.newsletter.model.GeneratedMessageHistory
+import com.condado.newsletter.model.TaskGenerationSource
 import com.condado.newsletter.repository.EntityTaskRepository
 import com.condado.newsletter.repository.GeneratedMessageHistoryRepository
 import org.springframework.stereotype.Service
@@ -16,7 +17,8 @@ import java.util.UUID
 class TaskGeneratedMessageService(
    private val generatedMessageHistoryRepository: GeneratedMessageHistoryRepository,
    private val entityTaskRepository: EntityTaskRepository,
-    private val llamaPreviewService: LlamaPreviewService
+    private val llamaPreviewService: LlamaPreviewService,
+    private val aiService: AiService
 ) {

    /** Lists persisted generated messages for a task. */
@@ -25,15 +27,19 @@ class TaskGeneratedMessageService(
            .findAllByTask_IdOrderByCreatedAtDesc(taskId)
            .map { GeneratedMessageHistoryResponseDto.from(it) }

-    /**
-     * Generates a new message using local Llama, persists it, and returns it.
-     */
+    /** Generates a new message with the task-selected provider, persists it, and returns it. */
    @Transactional
    fun generateAndSave(taskId: UUID, request: TaskPreviewGenerateRequestDto): GeneratedMessageHistoryResponseDto {
        val task = entityTaskRepository.findById(taskId)
            .orElseThrow { IllegalArgumentException("Task not found: $taskId") }
        val prompt = buildPrompt(request)
-        val generatedContent = llamaPreviewService.generate(prompt)
+        val generatedContent = when (task.generationSource) {
+            TaskGenerationSource.LLAMA -> llamaPreviewService.generate(prompt)
+            TaskGenerationSource.OPENAI -> {
+                val parsed = aiService.generate(prompt)
+                "SUBJECT: ${parsed.subject}\nBODY:\n${parsed.body}"
+            }
+        }
        val nextLabel = "Message #${generatedMessageHistoryRepository.countByTask_Id(taskId) + 1}"

        val saved = generatedMessageHistoryRepository.save(
--- a/backend/src/test/kotlin/com/condado/newsletter/controller/EntityTaskControllerTest.kt
+++ b/backend/src/test/kotlin/com/condado/newsletter/controller/EntityTaskControllerTest.kt
@@ -8,6 +8,7 @@ import com.condado.newsletter.scheduler.EntityScheduler
 import com.condado.newsletter.service.JwtService
 import com.ninjasquad.springmockk.MockkBean
 import jakarta.servlet.http.Cookie
+import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Test
 import org.springframework.beans.factory.annotation.Autowired
@@ -15,6 +16,7 @@ import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureMock
 import org.springframework.boot.test.context.SpringBootTest
 import org.springframework.http.MediaType
 import org.springframework.test.web.servlet.MockMvc
+import org.springframework.test.web.servlet.request.MockMvcRequestBuilders.put
 import org.springframework.test.web.servlet.request.MockMvcRequestBuilders.post
 import org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath
 import org.springframework.test.web.servlet.result.MockMvcResultMatchers.status
@@ -56,7 +58,8 @@ class EntityTaskControllerTest {
              "name": "Morning Blast",
              "prompt": "",
              "scheduleCron": "0 8 * * 1-5",
-              "emailLookback": "last_week"
+                            "emailLookback": "last_week",
+                            "generationSource": "openai"
            }
        """.trimIndent()

@@ -70,5 +73,96 @@ class EntityTaskControllerTest {
            .andExpect(jsonPath("$.entityId").value(entity.id.toString()))
            .andExpect(jsonPath("$.name").value("Morning Blast"))
            .andExpect(jsonPath("$.prompt").value(""))
+            .andExpect(jsonPath("$.generationSource").value("openai"))
+
+        val persisted = entityTaskRepository.findAll().first()
+        assertThat(persisted.generationSource.value).isEqualTo("openai")
+    }
+
+    @Test
+    fun should_updateTaskAndPersistGenerationSource_when_validRequestProvided() {
+        val entity = virtualEntityRepository.save(
+            VirtualEntity(
+                name = "Entity B",
+                email = "entity-b@condado.com",
+                jobTitle = "Ops"
+            )
+        )
+
+        val createdPayload = """
+            {
+              "entityId": "${entity.id}",
+              "name": "Task One",
+              "prompt": "Initial prompt",
+              "scheduleCron": "0 8 * * 1-5",
+              "emailLookback": "last_week",
+              "generationSource": "openai"
+            }
+        """.trimIndent()
+
+        val createdResult = mockMvc.perform(
+            post("/api/v1/tasks")
+                .cookie(authCookie())
+                .contentType(MediaType.APPLICATION_JSON)
+                .content(createdPayload)
+        )
+            .andExpect(status().isCreated)
+            .andReturn()
+
+        val taskId = com.jayway.jsonpath.JsonPath.read<String>(createdResult.response.contentAsString, "$.id")
+
+        val updatePayload = """
+            {
+              "entityId": "${entity.id}",
+              "name": "Task One Updated",
+              "prompt": "Updated prompt",
+              "scheduleCron": "0 10 * * 1-5",
+              "emailLookback": "last_day",
+              "generationSource": "llama"
+            }
+        """.trimIndent()
+
+        mockMvc.perform(
+            put("/api/v1/tasks/$taskId")
+                .cookie(authCookie())
+                .contentType(MediaType.APPLICATION_JSON)
+                .content(updatePayload)
+        )
+            .andExpect(status().isOk)
+            .andExpect(jsonPath("$.name").value("Task One Updated"))
+            .andExpect(jsonPath("$.generationSource").value("llama"))
+
+        val persisted = entityTaskRepository.findById(java.util.UUID.fromString(taskId)).orElseThrow()
+        assertThat(persisted.generationSource.value).isEqualTo("llama")
+    }
+
+    @Test
+    fun should_returnBadRequest_when_generationSourceIsInvalid() {
+        val entity = virtualEntityRepository.save(
+            VirtualEntity(
+                name = "Entity C",
+                email = "entity-c@condado.com",
+                jobTitle = "Ops"
+            )
+        )
+
+        val payload = """
+            {
+              "entityId": "${entity.id}",
+              "name": "Morning Blast",
+              "prompt": "Prompt",
+              "scheduleCron": "0 8 * * 1-5",
+              "emailLookback": "last_week",
+              "generationSource": "invalid-provider"
+            }
+        """.trimIndent()
+
+        mockMvc.perform(
+            post("/api/v1/tasks")
+                .cookie(authCookie())
+                .contentType(MediaType.APPLICATION_JSON)
+                .content(payload)
+        )
+            .andExpect(status().isBadRequest)
    }
 }
--- a/backend/src/test/kotlin/com/condado/newsletter/service/TaskGeneratedMessageServiceTest.kt
+++ b/backend/src/test/kotlin/com/condado/newsletter/service/TaskGeneratedMessageServiceTest.kt
@@ -5,6 +5,8 @@ import com.condado.newsletter.dto.TaskPreviewGenerateRequestDto
 import com.condado.newsletter.dto.TaskPreviewTaskDto
 import com.condado.newsletter.model.EntityTask
 import com.condado.newsletter.model.GeneratedMessageHistory
+import com.condado.newsletter.model.ParsedAiResponse
+import com.condado.newsletter.model.TaskGenerationSource
 import com.condado.newsletter.model.VirtualEntity
 import com.condado.newsletter.repository.EntityTaskRepository
 import com.condado.newsletter.repository.GeneratedMessageHistoryRepository
@@ -21,15 +23,17 @@ class TaskGeneratedMessageServiceTest {
    private val generatedMessageHistoryRepository: GeneratedMessageHistoryRepository = mockk()
    private val entityTaskRepository: EntityTaskRepository = mockk()
    private val llamaPreviewService: LlamaPreviewService = mockk()
+    private val aiService: AiService = mockk()

    private val service = TaskGeneratedMessageService(
        generatedMessageHistoryRepository = generatedMessageHistoryRepository,
        entityTaskRepository = entityTaskRepository,
-        llamaPreviewService = llamaPreviewService
+        llamaPreviewService = llamaPreviewService,
+        aiService = aiService
    )

    @Test
-    fun should_generateAndPersistMessage_when_generateAndSaveCalled() {
+    fun should_useLlamaProvider_when_taskGenerationSourceIsLlama() {
        val taskId = UUID.randomUUID()
        val entity = VirtualEntity(name = "Entity", email = "e@x.com", jobTitle = "Ops").apply { id = UUID.randomUUID() }
        val task = EntityTask(
@@ -37,7 +41,8 @@ class TaskGeneratedMessageServiceTest {
            name = "Task",
            prompt = "Prompt",
            scheduleCron = "0 9 * * 1",
-            emailLookback = "last_week"
+            emailLookback = "last_week",
+            generationSource = TaskGenerationSource.LLAMA
        ).apply { id = taskId }
        val captured = slot<GeneratedMessageHistory>()

@@ -59,9 +64,40 @@ class TaskGeneratedMessageServiceTest {
        assertThat(captured.captured.task.id).isEqualTo(taskId)

        verify(exactly = 1) { llamaPreviewService.generate(any()) }
+        verify(exactly = 0) { aiService.generate(any()) }
        verify(exactly = 1) { generatedMessageHistoryRepository.save(any()) }
    }

+    @Test
+    fun should_useOpenAiProvider_when_taskGenerationSourceIsOpenai() {
+        val taskId = UUID.randomUUID()
+        val entity = VirtualEntity(name = "Entity", email = "e@x.com", jobTitle = "Ops").apply { id = UUID.randomUUID() }
+        val task = EntityTask(
+            virtualEntity = entity,
+            name = "Task",
+            prompt = "Prompt",
+            scheduleCron = "0 9 * * 1",
+            emailLookback = "last_week",
+            generationSource = TaskGenerationSource.OPENAI
+        ).apply { id = taskId }
+        val captured = slot<GeneratedMessageHistory>()
+
+        every { aiService.generate(any()) } returns ParsedAiResponse(subject = "Open Subject", body = "Open Body")
+        every { entityTaskRepository.findById(taskId) } returns java.util.Optional.of(task)
+        every { generatedMessageHistoryRepository.countByTask_Id(taskId) } returns 0
+        every { generatedMessageHistoryRepository.save(capture(captured)) } answers {
+            captured.captured.apply {
+                id = UUID.fromString("00000000-0000-0000-0000-000000000001")
+            }
+        }
+
+        val response = service.generateAndSave(taskId, sampleRequest())
+
+        assertThat(response.content).isEqualTo("SUBJECT: Open Subject\nBODY:\nOpen Body")
+        verify(exactly = 1) { aiService.generate(any()) }
+        verify(exactly = 0) { llamaPreviewService.generate(any()) }
+    }
+
    private fun sampleRequest() = TaskPreviewGenerateRequestDto(
        entity = TaskPreviewEntityDto(
            id = UUID.randomUUID().toString(),
--- a/frontend/src/tests/api/tasksApi.test.ts
+++ b/frontend/src/tests/api/tasksApi.test.ts
@@ -39,6 +39,7 @@ const taskOne: EntityTaskResponse = {
  prompt: 'Summarize jokes',
  scheduleCron: '0 9 * * 1',
  emailLookback: 'last_week',
+  generationSource: 'openai',
  active: true,
  createdAt: '2026-03-26T10:00:00Z',
 }
@@ -50,6 +51,7 @@ const taskTwo: EntityTaskResponse = {
  prompt: 'Escalate sandwich policy',
  scheduleCron: '0 11 1 * *',
  emailLookback: 'last_month',
+  generationSource: 'llama',
  active: false,
  createdAt: '2026-03-26T11:00:00Z',
 }
@@ -72,6 +74,7 @@ const previewTask = {
  prompt: 'Draft an absurdly official update about disappearing crackers.',
  scheduleCron: '15 10 * * 2',
  emailLookback: 'last_week' as const,
+  generationSource: 'openai' as const,
 }

 describe('tasksApi', () => {
@@ -143,6 +146,7 @@ describe('tasksApi', () => {
      prompt: 'Ask about ceremonial coffee',
      scheduleCron: '0 8 * * 1-5',
      emailLookback: 'last_day',
+      generationSource: 'openai',
    })

    expect(createdTask).toEqual(
@@ -157,6 +161,7 @@ describe('tasksApi', () => {
      prompt: 'Ask about ceremonial coffee',
      scheduleCron: '0 8 * * 1-5',
      emailLookback: 'last_day',
+      generationSource: 'openai',
    })
  })

@@ -168,6 +173,7 @@ describe('tasksApi', () => {
        prompt: 'Ask about ceremonial coffee',
        scheduleCron: '0 8 * * 1-5',
        emailLookback: 'last_day',
+        generationSource: 'llama',
      },
    })

@@ -177,6 +183,7 @@ describe('tasksApi', () => {
      prompt: 'Ask about ceremonial coffee',
      scheduleCron: '0 8 * * 1-5',
      emailLookback: 'last_day',
+      generationSource: 'llama',
    })

    expect(updatedTask).toEqual({
@@ -185,6 +192,7 @@ describe('tasksApi', () => {
      prompt: 'Ask about ceremonial coffee',
      scheduleCron: '0 8 * * 1-5',
      emailLookback: 'last_day',
+      generationSource: 'llama',
    })
    expect(mockedApiClient.put).toHaveBeenCalledWith('/v1/tasks/task-1', {
      entityId: 'entity-1',
@@ -192,6 +200,7 @@ describe('tasksApi', () => {
      prompt: 'Ask about ceremonial coffee',
      scheduleCron: '0 8 * * 1-5',
      emailLookback: 'last_day',
+      generationSource: 'llama',
    })
  })

--- a/frontend/src/tests/pages/CreateTaskPage.test.tsx
+++ b/frontend/src/tests/pages/CreateTaskPage.test.tsx
@@ -50,6 +50,7 @@ describe('CreateTaskPage', () => {

    expect(screen.getByLabelText(/task name/i)).toBeInTheDocument()
    expect(screen.queryByLabelText(/task prompt/i)).not.toBeInTheDocument()
+    expect(screen.getByLabelText(/^Generation Source$/i)).toHaveValue('openai')
    expect(screen.getByLabelText(/^Email Period$/i)).toBeInTheDocument()
    expect(screen.getByLabelText(/^Minute$/i)).toBeInTheDocument()
    expect(screen.getByLabelText(/^Hour$/i)).toBeInTheDocument()
@@ -68,6 +69,7 @@ describe('CreateTaskPage', () => {
      prompt: '',
      scheduleCron: '0 8 * * 1-5',
      emailLookback: 'last_week',
+      generationSource: 'openai',
      active: false,
      createdAt: '2026-03-26T10:00:00Z',
    })
@@ -78,6 +80,7 @@ describe('CreateTaskPage', () => {
      prompt: '',
      scheduleCron: '0 8 * * 1-5',
      emailLookback: 'last_week',
+      generationSource: 'openai',
      active: false,
      createdAt: '2026-03-26T10:00:00Z',
    })
@@ -118,6 +121,7 @@ describe('CreateTaskPage', () => {
          prompt: '',
          scheduleCron: '0 8 * * 1-5',
          emailLookback: 'last_week',
+          generationSource: 'openai',
        })
      )
      expect(tasksApi.inactivateTask).toHaveBeenCalledWith('task-2')
--- a/frontend/src/tests/pages/EditTaskPage.test.tsx
+++ b/frontend/src/tests/pages/EditTaskPage.test.tsx
@@ -57,6 +57,7 @@ const mockTask = {
  prompt: 'Summarize jokes',
  scheduleCron: '0 9 * * 1',
  emailLookback: 'last_week' as const,
+  generationSource: 'openai' as const,
  active: true,
  createdAt: '2026-03-26T10:00:00Z',
 }
@@ -77,7 +78,7 @@ describe('EditTaskPage', () => {
    vi.mocked(tasksApi.getTask).mockResolvedValue(mockTask)
    vi.mocked(tasksApi.buildTaskPreviewPrompt).mockImplementation(
      (entity, task) =>
-        `PROMPT FOR ${entity.name}: ${task.name} | ${task.prompt} | ${task.scheduleCron} | ${task.emailLookback}`
+        `PROMPT FOR ${entity.name}: ${task.name} | ${task.prompt} | ${task.scheduleCron} | ${task.emailLookback} | ${task.generationSource}`
    )
    vi.mocked(tasksApi.activateTask).mockResolvedValue({ ...mockTask, active: true })
    vi.mocked(tasksApi.inactivateTask).mockResolvedValue({ ...mockTask, active: false })
@@ -97,6 +98,7 @@ describe('EditTaskPage', () => {
    expect(screen.getByRole('heading', { name: /edit task/i })).toBeInTheDocument()
    expect(screen.getByLabelText(/task name/i)).toHaveValue('Weekly Check-in')
    expect(screen.getByLabelText(/task prompt/i)).toHaveValue('Summarize jokes')
+    expect(screen.getByLabelText(/^Generation Source$/i)).toHaveValue('openai')
    expect(screen.getByLabelText(/^Email Period$/i)).toHaveValue('last_week')
    expect(screen.getByLabelText(/^Minute$/i)).toHaveValue('0')
    expect(screen.getByLabelText(/^Hour$/i)).toHaveValue('9')
@@ -124,6 +126,7 @@ describe('EditTaskPage', () => {
      prompt: 'Ask about ceremonial coffee',
      scheduleCron: '0 8 * * 1-5',
      emailLookback: 'last_day',
+      generationSource: 'llama',
    })

    const { queryClient } = renderPage()
@@ -137,13 +140,16 @@ describe('EditTaskPage', () => {
    fireEvent.change(screen.getByLabelText(/^Email Period$/i), {
      target: { value: 'last_day' },
    })
+    fireEvent.change(screen.getByLabelText(/^Generation Source$/i), {
+      target: { value: 'llama' },
+    })
    fireEvent.click(screen.getByRole('button', { name: /Weekdays/i }))
    fireEvent.change(screen.getByLabelText(/^Hour$/i), { target: { value: '8' } })

    expect(screen.getByText(/Final Prompt/i)).toBeInTheDocument()
    expect(
      screen.getByText(
-        'PROMPT FOR Entity A: Daily Check-in | Ask about ceremonial coffee | 0 8 * * 1-5 | last_day'
+        'PROMPT FOR Entity A: Daily Check-in | Ask about ceremonial coffee | 0 8 * * 1-5 | last_day | llama'
      )
    ).toBeInTheDocument()

@@ -159,6 +165,7 @@ describe('EditTaskPage', () => {
          prompt: 'Ask about ceremonial coffee',
          scheduleCron: '0 8 * * 1-5',
          emailLookback: 'last_day',
+          generationSource: 'llama',
        })
      )
      expect(vi.mocked(tasksApi.generateTaskPreview).mock.calls[0][0]).toEqual('task-1')
@@ -171,6 +178,7 @@ describe('EditTaskPage', () => {
            prompt: 'Ask about ceremonial coffee',
            scheduleCron: '0 8 * * 1-5',
            emailLookback: 'last_day',
+            generationSource: 'llama',
          }),
        })
      )
@@ -186,6 +194,7 @@ describe('EditTaskPage', () => {
        prompt: 'Ask about ceremonial coffee',
        scheduleCron: '0 8 * * 1-5',
        emailLookback: 'last_day',
+        generationSource: 'llama',
      })
      expect(invalidateQueriesSpy).toHaveBeenCalledWith({ queryKey: ['entity-tasks', 'entity-1'] })
      expect(invalidateQueriesSpy).toHaveBeenCalledWith({ queryKey: ['entity-task', 'task-1'] })
--- a/frontend/src/tests/pages/EntityDetailPage.test.tsx
+++ b/frontend/src/tests/pages/EntityDetailPage.test.tsx
@@ -120,6 +120,7 @@ describe('EntityDetailPage', () => {
        prompt: 'Summarize jokes',
        scheduleCron: '0 9 * * 1',
        emailLookback: 'last_week',
+        generationSource: 'openai',
        active: true,
        createdAt: '2026-03-26T10:00:00Z',
      },
@@ -165,6 +166,7 @@ describe('EntityDetailPage', () => {
        prompt: 'Archive the sandwich minutes',
        scheduleCron: '0 9 * * 1',
        emailLookback: 'last_week',
+        generationSource: 'llama',
        active: false,
        createdAt: '2026-03-26T10:00:00Z',
      },
--- a/frontend/src/api/tasksApi.ts
+++ b/frontend/src/api/tasksApi.ts
@@ -2,6 +2,7 @@ import type { VirtualEntityResponse } from './entitiesApi'
 import apiClient from './apiClient'

 export type EmailLookback = 'last_day' | 'last_week' | 'last_month'
+export type GenerationSource = 'openai' | 'llama'

 export interface EntityTaskResponse {
  id: string
@@ -10,6 +11,7 @@ export interface EntityTaskResponse {
  prompt: string
  scheduleCron: string
  emailLookback: EmailLookback
+  generationSource: GenerationSource
  active: boolean
  createdAt: string
 }
@@ -20,6 +22,7 @@ export interface EntityTaskCreateDto {
  prompt: string
  scheduleCron: string
  emailLookback: EmailLookback
+  generationSource: GenerationSource
 }

 export type EntityTaskUpdateDto = EntityTaskCreateDto
--- a/frontend/src/pages/CreateTaskPage.tsx
+++ b/frontend/src/pages/CreateTaskPage.tsx
@@ -6,12 +6,14 @@ import {
  createTask,
  inactivateTask,
  type EmailLookback,
+  type GenerationSource,
 } from '../api/tasksApi'

 interface TaskFormState {
  name: string
  scheduleCron: string
  emailLookback: EmailLookback
+  generationSource: GenerationSource
 }

 interface CronParts {
@@ -72,6 +74,7 @@ const DEFAULT_TASK_FORM: TaskFormState = {
  name: '',
  scheduleCron: buildCron(DEFAULT_CRON_PARTS),
  emailLookback: 'last_week',
+  generationSource: 'openai',
 }

 export default function CreateTaskPage() {
@@ -151,6 +154,7 @@ export default function CreateTaskPage() {
              prompt: '',
              scheduleCron: taskForm.scheduleCron,
              emailLookback: taskForm.emailLookback,
+              generationSource: taskForm.generationSource,
            })
          }}
        >
@@ -167,6 +171,26 @@ export default function CreateTaskPage() {
            />
          </div>

+          <div>
+            <label htmlFor="task-generation-source" className="text-sm font-medium text-slate-200">
+              Generation Source
+            </label>
+            <select
+              id="task-generation-source"
+              value={taskForm.generationSource}
+              onChange={(event) =>
+                setTaskForm((prev) => ({
+                  ...prev,
+                  generationSource: event.target.value as GenerationSource,
+                }))
+              }
+              className="mt-1 w-full rounded-md border border-slate-700 bg-slate-900 px-3 py-2 text-sm text-slate-100"
+            >
+              <option value="openai">OpenAI</option>
+              <option value="llama">Llama</option>
+            </select>
+          </div>
+
          <div>
            <label htmlFor="task-lookback" className="text-sm font-medium text-slate-200">
              Email Period
--- a/frontend/src/pages/EditTaskPage.tsx
+++ b/frontend/src/pages/EditTaskPage.tsx
@@ -13,6 +13,7 @@ import {
  inactivateTask,
  updateTask,
  type EmailLookback,
+  type GenerationSource,
 } from '../api/tasksApi'

 interface TaskFormState {
@@ -20,6 +21,7 @@ interface TaskFormState {
  prompt: string
  scheduleCron: string
  emailLookback: EmailLookback
+  generationSource: GenerationSource
 }

 interface CronParts {
@@ -97,6 +99,7 @@ const DEFAULT_TASK_FORM: TaskFormState = {
  prompt: '',
  scheduleCron: buildCron(DEFAULT_CRON_PARTS),
  emailLookback: 'last_week',
+  generationSource: 'openai',
 }

 async function invalidateTaskQueries(
@@ -164,6 +167,7 @@ export default function EditTaskPage() {
      prompt: task.prompt,
      scheduleCron: task.scheduleCron,
      emailLookback: task.emailLookback,
+      generationSource: task.generationSource,
    })
  }, [task])

@@ -175,6 +179,7 @@ export default function EditTaskPage() {
        prompt: data.prompt,
        scheduleCron: data.scheduleCron,
        emailLookback: data.emailLookback,
+        generationSource: data.generationSource,
      }),
    onSuccess: async () => {
      await invalidateTaskQueries(queryClient, entityId, taskId)
@@ -236,8 +241,16 @@ export default function EditTaskPage() {
      prompt: taskForm.prompt,
      scheduleCron: taskForm.scheduleCron,
      emailLookback: taskForm.emailLookback,
+      generationSource: taskForm.generationSource,
    }),
-    [entityId, taskForm.emailLookback, taskForm.name, taskForm.prompt, taskForm.scheduleCron]
+    [
+      entityId,
+      taskForm.emailLookback,
+      taskForm.generationSource,
+      taskForm.name,
+      taskForm.prompt,
+      taskForm.scheduleCron,
+    ]
  )

  const finalPrompt = useMemo(() => {
@@ -352,6 +365,26 @@ export default function EditTaskPage() {
            />
          </div>

+          <div>
+            <label htmlFor="task-generation-source" className="text-sm font-medium text-slate-200">
+              Generation Source
+            </label>
+            <select
+              id="task-generation-source"
+              value={taskForm.generationSource}
+              onChange={(event) =>
+                setTaskForm((prev) => ({
+                  ...prev,
+                  generationSource: event.target.value as GenerationSource,
+                }))
+              }
+              className="mt-1 w-full rounded-md border border-slate-700 bg-slate-900 px-3 py-2 text-sm text-slate-100"
+            >
+              <option value="openai">OpenAI</option>
+              <option value="llama">Llama</option>
+            </select>
+          </div>
+
          <div>
            <label htmlFor="task-lookback" className="text-sm font-medium text-slate-200">
              Email Period