From c5dabddb56f465d75a355357f2232039a3740447 Mon Sep 17 00:00:00 2001 From: Louis Beaumont Date: Fri, 24 Nov 2023 11:11:57 -0800 Subject: [PATCH] feat: docker compose shit + kr file storage v0 --- Makefile | 26 ++++++++ assistants-core/Cargo.toml | 5 +- assistants-core/src/assistant.rs | 65 +++++++----------- assistants-core/src/knowledge-retrieval.rs | 76 ++++++++++++++++++++++ assistants-core/src/migrations.sql | 3 + docker/docker-compose.yml | 30 +++++++++ 6 files changed, 161 insertions(+), 44 deletions(-) create mode 100644 Makefile create mode 100644 assistants-core/src/knowledge-retrieval.rs create mode 100644 docker/docker-compose.yml diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..7bc5445 --- /dev/null +++ b/Makefile @@ -0,0 +1,26 @@ +.PHONY: help docker clean + +.DEFAULT_GOAL := help + +## Colors +YELLOW := $(shell tput -Txterm setaf 3) +RESET := $(shell tput -Txterm sgr0) + +## Help documentation +help: + @echo "${YELLOW}Available commands:${RESET}" + @echo + @awk 'BEGIN {FS = ":.*##"; printf "\nUsage:\n make \033[36m\033[0m\n"} /^[a-zA-Z_-]+:.*?##/ { printf " \033[36m%-15s\033[0m %s\n", $$1, $$2 } /^##@/ { printf "\n\033[1m%s\033[0m\n", substr($$0, 5) } ' $(MAKEFILE_LIST) + +## Docker compose up +docker: ## Run docker compose up + docker-compose -f docker/docker-compose.yml up -d + while ! docker exec -it pg pg_isready -U postgres; do sleep 1; done + docker exec -it pg psql -U postgres -c "CREATE DATABASE mydatabase;" > /dev/null 2>&1 || echo "Database already exists" + docker exec -i pg psql -U postgres -d mydatabase < assistants-core/src/migrations.sql > /dev/null 2>&1 || echo "Migrations already applied" + +## Stop and remove containers +clean: ## Stop and remove docker-postgres-1, docker-redis-1, docker-minio-1 containers + docker stop $$(docker ps -a -q --filter name=pg --filter name=redis --filter name=minio1) > /dev/null 2>&1 || echo "No containers to stop" + docker rm $$(docker ps -a -q --filter name=pg --filter name=redis --filter name=minio1) > /dev/null 2>&1 || true + docker volume rm $$(docker volume ls -q --filter name=pg --filter name=redis --filter name=minio1) > /dev/null 2>&1 || true diff --git a/assistants-core/Cargo.toml b/assistants-core/Cargo.toml index 5efd97e..3a87dfe 100644 --- a/assistants-core/Cargo.toml +++ b/assistants-core/Cargo.toml @@ -20,12 +20,13 @@ serde_json = "1.0" serde = { version = "1.0", features = ["derive"] } redis = { version = "0.23.3", features = ["tokio-comp"] } assistants-extra = { path = "../assistants-extra" } - +rusoto_core = "0.46.0" +rusoto_s3 = "0.46.0" [build-dependencies] syn = "1" [dev-dependencies] dotenv = "0.15" - +tempfile = "3.2.0" diff --git a/assistants-core/src/assistant.rs b/assistants-core/src/assistant.rs index cbedf2d..754e96e 100644 --- a/assistants-core/src/assistant.rs +++ b/assistants-core/src/assistant.rs @@ -1,16 +1,28 @@ -// data storage -// init -// docker run --name pg -e POSTGRES_PASSWORD=secret -d -p 5432:5432 postgres -// docker exec -it pg psql -U postgres -c "CREATE DATABASE mydatabase;" +/* +data storage +init +docker run --name pg -e POSTGRES_PASSWORD=secret -d -p 5432:5432 postgres +docker exec -it pg psql -U postgres -c "CREATE DATABASE mydatabase;" -// migrations -// docker exec -i pg psql -U postgres -d mydatabase < assistants-core/src/migrations.sql +migrations +docker exec -i pg psql -U postgres -d mydatabase < assistants-core/src/migrations.sql -// checks -// docker exec -it pg psql -U postgres -d mydatabase -c "\dt" +checks +docker exec -it pg psql -U postgres -d mydatabase -c "\dt" -// queue -// docker run --name redis -d -p 6379:6379 redis +queue +docker run --name redis -d -p 6379:6379 redis + +MINIO + +docker run -d -p 9000:9000 -p 9001:9001 \ +--name minio1 \ +-e "MINIO_ROOT_USER=minioadmin" \ +-e "MINIO_ROOT_PASSWORD=minioadmin" \ +minio/minio server /data --console-address ":9001" + +check docker/docker-compose.yml +*/ use sqlx::PgPool; use serde_json; @@ -58,13 +70,7 @@ pub struct Run { pub instructions: String, pub status: String, } -fn from_sql_value<'de, D>(deserializer: D) -> Result, D::Error> -where - D: Deserializer<'de>, -{ - let value: serde_json::Value = Deserialize::deserialize(deserializer)?; - serde_json::from_value(value).map_err(serde::de::Error::custom) -} + pub struct Assistant { pub instructions: String, @@ -135,31 +141,6 @@ pub async fn list_messages(pool: &PgPool, thread_id: &str) -> Result Result, sqlx::Error> { // Change run_id type to i32 and return type to Option -// let row = sqlx::query!( -// r#" -// SELECT status FROM runs WHERE id = $1 -// "#, -// &run_id -// ) -// .fetch_one(pool) -// .await?; -// Ok(row.status) -// } - -// // Change the return type to Vec in display_assistant_response function -// pub async fn display_assistant_response(pool: &PgPool, thread_id: &str, user_id: &str) -> Result, sqlx::Error> { -// let rows = sqlx::query!( -// r#" -// SELECT content FROM messages WHERE thread_id = $1 AND user_id = $2 AND role = 'assistant' -// "#, -// &thread_id, &user_id -// ) -// .fetch_all(pool) -// .await?; -// Ok(rows) -// } - pub async fn create_assistant(pool: &PgPool, assistant: &Assistant) -> Result<(), sqlx::Error> { // Convert Vec<&str> to Vec diff --git a/assistants-core/src/knowledge-retrieval.rs b/assistants-core/src/knowledge-retrieval.rs new file mode 100644 index 0000000..e8ba35d --- /dev/null +++ b/assistants-core/src/knowledge-retrieval.rs @@ -0,0 +1,76 @@ + + + +use rusoto_core::Region; +use rusoto_s3::{S3Client, S3}; + +pub struct KnowledgeRetrieval { + s3_client: S3Client, +} + +impl KnowledgeRetrieval { + pub fn new() -> Self { + let s3_client = S3Client::new(Region::Custom { + name: "us-east-1".to_owned(), + endpoint: "http://localhost:9000".to_owned(), + }); + + Self { s3_client } + } + + pub async fn upload_file(&self, file_path: &Path) -> Result> { + let mut file = File::open(file_path)?; + let mut contents = Vec::new(); + file.read_to_end(&mut contents)?; + + let req = PutObjectRequest { + bucket: "my-bucket".to_owned(), + key: file_path.file_name().unwrap().to_str().unwrap().to_owned(), + body: Some(contents.into()), + ..Default::default() + }; + + self.s3_client.put_object(req).await?; + + Ok(file_path.file_name().unwrap().to_str().unwrap().to_owned()) + } +} + + +#[cfg(test)] +mod tests { + use super::*; + use std::fs::File; + use std::io::Write; + use std::path::PathBuf; + use tempfile::tempdir; + + #[tokio::test] + async fn test_upload_file() { + // Create a temporary directory. + let dir = tempdir().unwrap(); + + // Create a file path in the temporary directory. + let file_path = dir.path().join("test.txt"); + + // Write some data to the file. + let mut file = File::create(&file_path).unwrap(); + writeln!(file, "Hello, world!").unwrap(); + + // Create a new KnowledgeRetrieval instance. + let kr = KnowledgeRetrieval::new(); + + // Upload the file. + let result = kr.upload_file(&file_path).await; + + // Check that the upload was successful. + assert!(result.is_ok()); + + // Check that the returned key is correct. + assert_eq!(result.unwrap(), "test.txt"); + + // Clean up the temporary directory. + dir.close().unwrap(); + } +} + diff --git a/assistants-core/src/migrations.sql b/assistants-core/src/migrations.sql index 2f83830..97e395b 100644 --- a/assistants-core/src/migrations.sql +++ b/assistants-core/src/migrations.sql @@ -1,3 +1,6 @@ + +\c mydatabase; + -- Drop existing tables DROP TABLE IF EXISTS assistants; DROP TABLE IF EXISTS threads; diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml new file mode 100644 index 0000000..94f795d --- /dev/null +++ b/docker/docker-compose.yml @@ -0,0 +1,30 @@ +version: '3.8' +services: + postgres: + container_name: pg + image: postgres + restart: always + environment: + POSTGRES_PASSWORD: secret + POSTGRES_DB: mydatabase + ports: + - 5432:5432 + command: postgres + redis: + container_name: redis + image: redis + restart: always + ports: + - 6379:6379 + + minio: + container_name: minio1 + image: minio/minio + restart: always + ports: + - 9000:9000 + - 9001:9001 + environment: + MINIO_ROOT_USER: minioadmin + MINIO_ROOT_PASSWORD: minioadmin + command: server /data --console-address ":9001" \ No newline at end of file