poimen-memory/crates/mem-cli/src/handlers/unified_synthesis.rs

//! Unified Synthesis Endpoint (Phase 5.5)
//!
//! Single composable endpoint combining entity linking, inference, reasoning, summarization.

use actix_web::{web, HttpRequest, HttpResponse};
use serde::{Deserialize, Serialize};
use crate::query::{
    EntityLinker, InferenceEngine, QueryReasoner, Summarizer,
    SummarizationStrategy, MentionLink,
};
use crate::handlers::response_builder;
use tracing::{debug, info, error};

/// Unified synthesis request
#[derive(Debug, Deserialize)]
pub struct UnifiedSynthesisRequest {
    pub project: String,
    pub content: String,

    // Entity linking options
    #[serde(default)]
    pub link_entities: bool,
    #[serde(default)]
    pub detect_aliases: bool,

    // Inference options
    #[serde(default)]
    pub infer_facts: bool,
    #[serde(default)]
    pub transitive_closure: bool,

    // Reasoning options
    #[serde(default)]
    pub reason_query: bool,

    // Summarization options
    #[serde(default)]
    pub summarize: bool,
    #[serde(default = "default_max_length")]
    pub max_length: usize,
    #[serde(default = "default_strategy")]
    pub strategy: String,
}

fn default_max_length() -> usize { 200 }
fn default_strategy() -> String { "hybrid".to_string() }

/// Unified synthesis response
#[derive(Debug, Serialize)]
pub struct UnifiedSynthesisResponse {
    pub project: String,
    pub entity_linking: Option<EntityLinkingResult>,
    pub inference: Option<InferenceResult>,
    pub reasoning: Option<ReasoningResult>,
    pub summarization: Option<SummarizationResult>,
    pub process_time_ms: u128,
}

/// Entity linking result
#[derive(Debug, Serialize)]
pub struct EntityLinkingResult {
    pub mention_links: Vec<MentionLinkResponse>,
    pub alias_count: usize,
}

/// Mention link in response
#[derive(Debug, Serialize)]
pub struct MentionLinkResponse {
    pub mention: String,
    pub entity_id: String,
    pub confidence: f32,
}

/// Inference result
#[derive(Debug, Serialize)]
pub struct InferenceResult {
    pub inferred_facts: Vec<InferredFactResponse>,
    pub fact_count: usize,
}

/// Inferred fact in response
#[derive(Debug, Serialize)]
pub struct InferredFactResponse {
    pub source: String,
    pub relation: String,
    pub target: String,
    pub confidence: f32,
}

/// Reasoning result
#[derive(Debug, Serialize)]
pub struct ReasoningResult {
    pub question: String,
    pub answers: Vec<String>,
    pub confidence: f32,
    pub step_count: usize,
}

/// Summarization result
#[derive(Debug, Serialize)]
pub struct SummarizationResult {
    pub summary: String,
    pub compression_ratio: f32,
    pub coherence: f32,
    pub key_facts_count: usize,
}

/// POST /memory/synthesis - Unified synthesis endpoint
/// Delegates reasoning to Temporal workflows (activities persist to DB)
pub async fn unified_synthesis_handler(
    req: HttpRequest,
    body: web::Json<UnifiedSynthesisRequest>,
    state: web::Data<crate::AppState>,
) -> HttpResponse {
    let start_time = std::time::Instant::now();

    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
        &req, &state, "synthesis", 50
    ) {
        return response;
    }

    if body.content.is_empty() || body.content.len() > 100000 {
        return response_builder::bad_request("Content must be 1-100K characters");
    }

    // Check at least one operation requested
    if !body.link_entities && !body.infer_facts && !body.reason_query && !body.summarize {
        return response_builder::bad_request(
            "At least one operation must be requested (link_entities, infer_facts, reason_query, summarize)"
        );
    }

    // Extract JWT token for Temporal workflow calls
    let jwt = match crate::handlers::extract_jwt_token(&req) {
        Some(token) => token,
        None => {
            return response_builder::unauthorized("Bearer token required for synthesis");
        }
    };

    debug!(
        "Unified synthesis: linking={}, inferring={}, reasoning={}, summarizing={}",
        body.link_entities, body.infer_facts, body.reason_query, body.summarize
    );

    // Create Temporal workflow client
    let synthesis_client = crate::agent::client_sdk::SynthesisClient::new(
        "https://api.riotpiao.com".to_string(),
        jwt,
    );

    let mut entity_linking = None;
    let mut inference = None;
    let mut reasoning = None;
    let mut summarization = None;

    // Entity Linking
    if body.link_entities {
        let linker = EntityLinker::new(state.pool.clone());
        match linker.link_entities(&body.content) {
            Ok(links) => {
                let alias_count = links.iter().filter(|l| l.confidence > 0.85).count();
                entity_linking = Some(EntityLinkingResult {
                    mention_links: links.iter().map(|l| MentionLinkResponse {
                        mention: l.mention.clone(),
                        entity_id: l.entity_id.clone(),
                        confidence: l.confidence,
                    }).collect(),
                    alias_count,
                });
            }
            Err(e) => {
                error!("Entity linking failed: {}", e);
                return response_builder::internal_error("Entity linking failed");
            }
        }
    }

    // Inference
    if body.infer_facts {
        let engine = InferenceEngine::new(state.pool.clone());
        match engine.infer_facts(&body.content, 5, 0.6, &body.project) {
            Ok(facts) => {
                inference = Some(InferenceResult {
                    inferred_facts: facts.iter().map(|f| InferredFactResponse {
                        source: f.source.clone(),
                        relation: f.relation.clone(),
                        target: f.target.clone(),
                        confidence: f.confidence,
                    }).collect(),
                    fact_count: facts.len(),
                });
            }
            Err(e) => {
                error!("Inference failed: {}", e);
                return response_builder::internal_error("Inference failed");
            }
        }
    }

    // Reasoning via Temporal workflow
    // Temporal activity calls LLMInferenceActivity + persists results to memory_entity/memory_edge
    if body.reason_query {
        reasoning = match execute_reasoning_workflow(
            &synthesis_client,
            &body,
        ).await {
            Ok(result) => Some(result),
            Err(e) => {
                error!("Reasoning workflow failed: {}", e);
                return response_builder::internal_error(&e);
            }
        };
    }

    // Summarization
    if body.summarize {
        let summarizer = Summarizer::new();
        let strategy = match body.strategy.to_lowercase().as_str() {
            "extractive" => SummarizationStrategy::Extractive,
            "abstractive" => SummarizationStrategy::Abstractive,
            "hybrid" | _ => SummarizationStrategy::Hybrid,
        };

        match summarizer.summarize(&body.content, body.max_length, strategy) {
            Ok(summary) => {
                summarization = Some(SummarizationResult {
                    summary: summary.text,
                    compression_ratio: summary.compression_ratio,
                    coherence: summary.coherence,
                    key_facts_count: summary.key_facts.len(),
                });
            }
            Err(e) => {
                error!("Summarization failed: {}", e);
                return response_builder::internal_error("Summarization failed");
            }
        }
    }

    let elapsed = start_time.elapsed().as_millis();

    info!(
        "Unified synthesis completed in {}ms: linking={}, inference={}, reasoning={}, summary={}",
        elapsed,
        entity_linking.is_some(),
        inference.is_some(),
        reasoning.is_some(),
        summarization.is_some()
    );

    response_builder::success_response(UnifiedSynthesisResponse {
        project: body.project.clone(),
        entity_linking,
        inference,
        reasoning,
        summarization,
        process_time_ms: elapsed,
    })
}

/// Select LLM model based on operations complexity
fn select_llm_model(operations: &serde_json::Value) -> &'static str {
    let reason_query = operations.get("reason_query")
        .and_then(|v| v.as_bool())
        .unwrap_or(false);
    let summarize = operations.get("summarize")
        .and_then(|v| v.as_bool())
        .unwrap_or(false);

    match (reason_query, summarize) {
        (true, true) => "reasoning",      // Complex: extract + reason + summarize
        (true, false) => "ornith:35b",    // Medium: extract + reason
        (false, _) => "qwen2.5:3b",       // Quick: only linking/inference
    }
}

/// Build LLM system prompt for entity/fact extraction
fn build_system_prompt(operations: &serde_json::Value) -> String {
    let reason_query = operations.get("reason_query")
        .and_then(|v| v.as_bool())
        .unwrap_or(false);

    // Use Zep-inspired reasoning prompt for complex reasoning, entity extraction otherwise
    if reason_query {
        crate::handlers::llm_prompts::reasoning_system_prompt().to_string()
    } else {
        crate::handlers::llm_prompts::entity_extraction_system_prompt().to_string()
    }
}

/// Execute reasoning workflow via Temporal
/// Returns parsed ReasoningResult from workflow output
async fn execute_reasoning_workflow(
    client: &crate::agent::client_sdk::SynthesisClient,
    body: &UnifiedSynthesisRequest,
) -> Result<ReasoningResult, String> {
    // Prepare operations metadata
    let operations = serde_json::json!({
        "link_entities": body.link_entities,
        "infer_facts": body.infer_facts,
        "reason_query": body.reason_query,
        "summarize": body.summarize
    });

    // Select model based on complexity
    let model = select_llm_model(&operations);
    let system_prompt = build_system_prompt(&operations);

    // Build START_WORKFLOW request with LLMInferenceActivity inputs
    let workflow_input = serde_json::json!({
        // Workflow input
        "question": body.content,
        "project": body.project,
        "operations": operations,
        
        // LLMInferenceActivity inputs (passed to Temporal activity)
        "llm_activity": {
            "model": model,
            "system_prompt": system_prompt,
            "user_prompt": body.content,
            "temperature": 0.7,
            "max_tokens": 2048
        }
    });

    let workflow_builder = crate::handlers::WorkflowBuilder::new("ReasoningWorkflow")
        .with_input(workflow_input);
    let workflow_id = workflow_builder.workflow_id().to_string();
    let workflow_req = workflow_builder.build();

    debug!("Starting ReasoningWorkflow: {}", workflow_id);

    // Start workflow
    client.execute_workflow(workflow_req).await?;

    // Poll until completion
    let (_, result) = crate::handlers::poll_workflow_until_complete(
        client,
        &workflow_id,
        crate::handlers::PollConfig::default(),
    ).await?;

    // Extract and parse result
    let result = result.ok_or("Workflow returned no result".to_string())?;

    Ok(ReasoningResult {
        question: result.get("question")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string(),
        answers: result.get("answers")
            .and_then(|v| v.as_array())
            .map(|arr| arr.iter()
                .filter_map(|v| v.as_str().map(|s| s.to_string()))
                .collect())
            .unwrap_or_default(),
        confidence: result.get("confidence")
            .and_then(|v| v.as_f64())
            .unwrap_or(0.0) as f32,
        step_count: result.get("reasoning_steps")
            .and_then(|v| v.as_array())
            .map(|arr| arr.len())
            .unwrap_or(0),
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_unified_synthesis_request_structure() {
        let req = UnifiedSynthesisRequest {
            project: "poimen".to_string(),
            content: "Test content".to_string(),
            link_entities: true,
            detect_aliases: false,
            infer_facts: false,
            transitive_closure: false,
            reason_query: false,
            summarize: false,
            max_length: 200,
            strategy: "hybrid".to_string(),
        };
        assert!(req.link_entities);
    }

    #[test]
    fn test_default_max_length() {
        assert_eq!(default_max_length(), 200);
    }

    #[test]
    fn test_default_strategy() {
        assert_eq!(default_strategy(), "hybrid");
    }

    #[test]
    fn test_all_operations_enabled() {
        let req = UnifiedSynthesisRequest {
            project: "p".to_string(),
            content: "c".to_string(),
            link_entities: true,
            detect_aliases: true,
            infer_facts: true,
            transitive_closure: true,
            reason_query: true,
            summarize: true,
            max_length: 200,
            strategy: "hybrid".to_string(),
        };
        assert!(req.link_entities && req.infer_facts && req.reason_query && req.summarize);
    }

    #[test]
    fn test_entity_linking_result_structure() {
        let result = EntityLinkingResult {
            mention_links: vec![],
            alias_count: 0,
        };
        assert_eq!(result.alias_count, 0);
    }

    #[test]
    fn test_inference_result_structure() {
        let result = InferenceResult {
            inferred_facts: vec![],
            fact_count: 0,
        };
        assert_eq!(result.fact_count, 0);
    }

    #[test]
    fn test_reasoning_result_structure() {
        let result = ReasoningResult {
            question: "Test?".to_string(),
            answers: vec![],
            confidence: 0.8,
            step_count: 1,
        };
        assert_eq!(result.step_count, 1);
    }

    #[test]
    fn test_summarization_result_structure() {
        let result = SummarizationResult {
            summary: "Summary".to_string(),
            compression_ratio: 0.5,
            coherence: 0.8,
            key_facts_count: 3,
        };
        assert_eq!(result.key_facts_count, 3);
    }
}