Skip to content

Commit

Permalink
cleanup: properly return 206 response code for slim_chunk queries ins…
Browse files Browse the repository at this point in the history 
…tead of enum
  • Loading branch information
@skeptrunedev @cdxker
skeptrunedev authored and cdxker committed Apr 9, 2024
1 parent a896cd2 commit 8e49f6e
Show file tree
Hide file tree
Showing 5 changed files with 42 additions and 159 deletions.
4 changes: 2 additions & 2 deletions server/src/data/models.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -564,14 +564,14 @@ pub struct GroupScoreSlimChunks {
}

#[derive(Serialize, Deserialize, ToSchema)]
pub struct SearchGroupSlimChunksResult {
pub struct SearchWithinGroupSlimResults {
pub bookmarks: Vec<ScoreSlimChunks>,
pub group: ChunkGroup,
pub total_pages: i64,
}

#[derive(Serialize, Deserialize, ToSchema)]
pub struct SearchOverGroupsSlimChunksResponseBody {
pub struct SearchOverGroupsSlimResults {
pub group_chunks: Vec<GroupScoreSlimChunks>,
pub total_chunk_pages: i64,
}
Expand Down
53 changes: 6 additions & 47 deletions server/src/handlers/chunk_handler.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1258,7 +1258,8 @@ pub enum SearchChunkResponseTypes {
tag = "chunk",
request_body(content = SearchChunkData, description = "JSON request payload to semantically search for chunks (chunks)", content_type = "application/json"),
responses(
(status = 200, description = "Chunks with embedding vectors which are similar to those in the request body", body = SearchChunkResponseTypes),
(status = 200, description = "Chunks with embedding vectors which are similar to those in the request body if slim_chunks is false or not specified in the request body", body = SearchChunkQueryResponseBody),
(status = 206, description = "Chunks with embedding vectors which are similar to those in the request body if slim_chunks is true in the request body", body = SearchSlimChunkQueryResponseBody),
(status = 400, description = "Service error relating to searching", body = ErrorResponseBody),
),
Expand Down Expand Up @@ -1349,7 +1350,7 @@ pub async fn search_chunk(
total_chunk_pages: result_chunks.total_chunk_pages,
};

return Ok(HttpResponse::Ok()
return Ok(HttpResponse::PartialContent()
.insert_header((Timer::header_key(), timer.header_value()))
.json(res));
}
Expand Down Expand Up @@ -1449,49 +1450,6 @@ pub struct RecommendChunksRequest {
pub slim_chunks: Option<bool>,
}

#[derive(Serialize, Deserialize, Debug, ToSchema)]
pub struct RecommendChunkMetadata(Vec<ChunkMetadataWithScore>);
#[derive(Serialize, Deserialize, Debug, ToSchema)]
pub struct RecommendSlimChunkMetadata(Vec<SlimChunkMetadataWithScore>);

#[derive(Serialize, Deserialize, Debug, ToSchema)]
#[serde(untagged)]
pub enum RecommendChunksResponseTypes {
#[schema(example = json!([{
"id": "e3e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"content": "Hello, world!",
"link": "https://trieve.ai",
"qdrant_point_id": "e3e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"created_at": "2021-01-01T00:00:00",
"updated_at": "2021-01-01T00:00:00",
"tag_set": "tag1,tag2",
"chunk_html": "<p>Hello, world!</p>",
"metadata": {"key": "value"},
"tracking_id": "e3e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"time_stamp": "2021-01-01T00:00:00",
"dataset_id": "e3e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"weight": 0.5,
"score": 0.9,
}]))]
Chunks(RecommendChunkMetadata),
#[schema(example = json!([{
"id": "e3e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"link": "https://trieve.ai",
"qdrant_point_id": "e3e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"created_at": "2021-01-01T00:00:00",
"updated_at": "2021-01-01T00:00:00",
"tag_set": "tag1,tag2",
"metadata": {"key": "value"},
"tracking_id": "e3e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"time_stamp": "2021-01-01T00:00:00",
"dataset_id": "e3e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"weight": 0.5,
"score": 0.9,
}]))]
#[schema(title = "SlimChunkMetadata")]
SlimChunks(RecommendSlimChunkMetadata),
}

/// Get Recommended Chunks
///
/// Get recommendations of chunks similar to the chunks in the request. Think about this as a feature similar to the "add to playlist" recommendation feature on Spotify. This request pairs especially well with our groups endpoint.
Expand All @@ -1503,7 +1461,8 @@ pub enum RecommendChunksResponseTypes {
request_body(content = RecommendChunksRequest, description = "JSON request payload to get recommendations of chunks similar to the chunks in the request", content_type = "application/json"),
responses(
(status = 200, description = "Chunks with embedding vectors which are similar to those in the request body", body = RecommendChunksResponseTypes),
(status = 200, description = "Chunks with embedding vectors which are similar to positives and dissimilar to negatives if slim_chunks is false in the request", body = Vec<ChunkMetadataWithScore>),
(status = 206, description = "Chunks with embedding vectors which are similar to positives and dissimilar to negatives if slim_chunks is true in the request", body = Vec<SlimChunkMetadataWithScore>),
(status = 400, description = "Service error relating to to getting similar chunks", body = ErrorResponseBody),
),
params(
Expand Down Expand Up @@ -1677,7 +1636,7 @@ pub async fn get_recommended_chunks(
.map(|chunk| chunk.into())
.collect::<Vec<SlimChunkMetadataWithScore>>();

return Ok(HttpResponse::Ok().json(res));
return Ok(HttpResponse::PartialContent().json(res));
}

Ok(HttpResponse::Ok()
Expand Down
97 changes: 15 additions & 82 deletions server/src/handlers/group_handler.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@ use crate::{
data::models::{
ChunkGroup, ChunkGroupAndFile, ChunkGroupBookmark, ChunkMetadata,
DatasetAndOrgWithSubAndPlan, GroupScoreSlimChunks, Pool, ScoreSlimChunks,
SearchGroupSlimChunksResult, SearchOverGroupsSlimChunksResponseBody,
ServerDatasetConfiguration, UnifiedId,
SearchOverGroupsSlimResults, SearchWithinGroupSlimResults, ServerDatasetConfiguration,
UnifiedId,
},
errors::ServiceError,
operators::{
Expand All @@ -20,8 +20,7 @@ use crate::{
search_operator::{
full_text_search_over_groups, get_metadata_from_groups, hybrid_search_over_groups,
search_full_text_groups, search_hybrid_groups, search_semantic_groups,
semantic_search_over_groups, GroupScoreChunk, SearchOverGroupsQueryResult,
SearchOverGroupsResponseBody,
semantic_search_over_groups, SearchOverGroupsQueryResult,
},
},
};
Expand Down Expand Up @@ -895,61 +894,6 @@ pub struct RecommendGroupChunksRequest {
pub slim_chunks: Option<bool>,
}

#[derive(Debug, Serialize, Deserialize, Clone, ToSchema)]
pub struct RecommendGroupChunks(pub Vec<GroupScoreChunk>);
#[derive(Debug, Serialize, Deserialize, Clone, ToSchema)]
pub struct RecommendGroupSlimChunks(pub Vec<GroupScoreSlimChunks>);

#[derive(Serialize, Deserialize, Debug, ToSchema)]
#[serde(untagged)]
pub enum RecommendGroupChunkResponseTypes {
#[schema(example = json!([{
"group_id": "e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"metadata": [
{
"metadata": [
{
"id": "e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"content": "This is a test content",
"link": "https://www.google.com",
"tag_set": "test",
"metadata": {
"key": "value"
},
"tracking_id": "e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"time_stamp": "2021-01-01T00:00:00Z",
"weight": 1.0
}
],
"score": 0.5
}
]
}]))]
GroupSlimChunksDTO(RecommendGroupSlimChunks),
#[schema(example = json!({
"group_id": "e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"metadata": [
{
"metadata": [
{
"id": "e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"link": "https://www.google.com",
"tag_set": "test",
"metadata": {
"key": "value"
},
"tracking_id": "e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3",
"time_stamp": "2021-01-01T00:00:00Z",
"weight": 1.0
}
],
"score": 0.5
}
]
}))]
GroupScoreChunkDTO(RecommendGroupChunks),
}

/// Get Recommended Groups
///
/// Route to get recommended groups. This route will return groups which are similar to the groups in the request body.
Expand All @@ -960,7 +904,8 @@ pub enum RecommendGroupChunkResponseTypes {
tag = "chunk_group",
request_body(content = RecommendGroupChunksRequest, description = "JSON request payload to get recommendations of chunks similar to the chunks in the request", content_type = "application/json"),
responses(
(status = 200, description = "JSON body representing the groups which are similar to the groups in the request", body = RecommendGroupChunkResponseTypes),
(status = 200, description = "JSON body representing the groups which are similar to the positive groups and dissimilar to the negative ones if slim_chunks is false in the request", body = Vec<GroupScoreChunk>),
(status = 206, description = "JSON body representing the groups which are similar to the positive groups and dissimilar to the negative ones if slim_chunks is false in the request", body = Vec<GroupScoreSlimChunks>),
(status = 400, description = "Service error relating to to getting similar chunks", body = ErrorResponseBody),
),
params(
Expand Down Expand Up @@ -1179,19 +1124,12 @@ impl From<SearchWithinGroupData> for SearchChunkData {
}

#[derive(Serialize, Deserialize, ToSchema)]
pub struct SearchGroupsResult {
pub struct SearchWithinGroupResults {
pub bookmarks: Vec<ScoreChunkDTO>,
pub group: ChunkGroup,
pub total_pages: i64,
}

#[derive(Serialize, Deserialize, ToSchema)]
#[serde(untagged)]
pub enum SearchWithinGroupResponseTypes {
SearchGroupsResult(SearchGroupsResult),
SearchGroupSlimChunksResult(SearchGroupSlimChunksResult),
}

/// Search Within Group
///
/// This route allows you to search only within a group. This is useful for when you only want search results to contain chunks which are members of a specific group. If choosing hybrid search, the results will be re-ranked using BAAI/bge-reranker-large.
Expand All @@ -1202,7 +1140,8 @@ pub enum SearchWithinGroupResponseTypes {
tag = "chunk_group",
request_body(content = SearchWithinGroupData, description = "JSON request payload to semantically search a group", content_type = "application/json"),
responses(
(status = 200, description = "Group chunks which are similar to the embedding vector of the search query", body = SearchWithinGroupResponseTypes),
(status = 200, description = "Group chunks which are similar to the embedding vector of the search query if slim_chunks is false in the request", body = SearchWithinGroupResults),
(status = 206, description = "Group chunks which are similar to the embedding vector of the search query if slim_chunks is true in the request", body = SearchWithinGroupSlimResults),
(status = 400, description = "Service error relating to getting the groups that the chunk is in", body = ErrorResponseBody),
),
params(
Expand Down Expand Up @@ -1298,13 +1237,13 @@ pub async fn search_within_group(
.map(|metadata| metadata.into())
.collect::<Vec<ScoreSlimChunks>>();

let res = SearchGroupSlimChunksResult {
let res = SearchWithinGroupSlimResults {
bookmarks: ids,
group: result_chunks.group,
total_pages: result_chunks.total_pages,
};

return Ok(HttpResponse::Ok().json(res));
return Ok(HttpResponse::PartialContent().json(res));
}

Ok(HttpResponse::Ok().json(result_chunks))
Expand Down Expand Up @@ -1336,13 +1275,6 @@ pub struct SearchOverGroupsData {
pub slim_chunks: Option<bool>,
}

#[derive(Serialize, Deserialize, ToSchema)]
#[serde(untagged)]
pub enum SearchOverGroupsResponseTypes {
SearchOverGroupsResponseBody(SearchOverGroupsResponseBody),
SearchOverGroupsSlimChunksResponseBody(SearchOverGroupsSlimChunksResponseBody),
}

/// Search Over Groups
///
/// This route allows you to get groups as results instead of chunks. Each group returned will have the matching chunks sorted by similarity within the group. This is useful for when you want to get groups of chunks which are similar to the search query. If choosing hybrid search, the results will be re-ranked using BAAI/bge-reranker-large. Compatible with semantic, fulltext, or hybrid search modes.
Expand All @@ -1353,8 +1285,9 @@ pub enum SearchOverGroupsResponseTypes {
tag = "chunk_group",
request_body(content = SearchOverGroupsData, description = "JSON request payload to semantically search over groups", content_type = "application/json"),
responses(
(status = 200, description = "Group chunks which are similar to the embedding vector of the search query", body = SearchOverGroupsResponseTypes),
(status = 400, description = "Service error relating to getting the groups that the chunk is in", body = ErrorResponseBody),
(status = 200, description = "Group chunks which are similar to the embedding vector of the search query if slim_chunks is false or unspecified in the request body", body = SearchOverGroupsResults),
(status = 206, description = "Group chunks which are similar to the embedding vector of the search query if slim_chunks is true in the request body", body = SearchOverGroupsSlimResults),
(status = 400, description = "Service error relating to searching over groups", body = ErrorResponseBody),
),
params(
("TR-Dataset" = String, Header, description = "The dataset id to use for the request"),
Expand Down Expand Up @@ -1441,12 +1374,12 @@ pub async fn search_over_groups(
})
.collect::<Vec<GroupScoreSlimChunks>>();

let res = SearchOverGroupsSlimChunksResponseBody {
let res = SearchOverGroupsSlimResults {
group_chunks: ids,
total_chunk_pages: result_chunks.total_chunk_pages,
};

return Ok(HttpResponse::Ok()
return Ok(HttpResponse::PartialContent()
.insert_header((Timer::header_key(), timer.header_value()))
.json(res));
}
Expand Down
19 changes: 5 additions & 14 deletions server/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ impl Modify for SecurityAddon {
name = "BSL",
url = "https://github.com/devflowinc/trieve/blob/main/LICENSE.txt",
),
version = "0.6.1",
version = "0.6.2",
),
servers(
(url = "https://api.trieve.ai",
Expand Down Expand Up @@ -233,18 +233,9 @@ impl Modify for SecurityAddon {
handlers::chunk_handler::GenerateChunksRequest,
handlers::chunk_handler::SearchChunkData,
handlers::chunk_handler::ScoreChunkDTO,
handlers::chunk_handler::SearchChunkResponseTypes,
handlers::chunk_handler::RecommendChunksResponseTypes,
handlers::chunk_handler::RecommendChunkMetadata,
handlers::chunk_handler::RecommendSlimChunkMetadata,
handlers::group_handler::RecommendGroupChunks,
handlers::group_handler::RecommendGroupSlimChunks,
handlers::group_handler::SearchWithinGroupData,
handlers::group_handler::SearchOverGroupsData,
handlers::group_handler::SearchGroupsResult,
handlers::group_handler::SearchWithinGroupResponseTypes,
handlers::group_handler::RecommendGroupChunkResponseTypes,
handlers::group_handler::SearchOverGroupsResponseTypes,
handlers::group_handler::SearchWithinGroupResults,
handlers::chunk_handler::SearchChunkQueryResponseBody,
handlers::chunk_handler::ChunkFilter,
handlers::chunk_handler::FieldCondition,
Expand All @@ -270,7 +261,7 @@ impl Modify for SecurityAddon {
handlers::organization_handler::CreateOrganizationData,
handlers::organization_handler::UpdateOrganizationData,
operators::event_operator::EventReturn,
operators::search_operator::SearchOverGroupsResponseBody,
operators::search_operator::SearchOverGroupsResults,
operators::search_operator::GroupScoreChunk,
handlers::dataset_handler::CreateDatasetRequest,
handlers::dataset_handler::UpdateDatasetRequest,
Expand Down Expand Up @@ -300,8 +291,8 @@ impl Modify for SecurityAddon {
data::models::ScoreSlimChunks,
data::models::SlimChunkMetadata,
data::models::SlimChunkMetadataWithScore,
data::models::SearchGroupSlimChunksResult,
data::models::SearchOverGroupsSlimChunksResponseBody,
data::models::SearchWithinGroupSlimResults,
data::models::SearchOverGroupsSlimResults,
data::models::GroupScoreSlimChunks,
handlers::chunk_handler::RangeCondition,
errors::ErrorResponseBody,
Expand Down
Loading

0 comments on commit 8e49f6e

Please sign in to comment.