gitdataai/libs/service/chat/context.rs

use std::collections::HashSet;

use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};
use uuid::Uuid;

use crate::AppService;
use crate::error::AppError;
use models::projects::project_skill;
use models::repos::repo;

fn metadata_object(
    metadata: Option<&serde_json::Value>,
) -> Option<&serde_json::Map<String, serde_json::Value>> {
    metadata?.as_object()
}

fn slash_context_object(
    metadata: Option<&serde_json::Value>,
) -> Option<&serde_json::Map<String, serde_json::Value>> {
    metadata_object(metadata)?.get("slash_context")?.as_object()
}

fn stringify_text(value: &serde_json::Value) -> Option<String> {
    value
        .as_str()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToOwned::to_owned)
}

fn repo_ids_from_metadata(metadata: Option<&serde_json::Value>) -> Vec<Uuid> {
    let Some(context) = slash_context_object(metadata) else {
        return Vec::new();
    };

    let Some(repos) = context.get("repos").and_then(|value| value.as_array()) else {
        return Vec::new();
    };

    let mut seen = HashSet::new();
    let mut ids = Vec::new();

    for repo in repos {
        let Some(repo_id) = repo
            .as_object()
            .and_then(|value| value.get("id"))
            .and_then(stringify_text)
        else {
            continue;
        };

        let Ok(repo_uuid) = Uuid::parse_str(&repo_id) else {
            continue;
        };

        if seen.insert(repo_uuid) {
            ids.push(repo_uuid);
        }
    }

    ids
}

fn skill_ids_from_metadata(metadata: Option<&serde_json::Value>) -> Vec<i64> {
    let Some(context) = slash_context_object(metadata) else {
        return Vec::new();
    };

    let Some(skills) = context.get("skills").and_then(|value| value.as_array()) else {
        return Vec::new();
    };

    let mut seen = HashSet::new();
    let mut ids = Vec::new();

    for skill in skills {
        let Some(skill_id) = skill
            .as_object()
            .and_then(|value| value.get("id"))
            .and_then(stringify_text)
        else {
            continue;
        };

        let Ok(skill_id) = skill_id.parse::<i64>() else {
            continue;
        };

        if seen.insert(skill_id) {
            ids.push(skill_id);
        }
    }

    ids
}

impl AppService {
    pub async fn build_message_context_prompts(
        &self,
        project_id: Option<Uuid>,
        metadata: Option<&serde_json::Value>,
    ) -> Result<Vec<String>, AppError> {
        let mut prompts = Vec::new();

        let repo_ids = repo_ids_from_metadata(metadata);
        for repo_id in repo_ids {
            let mut query = repo::Entity::find().filter(repo::Column::Id.eq(repo_id));
            if let Some(project_id) = project_id {
                query = query.filter(repo::Column::Project.eq(project_id));
            }

            if let Some(repo) = query.one(self.db.reader()).await? {
                let mut parts = vec![
                    format!("Repository name: {}", repo.repo_name),
                    format!("Repository id: {}", repo.id),
                    format!("Default branch: {}", repo.default_branch),
                    format!(
                        "Visibility: {}",
                        if repo.is_private { "private" } else { "public" }
                    ),
                ];

                if let Some(description) = repo.description.as_deref() {
                    parts.push(format!("Description: {}", description));
                }

                prompts.push(format!(
                    "[Selected repository context]\n{}",
                    parts.join("\n")
                ));
            }
        }

        let skill_ids = skill_ids_from_metadata(metadata);
        if let Some(project_id) = project_id {
            for skill_id in skill_ids {
                if let Some(skill) = project_skill::Entity::find()
                    .filter(project_skill::Column::Id.eq(skill_id))
                    .filter(project_skill::Column::ProjectUuid.eq(project_id))
                    .filter(project_skill::Column::Enabled.eq(true))
                    .one(self.db.reader())
                    .await?
                {
                    let mut header = vec![
                        format!("Skill name: {}", skill.name),
                        format!("Skill slug: {}", skill.slug),
                        format!("Skill source: {}", skill.source),
                    ];

                    if let Some(description) = skill.description.as_deref() {
                        header.push(format!("Description: {}", description));
                    }

                    prompts.push(format!(
                        "[Selected skill context]\n{}\n\n{}",
                        header.join("\n"),
                        skill.content
                    ));
                }
            }
        }

        Ok(prompts)
    }
}