From b14f85022b2351d31c1fe701ed2bdc9e81c1a08c Mon Sep 17 00:00:00 2001 From: Stefano Pigozzi Date: Mon, 1 Aug 2022 00:49:44 +0200 Subject: [PATCH] Keep building the skeleton of the bot --- Cargo.toml | 1 + src/data/load.rs | 2 + src/data/map.rs | 14 +++++++ src/data/mod.rs | 1 + src/main.rs | 94 +++++++++++++++++++++++++++++++++------------ src/search/mod.rs | 1 + src/search/query.rs | 49 +++++++++++++++++++++++ 7 files changed, 137 insertions(+), 25 deletions(-) create mode 100644 src/data/map.rs create mode 100644 src/search/query.rs diff --git a/Cargo.toml b/Cargo.toml index 5d8f2da..f24f68a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,3 +15,4 @@ pretty_env_logger = "0.4.0" glob = "0.3.0" itertools = "0.10.2" tantivy = "0.18.0" +reqwest = "0.11.11" \ No newline at end of file diff --git a/src/data/load.rs b/src/data/load.rs index b462fe3..e223143 100644 --- a/src/data/load.rs +++ b/src/data/load.rs @@ -1,5 +1,6 @@ //! This module contains functions to load **Set Bundles** from [Data Dragon](https://developer.riotgames.com/docs/lol). +use log::*; use itertools::Itertools; use crate::data::schema::Card; @@ -13,6 +14,7 @@ enum LoadingError { /// Load a single Set Bundle and create a [Vec] with the cards contained in it. fn load_setbundle(path: std::path::PathBuf) -> Result, LoadingError> { + debug!("Loading Set Bundle {:?}", &path); let file = std::fs::File::open(path) .map_err(LoadingError::IO)?; let data = serde_json::de::from_reader::>(file) diff --git a/src/data/map.rs b/src/data/map.rs new file mode 100644 index 0000000..7101e8c --- /dev/null +++ b/src/data/map.rs @@ -0,0 +1,14 @@ +use std::collections::HashMap; +use crate::data::schema::Card; + + +/// Build a [HashMap] mapping card codes strings to [Card] structs. +pub fn build_card_code_hashmap(cards: &Vec) -> HashMap { + let mut map = HashMap::::new(); + + for card in cards { + map.insert(card.card_code.to_owned(), card.to_owned()); + } + + map +} diff --git a/src/data/mod.rs b/src/data/mod.rs index 2de4b23..605c8fc 100644 --- a/src/data/mod.rs +++ b/src/data/mod.rs @@ -1,2 +1,3 @@ pub mod schema; pub mod load; +pub mod map; diff --git a/src/main.rs b/src/main.rs index 21b90ce..0a9ff8a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,52 +1,96 @@ #[macro_use] extern crate tantivy; -#[macro_use] extern crate log; use teloxide::prelude::*; use teloxide::types::*; use log::*; -use tantivy::{IndexReader, ReloadPolicy}; +use reqwest::Url; +use itertools::Itertools; mod data; mod search; +/// Run the bot. #[tokio::main] async fn main() { pretty_env_logger::init(); debug!("Loading Set Bundles..."); - let cards = data::load::load_setbundles_infallible("./data/*/en_us/data/*-en_us.json"); + let cards = data::load::load_setbundles_infallible("./data/*/en_us/data/set*-en_us.json"); + debug!("Loaded {} cards!", &cards.len()); - debug!("Creating Tantivy index..."); + debug!("Creating Card index..."); let card_index = search::index::build_card_index(); + debug!("Card index created successfully!"); - debug!("Writing cards to the index..."); + debug!("Creating Card hashmap..."); + let card_map = data::map::build_card_code_hashmap(&cards); + debug!("Card hashmap contains {} cards!", &card_map.len()); + + debug!("Writing Cards to the Card index..."); search::index::write_cards_to_index(&card_index, &cards); + debug!("Wrote Cards to the Card index!"); - debug!("Creating Tantivy reader..."); + debug!("Creating Card query parser..."); + let card_query_parser = search::query::build_query_parser(&card_index); + debug!("Created Card query parser!"); + + debug!("Creating Card reader..."); let card_reader = search::index::build_reader(&card_index); + debug!("Created Card reader!"); - let inline_query_closure = |query: InlineQuery, bot: AutoSend| async move { - bot.answer_inline_query(&query.id, handle_query(&query, &card_reader)).await; - respond(()) - }; - - info!("patched-porobot is starting..."); - + debug!("Creating Telegram Bot with parameters from the environment..."); let bot = Bot::from_env().auto_send(); + let bot_me = bot.get_me().await.expect("Telegram Bot parameters to be valid"); + debug!("Created Telegram Bot @{}!", &bot_me.username.as_ref().unwrap()); - let handler = Update::filter_inline_query().branch(dptree::endpoint(inline_query_closure)); + let handler = Update::filter_inline_query().branch(dptree::endpoint(move |query: InlineQuery, bot: AutoSend| { + let card_schema = card_index.schema(); + let card_query_parser = &card_query_parser; + let card_reader = &card_reader; + let card_map = &card_map; + + let results = search::query::search_card( + &card_schema, + &card_query_parser, + &card_reader, + &card_map, + &query.query + ); + + async move { + let reply = results.iter() + /* + .map(|card| InlineQueryResult::Photo(InlineQueryResultPhoto { + id: card.card_code.to_owned(), + title: Some(card.name.to_owned()), + description: Some(card.description_raw.to_owned()), + caption: Some(card.description_raw.to_owned()), + photo_url: Url::parse( & card.assets.get(0).expect("card to have assets").game_absolute_path).expect("card to have a valid asset URL"), + thumb_url: Url::parse( & card.assets.get(0).expect("card to have assets").full_absolute_path).expect("card to have a valid asset URL"), + photo_width: None, + photo_height: None, + parse_mode: None, + caption_entities: None, + reply_markup: None, + input_message_content: None, + })) + */ + .map(|card| InlineQueryResult::Article(InlineQueryResultArticle::new( + card.card_code.to_owned(), + card.name.to_owned(), + InputMessageContent::Text(InputMessageContentText::new( + format!("{}\n\n{}", &card.name, &card.description_raw) + )) + ))) + .collect_vec(); + + if let Err(e) = bot.answer_inline_query(&query.id, reply).send().await { + error!("{:?}", e); + }; + respond(()) + } + })); Dispatcher::builder(bot, handler).enable_ctrlc_handler().build().dispatch().await; } - - -/// Handle a [InlineQuery] incoming from Telegram. -fn handle_query(query: &InlineQuery, reader: &IndexReader) -> Vec { - debug!("Creating Tantivy searcher..."); - let card_searcher = reader.searcher(); - - let result = InlineQueryResult::Article(InlineQueryResultArticle::new("test", "Test", InputMessageContent::Text(InputMessageContentText::new("Qui รจ dove metterei la mia carta, se solo ne avessi una!")))); - - vec![result] -} \ No newline at end of file diff --git a/src/search/mod.rs b/src/search/mod.rs index ec6b8d2..6ad594f 100644 --- a/src/search/mod.rs +++ b/src/search/mod.rs @@ -1,2 +1,3 @@ pub mod schema; pub mod index; +pub mod query; diff --git a/src/search/query.rs b/src/search/query.rs new file mode 100644 index 0000000..3e72a86 --- /dev/null +++ b/src/search/query.rs @@ -0,0 +1,49 @@ +use log::*; +use std::collections::HashMap; +use tantivy::Index; +use tantivy::IndexReader; +use tantivy::LeasedItem; +use tantivy::Searcher; +use tantivy::collector::TopDocs; +use tantivy::query::QueryParser; +use tantivy::schema::Schema; +use crate::data::schema::Card; +use itertools::Itertools; + + +pub fn build_query_parser(index: &Index) -> QueryParser { + let schema = index.schema(); + + let name = schema.get_field("name").unwrap(); + let description = schema.get_field("description").unwrap(); + let code = schema.get_field("code").unwrap(); + + QueryParser::for_index(&index, vec![name, description, code]) +} + + +pub fn build_searcher(reader: &IndexReader) -> LeasedItem { + reader.searcher() +} + + +pub fn search_card(schema: &Schema, parser: &QueryParser, reader: &IndexReader, map: &HashMap, q: &str) -> Vec { + debug!("Searching for `{}`...", &q); + + let code = schema.get_field("code").unwrap(); + + debug!("Building Card searcher..."); + let searcher = build_searcher(reader); + let query = parser.parse_query(q) + .expect("to be able to parse the query"); + let search = searcher.search(&*query, &TopDocs::with_limit(50)) + .expect("to be able to search for a card"); + debug!("Retrieved {} results!", &search.len()); + + search.iter().filter_map(|(_score, address)| searcher.doc(address.to_owned()).ok()) + .filter_map(|doc| doc.get_first(code).cloned()) + .filter_map(|field| field.as_text().map(String::from)) + .filter_map(|code| map.get(&*code)) + .cloned() + .collect_vec() +} \ No newline at end of file