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}