diff --git a/datalog-lsp/src/Datalog/LSP.hs b/datalog-lsp/src/Datalog/LSP.hs
new file mode 100644
index 0000000..6fe6131
--- /dev/null
+++ b/datalog-lsp/src/Datalog/LSP.hs
@@ -0,0 +1,43 @@
+module Datalog.LSP where
+
+import Control.Monad.IO.Class
+import Language.LSP.Protocol.Message
+import Language.LSP.Protocol.Types
+import Language.LSP.Server
+import Datalog.LSP.Hover (hoverHandler)
+import Datalog.LSP.Highlight (tokenHandler)
+
+handlers :: Handlers (LspM ())
+handlers =
+    mconcat
+        [ initHandler
+        , hoverHandler
+        , tokenHandler
+        ]
+
+initHandler :: Handlers (LspM ())
+initHandler = notificationHandler SMethod_Initialized $ \_ -> pure ()
+
+serverDefinition :: ServerDefinition ()
+serverDefinition = 
+        ServerDefinition
+            { parseConfig = const $ const $ Right ()
+            , onConfigChange = const $ pure ()
+            , defaultConfig = ()
+            , configSection = "demo"
+            , doInitialize = \env _req -> pure $ Right env
+            , staticHandlers = const handlers
+            , interpretHandler = \env -> Iso (runLspT env) liftIO
+            , options =
+                defaultOptions
+                    { optTextDocumentSync =
+                        Just $
+                            TextDocumentSyncOptions
+                                { _openClose = Just True
+                                , _change = Just TextDocumentSyncKind_Full
+                                , _willSave = Just False
+                                , _willSaveWaitUntil = Just False
+                                , _save = Just (InR (SaveOptions (Just False)))
+                                }
+                    }
+            }
diff --git a/datalog-lsp/src/Datalog/LSP/Highlight.hs b/datalog-lsp/src/Datalog/LSP/Highlight.hs
new file mode 100644
index 0000000..0a5890e
--- /dev/null
+++ b/datalog-lsp/src/Datalog/LSP/Highlight.hs
@@ -0,0 +1,28 @@
+module Datalog.LSP.Highlight where
+
+import Control.Lens ((^.), Lens')
+import Data.Either (fromRight)
+import Data.Maybe (fromJust)
+import Language.LSP.Protocol.Lens (params, textDocument, uri, HasTextDocument, HasParams, HasUri)
+import Language.LSP.Protocol.Message
+import Language.LSP.Protocol.Types
+import Language.LSP.Server
+
+tokenHandler :: Handlers (LspM ())
+tokenHandler = requestHandler SMethod_TextDocumentSemanticTokensFull $ \req responder -> do
+    let 
+        openUri = toNormalizedUri $ req ^. docUri
+        token =
+            SemanticTokenAbsolute
+                { _line = 0
+                , _startChar = 0
+                , _length = 5
+                , _tokenType = SemanticTokenTypes_Keyword
+                , _tokenModifiers = []
+                }
+        tokens = fromRight (error "failed") $ makeSemanticTokens defaultSemanticTokensLegend [token]
+    c <- fromJust <$> getVirtualFile openUri
+    responder (Right $ InL tokens)
+
+docUri :: (HasParams s s1,  HasTextDocument s1 s2, HasUri s2 a) => Lens' s a 
+docUri = params . textDocument . uri
diff --git a/datalog-lsp/src/Datalog/LSP/Hover.hs b/datalog-lsp/src/Datalog/LSP/Hover.hs
new file mode 100644
index 0000000..ef11663
--- /dev/null
+++ b/datalog-lsp/src/Datalog/LSP/Hover.hs
@@ -0,0 +1,23 @@
+module Datalog.LSP.Hover where
+
+import Control.Lens ((^.), Lens')
+import Data.Maybe (fromJust)
+import Language.LSP.Protocol.Lens (params, textDocument, uri, HasTextDocument, HasParams, HasUri)
+import Language.LSP.Protocol.Message
+import Language.LSP.Protocol.Types
+import Language.LSP.Server
+import Language.LSP.VFS (virtualFileText)
+
+hoverHandler :: Handlers (LspM ())
+hoverHandler = requestHandler SMethod_TextDocumentHover $ \req responder -> do
+    let
+        openUri = toNormalizedUri $ req ^. docUri
+        TRequestMessage _ _ _ (HoverParams _doc pos _workDone) = req
+        Position _l _c' = pos
+        rsp txt = Hover (InL . mkMarkdown $ txt) (Just range)
+        range = Range pos pos
+    c <- fromJust <$> getVirtualFile openUri
+    responder (Right . InL . rsp . virtualFileText $ c)
+
+docUri :: (HasParams s s1,  HasTextDocument s1 s2, HasUri s2 a) => Lens' s a 
+docUri = params . textDocument . uri
diff --git a/datalog-lsp/src/Main.hs b/datalog-lsp/src/Main.hs
index 7d7904b..ee612d0 100644
--- a/datalog-lsp/src/Main.hs
+++ b/datalog-lsp/src/Main.hs
@@ -1,65 +1,7 @@
 {-# LANGUAGE OverloadedStrings #-}
 
-import Control.Lens ((^.))
-import Control.Monad.IO.Class
-import Data.Either (fromRight)
-import Data.Maybe (fromJust)
-import Language.LSP.Protocol.Lens (params, textDocument, uri)
-import Language.LSP.Protocol.Message
-import Language.LSP.Protocol.Types
-import Language.LSP.Server
-import Language.LSP.VFS (virtualFileText)
-
-handlers :: Handlers (LspM ())
-handlers =
-    mconcat
-        [ notificationHandler SMethod_Initialized $ \_ -> pure ()
-        , requestHandler SMethod_TextDocumentHover $ \req responder -> do
-            let
-                docUri = toNormalizedUri $ req ^. (params . textDocument . uri)
-                TRequestMessage _ _ _ (HoverParams _doc pos _workDone) = req
-                Position _l _c' = pos
-                rsp txt = Hover (InL . mkMarkdown $ txt) (Just range)
-                range = Range pos pos
-            c <- fromJust <$> getVirtualFile docUri
-            responder (Right . InL . rsp . virtualFileText $ c)
-        , requestHandler SMethod_TextDocumentSemanticTokensFull $ \req responder -> do
-            let TRequestMessage _ _ _ (SemanticTokensParams _doc _workDone _partial) = req
-                docUri = toNormalizedUri $ req ^. (params . textDocument . uri)
-                token =
-                    SemanticTokenAbsolute
-                        { _line = 0
-                        , _startChar = 0
-                        , _length = 5
-                        , _tokenType = SemanticTokenTypes_Keyword
-                        , _tokenModifiers = []
-                        }
-                tokens = fromRight (error "failed") $ makeSemanticTokens defaultSemanticTokensLegend [token]
-            c <- fromJust <$> getVirtualFile docUri
-            responder (Right $ InL tokens)
-        ]
+import Datalog.LSP (serverDefinition)
+import Language.LSP.Server (runServer)
 
 main :: IO Int
-main =
-    runServer $
-        ServerDefinition
-            { parseConfig = const $ const $ Right ()
-            , onConfigChange = const $ pure ()
-            , defaultConfig = ()
-            , configSection = "demo"
-            , doInitialize = \env _req -> pure $ Right env
-            , staticHandlers = const handlers
-            , interpretHandler = \env -> Iso (runLspT env) liftIO
-            , options =
-                defaultOptions
-                    { optTextDocumentSync =
-                        Just $
-                            TextDocumentSyncOptions
-                                { _openClose = Just True
-                                , _change = Just TextDocumentSyncKind_Full
-                                , _willSave = Just False
-                                , _willSaveWaitUntil = Just False
-                                , _save = Just (InR (SaveOptions (Just False)))
-                                }
-                    }
-            }
+main = runServer serverDefinition
diff --git a/geolog-lsp.cabal b/geolog-lsp.cabal
index c4968c0..7bc6a4b 100644
--- a/geolog-lsp.cabal
+++ b/geolog-lsp.cabal
@@ -42,6 +42,12 @@ executable datalog-lsp
                       lsp,
                       text,
                       containers,
-                      lens
+                      lens,
+                      megaparsec,
+                      datalog-parser
+    other-modules:
+        Datalog.LSP
+        Datalog.LSP.Highlight
+        Datalog.LSP.Hover
     hs-source-dirs:   datalog-lsp/src
     default-language: GHC2024