gemini_rust/embedding/
model.rs

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