initial commit

.env.example

@@ -0,0 +1,9 @@
+homeserver_url="https://matrix.org"
+username="bot"
+password="bot"
+
+ollama_host="http://localhost"
+ollama_port=1111
+ollama_model="neural-chat:latest"
+
+prefix=".ask"

.gitignore

@@ -0,0 +1,3 @@
+/target
+result*
+.env

Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "celestial"
+version = "0.1.0"
+edition = "2021"
+license = "MIT"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow = "1.0.82"
+dotenv = "0.15.0"
+env_logger = "0.11.3"
+envy = "0.4.2"
+log = "0.4.21"
+matrix-sdk = "0.7.1"
+ollama-rs = "0.1.9"
+serde = "1.0.200"
+tokio = { version = "1.24.2", features = ["macros", "rt-multi-thread"] }
+tracing-subscriber = "0.3.15"

default.nix

@@ -0,0 +1,20 @@
+{ lib, rustPlatform, fetchFromGitea, pkgs, stdenv, ... }:
+
+rustPlatform.buildRustPackage rec {
+  pname = "celestial";
+  version = "1.1.1";
+
+  nativeBuildInputs = lib.optionals stdenv.isLinux [ pkgs.pkg-config ];
+  OPENSSL_NO_VENDOR = 1;
+  buildInputs = lib.optionals stdenv.isLinux [ pkgs.openssl ];
+
+  src = "./.";
+  # TODO 
+  #cargoHash = "sha256-dpN7hHpqSur6KjtEikVjQMqnF8PW27ax7ccpQNU5vdA=";
+
+  meta = with lib; {
+    description = "Celestial ai bot";
+    homepage = "https://git.4o1x5.dev/4o1x5/celestial";
+    license = licenses.mit;
+  };
+}

deny.toml

@@ -0,0 +1,4 @@
+[licenses]
+allow = [
+  "MIT"
+]

flake.lock

@@ -0,0 +1,121 @@
+{
+  "nodes": {
+    "advisory-db": {
+      "flake": false,
+      "locked": {
+        "lastModified": 1713971667,
+        "narHash": "sha256-VtlITecqZHZcm/Fzfa+0IP7I3gcFe5AS+AG0Tv1qIHk=",
+        "owner": "rustsec",
+        "repo": "advisory-db",
+        "rev": "1d1431ceb4d312c0d5c98be63d518b5d472a1149",
+        "type": "github"
+      },
+      "original": {
+        "owner": "rustsec",
+        "repo": "advisory-db",
+        "type": "github"
+      }
+    },
+    "crane": {
+      "inputs": {
+        "nixpkgs": [
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1713979152,
+        "narHash": "sha256-apdecPuh8SOQnkEET/kW/UcfjCRb8JbV5BKjoH+DcP4=",
+        "owner": "ipetkov",
+        "repo": "crane",
+        "rev": "a5eca68a2cf11adb32787fc141cddd29ac8eb79c",
+        "type": "github"
+      },
+      "original": {
+        "owner": "ipetkov",
+        "repo": "crane",
+        "type": "github"
+      }
+    },
+    "fenix": {
+      "inputs": {
+        "nixpkgs": [
+          "nixpkgs"
+        ],
+        "rust-analyzer-src": []
+      },
+      "locked": {
+        "lastModified": 1714026264,
+        "narHash": "sha256-rIRsxOZ/eUnWVHfbJlXXQtYriPICFgHGyao5jxm1FMQ=",
+        "owner": "nix-community",
+        "repo": "fenix",
+        "rev": "4e14e4f21fbd7afae40a492b7d937cbf16f76b11",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-community",
+        "repo": "fenix",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1710146030,
+        "narHash": "sha256-SZ5L6eA7HJ/nmkzGG7/ISclqe6oZdOZTNoesiInkXPQ=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "b1d9ab70662946ef0850d488da1c9019f3a9752a",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1713805509,
+        "narHash": "sha256-YgSEan4CcrjivCNO5ZNzhg7/8ViLkZ4CB/GrGBVSudo=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "1e1dc66fe68972a76679644a5577828b6a7e8be4",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "ref": "nixpkgs-unstable",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "advisory-db": "advisory-db",
+        "crane": "crane",
+        "fenix": "fenix",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,140 @@
+{
+  description = "Build a cargo project";
+
+  inputs = {
+    nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable";
+
+    crane = {
+      url = "github:ipetkov/crane";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
+
+    fenix = {
+      url = "github:nix-community/fenix";
+      inputs.nixpkgs.follows = "nixpkgs";
+      inputs.rust-analyzer-src.follows = "";
+    };
+
+    flake-utils.url = "github:numtide/flake-utils";
+
+    advisory-db = {
+      url = "github:rustsec/advisory-db";
+      flake = false;
+    };
+  };
+
+  outputs = { self, nixpkgs, crane, fenix, flake-utils, advisory-db, ... }:
+    flake-utils.lib.eachDefaultSystem (system:
+      let
+        pkgs = nixpkgs.legacyPackages.${system};
+
+        inherit (pkgs) lib;
+
+        craneLib = crane.lib.${system};
+        src = craneLib.cleanCargoSource (craneLib.path ./.);
+
+        commonArgs = {
+          inherit src;
+          strictDeps = true;
+
+          buildInputs = with pkgs; [
+            pkg-config
+            openssl
+          ] ++ lib.optionals pkgs.stdenv.isDarwin [
+            pkgs.libiconv
+          ];
+
+        };
+
+        craneLibLLvmTools = craneLib.overrideToolchain
+          (fenix.packages.${system}.complete.withComponents [
+            "cargo"
+            "llvm-tools"
+            "rustc"
+          ]);
+
+        # Build *just* the cargo dependencies, so we can reuse
+        # all of that work (e.g. via cachix) when running in CI
+        cargoArtifacts = craneLib.buildDepsOnly commonArgs;
+
+        # Build the actual crate itself, reusing the dependency
+        # artifacts from above.
+        celestial = craneLib.buildPackage (commonArgs // {
+          inherit cargoArtifacts;
+        });
+      in
+      {
+        checks = {
+          # Build the crate as part of `nix flake check` for convenience
+          inherit celestial;
+
+          # Run clippy (and deny all warnings) on the crate source,
+          # again, reusing the dependency artifacts from above.
+          #
+          # Note that this is done as a separate derivation so that
+          # we can block the CI if there are issues here, but not
+          # prevent downstream consumers from building our crate by itself.
+          celestial-clippy = craneLib.cargoClippy (commonArgs // {
+            inherit cargoArtifacts;
+            cargoClippyExtraArgs = "--all-targets -- --deny warnings";
+          });
+
+          celestial-doc = craneLib.cargoDoc (commonArgs // {
+            inherit cargoArtifacts;
+          });
+
+          # Check formatting
+          celestial-fmt = craneLib.cargoFmt {
+            inherit src;
+          };
+
+          # Audit dependencies
+          celestial-audit = craneLib.cargoAudit {
+            inherit src advisory-db;
+          };
+
+          # Audit licenses
+          celestial-deny = craneLib.cargoDeny {
+            inherit src;
+          };
+
+          # Run tests with cargo-nextest
+          # Consider setting `doCheck = false` on `celestial` if you do not want
+          # the tests to run twice
+          celestial-nextest = craneLib.cargoNextest (commonArgs // {
+            inherit cargoArtifacts;
+            partitions = 1;
+            partitionType = "count";
+          });
+        };
+
+        packages = {
+          default = celestial;
+        } // lib.optionalAttrs (!pkgs.stdenv.isDarwin) {
+          celestial-llvm-coverage = craneLibLLvmTools.cargoLlvmCov (commonArgs // {
+            inherit cargoArtifacts;
+          });
+        };
+
+        apps.default = flake-utils.lib.mkApp {
+          drv = celestial;
+        };
+
+        devShells.default = craneLib.devShell {
+          # Inherit inputs from checks.
+          checks = self.checks.${system};
+
+          LIBCLANG_PATH = "${pkgs.llvmPackages_17.libclang.lib}/lib";
+          RUST_BACKTRACE = 1;
+          RUST_SRC_PATH = "${pkgs.rust.packages.stable.rustPlatform.rustLibSrc}";
+
+          packages = with pkgs; [
+            rustfmt
+            rust-analyzer
+            clippy
+            openssl
+            sqlite
+          ];
+        };
+      });
+}

readme.md

@@ -0,0 +1,4 @@
+### Celestial
+
+This is just a funny project I slapped together in a span of around 30 minutes.
+I plan on expanding it with other features in the future.

src/main.rs

@@ -0,0 +1,188 @@
+use std::time::Instant;
+
+use matrix_sdk::{
+    config::SyncSettings,
+    event_handler::Ctx,
+    ruma::events::room::{
+        member::StrippedRoomMemberEvent,
+        message::{RoomMessageEventContent, SyncRoomMessageEvent},
+    },
+    Client, Room,
+};
+use ollama_rs::{generation::completion::request::GenerationRequest, Ollama};
+use serde::Deserialize;
+use tokio::time::{sleep, Duration};
+
+#[derive(Deserialize, Debug)]
+struct Config {
+    homeserver_url: String,
+    username: String,
+    password: String,
+
+    ollama_host: String,
+    ollama_port: i16,
+    ollama_model: String,
+
+    prefix: String,
+}
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    tracing_subscriber::fmt::init();
+    let _ = dotenv::dotenv();
+    log::info!("Bot has started");
+    let c = envy::from_env::<Config>().expect("Please provide variables");
+
+    login_and_sync(
+        c.homeserver_url.to_string(),
+        &c.username,
+        &c.password,
+        c.prefix,
+        c.ollama_host,
+        c.ollama_port,
+        c.ollama_model,
+    )
+    .await?;
+    Ok(())
+}
+
+// The core sync loop we have running.
+async fn login_and_sync(
+    homeserver_url: String,
+    username: &str,
+    password: &str,
+    prefix: String,
+    ollama_host: String,
+    ollama_port: i16,
+    ollama_model: String,
+) -> anyhow::Result<()> {
+    let client = Client::builder()
+        .homeserver_url(homeserver_url)
+        .build()
+        .await?;
+
+    client
+        .matrix_auth()
+        .login_username(username, password)
+        .initial_device_display_name("getting started bot")
+        .await?;
+
+    println!("logged in as {username}");
+
+    let ollama = Ollama::new(ollama_host, ollama_port as u16);
+
+    client.add_event_handler(on_stripped_state_member);
+    let sync_token = client
+        .sync_once(SyncSettings::default())
+        .await
+        .unwrap()
+        .next_batch;
+
+    #[derive(Debug, Clone)]
+    struct MyContext {
+        botId: String,
+        ai: Ollama,
+    }
+    client.add_event_handler_context(MyContext {
+        botId: client.clone().user_id().unwrap().to_string(),
+        ai: ollama.clone(),
+    });
+
+    client.add_event_handler(
+        |ev: SyncRoomMessageEvent, room: Room, context: Ctx<MyContext>| async move {
+            if ev.as_original().unwrap().sender.to_string() == context.botId {
+                return;
+            };
+
+            // add . prefix
+            if !ev
+                .as_original()
+                .unwrap()
+                .content
+                .body()
+                .to_string()
+                .contains(prefix.as_str())
+            {
+                return;
+            }
+
+            // Send seen
+            room.send_single_receipt(
+                matrix_sdk::ruma::api::client::receipt::create_receipt::v3::ReceiptType::Read,
+                matrix_sdk::ruma::events::receipt::ReceiptThread::Main,
+                ev.event_id().to_owned(),
+            )
+            .await;
+
+            room.typing_notice(true).await;
+
+            let prompt = format!(
+                "{} says: {}",
+                room.get_member(&ev.sender())
+                    .await
+                    .unwrap()
+                    .unwrap()
+                    .display_name()
+                    .unwrap(),
+                ev.as_original()
+                    .unwrap()
+                    .content
+                    .body()
+                    .replace(prefix.as_str(), "")
+            );
+            let now = Instant::now();
+            let res = ollama
+                .generate(GenerationRequest::new(ollama_model, prompt))
+                .await;
+
+            if let Ok(res) = res {
+                let mut asd = res.response.to_string();
+                asd.push_str(format!("\n prompt took {:?}", now.elapsed()).as_str());
+
+                let content = RoomMessageEventContent::text_plain(asd);
+                println!("Got res from ai: {}", res.clone().response);
+                room.send(content).await.unwrap();
+
+                room.typing_notice(false).await;
+            }
+        },
+    );
+
+    let settings = SyncSettings::default().token(sync_token);
+    client.sync(settings).await?;
+
+    Ok(())
+}
+
+async fn on_stripped_state_member(
+    room_member: StrippedRoomMemberEvent,
+    client: Client,
+    room: Room,
+) {
+    if room_member.state_key != client.user_id().unwrap() {
+        return;
+    }
+
+    tokio::spawn(async move {
+        println!("Autojoining room {}", room.room_id());
+        let mut delay = 2;
+
+        while let Err(err) = room.join().await {
+            // retry autojoin due to synapse sending invites, before the
+            // invited user can join for more information see
+            // https://github.com/matrix-org/synapse/issues/4345
+            eprintln!(
+                "Failed to join room {} ({err:?}), retrying in {delay}s",
+                room.room_id()
+            );
+
+            sleep(Duration::from_secs(delay)).await;
+            delay *= 2;
+
+            if delay > 3600 {
+                eprintln!("Can't join room {} ({err:?})", room.room_id());
+                break;
+            }
+        }
+        println!("Successfully joined room {}", room.room_id());
+    });
+}