index.rs

  1use client::Client;
  2use futures::channel::oneshot;
  3use gpui::App;
  4use http_client::HttpClientWithUrl;
  5use isahc_http_client::IsahcHttpClient;
  6use language::language_settings::AllLanguageSettings;
  7use project::Project;
  8use semantic_index::{OpenAiEmbeddingModel, OpenAiEmbeddingProvider, SemanticDb};
  9use settings::SettingsStore;
 10use std::{
 11    path::{Path, PathBuf},
 12    sync::Arc,
 13};
 14
 15fn main() {
 16    env_logger::init();
 17
 18    use clock::FakeSystemClock;
 19
 20    App::new().run(|cx| {
 21        let store = SettingsStore::test(cx);
 22        cx.set_global(store);
 23        language::init(cx);
 24        Project::init_settings(cx);
 25        SettingsStore::update(cx, |store, cx| {
 26            store.update_user_settings::<AllLanguageSettings>(cx, |_| {});
 27        });
 28
 29        let clock = Arc::new(FakeSystemClock::default());
 30
 31        let http = Arc::new(HttpClientWithUrl::new(
 32            IsahcHttpClient::new(None, None),
 33            "http://localhost:11434",
 34            None,
 35        ));
 36        let client = client::Client::new(clock, http.clone(), cx);
 37        Client::set_global(client.clone(), cx);
 38
 39        let args: Vec<String> = std::env::args().collect();
 40        if args.len() < 2 {
 41            eprintln!("Usage: cargo run --example index -p semantic_index -- <project_path>");
 42            cx.quit();
 43            return;
 44        }
 45
 46        // let embedding_provider = semantic_index::FakeEmbeddingProvider;
 47
 48        let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
 49
 50        let embedding_provider = Arc::new(OpenAiEmbeddingProvider::new(
 51            http.clone(),
 52            OpenAiEmbeddingModel::TextEmbedding3Small,
 53            open_ai::OPEN_AI_API_URL.to_string(),
 54            api_key,
 55        ));
 56
 57        cx.spawn(|mut cx| async move {
 58            let semantic_index = SemanticDb::new(
 59                PathBuf::from("/tmp/semantic-index-db.mdb"),
 60                embedding_provider,
 61                &mut cx,
 62            );
 63
 64            let mut semantic_index = semantic_index.await.unwrap();
 65
 66            let project_path = Path::new(&args[1]);
 67
 68            let project = Project::example([project_path], &mut cx).await;
 69
 70            cx.update(|cx| {
 71                let language_registry = project.read(cx).languages().clone();
 72                let node_runtime = project.read(cx).node_runtime().unwrap().clone();
 73                languages::init(language_registry, node_runtime, cx);
 74            })
 75            .unwrap();
 76
 77            let project_index = cx
 78                .update(|cx| semantic_index.project_index(project.clone(), cx))
 79                .unwrap()
 80                .unwrap();
 81
 82            let (tx, rx) = oneshot::channel();
 83            let mut tx = Some(tx);
 84            let subscription = cx.update(|cx| {
 85                cx.subscribe(&project_index, move |_, event, _| {
 86                    if let Some(tx) = tx.take() {
 87                        _ = tx.send(*event);
 88                    }
 89                })
 90            });
 91
 92            let index_start = std::time::Instant::now();
 93            rx.await.expect("no event emitted");
 94            drop(subscription);
 95            println!("Index time: {:?}", index_start.elapsed());
 96
 97            let results = cx
 98                .update(|cx| {
 99                    let project_index = project_index.read(cx);
100                    let query = "converting an anchor to a point";
101                    project_index.search(query.into(), 4, cx)
102                })
103                .unwrap()
104                .await
105                .unwrap();
106
107            for search_result in results {
108                let path = search_result.path.clone();
109
110                let content = cx
111                    .update(|cx| {
112                        let worktree = search_result.worktree.read(cx);
113                        let entry_abs_path = worktree.abs_path().join(search_result.path.clone());
114                        let fs = project.read(cx).fs().clone();
115                        cx.spawn(|_| async move { fs.load(&entry_abs_path).await.unwrap() })
116                    })
117                    .unwrap()
118                    .await;
119
120                let range = search_result.range.clone();
121                let content = content[search_result.range].to_owned();
122
123                println!(
124                    "✄✄✄✄✄✄✄✄✄✄✄✄✄✄ {:?} @ {} ✄✄✄✄✄✄✄✄✄✄✄✄✄✄",
125                    path, search_result.score
126                );
127                println!("{:?}:{:?}:{:?}", path, range.start, range.end);
128                println!("{}", content);
129            }
130
131            cx.background_executor()
132                .timer(std::time::Duration::from_secs(100000))
133                .await;
134
135            cx.update(|cx| cx.quit()).unwrap();
136        })
137        .detach();
138    });
139}