llm 1.3.8

A Rust library unifying multiple LLM backends.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use axum::{extract::State, http::HeaderMap, Json};
use uuid::Uuid;

use super::chain::handle_chain_request;
use super::helpers::{bad_request, internal_error, parse_model, unauthorized, ApiResult};
use crate::api::types::{ChatRequest, ChatResponse, Choice, Message};
use crate::api::ServerState;
use crate::chat::{ChatMessage, ChatRole, MessageType};

pub async fn handle_chat(
    State(state): State<ServerState>,
    headers: HeaderMap,
    Json(req): Json<ChatRequest>,
) -> ApiResult<Json<ChatResponse>> {
    validate_auth(&state, &headers)?;
    if !req.steps.is_empty() {
        return handle_chain_request(state, req).await;
    }

    let messages = build_messages(req.messages);
    let model = req.model.ok_or_else(|| bad_request("Model is required"))?;
    let (provider_id, model_name) = parse_model(&model)?;
    let provider = state
        .llms
        .get(&provider_id)
        .ok_or_else(|| bad_request(format!("Unknown provider: {provider_id}")))?;

    let response = provider
        .chat(&messages)
        .await
        .map_err(|e| internal_error(e.to_string()))?;

    Ok(Json(build_response(
        model_name,
        response.text().unwrap_or_default(),
    )))
}

fn validate_auth(state: &ServerState, headers: &HeaderMap) -> ApiResult<()> {
    let Some(key) = &state.auth_key else {
        return Ok(());
    };

    let auth_header = headers
        .get("Authorization")
        .ok_or_else(|| unauthorized("Missing authorization"))?;
    let auth_str = auth_header
        .to_str()
        .map_err(|_| unauthorized("Invalid authorization header"))?;

    if !auth_str.starts_with("Bearer ") || &auth_str[7..] != key {
        return Err(unauthorized("Invalid API key"));
    }

    Ok(())
}

fn build_messages(messages: Option<Vec<Message>>) -> Vec<ChatMessage> {
    messages
        .unwrap_or_default()
        .into_iter()
        .map(|msg| ChatMessage {
            role: parse_role(&msg.role),
            message_type: MessageType::Text,
            content: msg.content,
        })
        .collect()
}

fn parse_role(role: &str) -> ChatRole {
    match role {
        "user" => ChatRole::User,
        "assistant" => ChatRole::Assistant,
        _ => ChatRole::User,
    }
}

fn build_response(model: String, content: String) -> ChatResponse {
    ChatResponse {
        id: format!("chatcmpl-{}", Uuid::new_v4()),
        object: "chat.completion".to_string(),
        created: std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs(),
        model,
        choices: vec![Choice {
            index: 0,
            message: Message {
                role: "assistant".to_string(),
                content,
            },
            finish_reason: "stop".to_string(),
        }],
    }
}