Merge branch 'master' into caglacelik/tps-reporter

Author: cceliks

.env.example

@@ -28,6 +28,9 @@ OPENAI_API_KEY=
 ## Gemini (if used, required) ##
 GEMINI_API_KEY=
 
+## Open Router (if used, required) ##
+OPENROUTER_API_KEY=
+
 ## Ollama (if used, optional) ##
 # do not change this, it is used by Docker
 OLLAMA_HOST=http://host.docker.internal

Cargo.lock

@@ -7,7 +7,7 @@ default-members = ["compute"]
 
 [workspace.package]
 edition = "2021"
-version = "0.2.21"
+version = "0.2.22"
 license = "Apache-2.0"
 readme = "README.md"
 

Cargo.toml

@@ -1,4 +1,4 @@
-FROM --platform=$BUILDPLATFORM rust:1.75 as builder
+FROM --platform=$BUILDPLATFORM rust:1.75 AS builder
 
 # https://docs.docker.com/engine/reference/builder/#automatic-platform-args-in-the-global-scope
 #

Dockerfile

@@ -23,6 +23,9 @@ pub struct DKNMessage {
     ///
     /// NOTE: This can be obtained via Identify protocol version
     pub(crate) version: String,
+    /// Identity protocol string of the Dria Compute Node
+    #[serde(default)]
+    pub(crate) identity: String,
     /// The timestamp of the message, in nanoseconds
     ///
     /// NOTE: This can be obtained via DataTransform in GossipSub
@@ -46,6 +49,9 @@ impl DKNMessage {
             payload: BASE64_STANDARD.encode(data),
             topic: topic.to_string(),
             version: DRIA_COMPUTE_NODE_VERSION.to_string(),
+            identity: dkn_p2p::P2P_IDENTITY_PREFIX
+                .trim_end_matches('/')
+                .to_string(),
             timestamp: get_current_time_nanos(),
         }
     }

compute/src/utils/message.rs

@@ -254,6 +254,7 @@ impl DriaP2PClient {
                     log::warn!("Local node is listening on {}", address);
                 }
                 SwarmEvent::ExternalAddrConfirmed { address } => {
+                    // this is usually the external address via relay
                     log::info!("External address confirmed: {}", address);
                 }
                 event => log::trace!("Unhandled Swarm Event: {:?}", event),
@@ -275,6 +276,13 @@ impl DriaP2PClient {
                 info.protocol_version,
                 self.identity_protocol
             );
+
+            // blacklist peers with different protocol
+            self.swarm
+                .behaviour_mut()
+                .gossipsub
+                .blacklist_peer(&peer_id);
+
             return;
         }
 

p2p/src/client.rs

@@ -7,10 +7,10 @@ mod client;
 pub use client::DriaP2PClient;
 
 /// Prefix for Kademlia protocol, must start with `/`!
-pub(crate) const P2P_KADEMLIA_PREFIX: &str = "/dria/kad/";
+pub const P2P_KADEMLIA_PREFIX: &str = "/dria/kad/";
 
 /// Prefix for Identity protocol string.
-pub(crate) const P2P_IDENTITY_PREFIX: &str = "dria/";
+pub const P2P_IDENTITY_PREFIX: &str = "dria/";
 
 // re-exports
 pub use libp2p;

p2p/src/lib.rs

@@ -1,6 +1,6 @@
 use crate::{
     apis::{JinaConfig, SerperConfig},
-    providers::{GeminiConfig, OllamaConfig, OpenAIConfig},
+    providers::{GeminiConfig, OllamaConfig, OpenAIConfig, OpenRouterConfig},
     split_csv_line, Model, ModelProvider,
 };
 use eyre::{eyre, Result};
@@ -19,6 +19,9 @@ pub struct DriaWorkflowsConfig {
     /// Gemini configurations, e.g. API key, in case Gemini is used.
     /// Otherwise, can be ignored.
     pub gemini: GeminiConfig,
+    /// OpenRouter configurations, e.g. API key, in case OpenRouter is used.
+    /// Otherwise, can be ignored.
+    pub openrouter: OpenRouterConfig,
     /// Serper configurations, e.g. API key, in case Serper is given in environment.
     /// Otherwise, can be ignored.
     pub serper: SerperConfig,
@@ -39,6 +42,7 @@ impl DriaWorkflowsConfig {
             models: models_and_providers,
             ollama: OllamaConfig::new(),
             openai: OpenAIConfig::new(),
+            openrouter: OpenRouterConfig::new(),
             gemini: GeminiConfig::new(),
             serper: SerperConfig::new(),
             jina: JinaConfig::new(),
@@ -230,6 +234,18 @@ impl DriaWorkflowsConfig {
             );
         }
 
+        // if OpenRouter is a provider, check that the API key is set
+        if unique_providers.contains(&ModelProvider::OpenRouter) {
+            let provider_models = self.get_models_for_provider(ModelProvider::OpenRouter);
+            good_models.extend(
+                self.openrouter
+                    .check(provider_models)
+                    .await?
+                    .into_iter()
+                    .map(|m| (ModelProvider::OpenRouter, m)),
+            );
+        }
+
         // update good models
         if good_models.is_empty() {
             Err(eyre!("No good models found, please check logs for errors."))

workflows/src/config.rs

@@ -6,3 +6,6 @@ pub use openai::OpenAIConfig;
 
 mod gemini;
 pub use gemini::GeminiConfig;
+
+mod openrouter;
+pub use openrouter::OpenRouterConfig;