cim_summary/backend/src/models/AgenticRAGModels.ts

import db from '../config/database';
import { AgentExecution, AgenticRAGSession, QualityMetrics } from './agenticTypes';
import { logger } from '../utils/logger';

export class AgentExecutionModel {
  /**
   * Create a new agent execution record
   */
  static async create(execution: Omit<AgentExecution, 'id' | 'createdAt' | 'updatedAt'>): Promise<AgentExecution> {
    const query = `
      INSERT INTO agent_executions (
        document_id, session_id, agent_name, step_number, status, 
        input_data, output_data, validation_result, processing_time_ms, 
        error_message, retry_count
      ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
      RETURNING *
    `;

    const values = [
      execution.documentId,
      execution.sessionId,
      execution.agentName,
      execution.stepNumber,
      execution.status,
      execution.inputData,
      execution.outputData,
      execution.validationResult,
      execution.processingTimeMs,
      execution.errorMessage,
      execution.retryCount
    ];

    try {
      const result = await db.query(query, values);
      return this.mapRowToAgentExecution(result.rows[0]);
    } catch (error) {
      logger.error('Failed to create agent execution', { error, execution });
      throw error;
    }
  }

  /**
   * Update an agent execution record
   */
  static async update(id: string, updates: Partial<AgentExecution>): Promise<AgentExecution> {
    const setClauses: string[] = [];
    const values: any[] = [];
    let paramCount = 1;

    // Build dynamic update query
    if (updates.status !== undefined) {
      setClauses.push(`status = $${paramCount++}`);
      values.push(updates.status);
    }
    if (updates.outputData !== undefined) {
      setClauses.push(`output_data = $${paramCount++}`);
      values.push(updates.outputData);
    }
    if (updates.validationResult !== undefined) {
      setClauses.push(`validation_result = $${paramCount++}`);
      values.push(updates.validationResult);
    }
    if (updates.processingTimeMs !== undefined) {
      setClauses.push(`processing_time_ms = $${paramCount++}`);
      values.push(updates.processingTimeMs);
    }
    if (updates.errorMessage !== undefined) {
      setClauses.push(`error_message = $${paramCount++}`);
      values.push(updates.errorMessage);
    }
    if (updates.retryCount !== undefined) {
      setClauses.push(`retry_count = $${paramCount++}`);
      values.push(updates.retryCount);
    }

    if (setClauses.length === 0) {
      throw new Error('No updates provided');
    }

    values.push(id);
    const query = `
      UPDATE agent_executions 
      SET ${setClauses.join(', ')}, updated_at = NOW()
      WHERE id = $${paramCount}
      RETURNING *
    `;

    try {
      const result = await db.query(query, values);
      if (result.rows.length === 0) {
        throw new Error(`Agent execution with id ${id} not found`);
      }
      return this.mapRowToAgentExecution(result.rows[0]);
    } catch (error) {
      logger.error('Failed to update agent execution', { error, id, updates });
      throw error;
    }
  }

  /**
   * Get agent executions by session ID
   */
  static async getBySessionId(sessionId: string): Promise<AgentExecution[]> {
    const query = `
      SELECT * FROM agent_executions 
      WHERE session_id = $1 
      ORDER BY step_number ASC
    `;

    try {
      const result = await db.query(query, [sessionId]);
      return result.rows.map((row: any) => this.mapRowToAgentExecution(row));
    } catch (error) {
      logger.error('Failed to get agent executions by session ID', { error, sessionId });
      throw error;
    }
  }

  /**
   * Get agent execution by ID
   */
  static async getById(id: string): Promise<AgentExecution | null> {
    const query = 'SELECT * FROM agent_executions WHERE id = $1';

    try {
      const result = await db.query(query, [id]);
      return result.rows.length > 0 ? this.mapRowToAgentExecution(result.rows[0]) : null;
    } catch (error) {
      logger.error('Failed to get agent execution by ID', { error, id });
      throw error;
    }
  }

  private static mapRowToAgentExecution(row: any): AgentExecution {
    return {
      id: row.id,
      documentId: row.document_id,
      sessionId: row.session_id,
      agentName: row.agent_name,
      stepNumber: row.step_number,
      status: row.status,
      inputData: row.input_data,
      outputData: row.output_data,
      validationResult: row.validation_result,
      processingTimeMs: row.processing_time_ms,
      errorMessage: row.error_message,
      retryCount: row.retry_count,
      createdAt: new Date(row.created_at),
      updatedAt: new Date(row.updated_at)
    };
  }
}

export class AgenticRAGSessionModel {
  /**
   * Create a new agentic RAG session
   */
  static async create(session: Omit<AgenticRAGSession, 'id' | 'createdAt' | 'completedAt'>): Promise<AgenticRAGSession> {
    const query = `
      INSERT INTO agentic_rag_sessions (
        document_id, user_id, strategy, status, total_agents, 
        completed_agents, failed_agents, overall_validation_score,
        processing_time_ms, api_calls_count, total_cost, 
        reasoning_steps, final_result
      ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
      RETURNING *
    `;

    const values = [
      session.documentId,
      session.userId,
      session.strategy,
      session.status,
      session.totalAgents,
      session.completedAgents,
      session.failedAgents,
      session.overallValidationScore,
      session.processingTimeMs,
      session.apiCallsCount,
      session.totalCost,
      session.reasoningSteps,
      session.finalResult
    ];

    try {
      const result = await db.query(query, values);
      return this.mapRowToSession(result.rows[0]);
    } catch (error) {
      logger.error('Failed to create agentic RAG session', { error, session });
      throw error;
    }
  }

  /**
   * Update an agentic RAG session
   */
  static async update(id: string, updates: Partial<AgenticRAGSession>): Promise<AgenticRAGSession> {
    const setClauses: string[] = [];
    const values: any[] = [];
    let paramCount = 1;

    // Build dynamic update query
    if (updates.status !== undefined) {
      setClauses.push(`status = $${paramCount++}`);
      values.push(updates.status);
    }
    if (updates.completedAgents !== undefined) {
      setClauses.push(`completed_agents = $${paramCount++}`);
      values.push(updates.completedAgents);
    }
    if (updates.failedAgents !== undefined) {
      setClauses.push(`failed_agents = $${paramCount++}`);
      values.push(updates.failedAgents);
    }
    if (updates.overallValidationScore !== undefined) {
      setClauses.push(`overall_validation_score = $${paramCount++}`);
      values.push(updates.overallValidationScore);
    }
    if (updates.processingTimeMs !== undefined) {
      setClauses.push(`processing_time_ms = $${paramCount++}`);
      values.push(updates.processingTimeMs);
    }
    if (updates.apiCallsCount !== undefined) {
      setClauses.push(`api_calls_count = $${paramCount++}`);
      values.push(updates.apiCallsCount);
    }
    if (updates.totalCost !== undefined) {
      setClauses.push(`total_cost = $${paramCount++}`);
      values.push(updates.totalCost);
    }
    if (updates.reasoningSteps !== undefined) {
      setClauses.push(`reasoning_steps = $${paramCount++}`);
      values.push(updates.reasoningSteps);
    }
    if (updates.finalResult !== undefined) {
      setClauses.push(`final_result = $${paramCount++}`);
      values.push(updates.finalResult);
    }
    if (updates.completedAt !== undefined) {
      setClauses.push(`completed_at = $${paramCount++}`);
      values.push(updates.completedAt);
    }

    if (setClauses.length === 0) {
      throw new Error('No updates provided');
    }

    values.push(id);
    const query = `
      UPDATE agentic_rag_sessions 
      SET ${setClauses.join(', ')}
      WHERE id = $${paramCount}
      RETURNING *
    `;

    try {
      const result = await db.query(query, values);
      if (result.rows.length === 0) {
        throw new Error(`Session with id ${id} not found`);
      }
      return this.mapRowToSession(result.rows[0]);
    } catch (error) {
      logger.error('Failed to update agentic RAG session', { error, id, updates });
      throw error;
    }
  }

  /**
   * Get session by ID
   */
  static async getById(id: string): Promise<AgenticRAGSession | null> {
    const query = 'SELECT * FROM agentic_rag_sessions WHERE id = $1';

    try {
      const result = await db.query(query, [id]);
      return result.rows.length > 0 ? this.mapRowToSession(result.rows[0]) : null;
    } catch (error) {
      logger.error('Failed to get session by ID', { error, id });
      throw error;
    }
  }

  /**
   * Get sessions by document ID
   */
  static async getByDocumentId(documentId: string): Promise<AgenticRAGSession[]> {
    const query = `
      SELECT * FROM agentic_rag_sessions 
      WHERE document_id = $1 
      ORDER BY created_at DESC
    `;

    try {
      const result = await db.query(query, [documentId]);
      return result.rows.map((row: any) => this.mapRowToSession(row));
    } catch (error) {
      logger.error('Failed to get sessions by document ID', { error, documentId });
      throw error;
    }
  }

  /**
   * Get sessions by user ID
   */
  static async getByUserId(userId: string): Promise<AgenticRAGSession[]> {
    const query = `
      SELECT * FROM agentic_rag_sessions 
      WHERE user_id = $1 
      ORDER BY created_at DESC
    `;

    try {
      const result = await db.query(query, [userId]);
      return result.rows.map((row: any) => this.mapRowToSession(row));
    } catch (error) {
      logger.error('Failed to get sessions by user ID', { error, userId });
      throw error;
    }
  }

  private static mapRowToSession(row: any): AgenticRAGSession {
    return {
      id: row.id,
      documentId: row.document_id,
      userId: row.user_id,
      strategy: row.strategy,
      status: row.status,
      totalAgents: row.total_agents,
      completedAgents: row.completed_agents,
      failedAgents: row.failed_agents,
      overallValidationScore: row.overall_validation_score,
      processingTimeMs: row.processing_time_ms,
      apiCallsCount: row.api_calls_count,
      totalCost: row.total_cost,
      reasoningSteps: row.reasoning_steps || [],
      finalResult: row.final_result,
      createdAt: new Date(row.created_at),
      completedAt: row.completed_at ? new Date(row.completed_at) : undefined
    };
  }
}

export class QualityMetricsModel {
  /**
   * Create a new quality metric record
   */
  static async create(metric: Omit<QualityMetrics, 'id' | 'createdAt'>): Promise<QualityMetrics> {
    const query = `
      INSERT INTO processing_quality_metrics (
        document_id, session_id, metric_type, metric_value, metric_details
      ) VALUES ($1, $2, $3, $4, $5)
      RETURNING *
    `;

    const values = [
      metric.documentId,
      metric.sessionId,
      metric.metricType,
      metric.metricValue,
      metric.metricDetails
    ];

    try {
      const result = await db.query(query, values);
      return this.mapRowToQualityMetric(result.rows[0]);
    } catch (error) {
      logger.error('Failed to create quality metric', { error, metric });
      throw error;
    }
  }

  /**
   * Get quality metrics by session ID
   */
  static async getBySessionId(sessionId: string): Promise<QualityMetrics[]> {
    const query = `
      SELECT * FROM processing_quality_metrics 
      WHERE session_id = $1 
      ORDER BY created_at ASC
    `;

    try {
      const result = await db.query(query, [sessionId]);
      return result.rows.map((row: any) => this.mapRowToQualityMetric(row));
    } catch (error) {
      logger.error('Failed to get quality metrics by session ID', { error, sessionId });
      throw error;
    }
  }

  /**
   * Get quality metrics by document ID
   */
  static async getByDocumentId(documentId: string): Promise<QualityMetrics[]> {
    const query = `
      SELECT * FROM processing_quality_metrics 
      WHERE document_id = $1 
      ORDER BY created_at DESC
    `;

    try {
      const result = await db.query(query, [documentId]);
      return result.rows.map((row: any) => this.mapRowToQualityMetric(row));
    } catch (error) {
      logger.error('Failed to get quality metrics by document ID', { error, documentId });
      throw error;
    }
  }

  private static mapRowToQualityMetric(row: any): QualityMetrics {
    return {
      id: row.id,
      documentId: row.document_id,
      sessionId: row.session_id,
      metricType: row.metric_type,
      metricValue: parseFloat(row.metric_value),
      metricDetails: row.metric_details,
      createdAt: new Date(row.created_at)
    };
  }
}