summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--Omni/Agent/Engine.hs317
1 files changed, 317 insertions, 0 deletions
diff --git a/Omni/Agent/Engine.hs b/Omni/Agent/Engine.hs
new file mode 100644
index 0000000..ac6c517
--- /dev/null
+++ b/Omni/Agent/Engine.hs
@@ -0,0 +1,317 @@
+{-# 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
+module Omni.Agent.Engine
+ ( Tool (..),
+ LLM (..),
+ AgentConfig (..),
+ Message (..),
+ Role (..),
+ ToolCall (..),
+ FunctionCall (..),
+ ToolResult (..),
+ ChatCompletionRequest (..),
+ ChatCompletionResponse (..),
+ Choice (..),
+ defaultLLM,
+ defaultAgentConfig,
+ chat,
+ main,
+ test,
+ )
+where
+
+import Alpha
+import Data.Aeson ((.!=), (.:), (.:?), (.=))
+import qualified Data.Aeson as Aeson
+import qualified Data.ByteString.Lazy as BL
+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://api.openai.com",
+ Test.unit "defaultAgentConfig has sensible defaults" <| do
+ agentMaxIterations defaultAgentConfig Test.@=? 10
+ ]
+
+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
+ }
+ deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON LLM
+
+instance Aeson.FromJSON LLM
+
+defaultLLM :: LLM
+defaultLLM =
+ LLM
+ { llmBaseUrl = "https://api.openai.com",
+ llmApiKey = "",
+ llmModel = "gpt-4"
+ }
+
+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 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 ChatCompletionResponse = ChatCompletionResponse
+ { respId :: Text,
+ respChoices :: [Choice],
+ respModel :: Text
+ }
+ deriving (Show, Eq, Generic)
+
+instance Aeson.FromJSON ChatCompletionResponse where
+ parseJSON =
+ Aeson.withObject "ChatCompletionResponse" <| \v ->
+ (ChatCompletionResponse </ (v .: "id"))
+ <*> (v .: "choices")
+ <*> (v .: "model")
+
+chat :: LLM -> [Tool] -> [Message] -> IO (Either Text Message)
+chat llm tools messages = do
+ let url = Text.unpack (llmBaseUrl llm) <> "/v1/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
+ }
+ req =
+ HTTP.setRequestMethod "POST"
+ <| HTTP.setRequestHeader "Content-Type" ["application/json"]
+ <| HTTP.setRequestHeader "Authorization" ["Bearer " <> TE.encodeUtf8 (llmApiKey llm)]
+ <| HTTP.setRequestBodyLBS (Aeson.encode body)
+ <| req0
+
+ 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 (choiceMessage c))
+ [] -> 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))))
generated by cgit v1.2.3 (git 2.25.1) at 2025-12-16 16:59:11 +0000