Skip to content

Commit

Permalink
Demo (#9)
Browse files Browse the repository at this point in the history
  • Loading branch information
lukaszKielar authored Oct 22, 2024
1 parent 2c7dc70 commit 7be7ec2
Show file tree
Hide file tree
Showing 10 changed files with 57 additions and 104 deletions.
42 changes: 12 additions & 30 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,10 @@

LokAI is a local AI assistant in your terminal.

## Demo

![Demo](demo/lokai-demo.gif)

## Running app

Before we get started make sure you have following tools installed:
Expand All @@ -14,16 +18,21 @@ To run LokAI type:
cargo run # for more configuration flags see CLI section
```

### Persistence

Downloaded LLM models, logs, database and conversations are saved into `~/.lokai` directory.

### CLI

LokAI allow you to set some options through CLI.

- `database-url` [default: `sqlite::memory:`] - default value spins new in-memory instance that won't persist conversations between restarts. Example value for persistent database `sqlite://db.sqlite3`
- `database-url` - defines location of SQLite database. Example values: "sqlite::memory:" (in-memory), "sqlite://db.slite3" (persistent), "db.sqlite3" (persitent)
- `enable-transcription` - transcribes voice into prompt

To use one, many or all options type:

```bash
cargo run -- --database-url <DB_URL>
cargo run -- --database-url sqlite::memory: --enable-transcription
```

To print help type:
Expand All @@ -40,34 +49,7 @@ cargo run -- --help
| <kbd>Ctrl</kbd> + <kbd>n</kbd> | Add new conversation | Global |
| <kbd>Tab</kbd> | Next focus | Global |
| <kbd>Shift</kbd> + <kbd>Tab</kbd> | Previous focus | Global |
| <kbd>↑</kbd>/<kbd>↓</kbd> | Switch between conversations | Conversation sidebar |
| <kbd>Delete</kbd> | Delete selected conversation | Conversation sidebar |
| <kbd>↑</kbd>/<kbd>↓</kbd> | Switch between conversations | Conversation sidebar |
| <kbd>↑</kbd>/<kbd>↓</kbd> | Scroll up/down | Chat/Prompt |
| <kbd>Esc</kbd> | Cancel action | Popups |

## Roadmap

- [ ] ? Settings persistance - save TOML file in user's dir
- [ ] Better error handling - new Result and Error structs allowing for clear distinction between critical and non-critical errors
- [ ] If nothing is presented in Chat area print shortcuts and welcoming graphics (logo)
- [ ] Create logo
- [ ] Conversations
- [x] Adding new conversation - design dedicated pop up
- [x] Deleting conversation
- [ ] Add `session_path` column to `conversations` table - it should store local path to chat session `LOKAI_DIR/chats/{chat_id}`
- [ ] Chat
- [ ] Highlighting code snippets returned by LLM
- [ ] Ability to copy chat or selected messages to clipboard
- [ ] Save/load Kalosm Chat history to/from disk
- [ ] Create simple cache (or reuse some tool) to store Chat sessions to avoid constant reading/writing from/to disk
- [ ] Prompt
- [ ] Set prompt's border to different colors depending on the factors like: empty prompt, LLM still replying, error
- [ ] Improve prompt transcription process. Currently there is no way to turn off microphone, and the app constantly listens until its killed. I need to toggle it on/off on demand.
- [ ] Popup or presenting shortcuts
- [ ] Implement `AppState` for sharing things like DB pool, Whisper, Llama, app config, lokai dir (app config is actually dependent on lokai dir)
- [ ] Bar that presents sliding messages (iterator for a piece of text that moves from right to left)
- [ ] Tracing
- [ ] Tests
- [ ] Improve unit test coverage
- [ ] Create integration tests
- [ ] Documentation improvements
Binary file added demo/lokai-demo.gif
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
10 changes: 4 additions & 6 deletions src/app.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,6 @@ impl AppFocus {
}
}

// TODO: create shared AppState(SqlitePool)

pub struct App {
pub chat: Chat,
pub conversations: Conversations,
Expand All @@ -60,24 +58,24 @@ pub struct App {
inference_tx: Sender<Message>,
running: bool,
sqlite: SqlitePool,
_assistant: Assistant,
}

impl App {
pub fn new(sqlite: SqlitePool, event_tx: UnboundedSender<Event>, llama: Llama) -> Self {
let (inference_tx, inference_rx) = mpsc::channel::<Message>(10);
Assistant::run(llama, sqlite.clone(), inference_rx, event_tx.clone());

Self {
chat: Chat::new(sqlite.clone()),
conversations: Conversations::new(sqlite.clone()),
prompt: Default::default(),
new_conversation_popup: Default::default(),
delete_conversation_popup: Default::default(),
focus: Default::default(),
event_tx: event_tx.clone(),
event_tx,
inference_tx,
running: true,
sqlite: sqlite.clone(),
_assistant: Assistant::new(llama, sqlite, inference_rx, event_tx),
sqlite,
}
}

Expand Down
28 changes: 10 additions & 18 deletions src/assistant.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,7 @@
use futures::StreamExt;
use kalosm::language::{Chat, Llama};
use sqlx::SqlitePool;
use tokio::{
sync::mpsc::{Receiver, UnboundedSender},
task::JoinHandle,
};
use tokio::sync::mpsc::{Receiver, UnboundedSender};

use crate::{
db,
Expand All @@ -13,24 +10,18 @@ use crate::{
AppResult,
};

pub struct Assistant {
_join_handle: JoinHandle<()>,
}
pub struct Assistant;

impl Assistant {
pub fn new(
pub fn run(
llama: Llama,
sqlite: SqlitePool,
inference_rx: Receiver<Message>,
event_tx: UnboundedSender<Event>,
) -> Self {
let join_handle = tokio::spawn(async move {
) {
tokio::spawn(async move {
inference_stream(sqlite, inference_rx, event_tx, llama).await;
});

Self {
_join_handle: join_handle,
}
}
}

Expand All @@ -55,24 +46,25 @@ async fn inference_stream(
while let Some(chunk) = text_stream.next().await {
assistant_response.content.push_str(&chunk);

event_tx.send(Event::Inference(
// ignore send errors, I can at least wait until the end of assistant's response and save it to db
// if the channel is closed we probably paniced anyway
let _ = event_tx.send(Event::Inference(
assistant_response.clone(),
InferenceType::Streaming,
))?;
));
}

let assistant_response =
db::update_message(&sqlite, &assistant_response.content, assistant_response.id).await?;
chat.add_message(assistant_response.content);

// TODO: handle errors
tokio::spawn(async move {
match chat.save_session(&conversation.session_path).await {
Ok(_) => tracing::info!("session saved to disk"),
Err(err) => tracing::error!("Error while saving session: {}", err),
}
})
.await;
.await?;
}

Ok(())
Expand Down
1 change: 0 additions & 1 deletion src/chat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ use crate::{db, models::Message, AppResult};
const BORDER_SIZE: usize = 1;

// TODO: automatically scroll to the bottom when messages are loaded
// TODO: this has to be wrapper for kalosm::Chat
pub struct Chat {
messages: Vec<Message>,
pub vertical_scrollbar_state: ScrollbarState,
Expand Down
1 change: 0 additions & 1 deletion src/conversations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,6 @@ impl Conversations {
}
}

// TODO: add internal attribute that will define error style and message
pub struct NewConversationPopup {
text: Option<String>,
text_area: TextArea<'static>,
Expand Down
26 changes: 11 additions & 15 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use sqlx::{migrate::MigrateDatabase, sqlite::SqlitePoolOptions, Executor, Sqlite
use tokio::sync::{mpsc, RwLock};
use tracing::{info, Level};
use tracing_subscriber::EnvFilter;
use transcribe::Transcriber;
use transcribe::transcribe;

use crate::{app::App, event::EventHandler, tui::Tui};

Expand Down Expand Up @@ -83,20 +83,16 @@ async fn main() -> AppResult<()> {
Cache::new(kalosm_cache_dir)
};

let _transcriber = {
if cli_args.enable_transcription {
let whisper = Whisper::builder()
.with_cache(kalosm_cache.clone())
.with_source(WhisperSource::BaseEn)
.with_language(Some(WhisperLanguage::English))
.build()
.await?;

Some(Transcriber::new(event_tx.clone(), whisper))
} else {
None
}
};
if cli_args.enable_transcription {
let whisper = Whisper::builder()
.with_cache(kalosm_cache.clone())
.with_source(WhisperSource::BaseEn)
.with_language(Some(WhisperLanguage::English))
.build()
.await?;

transcribe(event_tx.clone(), whisper)
}

let llama = Llama::builder()
.with_source(LlamaSource::llama_3_1_8b_chat().with_cache(kalosm_cache))
Expand Down
3 changes: 0 additions & 3 deletions src/models.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,6 @@ pub struct Conversation {
pub created_at: DateTime<Utc>,
}

// TODO: implement load and save methods
impl Conversation {}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub enum Role {
#[serde(rename = "assistant")]
Expand Down
49 changes: 20 additions & 29 deletions src/transcribe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,42 +5,33 @@ use kalosm::sound::TextStream;
use kalosm_sound::{
DenoisedExt, MicInput, TranscribeChunkedAudioStreamExt, VoiceActivityStreamExt, Whisper,
};
use tokio::{sync::mpsc::UnboundedSender, task::JoinHandle};
use tokio::sync::mpsc::UnboundedSender;
use tracing::info;

use crate::event::Event;

pub struct Transcriber {
_join_handle: JoinHandle<()>,
}

impl Transcriber {
pub fn new(event_tx: UnboundedSender<Event>, whisper: Whisper) -> Self {
let mic_input = MicInput::default();
let mic_stream = mic_input.stream().expect("Cannot create MicStream");
pub fn transcribe(event_tx: UnboundedSender<Event>, whisper: Whisper) {
let mic_input = MicInput::default();
let mic_stream = mic_input.stream().expect("Cannot create MicStream");

let join_handle = tokio::spawn(async move {
let voice_stream = mic_stream
.denoise_and_detect_voice_activity()
.rechunk_voice_activity();
let text_stream = voice_stream.transcribe(whisper);
let mut sentences = text_stream.sentences();
tokio::spawn(async move {
let voice_stream = mic_stream
.denoise_and_detect_voice_activity()
.rechunk_voice_activity();
let text_stream = voice_stream.transcribe(whisper);
let mut sentences = text_stream.sentences();

// TODO: use tokio::select!
loop {
match sentences.next().await {
Some(sentence) => {
info!("sentence: {:?}", sentence);
// TODO: handle error
event_tx.send(Event::PromptTranscription(sentence.to_string()));
}
None => tokio::time::sleep(Duration::from_secs(2)).await,
// TODO: use tokio::select!
loop {
match sentences.next().await {
Some(sentence) => {
info!("sentence: {:?}", sentence);
// ignore send errors, I can at least wait until the end of assistant's response and save it to db
// if the channel is closed we probably paniced anyway
let _ = event_tx.send(Event::PromptTranscription(sentence.to_string()));
}
None => tokio::time::sleep(Duration::from_secs(2)).await,
}
});

Self {
_join_handle: join_handle,
}
}
});
}
1 change: 0 additions & 1 deletion src/ui.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,6 @@ pub fn render(app: &mut App, frame: &mut Frame) {
);
frame.render_widget(&*app.prompt, messages_layout[1]);

// TODO: dimm other components when popup is active
if app.new_conversation_popup.is_activated() {
let (popup_width, popup_height) = (50, 3);
let (popup_x, popup_y) =
Expand Down

0 comments on commit 7be7ec2

Please sign in to comment.