elizaos_plugin_xai/
lib.rs1#![allow(missing_docs)]
2#![warn(missing_docs)]
34
35pub mod actions;
36pub mod client;
37pub mod error;
38pub mod grok;
39pub mod models;
40pub mod services;
41pub mod types;
42
43use anyhow::Result as AnyhowResult;
44use std::collections::HashMap;
45use std::sync::Arc;
46
47pub use crate::actions::{PostAction, PostActionResult};
49pub use crate::client::TwitterClient;
50pub use crate::grok::{EmbeddingParams, GrokClient, GrokConfig, TextGenerationParams};
51pub use crate::models::{TextEmbeddingHandler, TextLargeHandler, TextSmallHandler};
52pub use crate::types::TwitterConfig;
53pub use crate::services::{XService, XServiceSettings};
54
55pub fn get_x_client() -> AnyhowResult<TwitterClient> {
60 let config = TwitterConfig::from_env()?;
61 Ok(TwitterClient::new(config)?)
62}
63
64pub fn get_grok_client() -> AnyhowResult<GrokClient> {
69 let config = GrokConfig::from_env()?;
70 Ok(GrokClient::new(config)?)
71}
72
73pub fn create_xai_elizaos_plugin() -> AnyhowResult<elizaos::types::Plugin> {
83 use elizaos::types::{Plugin, PluginDefinition};
84
85 let grok = Arc::new(get_grok_client()?);
86
87 let mut model_handlers: HashMap<String, elizaos::types::ModelHandlerFn> = HashMap::new();
88
89 let grok_small = Arc::clone(&grok);
91 model_handlers.insert(
92 "TEXT_SMALL".to_string(),
93 Box::new(move |params: serde_json::Value| {
94 let grok = Arc::clone(&grok_small);
95 Box::pin(async move {
96 let prompt = params
97 .get("prompt")
98 .and_then(|v| v.as_str())
99 .ok_or_else(|| anyhow::anyhow!("Missing prompt"))?;
100
101 let system = params.get("system").and_then(|v| v.as_str()).map(str::to_string);
102 let temperature = params
103 .get("temperature")
104 .and_then(|v| v.as_f64())
105 .unwrap_or(0.7) as f32;
106
107 let max_tokens = params
108 .get("max_tokens")
109 .and_then(|v| v.as_u64())
110 .or_else(|| params.get("maxTokens").and_then(|v| v.as_u64()))
111 .map(|v| v as u32);
112
113 let mut tg = TextGenerationParams::new(prompt.to_string()).temperature(temperature);
114 if let Some(max) = max_tokens {
115 tg = tg.max_tokens(max);
116 }
117 if let Some(sys) = system {
118 tg = tg.system(sys);
119 }
120
121 let result = grok.generate_text(&tg, false).await?;
122 Ok(result.text)
123 })
124 }),
125 );
126
127 let grok_large = Arc::clone(&grok);
129 model_handlers.insert(
130 "TEXT_LARGE".to_string(),
131 Box::new(move |params: serde_json::Value| {
132 let grok = Arc::clone(&grok_large);
133 Box::pin(async move {
134 let prompt = params
135 .get("prompt")
136 .and_then(|v| v.as_str())
137 .ok_or_else(|| anyhow::anyhow!("Missing prompt"))?;
138
139 let system = params.get("system").and_then(|v| v.as_str()).map(str::to_string);
140 let temperature = params
141 .get("temperature")
142 .and_then(|v| v.as_f64())
143 .unwrap_or(0.7) as f32;
144
145 let max_tokens = params
146 .get("max_tokens")
147 .and_then(|v| v.as_u64())
148 .or_else(|| params.get("maxTokens").and_then(|v| v.as_u64()))
149 .map(|v| v as u32);
150
151 let mut tg = TextGenerationParams::new(prompt.to_string()).temperature(temperature);
152 if let Some(max) = max_tokens {
153 tg = tg.max_tokens(max);
154 }
155 if let Some(sys) = system {
156 tg = tg.system(sys);
157 }
158
159 let result = grok.generate_text(&tg, true).await?;
160 Ok(result.text)
161 })
162 }),
163 );
164
165 let grok_embed = Arc::clone(&grok);
167 model_handlers.insert(
168 "TEXT_EMBEDDING".to_string(),
169 Box::new(move |params: serde_json::Value| {
170 let grok = Arc::clone(&grok_embed);
171 Box::pin(async move {
172 let text = params
173 .get("text")
174 .and_then(|v| v.as_str())
175 .or_else(|| params.get("input").and_then(|v| v.as_str()))
176 .ok_or_else(|| anyhow::anyhow!("Missing text"))?;
177
178 let embedding = grok.create_embedding(&EmbeddingParams::new(text)).await?;
179 let json = serde_json::to_string(&embedding)?;
180 Ok(json)
181 })
182 }),
183 );
184
185 Ok(Plugin {
186 definition: PluginDefinition {
187 name: "xai".to_string(),
188 description: "xAI Grok models and X (formerly Twitter) API integration".to_string(),
189 ..Default::default()
190 },
191 model_handlers,
192 ..Default::default()
193 })
194}
195
196pub async fn start_x_service(runtime: Arc<elizaos::AgentRuntime>) -> AnyhowResult<Arc<XService>> {
200 let settings = XServiceSettings::from_env()?;
201 let service = XService::start(Arc::clone(&runtime), settings).await?;
202 let service_dyn: Arc<dyn elizaos::runtime::Service> = service.clone();
203 runtime.register_service(XService::SERVICE_TYPE, service_dyn).await;
204 Ok(service)
205}