index.rs

  1use client::Client;
  2use futures::channel::oneshot;
  3use gpui::Application;
  4use http_client::HttpClientWithUrl;
  5use language::language_settings::AllLanguageSettings;
  6use project::Project;
  7use semantic_index::{OpenAiEmbeddingModel, OpenAiEmbeddingProvider, SemanticDb};
  8use settings::SettingsStore;
  9use std::{
 10    path::{Path, PathBuf},
 11    sync::Arc,
 12};
 13
 14fn main() {
 15    zlog::init();
 16
 17    use clock::FakeSystemClock;
 18
 19    Application::new().run(|cx| {
 20        let store = SettingsStore::test(cx);
 21        cx.set_global(store);
 22        language::init(cx);
 23        Project::init_settings(cx);
 24        SettingsStore::update(cx, |store, cx| {
 25            store.update_user_settings::<AllLanguageSettings>(cx, |_| {});
 26        });
 27
 28        let clock = Arc::new(FakeSystemClock::new());
 29
 30        let http = Arc::new(HttpClientWithUrl::new(
 31            Arc::new(
 32                reqwest_client::ReqwestClient::user_agent("Zed semantic index example").unwrap(),
 33            ),
 34            "http://localhost:11434",
 35            None,
 36        ));
 37        let client = client::Client::new(clock, http.clone(), cx);
 38        Client::set_global(client, cx);
 39
 40        let args: Vec<String> = std::env::args().collect();
 41        if args.len() < 2 {
 42            eprintln!("Usage: cargo run --example index -p semantic_index -- <project_path>");
 43            cx.quit();
 44            return;
 45        }
 46
 47        // let embedding_provider = semantic_index::FakeEmbeddingProvider;
 48
 49        let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
 50
 51        let embedding_provider = Arc::new(OpenAiEmbeddingProvider::new(
 52            http,
 53            OpenAiEmbeddingModel::TextEmbedding3Small,
 54            open_ai::OPEN_AI_API_URL.to_string(),
 55            api_key,
 56        ));
 57
 58        cx.spawn(async move |cx| {
 59            let semantic_index = SemanticDb::new(
 60                PathBuf::from("/tmp/semantic-index-db.mdb"),
 61                embedding_provider,
 62                cx,
 63            );
 64
 65            let mut semantic_index = semantic_index.await.unwrap();
 66
 67            let project_path = Path::new(&args[1]);
 68
 69            let project = Project::example([project_path], cx).await;
 70
 71            cx.update(|cx| {
 72                let language_registry = project.read(cx).languages().clone();
 73                let node_runtime = project.read(cx).node_runtime().unwrap().clone();
 74                languages::init(language_registry, node_runtime, cx);
 75            })
 76            .unwrap();
 77
 78            let project_index = cx
 79                .update(|cx| semantic_index.project_index(project.clone(), cx))
 80                .unwrap()
 81                .unwrap();
 82
 83            let (tx, rx) = oneshot::channel();
 84            let mut tx = Some(tx);
 85            let subscription = cx.update(|cx| {
 86                cx.subscribe(&project_index, move |_, event, _| {
 87                    if let Some(tx) = tx.take() {
 88                        _ = tx.send(*event);
 89                    }
 90                })
 91            });
 92
 93            let index_start = std::time::Instant::now();
 94            rx.await.expect("no event emitted");
 95            drop(subscription);
 96            println!("Index time: {:?}", index_start.elapsed());
 97
 98            let results = cx
 99                .update(|cx| {
100                    let project_index = project_index.read(cx);
101                    let query = "converting an anchor to a point";
102                    project_index.search(vec![query.into()], 4, cx)
103                })
104                .unwrap()
105                .await
106                .unwrap();
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}