katniss-ai/src/server/decisions.rs

use crate::openai;
use crate::server::AppState;
use axum::Json;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::response::Response;
use axum_extra::routing::TypedPath;
use serde::Deserialize;
use serde_with::DeserializeFromStr;
use serde_with::SerializeDisplay;
use std::collections::HashMap;
use std::hash::RandomState;
use strum::EnumString;
use strum::VariantArray;

#[derive(Deserialize, TypedPath)]
#[typed_path("/v1/decisions")]
pub struct V1DecisionsPath;

#[derive(
    Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
)]
#[strum(serialize_all = "lowercase")]
pub enum Shortcut {
    S,
    M,
}

pub struct DecisionsResponse(openai::decisions::DecisionResponse<String>);

impl IntoResponse for DecisionsResponse {
    fn into_response(self) -> Response {
        let Self(embeddings) = self;

        Json(embeddings).into_response()
    }
}

pub enum DecisionsError {
    NoneSucceeded,
}

impl IntoResponse for DecisionsError {
    fn into_response(self) -> Response {
        match self {
            DecisionsError::NoneSucceeded => StatusCode::SERVICE_UNAVAILABLE.into_response(),
        }
    }
}

#[tracing::instrument(skip(app_state, request))]
pub async fn create_decision(
    _: V1DecisionsPath,
    State(app_state): State<AppState>,
    Json(request): Json<openai::decisions::DecisionRequest<Shortcut>>,
) -> Result<DecisionsResponse, DecisionsError> {
    use super::system_one;

    let openai::decisions::DecisionRequest {
        model,
        input,
        questions,
    } = request;

    let model = match model {
        Shortcut::S => system_one::Shortcut::S,
        Shortcut::M => system_one::Shortcut::M,
    };

    let (state, images, files) = match input {
        openai::decisions::DecisionInput::Messages(messages) => {
            let mut images = vec![];
            let mut texts = vec![];

            for openai::decisions::DecisionInputMessage { role, content } in messages {
                match content {
                    openai::decisions::DecisionInputContent::Text(text) => texts.push(text),
                    openai::decisions::DecisionInputContent::Parts(decision_input_parts) => {
                        for decision_input_part in decision_input_parts {
                            match decision_input_part {
                                openai::decisions::DecisionInputPart::Image { image_url, detail } => images.push(image_url),
                                openai::decisions::DecisionInputPart::Text { text } => texts.push(text),
                            }
                        }
                    },
                }
            };

            let text = texts.join("\n\n");

            let files = images.clone();

            (openai::system_one::State::Text(text), Some(images), Some(files))
        }
        openai::decisions::DecisionInput::Text(text) => (openai::system_one::State::Text(text), None, None),
    };

    let questions = questions.into_iter().map(|question| match question {
        openai::decisions::DecisionQuestion::Predicate(predicate_question) => todo!(),
        openai::decisions::DecisionQuestion::Choice(choice_question) => todo!(),
        openai::decisions::DecisionQuestion::Score(score_question) => todo!(),
    } ).collect();

    let request = openai::system_one::SystemOneRequest {
        model,
        state,
        files,
        images,
        questions,
    };

    system_one::create_system_one(system_one::V1SystemOnePath, State(app_state), Json(request))
        .await
        .map(
            |system_one::SystemOneResponse(openai::system_one::SystemOneResponse {
                 model,
                 answers,
                 usage,
             })| {
                let answers = answers
                    .into_iter()
                    .map(|(id, answer)| match answer {
                        openai::system_one::Answer::Noul(openai::system_one::NoulAnswer {
                            noul,
                        }) => openai::decisions::DecisionAnswer::Predicate(
                            openai::decisions::PredicateAnswer {
                                name: Some(id),
                                probability: noul,
                            },
                        ),
                        openai::system_one::Answer::Choice(openai::system_one::ChoiceAnswer {
                            choice,
                            probabilities,
                            confidence,
                        }) => openai::decisions::DecisionAnswer::Choice(
                            openai::decisions::ChoiceAnswer {
                                name: Some(id),
                                choice: openai::decisions::Choice::String(choice),
                                confidence,
                                probabilities: probabilities
                                    .into_iter()
                                    .map(|(choice_id, probability)| {
                                        openai::decisions::ChoiceProbability {
                                            value: openai::decisions::Choice::String(choice_id),
                                            probability,
                                        }
                                    })
                                    .collect(),
                            },
                        ),
                        openai::system_one::Answer::Score(openai::system_one::ScoreAnswer {
                            score,
                            legend,
                            probabilities,
                            confidence,
                        }) => openai::decisions::DecisionAnswer::Score(
                            openai::decisions::ScoreAnswer {
                                name: Some(id),
                                score,
                                confidence,
                                probabilities: {
                                    let reverse_legend = HashMap::<_, _, RandomState>::from_iter(
                                        legend.into_iter().map(|(index, value)| {
                                            (
                                                value,
                                                index
                                                    .parse::<i32>()
                                                    .expect("this should be an integer, I believe"),
                                            )
                                        }),
                                    );

                                    probabilities
                                        .into_iter()
                                        .map(|(choice_id, probability)| {
                                            openai::decisions::ScoreProbability {
                                                value: reverse_legend
                                                    .get(&choice_id)
                                                    .expect("it should exist, I think")
                                                    .clone(),
                                                label: choice_id,
                                                probability,
                                            }
                                        })
                                        .collect()
                                },
                            },
                        ),
                    })
                    .collect();

                let openai::system_one::Usage {
                    input_tokens,
                    output_tokens,
                } = usage;
                let input_tokens = input_tokens as i32;
                let output_tokens = output_tokens as i32;
                let total_tokens = input_tokens + output_tokens;
                let input_tokens_details = openai::decisions::TokenUsageDetails {
                    cache_write_tokens: 0,
                    cached_tokens: 0,
                };
                let output_tokens_details = openai::decisions::OutputTokenDetails {
                    reasoning_tokens: 0,
                };
                let usage = openai::decisions::TokenUsage {
                    input_tokens,
                    input_tokens_details,
                    output_tokens,
                    output_tokens_details,
                    total_tokens,
                };

                DecisionsResponse(openai::decisions::DecisionResponse {
                    model,
                    answers,
                    usage,
                })
            },
        )
        .map_err(|_| DecisionsError::NoneSucceeded)
}