include argus workflow
This commit is contained in:
joyzhuo
278
GeminiVLMClient.swift
Normal file
278
GeminiVLMClient.swift
Normal file
@@ -0,0 +1,278 @@
|
||||
// GeminiVLMClient.swift — Native Swift Gemini Vision API client
|
||||
// Ports the Python argus VLM analysis (vlm.py) directly into Swift.
|
||||
// No subprocess required: screenshots go straight from ScreenCaptureKit → Gemini → UI.
|
||||
|
||||
import Foundation
|
||||
|
||||
struct GeminiVLMClient {
|
||||
|
||||
private static let apiBase = "https://generativelanguage.googleapis.com/v1beta/models"
|
||||
private static let model = "gemini-3.1-pro-preview"
|
||||
|
||||
let apiKey: String
|
||||
|
||||
// MARK: - Public
|
||||
|
||||
/// Analyze a sequence of JPEG frames and return a structured distraction analysis.
|
||||
/// - Parameters:
|
||||
/// - frames: JPEG screenshot frames, oldest first, newest last.
|
||||
/// - taskTitle: Current task title (empty if no session).
|
||||
/// - taskGoal: Task description / goal.
|
||||
/// - steps: Active step list for the current task.
|
||||
/// - windowTitle: Frontmost app name from NSWorkspace.
|
||||
/// - recentSummaries: Rolling summaries from previous analyses (temporal context).
|
||||
func analyze(
|
||||
frames: [Data],
|
||||
taskTitle: String,
|
||||
taskGoal: String,
|
||||
steps: [Step],
|
||||
windowTitle: String,
|
||||
recentSummaries: [String]
|
||||
) async throws -> DistractionAnalysisResponse {
|
||||
let prompt = buildPrompt(
|
||||
taskTitle: taskTitle,
|
||||
taskGoal: taskGoal,
|
||||
steps: steps,
|
||||
windowTitle: windowTitle,
|
||||
recentSummaries: recentSummaries
|
||||
)
|
||||
let raw = try await callGemini(prompt: prompt, frames: frames)
|
||||
return try parseResponse(raw)
|
||||
}
|
||||
|
||||
// MARK: - Prompt Builder (ported from vlm.py build_system_prompt)
|
||||
|
||||
private func buildPrompt(
|
||||
taskTitle: String,
|
||||
taskGoal: String,
|
||||
steps: [Step],
|
||||
windowTitle: String,
|
||||
recentSummaries: [String]
|
||||
) -> String {
|
||||
let stepsText: String
|
||||
if steps.isEmpty {
|
||||
stepsText = " (no steps defined)"
|
||||
} else {
|
||||
stepsText = steps.map { s in
|
||||
let marker: String
|
||||
switch s.status {
|
||||
case "pending": marker = "○"
|
||||
case "in_progress": marker = "►"
|
||||
case "done": marker = "✓"
|
||||
default: marker = "?"
|
||||
}
|
||||
var line = " \(marker) [\(s.status)] (id=\(s.id)) \(s.sortOrder). \(s.title)"
|
||||
if let note = s.checkpointNote { line += " — checkpoint: \(note)" }
|
||||
return line
|
||||
}.joined(separator: "\n")
|
||||
}
|
||||
|
||||
let historyText: String
|
||||
if recentSummaries.isEmpty {
|
||||
historyText = " (no previous frames)"
|
||||
} else {
|
||||
historyText = recentSummaries.enumerated()
|
||||
.map { i, s in " [frame \(i + 1)] \(s)" }
|
||||
.joined(separator: "\n")
|
||||
}
|
||||
|
||||
return """
|
||||
You are a proactive focus assistant analyzing a TIME SEQUENCE of screenshots.
|
||||
|
||||
## How to read the screenshots
|
||||
|
||||
You receive screenshots in chronological order (oldest first, newest last).
|
||||
Each frame is ~5 seconds apart. This means:
|
||||
- 2 unchanged frames = ~10 seconds idle — significant.
|
||||
- 3 unchanged frames = ~15 seconds idle — user is stuck or distracted.
|
||||
- If ALL frames are identical, the user has been idle for 15+ seconds — flag it.
|
||||
|
||||
Your PRIMARY signal is the DIFFERENCES between consecutive frames.
|
||||
Where the screen CHANGED = where attention is. Static areas = ignore.
|
||||
|
||||
Diff signals and what they mean:
|
||||
- New text appearing / cursor advancing → user is actively typing (this IS their task)
|
||||
- Window or tab switch → context change, could be reference or distraction
|
||||
- Same content, no pixel changes → stalled, idle, or reading
|
||||
- Repeated switching between same 2-3 apps → repetitive loop (manual data transfer)
|
||||
- Error message that APPEARED between frames → user just triggered it, relevant
|
||||
- Error message already in ALL frames → stale, ignore
|
||||
|
||||
CRITICAL — looking at something ≠ working on something:
|
||||
- User switches to browser/another app and just LOOKS → distraction or quick reference.
|
||||
- User switches and starts TYPING/EDITING → might be a new task.
|
||||
- If the user has an active session and switches away WITHOUT typing in the new app,
|
||||
they are DISTRACTED from their session, not starting a new task.
|
||||
- A single app switch is NEVER enough to infer a new task. Wait for active work.
|
||||
|
||||
## Current task context
|
||||
|
||||
Task: \(taskTitle.isEmpty ? "(no active task)" : taskTitle)
|
||||
Goal: \(taskGoal.isEmpty ? taskTitle : taskGoal)
|
||||
Steps:
|
||||
\(stepsText)
|
||||
Window title (OS): \(windowTitle.isEmpty ? "(unknown)" : windowTitle)
|
||||
|
||||
## Recent screen history (for temporal context)
|
||||
\(historyText)
|
||||
|
||||
## What to output
|
||||
|
||||
Analyze the screenshots and return JSON with EXACTLY this structure (no extra fields, no markdown):
|
||||
{
|
||||
"on_task": true,
|
||||
"current_step_id": "step UUID or null",
|
||||
"inferred_task": "what the user is actually working on based on screen diffs",
|
||||
"checkpoint_note_update": "what specifically changed across these frames",
|
||||
"steps_completed": [],
|
||||
"friction": {
|
||||
"type": "repetitive_loop | stalled | tedious_manual | context_overhead | task_resumption | none",
|
||||
"confidence": 0.0,
|
||||
"description": "what the user is struggling with",
|
||||
"proposed_actions": [
|
||||
{
|
||||
"label": "specific verb phrase: what to do",
|
||||
"action_type": "auto_extract | brain_dump | other",
|
||||
"details": "natural language spec for what action to take"
|
||||
}
|
||||
],
|
||||
"source_context": "filename or app name, or null",
|
||||
"target_context": "filename or app name, or null"
|
||||
},
|
||||
"session_action": {
|
||||
"type": "none",
|
||||
"session_id": null,
|
||||
"reason": ""
|
||||
},
|
||||
"intent": "skimming | engaged | unclear | null",
|
||||
"distraction_type": "app_switch | browsing | idle | null",
|
||||
"app_name": "primary visible application",
|
||||
"confidence": 0.8,
|
||||
"gentle_nudge": "short nudge message if distracted but no friction action applies, otherwise null",
|
||||
"vlm_summary": "1-sentence description of what CHANGED across the frames (not what is static)"
|
||||
}
|
||||
|
||||
FRICTION DETECTION rules:
|
||||
- REPETITIVE_LOOP: Switching between same 2-3 windows (copying data manually)
|
||||
- STALLED: No meaningful pixel changes across 2+ frames; or user wrote then deleted
|
||||
- TEDIOUS_MANUAL: Doing automatable work (filling forms, transcribing, copying by hand)
|
||||
- CONTEXT_OVERHEAD: Many windows open, visibly searching across them
|
||||
- TASK_RESUMPTION: User just returned to a task they were working on earlier
|
||||
|
||||
If friction confidence < 0.5, set type to "none".
|
||||
Only set gentle_nudge when user is off-task AND no actionable friction applies.
|
||||
"""
|
||||
}
|
||||
|
||||
// MARK: - Action Executor
|
||||
|
||||
/// Execute a user-approved proactive action and return a plain-text result.
|
||||
func executeAction(
|
||||
label: String,
|
||||
actionType: String,
|
||||
details: String,
|
||||
screenshot: Data?
|
||||
) async throws -> String {
|
||||
let taskInstruction: String
|
||||
switch actionType {
|
||||
case "auto_extract":
|
||||
taskInstruction = "Extract the relevant data from the screenshot and present it concisely as plain text."
|
||||
case "brain_dump":
|
||||
taskInstruction = "Format this as a short brain-dump note the user should add to their task list."
|
||||
default:
|
||||
taskInstruction = "Provide 2–3 concrete next steps the user can take right now."
|
||||
}
|
||||
let prompt = """
|
||||
You are a productivity assistant. The user approved this action: "\(label)"
|
||||
Details: \(details.isEmpty ? "(none)" : details)
|
||||
\(taskInstruction)
|
||||
Be specific and brief (3–5 sentences max). No markdown, no preamble, plain text only.
|
||||
"""
|
||||
let frames: [Data] = screenshot.map { [$0] } ?? []
|
||||
return try await callGemini(prompt: prompt, frames: frames)
|
||||
}
|
||||
|
||||
// MARK: - Gemini REST API Call
|
||||
|
||||
private func callGemini(prompt: String, frames: [Data]) async throws -> String {
|
||||
let urlStr = "\(Self.apiBase)/\(Self.model):generateContent?key=\(apiKey)"
|
||||
guard let url = URL(string: urlStr) else { throw URLError(.badURL) }
|
||||
|
||||
// Build content parts: label + image for each frame, then instruction
|
||||
var parts: [[String: Any]] = []
|
||||
let total = frames.count
|
||||
for (i, frame) in frames.enumerated() {
|
||||
parts.append(["text": "[Screenshot \(i + 1)/\(total) — \((total - i) * 5)s ago]"])
|
||||
parts.append([
|
||||
"inlineData": [
|
||||
"mimeType": "image/jpeg",
|
||||
"data": frame.base64EncodedString()
|
||||
]
|
||||
])
|
||||
}
|
||||
parts.append(["text": "Analyze this screenshot sequence now. Reply with ONLY valid JSON — no markdown, no code fences."])
|
||||
|
||||
let body: [String: Any] = [
|
||||
"systemInstruction": ["parts": [["text": prompt]]],
|
||||
"contents": [["parts": parts]],
|
||||
"generationConfig": [
|
||||
"temperature": 0.2,
|
||||
"maxOutputTokens": 1024
|
||||
]
|
||||
]
|
||||
|
||||
var request = URLRequest(url: url)
|
||||
request.httpMethod = "POST"
|
||||
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
|
||||
request.httpBody = try JSONSerialization.data(withJSONObject: body)
|
||||
request.timeoutInterval = 60
|
||||
|
||||
let (data, response) = try await URLSession.shared.data(for: request)
|
||||
|
||||
if let http = response as? HTTPURLResponse, http.statusCode != 200 {
|
||||
let msg = String(data: data, encoding: .utf8) ?? "HTTP \(http.statusCode)"
|
||||
print("[GeminiVLM] API error \(http.statusCode): \(msg)")
|
||||
throw URLError(.badServerResponse)
|
||||
}
|
||||
|
||||
guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
|
||||
let candidates = json["candidates"] as? [[String: Any]],
|
||||
let first = candidates.first,
|
||||
let content = first["content"] as? [String: Any],
|
||||
let contentParts = content["parts"] as? [[String: Any]],
|
||||
let text = contentParts.first?["text"] as? String
|
||||
else {
|
||||
let raw = String(data: data, encoding: .utf8) ?? ""
|
||||
print("[GeminiVLM] Unexpected response shape: \(raw.prefix(300))")
|
||||
throw URLError(.cannotParseResponse)
|
||||
}
|
||||
|
||||
print("[GeminiVLM] Response (\(text.count) chars): \(text.prefix(200))")
|
||||
return text
|
||||
}
|
||||
|
||||
// MARK: - Response Parsing
|
||||
|
||||
private func parseResponse(_ text: String) throws -> DistractionAnalysisResponse {
|
||||
var cleaned = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
// Strip ```json ... ``` or ``` ... ``` fences
|
||||
if cleaned.hasPrefix("```") {
|
||||
let lines = cleaned.components(separatedBy: "\n")
|
||||
cleaned = lines.dropFirst().joined(separator: "\n")
|
||||
if let backtickRange = cleaned.range(of: "```") {
|
||||
cleaned = String(cleaned[..<backtickRange.lowerBound])
|
||||
}
|
||||
cleaned = cleaned.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
}
|
||||
// Find JSON object boundaries robustly
|
||||
guard let start = cleaned.firstIndex(of: "{"),
|
||||
let end = cleaned.lastIndex(of: "}") else {
|
||||
throw URLError(.cannotParseResponse)
|
||||
}
|
||||
let jsonStr = String(cleaned[start...end])
|
||||
guard let jsonData = jsonStr.data(using: .utf8) else {
|
||||
throw URLError(.cannotParseResponse)
|
||||
}
|
||||
return try JSONDecoder().decode(DistractionAnalysisResponse.self, from: jsonData)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user