Skip to main content

threatflux_vertex_rust_sdk/
lib.rs

1//! # `ThreatFlux` Vertex Rust SDK
2//!
3//! A Rust SDK for Google Cloud Vertex AI API, providing access to Gemini models and other AI services.
4//!
5//! ## Features
6//!
7//! - **Authentication**: `OAuth2`, Service Account, and Application Default Credentials
8//! - **Gemini Models**: Content generation with streaming and non-streaming support
9//! - **Function Calling**: Tool/function calling capabilities
10//! - **Token Counting**: Count tokens in content
11//! - **Chat Completions**: Multi-turn conversations
12//!
13//! ## Quick Start
14//!
15//! ```rust,no_run
16//! use threatflux_vertex_rust_sdk::{config::Config, GenerateContentRequest, VertexClient};
17//!
18//! #[tokio::main]
19//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
20//!     let config = Config {
21//!         project_id: "your-project-id".into(),
22//!         region: "us-central1".into(),
23//!         ..Config::default()
24//!     };
25//!     let client = VertexClient::new(config).await?;
26//!
27//!     let request = GenerateContentRequest::new("Why is the sky blue?");
28//!     let response = client.generate_content("gemini-2.5-flash", &request).await?;
29//!
30//!     if let Some(text) = response.text() {
31//!         println!("Response: {}", text);
32//!     }
33//!     Ok(())
34//! }
35//! ```
36
37pub mod api;
38pub mod auth;
39pub mod builders;
40pub mod cache;
41pub mod chat_core;
42pub mod claude;
43pub mod client;
44pub mod config;
45pub mod error;
46pub mod media;
47pub mod model_descriptor;
48pub mod model_info;
49pub mod models;
50pub mod streaming;
51pub mod streaming_support;
52pub mod types;
53
54// Re-export main types for convenience
55pub use api::chat::ChatConversation;
56pub use api::embeddings::{
57    EmbeddingInstance, EmbeddingParameters, EmbeddingPrediction, EmbeddingRequest,
58    EmbeddingResponse, EmbeddingTaskType, EmbeddingValues, EmbeddingsApi,
59};
60pub use api::generate::GenerateApi;
61pub use api::models::{ListLocationsResponse, ListModelsResponse, Location, Model, ModelsApi};
62pub use auth::{
63    from_env, ApplicationDefaultCredentials, AuthProvider, EnvAuth, ServiceAccountAuth,
64};
65pub use builders::{ContentRequestBuilder, FunctionBuilder, TokenCountBuilder};
66pub use cache::{
67    CacheApi, CacheUsageMetadata, CachedContent, CachedContentRef, CreateCachedContentRequest,
68    ListCachedContentsResponse, UpdateCachedContentRequest,
69};
70pub use client::{VertexClient, VertexClientBuilder};
71pub use error::{Result, VertexError};
72pub use media::{classify_inline_data, InlineDataClassification, InlineDataKind};
73pub use model_descriptor::ModelDescriptor;
74pub use model_info::{get_model_info, ModelInfo};
75pub use models::*;
76pub use streaming::{ChatStream, ChatStreamChunk, SseParser};
77pub use types::*;
78
79pub use crate as threatflux_vertex_rust_sdk;
80
81/// The default API endpoint for Vertex AI
82pub const DEFAULT_ENDPOINT: &str = "https://aiplatform.googleapis.com";
83
84/// The current version of the SDK
85pub const VERSION: &str = env!("CARGO_PKG_VERSION");
86
87/// User agent string for HTTP requests
88#[must_use]
89pub fn user_agent() -> String {
90    format!("vertex-rust-sdk/{VERSION}")
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn test_user_agent() {
99        let ua = user_agent();
100        assert!(ua.starts_with("vertex-rust-sdk/"));
101    }
102}