11use anyhow:: Result ;
22use clap:: { Parser , Subcommand } ;
33use s3_vectors:: {
4- rag:: { RagConfig , RagPipeline , rag_query } ,
4+ rag:: { rag_query , RagConfig , RagPipeline } ,
55 S3VectorsClient ,
66} ;
77use std:: path:: PathBuf ;
@@ -12,15 +12,15 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
1212struct Cli {
1313 #[ command( subcommand) ]
1414 command : Commands ,
15-
15+
1616 /// AWS region for S3 Vectors
1717 #[ arg( short, long, default_value = "us-east-1" ) ]
1818 region : String ,
19-
19+
2020 /// S3 Vectors bucket name
2121 #[ arg( short, long, default_value = "rag-demo-default" ) ]
2222 bucket : String ,
23-
23+
2424 /// S3 Vectors index name
2525 #[ arg( short, long, default_value = "documents-default" ) ]
2626 index : String ,
@@ -30,25 +30,25 @@ struct Cli {
3030enum Commands {
3131 /// Initialize the RAG pipeline (create bucket and index)
3232 Init ,
33-
33+
3434 /// Ingest documents from a directory
3535 Ingest {
3636 /// Directory containing documents to ingest
3737 #[ arg( short, long) ]
3838 directory : PathBuf ,
3939 } ,
40-
40+
4141 /// Query the RAG system
4242 Query {
4343 /// Query text
4444 #[ arg( short, long) ]
4545 query : String ,
46-
46+
4747 /// Number of results to return
4848 #[ arg( short, long, default_value = "5" ) ]
4949 top_k : u32 ,
5050 } ,
51-
51+
5252 /// Interactive query mode
5353 Interactive ,
5454}
@@ -58,27 +58,26 @@ async fn main() -> Result<()> {
5858 // Initialize tracing
5959 tracing_subscriber:: registry ( )
6060 . with (
61- tracing_subscriber:: EnvFilter :: try_from_default_env ( )
62- . unwrap_or_else ( |_| "info" . into ( ) ) ,
61+ tracing_subscriber:: EnvFilter :: try_from_default_env ( ) . unwrap_or_else ( |_| "info" . into ( ) ) ,
6362 )
6463 . with ( tracing_subscriber:: fmt:: layer ( ) )
6564 . init ( ) ;
66-
65+
6766 let cli = Cli :: parse ( ) ;
68-
67+
6968 // Create S3 Vectors client
7069 let client = S3VectorsClient :: from_env ( ) ?;
71-
70+
7271 // Create RAG config
7372 let config = RagConfig {
7473 bucket_name : cli. bucket . clone ( ) ,
7574 index_name : cli. index . clone ( ) ,
7675 ..Default :: default ( )
7776 } ;
78-
77+
7978 // Create RAG pipeline
8079 let pipeline = RagPipeline :: new ( config, client) ;
81-
80+
8281 match cli. command {
8382 Commands :: Init => {
8483 println ! ( "🚀 Initializing RAG pipeline..." ) ;
@@ -88,56 +87,56 @@ async fn main() -> Result<()> {
8887 println ! ( " Index: {}" , cli. index) ;
8988 println ! ( " Region: {}" , cli. region) ;
9089 }
91-
90+
9291 Commands :: Ingest { directory } => {
9392 println ! ( "📄 Ingesting documents from: {}" , directory. display( ) ) ;
94-
93+
9594 if !directory. exists ( ) {
9695 eprintln ! ( "❌ Directory does not exist: {}" , directory. display( ) ) ;
9796 std:: process:: exit ( 1 ) ;
9897 }
99-
98+
10099 let start = std:: time:: Instant :: now ( ) ;
101100 pipeline. ingest_documents ( & directory) . await ?;
102101 let elapsed = start. elapsed ( ) ;
103-
102+
104103 println ! ( "✅ Document ingestion completed in {:?}" , elapsed) ;
105104 }
106-
105+
107106 Commands :: Query { query, top_k } => {
108107 println ! ( "🔍 Searching for: {}" , query) ;
109108 println ! ( ) ;
110-
109+
111110 let response = rag_query ( & pipeline, & query, top_k) . await ?;
112111 println ! ( "{}" , response) ;
113112 }
114-
113+
115114 Commands :: Interactive => {
116115 println ! ( "🤖 Interactive RAG Query Mode" ) ;
117116 println ! ( "Type 'exit' or 'quit' to stop" ) ;
118117 println ! ( ) ;
119-
118+
120119 let stdin = std:: io:: stdin ( ) ;
121120 let mut input = String :: new ( ) ;
122-
121+
123122 loop {
124123 print ! ( "> " ) ;
125124 std:: io:: Write :: flush ( & mut std:: io:: stdout ( ) ) ?;
126-
125+
127126 input. clear ( ) ;
128127 stdin. read_line ( & mut input) ?;
129-
128+
130129 let query = input. trim ( ) ;
131-
130+
132131 if query. is_empty ( ) {
133132 continue ;
134133 }
135-
134+
136135 if query == "exit" || query == "quit" {
137136 println ! ( "👋 Goodbye!" ) ;
138137 break ;
139138 }
140-
139+
141140 match rag_query ( & pipeline, query, 5 ) . await {
142141 Ok ( response) => {
143142 println ! ( ) ;
@@ -151,17 +150,17 @@ async fn main() -> Result<()> {
151150 }
152151 }
153152 }
154-
153+
155154 Ok ( ( ) )
156155}
157156
158157#[ cfg( test) ]
159158mod tests {
160159 use super :: * ;
161-
160+
162161 #[ test]
163162 fn verify_cli ( ) {
164163 use clap:: CommandFactory ;
165164 Cli :: command ( ) . debug_assert ( ) ;
166165 }
167- }
166+ }
0 commit comments