chronicle-proxy 0.4.3

LLM Provider Abstraction and Logging
Documentation
use chrono::{TimeZone, Utc};
use serde_json::json;
use uuid::Uuid;

use crate::{
    database::logging::{ProxyLogEntry, ProxyLogEvent},
    workflow_events::{
        ErrorData, EventPayload, RunStartEvent, RunUpdateEvent, StepEndData, StepEventData,
        StepStartData, WorkflowEvent,
    },
};

pub const TEST_STEP1_ID: Uuid = Uuid::from_u128(1);
pub const TEST_STEP2_ID: Uuid = Uuid::from_u128(2);
pub const TEST_RUN_ID: Uuid = Uuid::from_u128(100);
pub const TEST_EVENT1_ID: Uuid = Uuid::from_u128(5);

pub fn test_events() -> Vec<ProxyLogEntry> {
    vec![
        ProxyLogEntry::Workflow(WorkflowEvent::RunStart(RunStartEvent {
            id: TEST_RUN_ID,
            name: "test run".to_string(),
            description: Some("test description".to_string()),
            application: Some("test application".to_string()),
            environment: Some("test environment".to_string()),
            status: None,
            input: Some(json!({"query":"abc"})),
            trace_id: Some("0123456789abcdef".to_string()),
            span_id: Some("12345678".to_string()),
            tags: vec!["tag1".to_string(), "tag2".to_string()],
            info: Some(json!({
                "info1": "value1",
                "info2": "value2"
            })),
            time: Some(Utc.timestamp_opt(1, 0).unwrap()),
        })),
        ProxyLogEntry::Workflow(WorkflowEvent::StepStart(StepEventData {
            step_id: TEST_STEP1_ID,
            run_id: TEST_RUN_ID,
            time: Some(Utc.timestamp_opt(2, 0).unwrap()),
            data: StepStartData {
                name: Some("source_node1".to_string()),
                typ: "step_type".to_string(),
                parent_step: None,
                span_id: Some("11111111".to_string()),
                info: Some(json!({ "model": "a_model" })),
                tags: vec!["dag".to_string(), "node".to_string()],
                input: json!({ "task_param": "value" }),
            },
        })),
        ProxyLogEntry::Workflow(WorkflowEvent::StepStart(StepEventData {
            step_id: TEST_STEP2_ID,
            run_id: TEST_RUN_ID,
            time: Some(Utc.timestamp_opt(3, 0).unwrap()),
            data: StepStartData {
                name: Some("source_node2".to_string()),
                typ: "llm".to_string(),
                parent_step: Some(TEST_STEP1_ID),
                span_id: Some("22222222".to_string()),
                info: Some(json!({ "model": "a_model" })),
                tags: vec![],
                input: json!({ "task_param2": "value" }),
            },
        })),
        ProxyLogEntry::Proxied(Box::new(ProxyLogEvent {
            id: TEST_EVENT1_ID,
            event_type: std::borrow::Cow::Borrowed("query"),
            timestamp: Utc.timestamp_opt(4, 0).unwrap(),
            request: None,
            response: None,
            latency: None,
            total_latency: None,
            was_rate_limited: Some(false),
            num_retries: Some(0),
            error: None,
            options: crate::ProxyRequestOptions {
                metadata: crate::ProxyRequestMetadata {
                    step_id: Some(TEST_STEP2_ID),
                    run_id: Some(TEST_RUN_ID),
                    extra: Some(
                        json!({
                            "some_key": "some_value",
                        })
                        .as_object()
                        .unwrap()
                        .clone(),
                    ),
                    ..Default::default()
                },
                ..Default::default()
            },
        })),
        ProxyLogEntry::Workflow(WorkflowEvent::Event(EventPayload {
            typ: "an_event".to_string(),
            data: Some(json!({
                "key": "value",
            })),
            error: Some(json!({
                "message": "something went wrong"
            })),
            step_id: TEST_STEP2_ID,
            run_id: TEST_RUN_ID,
            time: Some(Utc.timestamp_opt(5, 0).unwrap()),
            internal_metadata: None,
        })),
        ProxyLogEntry::Workflow(WorkflowEvent::StepError(StepEventData {
            step_id: TEST_STEP2_ID,
            run_id: TEST_RUN_ID,
            time: Some(Utc.timestamp_opt(5, 0).unwrap()),
            data: ErrorData {
                error: json!({"message": "an error"}),
            },
        })),
        ProxyLogEntry::Workflow(WorkflowEvent::StepEnd(StepEventData {
            step_id: TEST_STEP1_ID,
            run_id: TEST_RUN_ID,
            time: Some(Utc.timestamp_opt(5, 0).unwrap()),
            data: StepEndData {
                output: json!({ "result": "success" }),
                info: Some(json!({ "info3": "value3" })),
            },
        })),
        ProxyLogEntry::Workflow(WorkflowEvent::RunUpdate(RunUpdateEvent {
            id: TEST_RUN_ID,
            status: Some("finished".to_string()),
            output: Some(json!({ "result": "success" })),
            info: Some(json!({ "info2": "new_value", "info3": "value3"})),
            time: Some(Utc.timestamp_opt(5, 0).unwrap()),
        })),
    ]
}