gemini_rust/embedding/
model.rs

1use serde::{Deserialize, Serialize};
2use strum_macros::AsRefStr;
3
4use crate::{Content, Model};
5
6/// Text embedding representation
7#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
8pub struct ContentEmbedding {
9    /// The values generated
10    pub values: Vec<f32>, //Maybe Quantize this
11}
12
13/// Response for single embedding request
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct ContentEmbeddingResponse {
16    /// The embeddings generated
17    pub embedding: ContentEmbedding,
18}
19
20/// Response for batch embedding request
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct BatchContentEmbeddingResponse {
23    /// The embeddings generated
24    pub embeddings: Vec<ContentEmbedding>,
25}
26
27/// Request to generate embeddings for content
28#[derive(Debug, Clone, Serialize, Deserialize)]
29#[serde(rename_all = "camelCase")]
30pub struct EmbedContentRequest {
31    /// The specified embedding model
32    pub model: Model,
33    /// The chunks content to generate embeddings
34    pub content: Content,
35    /// The embedding task type (optional)
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub task_type: Option<TaskType>,
38    /// The title of the document (optional)
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub title: Option<String>,
41    /// Optional reduced dimension count for the output embedding
42    #[serde(skip_serializing_if = "Option::is_none")]
43    pub output_dimensionality: Option<i32>,
44}
45
46/// Request to generate embeddings for multiple contents
47#[derive(Debug, Clone, Serialize, Deserialize)]
48#[serde(rename_all = "camelCase")]
49pub struct BatchEmbedContentsRequest {
50    /// The list of embed requests
51    pub requests: Vec<EmbedContentRequest>,
52}
53
54/// Task types for embedding optimization
55#[derive(Debug, Clone, Serialize, Deserialize, AsRefStr)]
56#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
57#[strum(serialize_all = "SCREAMING_SNAKE_CASE")]
58pub enum TaskType {
59    /// Used to generate embeddings that are optimized to assess text similarity
60    SemanticSimilarity,
61    /// Used to generate embeddings that are optimized to classify texts according to preset labels
62    Classification,
63    /// Used to generate embeddings that are optimized to cluster texts based on their similarities
64    Clustering,
65
66    /// Used to generate embeddings that are optimized for document search or information retrieval
67    RetrievalDocument,
68    /// Used to generate embeddings optimized for search queries
69    RetrievalQuery,
70    /// Used to generate embeddings optimized for question answering tasks
71    QuestionAnswering,
72    /// Used to generate embeddings optimized for fact verification
73    FactVerification,
74
75    /// Used to retrieve a code block based on a natural language query, such as sort an array or reverse a linked list.
76    /// Embeddings of the code blocks are computed using RETRIEVAL_DOCUMENT.
77    CodeRetrievalQuery,
78}
gemini_rust/embedding/model.rs

gemini_rust/embedding/
model.rs
