1 files changed, 582 insertions, 0 deletions
diff --git a/Omni/Agent/Engine.hs b/Omni/Agent/Engine.hs
new file mode 100644
index 0000000..e019341
--- /dev/null
+++ b/Omni/Agent/Engine.hs
@@ -0,0 +1,582 @@
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+
+-- | LLM Agent Engine - Tool protocol and LLM provider abstraction.
+--
+-- This module provides the core abstractions for building LLM-powered agents:
+-- - Tool: Defines tools that agents can use
+-- - LLM: OpenAI-compatible chat completions API provider
+-- - AgentConfig: Configuration for running agents
+--
+-- : out omni-agent-engine
+-- : dep http-conduit
+-- : dep aeson
+-- : dep case-insensitive
+module Omni.Agent.Engine
+  ( Tool (..),
+    LLM (..),
+    EngineConfig (..),
+    AgentConfig (..),
+    AgentResult (..),
+    Message (..),
+    Role (..),
+    ToolCall (..),
+    FunctionCall (..),
+    ToolResult (..),
+    ChatCompletionRequest (..),
+    ChatCompletionResponse (..),
+    Choice (..),
+    Usage (..),
+    defaultLLM,
+    defaultEngineConfig,
+    defaultAgentConfig,
+    chat,
+    runAgent,
+    main,
+    test,
+  )
+where
+
+import Alpha
+import Data.Aeson ((.!=), (.:), (.:?), (.=))
+import qualified Data.Aeson as Aeson
+import qualified Data.ByteString.Lazy as BL
+import qualified Data.CaseInsensitive as CI
+import qualified Data.Map.Strict as Map
+import qualified Data.Text as Text
+import qualified Data.Text.Encoding as TE
+import qualified Network.HTTP.Simple as HTTP
+import qualified Omni.Test as Test
+
+main :: IO ()
+main = Test.run test
+
+test :: Test.Tree
+test =
+  Test.group
+    "Omni.Agent.Engine"
+    [ Test.unit "Tool JSON roundtrip" <| do
+        let tool =
+              Tool
+                { toolName = "get_weather",
+                  toolDescription = "Get weather for a location",
+                  toolJsonSchema = Aeson.object ["type" .= ("object" :: Text), "properties" .= Aeson.object []],
+                  toolExecute = \_ -> pure (Aeson.String "sunny")
+                }
+        let encoded = encodeToolForApi tool
+        case Aeson.decode (Aeson.encode encoded) of
+          Nothing -> Test.assertFailure "Failed to decode tool"
+          Just decoded -> toolName tool Test.@=? toolApiName decoded,
+      Test.unit "Message JSON roundtrip" <| do
+        let msg = Message User "Hello" Nothing Nothing
+        case Aeson.decode (Aeson.encode msg) of
+          Nothing -> Test.assertFailure "Failed to decode message"
+          Just decoded -> msgContent msg Test.@=? msgContent decoded,
+      Test.unit "defaultLLM has correct endpoint" <| do
+        llmBaseUrl defaultLLM Test.@=? "https://openrouter.ai/api/v1",
+      Test.unit "defaultLLM has OpenRouter headers" <| do
+        length (llmExtraHeaders defaultLLM) Test.@=? 2
+        llmModel defaultLLM Test.@=? "anthropic/claude-sonnet-4.5",
+      Test.unit "defaultAgentConfig has sensible defaults" <| do
+        agentMaxIterations defaultAgentConfig Test.@=? 10,
+      Test.unit "defaultEngineConfig has no-op callbacks" <| do
+        engineOnCost defaultEngineConfig 100 5
+        engineOnActivity defaultEngineConfig "test"
+        engineOnToolCall defaultEngineConfig "tool" "result"
+        True Test.@=? True,
+      Test.unit "buildToolMap creates correct map" <| do
+        let tool1 =
+              Tool
+                { toolName = "tool1",
+                  toolDescription = "First tool",
+                  toolJsonSchema = Aeson.object [],
+                  toolExecute = \_ -> pure Aeson.Null
+                }
+            tool2 =
+              Tool
+                { toolName = "tool2",
+                  toolDescription = "Second tool",
+                  toolJsonSchema = Aeson.object [],
+                  toolExecute = \_ -> pure Aeson.Null
+                }
+            toolMap = buildToolMap [tool1, tool2]
+        Map.size toolMap Test.@=? 2
+        Map.member "tool1" toolMap Test.@=? True
+        Map.member "tool2" toolMap Test.@=? True,
+      Test.unit "Usage JSON parsing" <| do
+        let json = "{\"prompt_tokens\":100,\"completion_tokens\":50,\"total_tokens\":150}"
+        case Aeson.decode json of
+          Nothing -> Test.assertFailure "Failed to decode usage"
+          Just usage -> do
+            usagePromptTokens usage Test.@=? 100
+            usageCompletionTokens usage Test.@=? 50
+            usageTotalTokens usage Test.@=? 150,
+      Test.unit "AgentResult JSON roundtrip" <| do
+        let result =
+              AgentResult
+                { resultFinalMessage = "Done",
+                  resultToolCallCount = 3,
+                  resultIterations = 2,
+                  resultTotalCost = 50,
+                  resultTotalTokens = 1500
+                }
+        case Aeson.decode (Aeson.encode result) of
+          Nothing -> Test.assertFailure "Failed to decode AgentResult"
+          Just decoded -> do
+            resultFinalMessage decoded Test.@=? "Done"
+            resultToolCallCount decoded Test.@=? 3
+            resultIterations decoded Test.@=? 2,
+      Test.unit "estimateCost calculates correctly" <| do
+        let gpt4oCost = estimateCost "gpt-4o" 1000
+            gpt4oMiniCost = estimateCost "gpt-4o-mini" 1000
+        (gpt4oCost >= gpt4oMiniCost) Test.@=? True,
+      Test.unit "ToolCall JSON roundtrip" <| do
+        let tc =
+              ToolCall
+                { tcId = "call_123",
+                  tcType = "function",
+                  tcFunction = FunctionCall "read_file" "{\"path\":\"/tmp/test\"}"
+                }
+        case Aeson.decode (Aeson.encode tc) of
+          Nothing -> Test.assertFailure "Failed to decode ToolCall"
+          Just decoded -> tcId decoded Test.@=? "call_123",
+      Test.unit "FunctionCall JSON roundtrip" <| do
+        let fc = FunctionCall "test_func" "{\"arg\":\"value\"}"
+        case Aeson.decode (Aeson.encode fc) of
+          Nothing -> Test.assertFailure "Failed to decode FunctionCall"
+          Just decoded -> do
+            fcName decoded Test.@=? "test_func"
+            fcArguments decoded Test.@=? "{\"arg\":\"value\"}",
+      Test.unit "Role JSON roundtrip for all roles" <| do
+        let roles = [System, User, Assistant, ToolRole]
+        forM_ roles <| \role ->
+          case Aeson.decode (Aeson.encode role) of
+            Nothing -> Test.assertFailure ("Failed to decode Role: " <> show role)
+            Just decoded -> decoded Test.@=? role
+    ]
+
+data Tool = Tool
+  { toolName :: Text,
+    toolDescription :: Text,
+    toolJsonSchema :: Aeson.Value,
+    toolExecute :: Aeson.Value -> IO Aeson.Value
+  }
+
+data ToolApi = ToolApi
+  { toolApiName :: Text,
+    toolApiDescription :: Text,
+    toolApiParameters :: Aeson.Value
+  }
+  deriving (Generic)
+
+instance Aeson.ToJSON ToolApi where
+  toJSON t =
+    Aeson.object
+      [ "type" .= ("function" :: Text),
+        "function"
+          .= Aeson.object
+            [ "name" .= toolApiName t,
+              "description" .= toolApiDescription t,
+              "parameters" .= toolApiParameters t
+            ]
+      ]
+
+instance Aeson.FromJSON ToolApi where
+  parseJSON =
+    Aeson.withObject "ToolApi" <| \v -> do
+      fn <- v .: "function"
+      (ToolApi </ (fn .: "name"))
+        <*> (fn .: "description")
+        <*> (fn .: "parameters")
+
+encodeToolForApi :: Tool -> ToolApi
+encodeToolForApi t =
+  ToolApi
+    { toolApiName = toolName t,
+      toolApiDescription = toolDescription t,
+      toolApiParameters = toolJsonSchema t
+    }
+
+data LLM = LLM
+  { llmBaseUrl :: Text,
+    llmApiKey :: Text,
+    llmModel :: Text,
+    llmExtraHeaders :: [(ByteString, ByteString)]
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON LLM where
+  toJSON l =
+    Aeson.object
+      [ "llmBaseUrl" .= llmBaseUrl l,
+        "llmApiKey" .= llmApiKey l,
+        "llmModel" .= llmModel l
+      ]
+
+instance Aeson.FromJSON LLM where
+  parseJSON =
+    Aeson.withObject "LLM" <| \v ->
+      (LLM </ (v .: "llmBaseUrl"))
+        <*> (v .: "llmApiKey")
+        <*> (v .: "llmModel")
+        <*> pure []
+
+defaultLLM :: LLM
+defaultLLM =
+  LLM
+    { llmBaseUrl = "https://openrouter.ai/api/v1",
+      llmApiKey = "",
+      llmModel = "anthropic/claude-sonnet-4.5",
+      llmExtraHeaders =
+        [ ("HTTP-Referer", "https://omni.dev"),
+          ("X-Title", "Omni Agent")
+        ]
+    }
+
+data AgentConfig = AgentConfig
+  { agentModel :: Text,
+    agentTools :: [Tool],
+    agentSystemPrompt :: Text,
+    agentMaxIterations :: Int
+  }
+
+defaultAgentConfig :: AgentConfig
+defaultAgentConfig =
+  AgentConfig
+    { agentModel = "gpt-4",
+      agentTools = [],
+      agentSystemPrompt = "You are a helpful assistant.",
+      agentMaxIterations = 10
+    }
+
+data EngineConfig = EngineConfig
+  { engineLLM :: LLM,
+    engineOnCost :: Int -> Int -> IO (),
+    engineOnActivity :: Text -> IO (),
+    engineOnToolCall :: Text -> Text -> IO ()
+  }
+
+defaultEngineConfig :: EngineConfig
+defaultEngineConfig =
+  EngineConfig
+    { engineLLM = defaultLLM,
+      engineOnCost = \_ _ -> pure (),
+      engineOnActivity = \_ -> pure (),
+      engineOnToolCall = \_ _ -> pure ()
+    }
+
+data AgentResult = AgentResult
+  { resultFinalMessage :: Text,
+    resultToolCallCount :: Int,
+    resultIterations :: Int,
+    resultTotalCost :: Int,
+    resultTotalTokens :: Int
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON AgentResult
+
+instance Aeson.FromJSON AgentResult
+
+data Role = System | User | Assistant | ToolRole
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON Role where
+  toJSON System = Aeson.String "system"
+  toJSON User = Aeson.String "user"
+  toJSON Assistant = Aeson.String "assistant"
+  toJSON ToolRole = Aeson.String "tool"
+
+instance Aeson.FromJSON Role where
+  parseJSON = Aeson.withText "Role" parseRole
+    where
+      parseRole "system" = pure System
+      parseRole "user" = pure User
+      parseRole "assistant" = pure Assistant
+      parseRole "tool" = pure ToolRole
+      parseRole _ = empty
+
+data Message = Message
+  { msgRole :: Role,
+    msgContent :: Text,
+    msgToolCalls :: Maybe [ToolCall],
+    msgToolCallId :: Maybe Text
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON Message where
+  toJSON m =
+    Aeson.object
+      <| catMaybes
+        [ Just ("role" .= msgRole m),
+          Just ("content" .= msgContent m),
+          ("tool_calls" .=) </ msgToolCalls m,
+          ("tool_call_id" .=) </ msgToolCallId m
+        ]
+
+instance Aeson.FromJSON Message where
+  parseJSON =
+    Aeson.withObject "Message" <| \v ->
+      (Message </ (v .: "role"))
+        <*> (v .:? "content" .!= "")
+        <*> (v .:? "tool_calls")
+        <*> (v .:? "tool_call_id")
+
+data ToolCall = ToolCall
+  { tcId :: Text,
+    tcType :: Text,
+    tcFunction :: FunctionCall
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON ToolCall where
+  toJSON tc =
+    Aeson.object
+      [ "id" .= tcId tc,
+        "type" .= tcType tc,
+        "function" .= tcFunction tc
+      ]
+
+instance Aeson.FromJSON ToolCall where
+  parseJSON =
+    Aeson.withObject "ToolCall" <| \v ->
+      (ToolCall </ (v .: "id"))
+        <*> (v .:? "type" .!= "function")
+        <*> (v .: "function")
+
+data FunctionCall = FunctionCall
+  { fcName :: Text,
+    fcArguments :: Text
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON FunctionCall where
+  toJSON fc =
+    Aeson.object
+      [ "name" .= fcName fc,
+        "arguments" .= fcArguments fc
+      ]
+
+instance Aeson.FromJSON FunctionCall where
+  parseJSON =
+    Aeson.withObject "FunctionCall" <| \v ->
+      (FunctionCall </ (v .: "name"))
+        <*> (v .: "arguments")
+
+data ToolResult = ToolResult
+  { trToolCallId :: Text,
+    trContent :: Text
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON ToolResult
+
+instance Aeson.FromJSON ToolResult
+
+data ChatCompletionRequest = ChatCompletionRequest
+  { reqModel :: Text,
+    reqMessages :: [Message],
+    reqTools :: Maybe [ToolApi]
+  }
+  deriving (Generic)
+
+instance Aeson.ToJSON ChatCompletionRequest where
+  toJSON r =
+    Aeson.object
+      <| catMaybes
+        [ Just ("model" .= reqModel r),
+          Just ("messages" .= reqMessages r),
+          ("tools" .=) </ reqTools r
+        ]
+
+data Choice = Choice
+  { choiceIndex :: Int,
+    choiceMessage :: Message,
+    choiceFinishReason :: Maybe Text
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON Choice where
+  parseJSON =
+    Aeson.withObject "Choice" <| \v ->
+      (Choice </ (v .: "index"))
+        <*> (v .: "message")
+        <*> (v .:? "finish_reason")
+
+data Usage = Usage
+  { usagePromptTokens :: Int,
+    usageCompletionTokens :: Int,
+    usageTotalTokens :: Int
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON Usage where
+  parseJSON =
+    Aeson.withObject "Usage" <| \v ->
+      (Usage </ (v .: "prompt_tokens"))
+        <*> (v .: "completion_tokens")
+        <*> (v .: "total_tokens")
+
+data ChatCompletionResponse = ChatCompletionResponse
+  { respId :: Text,
+    respChoices :: [Choice],
+    respModel :: Text,
+    respUsage :: Maybe Usage
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON ChatCompletionResponse where
+  parseJSON =
+    Aeson.withObject "ChatCompletionResponse" <| \v ->
+      (ChatCompletionResponse </ (v .: "id"))
+        <*> (v .: "choices")
+        <*> (v .: "model")
+        <*> (v .:? "usage")
+
+data ChatResult = ChatResult
+  { chatMessage :: Message,
+    chatUsage :: Maybe Usage
+  }
+  deriving (Show, Eq)
+
+chatWithUsage :: LLM -> [Tool] -> [Message] -> IO (Either Text ChatResult)
+chatWithUsage llm tools messages = do
+  let url = Text.unpack (llmBaseUrl llm) <> "/chat/completions"
+  req0 <- HTTP.parseRequest url
+  let toolApis = [encodeToolForApi t | not (null tools), t <- tools]
+      body =
+        ChatCompletionRequest
+          { reqModel = llmModel llm,
+            reqMessages = messages,
+            reqTools = if null toolApis then Nothing else Just toolApis
+          }
+      baseReq =
+        HTTP.setRequestMethod "POST"
+          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
+          <| HTTP.setRequestHeader "Authorization" ["Bearer " <> TE.encodeUtf8 (llmApiKey llm)]
+          <| HTTP.setRequestBodyLBS (Aeson.encode body)
+          <| req0
+      req = foldr addHeader baseReq (llmExtraHeaders llm)
+      addHeader (name, value) = HTTP.addRequestHeader (CI.mk name) value
+
+  response <- HTTP.httpLBS req
+  let status = HTTP.getResponseStatusCode response
+  if status >= 200 && status < 300
+    then case Aeson.decode (HTTP.getResponseBody response) of
+      Just resp ->
+        case respChoices resp of
+          (c : _) -> pure (Right (ChatResult (choiceMessage c) (respUsage resp)))
+          [] -> pure (Left "No choices in response")
+      Nothing -> pure (Left "Failed to parse response")
+    else pure (Left ("HTTP error: " <> tshow status <> " - " <> TE.decodeUtf8 (BL.toStrict (HTTP.getResponseBody response))))
+
+chat :: LLM -> [Tool] -> [Message] -> IO (Either Text Message)
+chat llm tools messages = do
+  result <- chatWithUsage llm tools messages
+  pure (chatMessage </ result)
+
+runAgent :: EngineConfig -> AgentConfig -> Text -> IO (Either Text AgentResult)
+runAgent engineCfg agentCfg userPrompt = do
+  let llm =
+        (engineLLM engineCfg)
+          { llmModel = agentModel agentCfg
+          }
+      tools = agentTools agentCfg
+      toolMap = buildToolMap tools
+      systemMsg = Message System (agentSystemPrompt agentCfg) Nothing Nothing
+      userMsg = Message User userPrompt Nothing Nothing
+      initialMessages = [systemMsg, userMsg]
+
+  engineOnActivity engineCfg "Starting agent loop"
+  loop llm tools toolMap initialMessages 0 0 0
+  where
+    maxIter = agentMaxIterations agentCfg
+
+    loop :: LLM -> [Tool] -> Map.Map Text Tool -> [Message] -> Int -> Int -> Int -> IO (Either Text AgentResult)
+    loop llm tools' toolMap msgs iteration totalCalls totalTokens
+      | iteration >= maxIter =
+          pure
+            <| Left
+            <| "Max iterations ("
+            <> tshow maxIter
+            <> ") reached"
+      | otherwise = do
+          engineOnActivity engineCfg <| "Iteration " <> tshow (iteration + 1)
+          result <- chatWithUsage llm tools' msgs
+          case result of
+            Left err -> pure (Left err)
+            Right chatRes -> do
+              let msg = chatMessage chatRes
+                  tokens = maybe 0 usageTotalTokens (chatUsage chatRes)
+                  cost = estimateCost (llmModel llm) tokens
+              engineOnCost engineCfg tokens cost
+              let newTokens = totalTokens + tokens
+              case msgToolCalls msg of
+                Nothing -> do
+                  engineOnActivity engineCfg "Agent completed"
+                  pure
+                    <| Right
+                    <| AgentResult
+                      { resultFinalMessage = msgContent msg,
+                        resultToolCallCount = totalCalls,
+                        resultIterations = iteration + 1,
+                        resultTotalCost = estimateTotalCost (llmModel llm) newTokens,
+                        resultTotalTokens = newTokens
+                      }
+                Just [] -> do
+                  engineOnActivity engineCfg "Agent completed (empty tool calls)"
+                  pure
+                    <| Right
+                    <| AgentResult
+                      { resultFinalMessage = msgContent msg,
+                        resultToolCallCount = totalCalls,
+                        resultIterations = iteration + 1,
+                        resultTotalCost = estimateTotalCost (llmModel llm) newTokens,
+                        resultTotalTokens = newTokens
+                      }
+                Just tcs -> do
+                  toolResults <- executeToolCalls engineCfg toolMap tcs
+                  let newMsgs = msgs <> [msg] <> toolResults
+                      newCalls = totalCalls + length tcs
+                  loop llm tools' toolMap newMsgs (iteration + 1) newCalls newTokens
+
+buildToolMap :: [Tool] -> Map.Map Text Tool
+buildToolMap = Map.fromList <. map (\t -> (toolName t, t))
+
+executeToolCalls :: EngineConfig -> Map.Map Text Tool -> [ToolCall] -> IO [Message]
+executeToolCalls engineCfg toolMap = traverse executeSingle
+  where
+    executeSingle tc = do
+      let name = fcName (tcFunction tc)
+          argsText = fcArguments (tcFunction tc)
+          callId = tcId tc
+      engineOnActivity engineCfg <| "Executing tool: " <> name
+      case Map.lookup name toolMap of
+        Nothing -> do
+          let errMsg = "Tool not found: " <> name
+          engineOnToolCall engineCfg name errMsg
+          pure <| Message ToolRole errMsg Nothing (Just callId)
+        Just tool -> do
+          case Aeson.decode (BL.fromStrict (TE.encodeUtf8 argsText)) of
+            Nothing -> do
+              let errMsg = "Invalid JSON arguments: " <> argsText
+              engineOnToolCall engineCfg name errMsg
+              pure <| Message ToolRole errMsg Nothing (Just callId)
+            Just args -> do
+              resultValue <- toolExecute tool args
+              let resultText = TE.decodeUtf8 (BL.toStrict (Aeson.encode resultValue))
+                  summary = Text.take 100 resultText
+              engineOnToolCall engineCfg name summary
+              pure <| Message ToolRole resultText Nothing (Just callId)
+
+estimateCost :: Text -> Int -> Int
+estimateCost model tokens
+  | "gpt-4o-mini" `Text.isInfixOf` model = tokens * 15 `div` 1000000
+  | "gpt-4o" `Text.isInfixOf` model = tokens * 250 `div` 100000
+  | "gpt-4" `Text.isInfixOf` model = tokens * 3 `div` 100000
+  | "claude" `Text.isInfixOf` model = tokens * 3 `div` 100000
+  | otherwise = tokens `div` 100000
+
+estimateTotalCost :: Text -> Int -> Int
+estimateTotalCost = estimateCost