Skip to content

Commit a369687

Browse files
committed
Fix regen
1 parent bd9db97 commit a369687

File tree

4 files changed

+99
-30
lines changed

4 files changed

+99
-30
lines changed

src/client/points.rs

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,12 @@ use crate::qdrant::{
1212
shard_key, ClearPayloadPoints, CountPoints, CountResponse, CreateFieldIndexCollection,
1313
DeleteFieldIndexCollection, DeletePayloadPoints, DeletePointVectors, DeletePoints,
1414
DiscoverBatchPoints, DiscoverBatchResponse, DiscoverPoints, DiscoverResponse, FieldType,
15-
GetPoints, GetResponse, HardwareUsage, PayloadIndexParams, PointId, PointVectors,
16-
PointsOperationResponse, PointsSelector, PointsUpdateOperation, ReadConsistency,
17-
RecommendBatchPoints, RecommendBatchResponse, RecommendGroupsResponse, RecommendPointGroups,
18-
RecommendPoints, RecommendResponse, ScrollPoints, ScrollResponse, SearchBatchPoints,
19-
SearchBatchResponse, SearchGroupsResponse, SearchPointGroups, SearchResponse, SetPayloadPoints,
20-
ShardKeySelector, UpdateBatchPoints, UpdateBatchResponse, UpdatePointVectors, UpsertPoints,
15+
GetPoints, GetResponse, PayloadIndexParams, PointId, PointVectors, PointsOperationResponse,
16+
PointsSelector, PointsUpdateOperation, ReadConsistency, RecommendBatchPoints,
17+
RecommendBatchResponse, RecommendGroupsResponse, RecommendPointGroups, RecommendPoints,
18+
RecommendResponse, ScrollPoints, ScrollResponse, SearchBatchPoints, SearchBatchResponse,
19+
SearchGroupsResponse, SearchPointGroups, SearchResponse, SetPayloadPoints, ShardKeySelector,
20+
UpdateBatchPoints, UpdateBatchResponse, UpdatePointVectors, UpsertPoints, Usage,
2121
VectorsSelector, WithPayloadSelector, WithVectorsSelector, WriteOrdering,
2222
};
2323

@@ -283,7 +283,7 @@ impl QdrantClient {
283283
.into_inner();
284284
resp.result = result;
285285
resp.time += time;
286-
HardwareUsage::aggregate_opts(resp.usage, usage);
286+
resp.usage = Usage::aggregate_opts(resp.usage, usage);
287287
}
288288
Ok(resp)
289289
})

src/grpc_conversions/extensions.rs

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
use std::fmt::{Display, Formatter};
22
use std::hash::{Hash, Hasher};
3+
use std::ops::AddAssign;
34

45
use crate::client::Payload;
56
#[allow(deprecated)]
67
use crate::error::NotA;
78
use crate::prelude::{PointStruct, Value};
89
use crate::qdrant::value::Kind;
910
use crate::qdrant::{
10-
HardwareUsage, ListValue, PointId, RetrievedPoint, ScoredPoint, Struct, Vectors,
11+
HardwareUsage, InferenceUsage, ListValue, ModelUsage, PointId, RetrievedPoint, ScoredPoint,
12+
Struct, Usage, Vectors,
1113
};
1214

1315
/// Null value
@@ -320,6 +322,51 @@ impl Hash for RetrievedPoint {
320322
}
321323
}
322324

325+
impl Usage {
326+
pub(crate) fn aggregate_opts(this: Option<Self>, other: Option<Self>) -> Option<Self> {
327+
match (this, other) {
328+
(Some(this), Some(other)) => Some(this.aggregate(other)),
329+
(Some(this), None) => Some(this),
330+
(None, Some(other)) => Some(other),
331+
(None, None) => None,
332+
}
333+
}
334+
335+
pub(crate) fn aggregate(self, other: Self) -> Self {
336+
Self {
337+
hardware: HardwareUsage::aggregate_opts(self.hardware, other.hardware),
338+
inference: InferenceUsage::aggregate_opts(self.inference, other.inference),
339+
}
340+
}
341+
}
342+
343+
impl InferenceUsage {
344+
pub(crate) fn aggregate_opts(this: Option<Self>, other: Option<Self>) -> Option<Self> {
345+
match (this, other) {
346+
(Some(this), Some(other)) => Some(this.aggregate(other)),
347+
(Some(this), None) => Some(this),
348+
(None, Some(other)) => Some(other),
349+
(None, None) => None,
350+
}
351+
}
352+
353+
pub(crate) fn aggregate(mut self, other: Self) -> Self {
354+
for (k, v) in other.models {
355+
self.models.entry(k).and_modify(|e| *e += v).or_insert(v);
356+
}
357+
358+
Self {
359+
models: self.models,
360+
}
361+
}
362+
}
363+
364+
impl AddAssign for ModelUsage {
365+
fn add_assign(&mut self, rhs: Self) {
366+
self.tokens += rhs.tokens;
367+
}
368+
}
369+
323370
impl HardwareUsage {
324371
pub(crate) fn aggregate_opts(this: Option<Self>, other: Option<Self>) -> Option<Self> {
325372
match (this, other) {

src/qdrant.rs

Lines changed: 41 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -4496,15 +4496,15 @@ pub struct UpdateBatchPoints {
44964496
#[prost(message, optional, tag = "4")]
44974497
pub ordering: ::core::option::Option<WriteOrdering>,
44984498
}
4499-
#[derive(Clone, Copy, PartialEq, ::prost::Message)]
4499+
#[derive(Clone, PartialEq, ::prost::Message)]
45004500
pub struct PointsOperationResponse {
45014501
#[prost(message, optional, tag = "1")]
45024502
pub result: ::core::option::Option<UpdateResult>,
45034503
/// Time spent to process
45044504
#[prost(double, tag = "2")]
45054505
pub time: f64,
45064506
#[prost(message, optional, tag = "3")]
4507-
pub usage: ::core::option::Option<HardwareUsage>,
4507+
pub usage: ::core::option::Option<Usage>,
45084508
}
45094509
#[derive(Clone, Copy, PartialEq, ::prost::Message)]
45104510
pub struct UpdateResult {
@@ -4600,7 +4600,7 @@ pub struct SearchResponse {
46004600
#[prost(double, tag = "2")]
46014601
pub time: f64,
46024602
#[prost(message, optional, tag = "3")]
4603-
pub usage: ::core::option::Option<HardwareUsage>,
4603+
pub usage: ::core::option::Option<Usage>,
46044604
}
46054605
#[derive(Clone, PartialEq, ::prost::Message)]
46064606
pub struct QueryResponse {
@@ -4610,7 +4610,7 @@ pub struct QueryResponse {
46104610
#[prost(double, tag = "2")]
46114611
pub time: f64,
46124612
#[prost(message, optional, tag = "3")]
4613-
pub usage: ::core::option::Option<HardwareUsage>,
4613+
pub usage: ::core::option::Option<Usage>,
46144614
}
46154615
#[derive(Clone, PartialEq, ::prost::Message)]
46164616
pub struct QueryBatchResponse {
@@ -4620,7 +4620,7 @@ pub struct QueryBatchResponse {
46204620
#[prost(double, tag = "2")]
46214621
pub time: f64,
46224622
#[prost(message, optional, tag = "3")]
4623-
pub usage: ::core::option::Option<HardwareUsage>,
4623+
pub usage: ::core::option::Option<Usage>,
46244624
}
46254625
#[derive(Clone, PartialEq, ::prost::Message)]
46264626
pub struct QueryGroupsResponse {
@@ -4630,7 +4630,7 @@ pub struct QueryGroupsResponse {
46304630
#[prost(double, tag = "2")]
46314631
pub time: f64,
46324632
#[prost(message, optional, tag = "3")]
4633-
pub usage: ::core::option::Option<HardwareUsage>,
4633+
pub usage: ::core::option::Option<Usage>,
46344634
}
46354635
#[derive(Clone, PartialEq, ::prost::Message)]
46364636
pub struct BatchResult {
@@ -4645,7 +4645,7 @@ pub struct SearchBatchResponse {
46454645
#[prost(double, tag = "2")]
46464646
pub time: f64,
46474647
#[prost(message, optional, tag = "3")]
4648-
pub usage: ::core::option::Option<HardwareUsage>,
4648+
pub usage: ::core::option::Option<Usage>,
46494649
}
46504650
#[derive(Clone, PartialEq, ::prost::Message)]
46514651
pub struct SearchGroupsResponse {
@@ -4655,17 +4655,17 @@ pub struct SearchGroupsResponse {
46554655
#[prost(double, tag = "2")]
46564656
pub time: f64,
46574657
#[prost(message, optional, tag = "3")]
4658-
pub usage: ::core::option::Option<HardwareUsage>,
4658+
pub usage: ::core::option::Option<Usage>,
46594659
}
4660-
#[derive(Clone, Copy, PartialEq, ::prost::Message)]
4660+
#[derive(Clone, PartialEq, ::prost::Message)]
46614661
pub struct CountResponse {
46624662
#[prost(message, optional, tag = "1")]
46634663
pub result: ::core::option::Option<CountResult>,
46644664
/// Time spent to process
46654665
#[prost(double, tag = "2")]
46664666
pub time: f64,
46674667
#[prost(message, optional, tag = "3")]
4668-
pub usage: ::core::option::Option<HardwareUsage>,
4668+
pub usage: ::core::option::Option<Usage>,
46694669
}
46704670
#[derive(Clone, PartialEq, ::prost::Message)]
46714671
pub struct ScrollResponse {
@@ -4678,7 +4678,7 @@ pub struct ScrollResponse {
46784678
#[prost(double, tag = "3")]
46794679
pub time: f64,
46804680
#[prost(message, optional, tag = "4")]
4681-
pub usage: ::core::option::Option<HardwareUsage>,
4681+
pub usage: ::core::option::Option<Usage>,
46824682
}
46834683
#[derive(Clone, Copy, PartialEq, ::prost::Message)]
46844684
pub struct CountResult {
@@ -4708,7 +4708,7 @@ pub struct GetResponse {
47084708
#[prost(double, tag = "2")]
47094709
pub time: f64,
47104710
#[prost(message, optional, tag = "3")]
4711-
pub usage: ::core::option::Option<HardwareUsage>,
4711+
pub usage: ::core::option::Option<Usage>,
47124712
}
47134713
#[derive(Clone, PartialEq, ::prost::Message)]
47144714
pub struct RecommendResponse {
@@ -4718,7 +4718,7 @@ pub struct RecommendResponse {
47184718
#[prost(double, tag = "2")]
47194719
pub time: f64,
47204720
#[prost(message, optional, tag = "3")]
4721-
pub usage: ::core::option::Option<HardwareUsage>,
4721+
pub usage: ::core::option::Option<Usage>,
47224722
}
47234723
#[derive(Clone, PartialEq, ::prost::Message)]
47244724
pub struct RecommendBatchResponse {
@@ -4728,7 +4728,7 @@ pub struct RecommendBatchResponse {
47284728
#[prost(double, tag = "2")]
47294729
pub time: f64,
47304730
#[prost(message, optional, tag = "3")]
4731-
pub usage: ::core::option::Option<HardwareUsage>,
4731+
pub usage: ::core::option::Option<Usage>,
47324732
}
47334733
#[derive(Clone, PartialEq, ::prost::Message)]
47344734
pub struct DiscoverResponse {
@@ -4738,7 +4738,7 @@ pub struct DiscoverResponse {
47384738
#[prost(double, tag = "2")]
47394739
pub time: f64,
47404740
#[prost(message, optional, tag = "3")]
4741-
pub usage: ::core::option::Option<HardwareUsage>,
4741+
pub usage: ::core::option::Option<Usage>,
47424742
}
47434743
#[derive(Clone, PartialEq, ::prost::Message)]
47444744
pub struct DiscoverBatchResponse {
@@ -4748,7 +4748,7 @@ pub struct DiscoverBatchResponse {
47484748
#[prost(double, tag = "2")]
47494749
pub time: f64,
47504750
#[prost(message, optional, tag = "3")]
4751-
pub usage: ::core::option::Option<HardwareUsage>,
4751+
pub usage: ::core::option::Option<Usage>,
47524752
}
47534753
#[derive(Clone, PartialEq, ::prost::Message)]
47544754
pub struct RecommendGroupsResponse {
@@ -4758,7 +4758,7 @@ pub struct RecommendGroupsResponse {
47584758
#[prost(double, tag = "2")]
47594759
pub time: f64,
47604760
#[prost(message, optional, tag = "3")]
4761-
pub usage: ::core::option::Option<HardwareUsage>,
4761+
pub usage: ::core::option::Option<Usage>,
47624762
}
47634763
#[derive(Clone, PartialEq, ::prost::Message)]
47644764
pub struct UpdateBatchResponse {
@@ -4767,6 +4767,8 @@ pub struct UpdateBatchResponse {
47674767
/// Time spent to process
47684768
#[prost(double, tag = "2")]
47694769
pub time: f64,
4770+
#[prost(message, optional, tag = "3")]
4771+
pub usage: ::core::option::Option<Usage>,
47704772
}
47714773
#[derive(Clone, PartialEq, ::prost::Message)]
47724774
pub struct FacetResponse {
@@ -4784,7 +4786,7 @@ pub struct SearchMatrixPairsResponse {
47844786
#[prost(double, tag = "2")]
47854787
pub time: f64,
47864788
#[prost(message, optional, tag = "3")]
4787-
pub usage: ::core::option::Option<HardwareUsage>,
4789+
pub usage: ::core::option::Option<Usage>,
47884790
}
47894791
#[derive(Clone, PartialEq, ::prost::Message)]
47904792
pub struct SearchMatrixOffsetsResponse {
@@ -4794,7 +4796,7 @@ pub struct SearchMatrixOffsetsResponse {
47944796
#[prost(double, tag = "2")]
47954797
pub time: f64,
47964798
#[prost(message, optional, tag = "3")]
4797-
pub usage: ::core::option::Option<HardwareUsage>,
4799+
pub usage: ::core::option::Option<Usage>,
47984800
}
47994801
#[derive(Clone, PartialEq, ::prost::Message)]
48004802
pub struct Filter {
@@ -5055,6 +5057,26 @@ pub struct GeoPoint {
50555057
#[prost(double, tag = "2")]
50565058
pub lat: f64,
50575059
}
5060+
/// ---------------------------------------------
5061+
/// ----------- Measurements collector ----------
5062+
/// ---------------------------------------------
5063+
#[derive(Clone, PartialEq, ::prost::Message)]
5064+
pub struct Usage {
5065+
#[prost(message, optional, tag = "1")]
5066+
pub hardware: ::core::option::Option<HardwareUsage>,
5067+
#[prost(message, optional, tag = "2")]
5068+
pub inference: ::core::option::Option<InferenceUsage>,
5069+
}
5070+
#[derive(Clone, PartialEq, ::prost::Message)]
5071+
pub struct InferenceUsage {
5072+
#[prost(map = "string, message", tag = "1")]
5073+
pub models: ::std::collections::HashMap<::prost::alloc::string::String, ModelUsage>,
5074+
}
5075+
#[derive(Clone, Copy, PartialEq, ::prost::Message)]
5076+
pub struct ModelUsage {
5077+
#[prost(uint64, tag = "1")]
5078+
pub tokens: u64,
5079+
}
50585080
#[derive(Clone, Copy, PartialEq, ::prost::Message)]
50595081
pub struct HardwareUsage {
50605082
#[prost(uint64, tag = "1")]

src/qdrant_client/points.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@ use crate::auth::TokenInterceptor;
88
use crate::qdrant::points_client::PointsClient;
99
use crate::qdrant::{
1010
CountPoints, CountResponse, DeletePointVectors, DeletePoints, FacetCounts, FacetResponse,
11-
GetPoints, GetResponse, HardwareUsage, PointsOperationResponse, ScrollPoints, ScrollResponse,
11+
GetPoints, GetResponse, PointsOperationResponse, ScrollPoints, ScrollResponse,
1212
SearchMatrixOffsetsResponse, SearchMatrixPairsResponse, SearchMatrixPoints, UpdateBatchPoints,
13-
UpdateBatchResponse, UpdatePointVectors, UpsertPoints,
13+
UpdateBatchResponse, UpdatePointVectors, UpsertPoints, Usage,
1414
};
1515
use crate::qdrant_client::{Qdrant, QdrantResult};
1616

@@ -148,7 +148,7 @@ impl Qdrant {
148148

149149
resp.result = result;
150150
resp.time += time;
151-
HardwareUsage::aggregate_opts(resp.usage, usage);
151+
resp.usage = Usage::aggregate_opts(resp.usage, usage);
152152
}
153153

154154
Ok(resp)

0 commit comments

Comments
 (0)