Skip to content

Commit 6aeecfe

Browse files
authored
feat: add quick AI access to search mode (#556)
* feat: add quick AI access to search mode * feat: add aI assistant quick access * refactor: adjusting lodash-es import location to optimize code structure * docs: update changelog * fix: fix the logic of assigning serverId in AskAi component * refactor: optimized layout * refactor: optimized some issues
1 parent 334e29d commit 6aeecfe

34 files changed

Lines changed: 997 additions & 188 deletions

File tree

docs/content.en/docs/release-notes/_index.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ Information about release notes of Coco Server is provided here.
2626
- feat: add option to allow self-signed certificates #509
2727
- feat: add AI summary component #518
2828
- feat: dynamic log level via env var COCO_LOG #535
29+
- feat: add quick AI access to search mode #556
2930

3031
### 🐛 Bug fix
3132

src-tauri/src/assistant/mod.rs

Lines changed: 57 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ use crate::common;
22
use crate::common::assistant::ChatRequestMessage;
33
use crate::common::http::GetResponse;
44
use crate::server::http_client::HttpClient;
5+
use futures_util::TryStreamExt;
6+
use http::Method;
57
use serde_json::Value;
68
use std::collections::HashMap;
7-
use tauri::{AppHandle, Runtime};
9+
use tauri::{AppHandle, Emitter, Runtime};
10+
use tokio::io::AsyncBufReadExt;
811

912
#[tauri::command]
1013
pub async fn chat_history<R: Runtime>(
@@ -143,8 +146,8 @@ pub async fn new_chat<R: Runtime>(
143146

144147
log::debug!("New chat response: {}", &body_text);
145148

146-
let chat_response: GetResponse =
147-
serde_json::from_str(&body_text).map_err(|e| format!("Failed to parse response JSON: {}", e))?;
149+
let chat_response: GetResponse = serde_json::from_str(&body_text)
150+
.map_err(|e| format!("Failed to parse response JSON: {}", e))?;
148151

149152
if chat_response.result != "created" {
150153
return Err(format!("Unexpected result: {}", chat_response.result));
@@ -178,8 +181,8 @@ pub async fn send_message<R: Runtime>(
178181
query_params,
179182
Some(body),
180183
)
181-
.await
182-
.map_err(|e| format!("Error cancel session: {}", e))?;
184+
.await
185+
.map_err(|e| format!("Error cancel session: {}", e))?;
183186

184187
common::http::get_response_body_text(response).await
185188
}
@@ -221,8 +224,8 @@ pub async fn update_session_chat(
221224
None,
222225
Some(reqwest::Body::from(serde_json::to_string(&body).unwrap())),
223226
)
224-
.await
225-
.map_err(|e| format!("Error updating session: {}", e))?;
227+
.await
228+
.map_err(|e| format!("Error updating session: {}", e))?;
226229

227230
Ok(response.status().is_success())
228231
}
@@ -250,11 +253,56 @@ pub async fn assistant_search<R: Runtime>(
250253
None,
251254
Some(reqwest::Body::from(body.to_string())),
252255
)
253-
.await
254-
.map_err(|e| format!("Error searching assistants: {}", e))?;
256+
.await
257+
.map_err(|e| format!("Error searching assistants: {}", e))?;
255258

256259
response
257260
.json::<Value>()
258261
.await
259262
.map_err(|err| err.to_string())
260263
}
264+
265+
#[tauri::command]
266+
pub async fn ask_ai<R: Runtime>(
267+
app_handle: AppHandle<R>,
268+
message: String,
269+
server_id: String,
270+
assistant_id: String,
271+
client_id: String,
272+
) -> Result<(), String> {
273+
let body = serde_json::json!({ "message": message });
274+
275+
let path = format!("/assistant/{}/_ask", assistant_id);
276+
277+
println!("Sending request to {}", &path);
278+
279+
let response = HttpClient::send_request(
280+
server_id.as_str(),
281+
Method::POST,
282+
path.as_str(),
283+
None,
284+
None,
285+
Some(reqwest::Body::from(body.to_string())),
286+
)
287+
.await?;
288+
289+
if !response.status().is_success() {
290+
return Err(format!("Request Failed: {}", response.status()));
291+
}
292+
293+
let stream = response.bytes_stream();
294+
let reader = tokio_util::io::StreamReader::new(
295+
stream.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)),
296+
);
297+
let mut lines = tokio::io::BufReader::new(reader).lines();
298+
299+
while let Ok(Some(line)) = lines.next_line().await {
300+
dbg!("Received line: {}", &line);
301+
302+
let _ = app_handle.emit(&client_id, line).map_err(|err| {
303+
println!("Failed to emit: {:?}", err);
304+
});
305+
}
306+
307+
Ok(())
308+
}

src-tauri/src/common/assistant.rs

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,18 +9,25 @@ pub struct ChatRequestMessage {
99
#[allow(dead_code)]
1010
pub struct NewChatResponse {
1111
pub _id: String,
12-
pub _source: Source,
12+
pub _source: Session,
1313
pub result: String,
1414
pub payload: Option<Value>,
1515
}
1616

1717
#[derive(Debug, Serialize, Deserialize)]
18-
pub struct Source {
18+
pub struct Session {
1919
pub id: String,
2020
pub created: String,
2121
pub updated: String,
2222
pub status: String,
2323
pub title: Option<String>,
2424
pub summary: Option<String>,
2525
pub manually_renamed_title: bool,
26+
pub visible: Option<bool>,
27+
pub context: Option<SessionContext>,
2628
}
29+
30+
#[derive(Debug, Serialize, Deserialize)]
31+
pub struct SessionContext {
32+
pub attachments: Option<Vec<String>>,
33+
}

src-tauri/src/common/search.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,9 @@ pub struct Total {
3636

3737
#[derive(Debug, Serialize, Deserialize)]
3838
pub struct SearchHit<T> {
39-
pub _index: String,
40-
pub _type: String,
41-
pub _id: String,
39+
pub _index: Option<String>,
40+
pub _type: Option<String>,
41+
pub _id: Option<String>,
4242
pub _score: Option<f64>,
4343
pub _source: T, // This will hold the type we pass in (e.g., DataSource)
4444
}

src-tauri/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,7 @@ pub fn run() {
158158
local::application::remove_app_search_path,
159159
settings::set_allow_self_signature,
160160
settings::get_allow_self_signature,
161+
assistant::ask_ai
161162
])
162163
.setup(|app| {
163164
let app_handle = app.handle().clone();

src-tauri/src/search/mod.rs

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,12 @@ pub async fn query_coco_fusion<R: Runtime>(
3131
// Time limit for each query
3232
let timeout_duration = Duration::from_millis(query_timeout);
3333

34+
log::debug!(
35+
"query_coco_fusion: {:?}, timeout: {:?}",
36+
query_strings,
37+
timeout_duration
38+
);
39+
3440
// Push all queries into futures
3541
for query_source in sources_list {
3642
let query_source_type = query_source.get_type().clone();
@@ -68,6 +74,8 @@ pub async fn query_coco_fusion<R: Runtime>(
6874
let source_id = response.source.id.clone();
6975

7076
for (doc, score) in response.hits {
77+
log::debug!("doc: {}, {:?}, {}", doc.id, doc.title, score);
78+
7179
let query_hit = QueryHits {
7280
source: Some(response.source.clone()),
7381
score,
@@ -83,7 +91,7 @@ pub async fn query_coco_fusion<R: Runtime>(
8391
}
8492
}
8593
Ok(Ok(Err(err))) => {
86-
log::debug!("Query error: {:?}", err);
94+
log::error!("{}", err);
8795
failed_requests.push(FailedRequest {
8896
source: QuerySource {
8997
r#type: "N/A".into(),
@@ -95,28 +103,19 @@ pub async fn query_coco_fusion<R: Runtime>(
95103
reason: None,
96104
});
97105
}
98-
Ok(Err(_elapsed)) => {
99-
log::debug!("Query timeout reached, skipping request");
106+
Ok(Err(err)) => {
107+
log::error!("{}", err);
100108
}
101109
// Timeout reached, skip this request
102110
_ => {
103-
failed_requests.push(FailedRequest {
104-
source: QuerySource {
105-
r#type: "N/A".into(),
106-
name: "N/A".into(),
107-
id: "N/A".into(),
108-
},
109-
status: 0,
110-
error: Some(format!("{:?}", &result)),
111-
reason: None,
112-
});
111+
log::debug!("timeout reached, skip this request");
113112
}
114113
}
115114
}
116115

117116
// Sort hits within each source by score (descending)
118117
for hits in hits_per_source.values_mut() {
119-
hits.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
118+
hits.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Greater));
120119
}
121120

122121
let total_sources = hits_per_source.len();
@@ -132,14 +131,22 @@ pub async fn query_coco_fusion<R: Runtime>(
132131
// Distribute hits fairly across sources
133132
for (_source_id, hits) in &mut hits_per_source {
134133
let take_count = hits.len().min(max_hits_per_source);
135-
for (doc, _) in hits.drain(0..take_count) {
134+
for (doc, score) in hits.drain(0..take_count) {
136135
if !seen_docs.contains(&doc.document.id) {
137136
seen_docs.insert(doc.document.id.clone());
137+
log::debug!(
138+
"collect doc: {}, {:?}, {}",
139+
doc.document.id,
140+
doc.document.title,
141+
score
142+
);
138143
final_hits.push(doc);
139144
}
140145
}
141146
}
142147

148+
log::debug!("final hits: {:?}", final_hits.len());
149+
143150
// If we still need more hits, take the highest-scoring remaining ones
144151
if final_hits.len() < size as usize {
145152
let remaining_needed = size as usize - final_hits.len();

src-tauri/src/server/http_client.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,7 @@ impl HttpClient {
212212
query_params,
213213
body,
214214
)
215-
.await
215+
.await
216216
}
217217

218218
// Convenience method for PUT requests
@@ -232,7 +232,7 @@ impl HttpClient {
232232
query_params,
233233
body,
234234
)
235-
.await
235+
.await
236236
}
237237

238238
// Convenience method for DELETE requests
@@ -251,6 +251,6 @@ impl HttpClient {
251251
query_params,
252252
None,
253253
)
254-
.await
254+
.await
255255
}
256256
}

src-tauri/src/server/servers.rs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -484,8 +484,6 @@ pub async fn mark_server_as_offline<R: Runtime>(
484484

485485
#[tauri::command]
486486
pub async fn disable_server<R: Runtime>(app_handle: AppHandle<R>, id: String) -> Result<(), ()> {
487-
println!("disable_server: {}", id);
488-
489487
let server = get_server_by_id(id.as_str());
490488
if let Some(mut server) = server {
491489
server.enabled = false;

src/components/Assistant/AssistantFetcher.tsx

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,17 +28,18 @@ export const AssistantFetcher = ({
2828
const fetchAssistant = async (params: {
2929
current: number;
3030
pageSize: number;
31+
serverId?: string;
3132
}) => {
3233
try {
33-
const { pageSize, current } = params;
34+
const { pageSize, current, serverId = currentService?.id } = params;
3435

3536
const from = (current - 1) * pageSize;
3637
const size = pageSize;
3738

3839
let response: any;
3940

4041
const body: Record<string, any> = {
41-
serverId: currentService?.id,
42+
serverId,
4243
from,
4344
size,
4445
};

src/components/Assistant/Chat.tsx

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import ConnectPrompt from "./ConnectPrompt";
2222
import type { Chat } from "@/types/chat";
2323
import PrevSuggestion from "@/components/ChatMessage/PrevSuggestion";
2424
import { useAppStore } from "@/stores/appStore";
25+
import { useSearchStore } from "@/stores/searchStore";
2526
// import ReadAloud from "./ReadAloud";
2627
import { useAuthStore } from "@/stores/authStore";
2728

@@ -91,11 +92,31 @@ const ChatAI = memo(
9192

9293
const [isSidebarOpenChat, setIsSidebarOpenChat] = useState(isSidebarOpen);
9394
const [chats, setChats] = useState<Chat[]>([]);
95+
const askAiSessionId = useSearchStore((state) => state.askAiSessionId);
96+
const setAskAiSessionId = useSearchStore(
97+
(state) => state.setAskAiSessionId
98+
);
9499

95100
useEffect(() => {
96101
activeChatProp && setActiveChat(activeChatProp);
97102
}, [activeChatProp]);
98103

104+
useEffect(() => {
105+
console.log("chats", chats);
106+
console.log("askAiSessionId", askAiSessionId);
107+
if (chats.length === 0 || !askAiSessionId) return;
108+
109+
const matched = chats.find((item) => item._id === askAiSessionId);
110+
111+
console.log("matched", matched);
112+
113+
if (matched) {
114+
onSelectChat(matched);
115+
116+
setAskAiSessionId(void 0);
117+
}
118+
}, [chats, chats]);
119+
99120
const [Question, setQuestion] = useState<string>("");
100121

101122
const [websocketSessionId, setWebsocketSessionId] = useState("");

0 commit comments

Comments
 (0)