summaryrefslogtreecommitdiff
path: root/Omni/Agent/Worker.hs
diff options
context:
space:
mode:
Diffstat (limited to 'Omni/Agent/Worker.hs')
-rw-r--r--Omni/Agent/Worker.hs186
1 files changed, 178 insertions, 8 deletions
diff --git a/Omni/Agent/Worker.hs b/Omni/Agent/Worker.hs
index 424a838..dafd0b2 100644
--- a/Omni/Agent/Worker.hs
+++ b/Omni/Agent/Worker.hs
@@ -13,10 +13,13 @@ import qualified Data.Text.Encoding as TE
import qualified Data.Text.IO as TIO
import qualified Data.Time
import qualified Omni.Agent.Core as Core
+import qualified Omni.Agent.Engine as Engine
import qualified Omni.Agent.Log as AgentLog
+import qualified Omni.Agent.Tools as Tools
import qualified Omni.Fact as Fact
import qualified Omni.Task.Core as TaskCore
import qualified System.Directory as Directory
+import qualified System.Environment as Env
import qualified System.Exit as Exit
import System.FilePath ((</>))
import qualified System.IO as IO
@@ -86,18 +89,35 @@ processTask worker task = do
TaskCore.updateTaskStatus tid TaskCore.InProgress []
say "[worker] Status -> InProgress"
- -- Run Amp with timing
- say "[worker] Starting amp..."
+ -- Check if we should use native engine
+ useEngine <- shouldUseEngine
+
+ -- Run agent with timing
startTime <- Data.Time.getCurrentTime
activityId <- TaskCore.logActivityWithMetrics tid TaskCore.Running Nothing Nothing (Just startTime) Nothing Nothing Nothing
- (exitCode, output) <- runAmp repo task
+
+ (exitCode, output, maybeCost) <-
+ if useEngine
+ then do
+ say "[worker] Starting native engine..."
+ (code, out, cost) <- runWithEngine repo task
+ pure (code, out, Just cost)
+ else do
+ say "[worker] Starting amp..."
+ (code, out) <- runAmp repo task
+ pure (code, out, Nothing)
+
endTime <- Data.Time.getCurrentTime
- say ("[worker] Amp exited with: " <> tshow exitCode)
+ say ("[worker] Agent exited with: " <> tshow exitCode)
- -- Capture metrics from agent log (thread URL, credits)
- status <- AgentLog.getStatus
- let threadUrl = ("https://ampcode.com/threads/" <>) </ AgentLog.statusThread status
- let costCents = Just <| floor (AgentLog.statusCredits status * 100)
+ -- Capture metrics - from engine result or agent log
+ (threadUrl, costCents) <- case maybeCost of
+ Just engineCost -> pure (Nothing, Just engineCost)
+ Nothing -> do
+ status <- AgentLog.getStatus
+ let url = ("https://ampcode.com/threads/" <>) </ AgentLog.statusThread status
+ let cost = Just <| floor (AgentLog.statusCredits status * 100)
+ pure (url, cost)
-- Update the activity record with metrics
TaskCore.updateActivityMetrics activityId threadUrl (Just endTime) costCents Nothing
@@ -300,6 +320,156 @@ runAmp repo task = do
killThread tid
pure (exitCode, output)
+-- | Check if we should use native engine instead of amp subprocess
+shouldUseEngine :: IO Bool
+shouldUseEngine = do
+ env <- Env.lookupEnv "JR_USE_ENGINE"
+ pure <| env == Just "1"
+
+-- | Run task using native Engine instead of amp subprocess
+-- Returns (ExitCode, output text, cost in cents)
+runWithEngine :: FilePath -> TaskCore.Task -> IO (Exit.ExitCode, Text, Int)
+runWithEngine repo task = do
+ -- Read API key from environment
+ maybeApiKey <- Env.lookupEnv "OPENROUTER_API_KEY"
+ case maybeApiKey of
+ Nothing -> pure (Exit.ExitFailure 1, "OPENROUTER_API_KEY not set", 0)
+ Just apiKey -> do
+ -- Check for retry context
+ maybeRetry <- TaskCore.getRetryContext (TaskCore.taskId task)
+
+ -- Build the full prompt (same as runAmp)
+ let ns = fromMaybe "." (TaskCore.taskNamespace task)
+ let basePrompt = buildBasePrompt task ns repo
+
+ -- Add retry context if present
+ let retryPrompt = buildRetryPrompt maybeRetry
+
+ let prompt = basePrompt <> retryPrompt
+
+ -- Read AGENTS.md
+ agentsMd <-
+ fmap (fromMaybe "") <| do
+ exists <- Directory.doesFileExist (repo </> "AGENTS.md")
+ if exists
+ then Just </ readFile (repo </> "AGENTS.md")
+ else pure Nothing
+
+ -- Get relevant facts from the knowledge base
+ relevantFacts <- getRelevantFacts task
+ let factsSection = formatFacts relevantFacts
+
+ -- Build system prompt
+ let systemPrompt =
+ prompt
+ <> "\n\nREPOSITORY GUIDELINES (AGENTS.md):\n"
+ <> agentsMd
+ <> factsSection
+
+ -- Build user prompt from task comments
+ let userPrompt = formatTask task
+
+ -- Select model based on task complexity (simple heuristic)
+ let model = selectModel task
+
+ -- Build Engine config with callbacks
+ totalCostRef <- newIORef (0 :: Int)
+ let engineCfg =
+ Engine.EngineConfig
+ { Engine.engineLLM =
+ Engine.defaultLLM
+ { Engine.llmApiKey = Text.pack apiKey
+ },
+ Engine.engineOnCost = \tokens cost -> do
+ modifyIORef' totalCostRef (+ cost)
+ AgentLog.log <| "Cost: " <> tshow cost <> " cents (" <> tshow tokens <> " tokens)",
+ Engine.engineOnActivity = \activity -> do
+ AgentLog.log <| "[engine] " <> activity,
+ Engine.engineOnToolCall = \toolName result -> do
+ AgentLog.log <| "[tool] " <> toolName <> ": " <> Text.take 100 result
+ }
+
+ -- Build Agent config
+ let agentCfg =
+ Engine.AgentConfig
+ { Engine.agentModel = model,
+ Engine.agentTools = Tools.allTools,
+ Engine.agentSystemPrompt = systemPrompt,
+ Engine.agentMaxIterations = 20
+ }
+
+ -- Run the agent
+ result <- Engine.runAgent engineCfg agentCfg userPrompt
+ totalCost <- readIORef totalCostRef
+
+ case result of
+ Left err -> pure (Exit.ExitFailure 1, "Engine error: " <> err, totalCost)
+ Right agentResult -> do
+ let output = Engine.resultFinalMessage agentResult
+ pure (Exit.ExitSuccess, output, totalCost)
+
+-- | Build the base prompt for the agent
+buildBasePrompt :: TaskCore.Task -> Text -> FilePath -> Text
+buildBasePrompt task ns repo =
+ "You are a Worker Agent.\n"
+ <> "Your goal is to implement the following task:\n\n"
+ <> formatTask task
+ <> "\n\nCRITICAL INSTRUCTIONS:\n"
+ <> "1. Analyze the codebase to understand where to make changes.\n"
+ <> "2. Implement the changes by editing files.\n"
+ <> "3. BEFORE finishing, you MUST run: bild --test "
+ <> ns
+ <> "\n"
+ <> "4. Fix ALL errors from bild --test (including hlint suggestions).\n"
+ <> "5. Keep running bild --test until it passes with no errors.\n"
+ <> "6. Do NOT update task status or manage git.\n"
+ <> "7. Only exit after bild --test passes.\n\n"
+ <> "IMPORTANT: The git commit will fail if hlint finds issues.\n"
+ <> "You must fix hlint suggestions like:\n"
+ <> "- 'Use list comprehension' -> use [x | cond] instead of if/else\n"
+ <> "- 'Avoid lambda' -> use function composition\n"
+ <> "- 'Redundant bracket' -> remove unnecessary parens\n\n"
+ <> "Context:\n"
+ <> "- Working directory: "
+ <> Text.pack repo
+ <> "\n"
+ <> "- Namespace: "
+ <> ns
+ <> "\n"
+
+-- | Build retry context prompt
+buildRetryPrompt :: Maybe TaskCore.RetryContext -> Text
+buildRetryPrompt Nothing = ""
+buildRetryPrompt (Just ctx) =
+ "\n\n## RETRY CONTEXT (IMPORTANT)\n\n"
+ <> "This task was previously attempted but failed. Attempt: "
+ <> tshow (TaskCore.retryAttempt ctx)
+ <> "/3\n"
+ <> "Reason: "
+ <> TaskCore.retryReason ctx
+ <> "\n\n"
+ <> ( if null (TaskCore.retryConflictFiles ctx)
+ then ""
+ else
+ "Conflicting files from previous attempt:\n"
+ <> Text.unlines (map (" - " <>) (TaskCore.retryConflictFiles ctx))
+ <> "\n"
+ )
+ <> "Original commit: "
+ <> TaskCore.retryOriginalCommit ctx
+ <> "\n\n"
+ <> maybe "" (\notes -> "## HUMAN NOTES/GUIDANCE\n\n" <> notes <> "\n\n") (TaskCore.retryNotes ctx)
+ <> "INSTRUCTIONS FOR RETRY:\n"
+ <> "- The codebase has changed since your last attempt\n"
+ <> "- Re-implement this task on top of the CURRENT codebase\n"
+ <> "- If there were merge conflicts, the conflicting files may have been modified by others\n"
+ <> "- Review the current state of those files before making changes\n"
+
+-- | Select model based on task complexity
+-- Currently always uses claude-sonnet-4, but can be extended for model selection
+selectModel :: TaskCore.Task -> Text
+selectModel _ = "anthropic/claude-sonnet-4-20250514"
+
formatTask :: TaskCore.Task -> Text
formatTask t =
"Task: "
generated by cgit v1.2.3 (git 2.25.1) at 2025-12-16 17:00:02 +0000