index.rs

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