potato_type/openai/
embedding.rs

1use crate::TypeError;
2use pyo3::{prelude::*, IntoPyObjectExt};
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
6#[pyclass]
7pub struct OpenAIEmbeddingConfig {
8    #[pyo3(get)]
9    pub model: String,
10
11    #[serde(skip_serializing_if = "Option::is_none")]
12    #[pyo3(get)]
13    pub dimensions: Option<i32>,
14
15    #[serde(skip_serializing_if = "Option::is_none")]
16    #[pyo3(get)]
17    pub encoding_format: Option<String>,
18
19    #[serde(skip_serializing_if = "Option::is_none")]
20    #[pyo3(get)]
21    pub user: Option<String>,
22}
23
24#[pymethods]
25impl OpenAIEmbeddingConfig {
26    #[new]
27    #[pyo3(signature = (model, dimensions=None, encoding_format=None, user=None))]
28    fn new(
29        model: String,
30        dimensions: Option<i32>,
31        encoding_format: Option<String>,
32        user: Option<String>,
33    ) -> Self {
34        OpenAIEmbeddingConfig {
35            model,
36            dimensions,
37            encoding_format,
38            user,
39        }
40    }
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
44#[pyclass]
45pub struct EmbeddingObject {
46    #[pyo3(get)]
47    pub object: String,
48    pub embedding: Vec<f32>,
49    #[pyo3(get)]
50    pub index: u32,
51}
52
53#[pymethods]
54impl EmbeddingObject {
55    #[getter]
56    pub fn embedding(&self) -> &Vec<f32> {
57        &self.embedding
58    }
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
62#[pyclass]
63pub struct UsageObject {
64    #[pyo3(get)]
65    pub prompt_tokens: u32,
66    #[pyo3(get)]
67    pub total_tokens: u32,
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
71#[pyclass]
72pub struct OpenAIEmbeddingResponse {
73    #[pyo3(get)]
74    pub object: String,
75    #[pyo3(get)]
76    pub data: Vec<EmbeddingObject>,
77    #[pyo3(get)]
78    #[serde(default)]
79    pub usage: UsageObject,
80}
81
82impl OpenAIEmbeddingResponse {
83    pub fn into_py_bound_any<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
84        let bound = Py::new(py, self.clone())?;
85        Ok(bound.into_bound_py_any(py)?)
86    }
87}
potato_type/openai/embedding.rs

potato_type/openai/
embedding.rs
