use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use serde::{Deserialize, Serialize};
use tokio::sync::{mpsc, oneshot};
use crate::llm::actions::ActionDefinition;
use crate::llm::ollama_client::Message;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BridgeRequestKind {
Generate,
Chat,
}
#[derive(Debug, Clone, Serialize)]
pub struct BridgeEventContext {
pub token: String,
pub server_id: u32,
pub connection_id: Option<u32>,
pub protocol: String,
pub event_type: String,
pub data: serde_json::Value,
}
#[derive(Debug, Serialize)]
pub struct BridgeRequest {
pub id: u64,
pub kind: BridgeRequestKind,
pub model: String,
pub messages: Vec<Message>,
pub tools: Vec<serde_json::Value>,
pub actions: Vec<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub event: Option<BridgeEventContext>,
#[serde(skip)]
pub reply: oneshot::Sender<Result<BridgeReply, String>>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct BridgeToolCall {
#[serde(default)]
pub id: Option<String>,
#[serde(default, alias = "function_name")]
pub name: String,
#[serde(default)]
pub arguments: serde_json::Value,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct BridgeReply {
#[serde(default)]
pub content: Option<String>,
#[serde(default)]
pub tool_calls: Vec<BridgeToolCall>,
#[serde(default, alias = "reasoning_content", alias = "thinking")]
pub reasoning: Option<String>,
#[serde(default)]
pub prompt_tokens: u64,
#[serde(default)]
pub completion_tokens: u64,
}
pub struct LlmBridge {
tx: mpsc::UnboundedSender<BridgeRequest>,
next_id: AtomicU64,
models: Mutex<Vec<String>>,
}
impl LlmBridge {
pub fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<BridgeRequest>) {
let (tx, rx) = mpsc::unbounded_channel();
(
Arc::new(Self {
tx,
next_id: AtomicU64::new(1),
models: Mutex::new(Vec::new()),
}),
rx,
)
}
pub fn submit(
&self,
kind: BridgeRequestKind,
model: String,
messages: Vec<Message>,
tools: Vec<serde_json::Value>,
offered: &[ActionDefinition],
) -> (u64, oneshot::Receiver<Result<BridgeReply, String>>) {
self.submit_with_event(kind, model, messages, tools, offered, None)
}
pub(crate) fn submit_with_event(
&self,
kind: BridgeRequestKind,
model: String,
messages: Vec<Message>,
tools: Vec<serde_json::Value>,
offered: &[ActionDefinition],
event: Option<BridgeEventContext>,
) -> (u64, oneshot::Receiver<Result<BridgeReply, String>>) {
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
let (reply, rx) = oneshot::channel();
let request = BridgeRequest {
id,
kind,
model,
messages,
tools,
actions: offered.iter().map(offered_action).collect(),
event,
reply,
};
let _ = self.tx.send(request);
(id, rx)
}
pub fn set_models(&self, models: Vec<String>) {
*self.models.lock().expect("bridge model list") = models;
}
pub fn models(&self) -> Vec<String> {
self.models.lock().expect("bridge model list").clone()
}
}
static GENERIC_ACTIONS: std::sync::LazyLock<std::collections::HashSet<String>> =
std::sync::LazyLock::new(|| {
let mut names: std::collections::HashSet<String> =
crate::llm::actions::common::get_network_event_common_actions()
.into_iter()
.map(|a| a.name)
.collect();
names.insert(crate::llm::actions::common::provide_feedback_action().name);
names
});
pub fn offered_action(action: &ActionDefinition) -> serde_json::Value {
let parameters: Vec<serde_json::Value> = action
.parameters
.iter()
.map(|p| {
let mut v = serde_json::json!({
"name": p.name,
"type": p.type_hint,
"description": p.description,
"required": p.required,
});
if let Some(choices) = p.choices() {
v["choices"] = serde_json::json!(choices);
}
v
})
.collect();
let schema = action
.to_tool_schema()
.pointer("/function/parameters")
.cloned()
.unwrap_or(serde_json::Value::Null);
serde_json::json!({
"name": action.name,
"description": action.description,
"tool": action.is_tool(),
"generic": GENERIC_ACTIONS.contains(action.name.as_str()),
"parameters": parameters,
"example": action.example,
"schema": schema,
})
}