Add base code and examples for the Vertex API

- Add base code for the Vertex API with types and a client with helper
  methods for prompting text and conversation.
- Add examples for prompting text and conversation.
This commit is contained in:
2024-02-06 20:50:36 +00:00
parent 995d2537b9
commit 67c6ba878b
10 changed files with 530 additions and 9 deletions

130
src/types.rs Normal file
View File

@@ -0,0 +1,130 @@
use std::collections::HashMap;
use serde::{Deserialize, Serialize};
#[derive(Debug, Serialize, Deserialize)]
pub struct CountTokensRequest {
pub contents: Content,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CountTokensResponse {
pub total_tokens: i32,
}
#[derive(Serialize, Deserialize)]
pub struct GenerateContentRequest {
pub contents: Vec<Content>,
pub generation_config: Option<GenerationConfig>,
pub tools: Option<Vec<Tools>>,
}
#[derive(Serialize, Deserialize)]
pub struct Tools {
pub function_declarations: Option<Vec<FunctionDeclaration>>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct Content {
pub role: String,
pub parts: Vec<Part>,
}
#[derive(Clone, Debug, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct GenerationConfig {
pub max_output_tokens: Option<i32>,
pub temperature: Option<f32>,
pub top_p: Option<f32>,
pub top_k: Option<i32>,
pub stop_sequences: Option<Vec<String>>,
pub candidate_count: Option<u32>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum Part {
Text(String),
InlineData {
mime_type: String,
data: String,
},
FileData {
mime_type: String,
file_uri: String,
},
FunctionCall {
name: String,
args: HashMap<String, String>,
},
}
#[derive(Debug, Serialize, Deserialize)]
pub struct GenerateContentResponse(pub Vec<ResponseStreamChunk>);
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ResponseStreamChunk {
pub candidates: Vec<Candidate>,
pub usage_metadata: Option<UsageMetadata>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Candidate {
pub content: Content,
pub citation_metadata: Option<CitationMetadata>,
pub safety_ratings: Vec<SafetyRating>,
pub finish_reason: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SafetyRating {
pub category: String,
pub probability: String,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Citation {
pub start_index: i32,
pub end_index: i32,
pub uri: String,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct CitationMetadata {
pub citations: Vec<Citation>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageMetadata {
pub candidates_token_count: Option<i32>,
pub prompt_token_count: i32,
pub total_token_count: i32,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FunctionDeclaration {
pub name: String,
pub description: String,
pub parameters: FunctionParameters,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FunctionParameters {
pub r#type: String,
pub properties: HashMap<String, FunctionParametersProperty>,
pub required: Vec<String>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FunctionParametersProperty {
pub r#type: String,
pub description: String,
}