|
| 1 | +import type { Bindings } from '@/types' |
| 2 | +import type { ProcessedTicket, TicketData } from '@/types/ticket' |
| 3 | +import type { LanguageModelV1 } from 'ai' |
| 4 | +import { ENV } from '@/config/environment' |
| 5 | +import { generateText } from 'ai' |
| 6 | +import { createWorkersAI } from 'workers-ai-provider' |
| 7 | + |
| 8 | +export class AIProcessor { |
| 9 | + private model: LanguageModelV1 |
| 10 | + private systemPrompt: string |
| 11 | + |
| 12 | + constructor(env: Bindings) { |
| 13 | + const workersai = createWorkersAI({ binding: env.AI }) |
| 14 | + this.model = workersai(ENV.MODEL_NAME, { |
| 15 | + safePrompt: true, |
| 16 | + }) |
| 17 | + this.systemPrompt = ` |
| 18 | + You are a ticket classification system. Classify tickets as either Bug, Story, Task, or Spike. |
| 19 | +
|
| 20 | + - Bugs: Issues or unexpected behavior |
| 21 | + - Stories: User-facing features |
| 22 | + - Tasks: Technical work items |
| 23 | + - Spikes: Research or exploration items |
| 24 | + |
| 25 | + Respond only with the classification label. |
| 26 | + ` |
| 27 | + } |
| 28 | + |
| 29 | + public async classifyTicket(ticket: TicketData): Promise<ProcessedTicket | Error> { |
| 30 | + const startTime = Date.now() |
| 31 | + const userContent = `Title: ${ticket.title}\nDescription: ${ticket.description}` |
| 32 | + |
| 33 | + const { text } = await generateText({ |
| 34 | + model: this.model, |
| 35 | + messages: [{ |
| 36 | + role: 'system', |
| 37 | + content: this.systemPrompt, |
| 38 | + }, { |
| 39 | + role: 'user', |
| 40 | + content: userContent, |
| 41 | + }], |
| 42 | + }) |
| 43 | + |
| 44 | + const result: ProcessedTicket = { |
| 45 | + ...ticket, |
| 46 | + predictedLabel: this.parseResponse(text), |
| 47 | + processingTime: Date.now() - startTime, |
| 48 | + } |
| 49 | + |
| 50 | + return result |
| 51 | + } |
| 52 | + |
| 53 | + private parseResponse(response: string): ProcessedTicket['predictedLabel'] { |
| 54 | + const normalized = response.trim().toLowerCase() |
| 55 | + if (normalized.includes('bug')) |
| 56 | + return 'Bug' |
| 57 | + if (normalized.includes('story')) |
| 58 | + return 'Story' |
| 59 | + if (normalized.includes('spike')) |
| 60 | + return 'Spike' |
| 61 | + return 'Task' |
| 62 | + } |
| 63 | +} |
0 commit comments