modelsocket 0.4.0

A Rust library for ModelSocket, a protocol for efficiently integrating with LLMs
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

use anyhow::{Context, Result};
use clap::Parser;
use modelsocket::client::{AppendOpts, GenOpts, ModelSocket};

mod utils;

#[derive(Parser, Debug)]
struct Args {
    #[clap(short, long, default_value = "wss://models.mixlayer.ai/ws")]
    url: String,
    #[clap(short, long, env = "MODELSOCKET_API_KEY")]
    api_key: String,
    #[clap(long, long, default_value_t = false)]
    hidden: bool,
    #[clap(long, long, default_value_t = true)]
    colors: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    let args = Args::parse();

    let socket = ModelSocket::connect(&args.url, Some(&args.api_key))
        .await
        .context("websocket connection failed")?;

    let seq = socket
        .open("meta/llama3.1-8b-instruct-free", None)
        .await
        .context("seq open failed")?;

    seq.append("Hello there.", AppendOpts::user()).await?;

    let stream = seq
        .generate(Some(GenOpts::assistant()))
        .await
        .context("generate failed")?;

    utils::print_stream(stream, args.hidden, args.colors).await?;

    Ok(())
}