ryedida.me / zotero-rag
Another remote for gh:zotero-rag/zotero-rag
0
fork

Configure Feed

Select the types of activity you want to include in your feed.

refactor: move roles into enum

Rahul Yedida 54524124 c30c4a34

+97 -74
+14 -13
zqa-rag/src/llm/anthropic.rs
··· 15 15 }; 16 16 use crate::http_client::HttpClient; 17 17 use crate::llm::base::{ 18 - ChatHistoryContent, ContentType, ReasoningConfig, ToolCallRequest, ToolCallResponse, USER_ROLE, 18 + ChatHistoryContent, ContentType, MessageRole, ReasoningConfig, ToolCallRequest, 19 + ToolCallResponse, 19 20 }; 20 21 use crate::llm::tools::{ 21 22 ANTHROPIC_SCHEMA_KEY, SerializedTool, get_owned_tools, process_tool_calls, ··· 27 28 #[derive(Clone, Serialize)] 28 29 pub(crate) struct AnthropicChatHistoryItem { 29 30 /// Either "user" or "assistant". 30 - pub role: String, 31 + pub role: MessageRole, 31 32 /// The contents of this item. 32 33 pub content: Vec<AnthropicResponseContent>, 33 34 } 34 35 35 36 impl From<ChatHistoryItem> for AnthropicChatHistoryItem { 36 37 fn from(value: ChatHistoryItem) -> Self { 37 - let role = value.role.clone(); 38 + let role = value.role; 38 39 Self { 39 40 content: value 40 41 .content ··· 116 117 .collect(); 117 118 118 119 messages.push(AnthropicChatHistoryItem { 119 - role: USER_ROLE.to_owned(), 120 + role: MessageRole::User, 120 121 content: vec![req.message.clone().into()], 121 122 }); 122 123 ··· 239 240 /// The model that generated the response 240 241 pub(crate) model: String, 241 242 /// The role of the message (usually "assistant") 242 - pub(crate) role: String, 243 + pub(crate) role: MessageRole, 243 244 /// Why the model stopped generating (e.g., "end_turn") 244 245 pub(crate) stop_reason: String, 245 246 /// The stop sequence that caused generation to end, if any ··· 373 374 374 375 // Append the contents 375 376 chat_history.push(AnthropicChatHistoryItem { 376 - role: "assistant".into(), 377 + role: MessageRole::Assistant, 377 378 content: response.content.clone(), 378 379 }); 379 380 ··· 408 409 409 410 // Append the new response to chat history 410 411 chat_history.push(AnthropicChatHistoryItem { 411 - role: "assistant".into(), 412 + role: MessageRole::Assistant, 412 413 content: response.content.clone(), 413 414 }); 414 415 ··· 460 461 use crate::http_client::{MockHttpClient, ReqwestClient, SequentialMockHttpClient}; 461 462 use crate::llm::anthropic::{AnthropicTextResponseContent, DEFAULT_CLAUDE_MODEL}; 462 463 use crate::llm::base::{ 463 - ApiClient, ChatHistoryContent, ChatHistoryItem, ChatRequest, ContentType, ToolCallResponse, 464 - USER_ROLE, 464 + ApiClient, ChatHistoryContent, ChatHistoryItem, ChatRequest, ContentType, MessageRole, 465 + ToolCallResponse, 465 466 }; 466 467 use crate::llm::tools::test_utils::MockTool; 467 468 ··· 494 495 let mock_response = AnthropicResponse { 495 496 id: "mock-id".to_string(), 496 497 model: DEFAULT_CLAUDE_MODEL.to_string(), 497 - role: "assistant".to_string(), 498 + role: MessageRole::Assistant, 498 499 stop_reason: "end_turn".to_string(), 499 500 stop_sequence: None, 500 501 usage: AnthropicUsageStats { ··· 573 574 let tool_call_response = AnthropicResponse { 574 575 id: "msg-1".into(), 575 576 model: DEFAULT_CLAUDE_MODEL.into(), 576 - role: "assistant".into(), 577 + role: MessageRole::Assistant, 577 578 stop_reason: "tool_use".into(), 578 579 stop_sequence: None, 579 580 usage: AnthropicUsageStats { ··· 596 597 let text_response = AnthropicResponse { 597 598 id: "msg-2".into(), 598 599 model: DEFAULT_CLAUDE_MODEL.into(), 599 - role: "assistant".into(), 600 + role: MessageRole::Assistant, 600 601 stop_reason: "end_turn".into(), 601 602 stop_sequence: None, 602 603 usage: AnthropicUsageStats { ··· 702 703 let second_message = ChatRequest { 703 704 chat_history: vec![ 704 705 ChatHistoryItem { 705 - role: USER_ROLE.into(), 706 + role: MessageRole::User, 706 707 content: vec![ChatHistoryContent::Text(first_message.message.clone())], 707 708 }, 708 709 chat_history_contents,
+13 -6
zqa-rag/src/llm/base.rs
··· 8 8 use super::errors::LLMError; 9 9 use crate::llm::tools::{CallbackFn, Tool}; 10 10 11 - /// The user role. 12 - pub const USER_ROLE: &str = "user"; 13 - /// The assistant role. 14 - pub const ASSISTANT_ROLE: &str = "assistant"; 11 + /// Roles for messages 12 + #[derive(Copy, Clone, Debug, PartialEq, Serialize, Deserialize)] 13 + #[serde(rename_all = "lowercase")] 14 + pub enum MessageRole { 15 + /// A user message/response. Some providers fold tool responses into this. 16 + User, 17 + /// An assistant response. 18 + Assistant, 19 + /// Some providers, like OpenRouter, use a separate role for tool responses. 20 + Tool, 21 + } 15 22 16 23 /// A user-facing, generic, tool call request. This contains the tool name and the parameters 17 24 /// passed to that tool. It also includes the tool call id, which most providers use to ··· 64 71 pub struct ChatHistoryItem { 65 72 /// The chat role. Every supported provider has a "user" (`USER_ROLE`) and an "assistant" 66 73 /// (`ASSISTANT_ROLE`) role. 67 - pub role: String, 74 + pub role: MessageRole, 68 75 /// The contents of this item. This is a `Vec` because some APIs expect tool call requests to 69 76 /// be bundled together with the text preceding it. 70 77 pub content: Vec<ChatHistoryContent>, ··· 91 98 .collect(); 92 99 93 100 Self { 94 - role: ASSISTANT_ROLE.into(), 101 + role: MessageRole::Assistant, 95 102 content, 96 103 } 97 104 }
+24 -11
zqa-rag/src/llm/gemini.rs
··· 15 15 }; 16 16 use crate::http_client::HttpClient; 17 17 use crate::llm::base::{ 18 - ChatHistoryContent, ChatHistoryItem, ContentType, ReasoningConfig, ToolCallRequest, USER_ROLE, 18 + ChatHistoryContent, ChatHistoryItem, ContentType, MessageRole, ReasoningConfig, ToolCallRequest, 19 19 }; 20 20 use crate::llm::tools::{GEMINI_SCHEMA_KEY, SerializedTool, get_owned_tools, process_tool_calls}; 21 21 ··· 63 63 }, 64 64 } 65 65 66 + /// Gemini uses "model", not "assistant", so we need our own enum. 67 + #[derive(Copy, Clone, PartialEq, Serialize, Deserialize)] 68 + #[serde(rename_all = "lowercase")] 69 + enum GeminiMessageRole { 70 + /// The model response. 71 + Model, 72 + /// The user message or response. 73 + User, 74 + } 75 + 66 76 /// Content for requests to the Gemini API 67 77 #[derive(Serialize, Deserialize, Clone)] 68 78 #[serde(rename_all = "camelCase")] 69 79 struct GeminiContent { 70 - role: String, 80 + role: GeminiMessageRole, 71 81 parts: Vec<GeminiPart>, 72 82 } 73 83 74 84 impl From<ChatHistoryItem> for GeminiContent { 75 85 fn from(value: ChatHistoryItem) -> Self { 76 86 Self { 77 - role: value.role, 87 + role: match value.role { 88 + MessageRole::User | MessageRole::Tool => GeminiMessageRole::User, 89 + MessageRole::Assistant => GeminiMessageRole::Model, 90 + }, 78 91 parts: value 79 92 .content 80 93 .into_iter() ··· 183 196 req.chat_history.iter().cloned().map(Into::into).collect(); 184 197 185 198 contents.push(GeminiContent { 186 - role: USER_ROLE.to_string(), 199 + role: GeminiMessageRole::User, 187 200 parts: vec![GeminiPart::Text { 188 201 text: req.message.clone(), 189 202 thought_signature: None, ··· 353 366 354 367 // Append the contents 355 368 chat_history.push(GeminiContent { 356 - role: "assistant".into(), 369 + role: GeminiMessageRole::Model, 357 370 parts: response.content.parts.clone(), 358 371 }); 359 372 ··· 397 410 398 411 // Append the new response to chat history 399 412 chat_history.push(GeminiContent { 400 - role: "assistant".into(), 413 + role: GeminiMessageRole::Model, 401 414 parts: response.content.parts.clone(), 402 415 }); 403 416 ··· 452 465 let mock_response = GeminiResponseBody { 453 466 candidates: vec![GeminiResponseCandidate { 454 467 content: GeminiContent { 455 - role: "model".into(), 468 + role: GeminiMessageRole::Model, 456 469 parts: vec![GeminiPart::Text { 457 470 text: "Hello from Gemini!".into(), 458 471 thought_signature: None, ··· 478 491 let request = ChatRequest { 479 492 message: "foo".into(), 480 493 chat_history: vec![ChatHistoryItem { 481 - role: "assistant".into(), 494 + role: MessageRole::Assistant, 482 495 content: vec![ChatHistoryContent::Text("Prior".into())], 483 496 }], 484 497 max_tokens: Some(256), ··· 612 625 let tool_call_response = GeminiResponseBody { 613 626 candidates: vec![GeminiResponseCandidate { 614 627 content: GeminiContent { 615 - role: "model".into(), 628 + role: GeminiMessageRole::Model, 616 629 parts: vec![GeminiPart::FunctionCall { 617 630 function_call: GeminiFunctionCall { 618 631 id: Some("call-1".into()), ··· 635 648 let text_response = GeminiResponseBody { 636 649 candidates: vec![GeminiResponseCandidate { 637 650 content: GeminiContent { 638 - role: "model".into(), 651 + role: GeminiMessageRole::Model, 639 652 parts: vec![GeminiPart::Text { 640 653 text: "Done!".into(), 641 654 thought_signature: None, ··· 708 721 let second_message = ChatRequest { 709 722 chat_history: vec![ 710 723 ChatHistoryItem { 711 - role: USER_ROLE.into(), 724 + role: MessageRole::User, 712 725 content: vec![ChatHistoryContent::Text(first_message.message.clone())], 713 726 }, 714 727 chat_history_contents,
+7 -7
zqa-rag/src/llm/ollama.rs
··· 17 17 AnthropicChatHistoryItem, AnthropicRequest, AnthropicResponse, AnthropicResponseContent, 18 18 build_anthropic_messages_and_tools, map_response_to_chat_contents, 19 19 }; 20 - use crate::llm::base::ContentType; 20 + use crate::llm::base::{ContentType, MessageRole}; 21 21 use crate::llm::tools::process_tool_calls; 22 22 23 23 /// Ollama supports the Anthropic Messages API, so we can reuse structs. ··· 105 105 106 106 // Append the contents 107 107 chat_history.push(AnthropicChatHistoryItem { 108 - role: "assistant".into(), 108 + role: MessageRole::Assistant, 109 109 content: response.content.clone(), 110 110 }); 111 111 ··· 140 140 141 141 // Append the new response to chat history 142 142 chat_history.push(AnthropicChatHistoryItem { 143 - role: "assistant".into(), 143 + role: MessageRole::Assistant, 144 144 content: response.content.clone(), 145 145 }); 146 146 ··· 191 191 AnthropicTextResponseContent, AnthropicToolUseResponseContent, AnthropicUsageStats, 192 192 }; 193 193 use crate::llm::base::{ 194 - ApiClient, ChatHistoryContent, ChatHistoryItem, ChatRequest, USER_ROLE, 194 + ApiClient, ChatHistoryContent, ChatHistoryItem, ChatRequest, MessageRole, 195 195 }; 196 196 use crate::llm::tools::test_utils::MockTool; 197 197 ··· 259 259 let tool_call_response = AnthropicResponse { 260 260 id: "msg-1".into(), 261 261 model: DEFAULT_OLLAMA_MODEL.into(), 262 - role: "assistant".into(), 262 + role: MessageRole::Assistant, 263 263 stop_reason: "tool_use".into(), 264 264 stop_sequence: None, 265 265 usage: AnthropicUsageStats { ··· 282 282 let text_response = AnthropicResponse { 283 283 id: "msg-2".into(), 284 284 model: DEFAULT_OLLAMA_MODEL.into(), 285 - role: "assistant".into(), 285 + role: MessageRole::Assistant, 286 286 stop_reason: "end_turn".into(), 287 287 stop_sequence: None, 288 288 usage: AnthropicUsageStats { ··· 355 355 let second_message = ChatRequest { 356 356 chat_history: vec![ 357 357 ChatHistoryItem { 358 - role: USER_ROLE.into(), 358 + role: MessageRole::User, 359 359 content: vec![ChatHistoryContent::Text(first_message.message.clone())], 360 360 }, 361 361 chat_history_contents,
+12 -10
zqa-rag/src/llm/openai.rs
··· 25 25 DEFAULT_OPENAI_REASONING_EFFORT, 26 26 }; 27 27 use crate::http_client::HttpClient; 28 - use crate::llm::base::{ChatHistoryContent, ContentType, ReasoningConfig, ToolUseStats, USER_ROLE}; 28 + use crate::llm::base::{ 29 + ChatHistoryContent, ContentType, MessageRole, ReasoningConfig, ToolUseStats, 30 + }; 29 31 use crate::llm::tools::{CallbackFn, OPENAI_SCHEMA_KEY, SerializedTool, get_owned_tools}; 30 32 31 33 /// OpenAI-specific tool wrapper that adds the `type` and `strict` fields ··· 108 110 #[derive(Clone, Serialize)] 109 111 pub(crate) struct OpenAIChatHistoryItem { 110 112 /// Either USER_ROLE or "assistant". 111 - pub role: String, 113 + pub role: MessageRole, 112 114 /// The content item for this role. 113 115 pub content: OpenAIRequestInputItem, 114 116 /// Always "message". ··· 130 132 .iter() 131 133 .map(|c| match c { 132 134 ChatHistoryContent::Text(s) => OpenAIRequestInput::Message(OpenAIChatHistoryItem { 133 - role: value.role.clone(), 135 + role: value.role, 134 136 r#type: "message".into(), 135 137 content: OpenAIRequestInputItem::Text(s.clone()), 136 138 }), ··· 217 219 .collect::<Vec<_>>(); 218 220 219 221 messages.push(OpenAIRequestInput::Message(OpenAIChatHistoryItem { 220 - role: USER_ROLE.to_owned(), 222 + role: MessageRole::User, 221 223 r#type: "message".into(), 222 224 content: OpenAIRequestInputItem::Text(req.message.clone()), 223 225 })); ··· 272 274 id: String, 273 275 status: String, 274 276 content: Vec<OpenAIContent>, 275 - role: String, 277 + role: MessageRole, 276 278 }, 277 279 /// A function call request to invoke a tool. 278 280 #[serde(rename = "function_call")] ··· 466 468 OpenAIOutput::Reasoning { .. } => None, 467 469 OpenAIOutput::Message { content, .. } => { 468 470 Some(OpenAIRequestInput::Message(OpenAIChatHistoryItem { 469 - role: "assistant".into(), 471 + role: MessageRole::Assistant, 470 472 r#type: "message".into(), 471 473 content: OpenAIRequestInputItem::Text( 472 474 content.first()?.text.clone().unwrap_or_else(String::new), ··· 685 687 use crate::constants::DEFAULT_OPENAI_EMBEDDING_DIM; 686 688 use crate::http_client::{MockHttpClient, ReqwestClient}; 687 689 use crate::llm::base::{ 688 - ApiClient, ChatHistoryContent, ChatHistoryItem, ChatRequest, ContentType, USER_ROLE, 690 + ApiClient, ChatHistoryContent, ChatHistoryItem, ChatRequest, ContentType, MessageRole, 689 691 }; 690 692 use crate::llm::tools::OPENAI_SCHEMA_KEY; 691 693 use crate::llm::tools::test_utils::MockTool; ··· 728 730 output: vec![OpenAIOutput::Message { 729 731 id: "msg_id".into(), 730 732 status: "completed".into(), 731 - role: "assistant".into(), 733 + role: MessageRole::Assistant, 732 734 content: vec![OpenAIContent { 733 735 r#type: "output_text".into(), 734 736 text: Some("Hi there!".into()), ··· 928 930 output: vec![OpenAIOutput::Message { 929 931 id: "msg-1".into(), 930 932 status: "completed".into(), 931 - role: "assistant".into(), 933 + role: MessageRole::Assistant, 932 934 content: vec![OpenAIContent { 933 935 r#type: "output_text".into(), 934 936 text: Some("Done!".into()), ··· 992 994 let second_message = ChatRequest { 993 995 chat_history: vec![ 994 996 ChatHistoryItem { 995 - role: USER_ROLE.into(), 997 + role: MessageRole::User, 996 998 content: vec![ChatHistoryContent::Text(first_message.message.clone())], 997 999 }, 998 1000 chat_history_contents,
+12 -12
zqa-rag/src/llm/openrouter.rs
··· 13 13 use crate::common::request_with_backoff; 14 14 use crate::http_client::HttpClient; 15 15 use crate::llm::base::{ 16 - ChatHistoryContent, ContentType, ReasoningConfig, ToolCallRequest, USER_ROLE, 16 + ChatHistoryContent, ContentType, MessageRole, ReasoningConfig, ToolCallRequest, 17 17 }; 18 18 use crate::llm::tools::{ 19 19 OPENROUTER_SCHEMA_KEY, SerializedTool, get_owned_tools, process_tool_calls, ··· 24 24 /// OpenRouter-specific message format 25 25 #[derive(Clone, Debug, Serialize)] 26 26 struct OpenRouterMessage { 27 - role: String, 27 + role: MessageRole, 28 28 #[serde(skip_serializing_if = "Option::is_none")] 29 29 content: Option<String>, 30 30 #[serde(skip_serializing_if = "Option::is_none")] ··· 119 119 ChatHistoryContent::ToolCallResponse(res) => { 120 120 // Tool responses in OpenRouter are sent as a separate message with role "tool" 121 121 messages.push(OpenRouterMessage { 122 - role: "tool".into(), 122 + role: MessageRole::Tool, 123 123 content: Some(serde_json::to_string(&res.result).unwrap_or_default()), 124 124 tool_calls: None, 125 125 tool_call_id: Some(res.id.clone()), ··· 133 133 messages.insert( 134 134 0, 135 135 OpenRouterMessage { 136 - role: item.role.clone(), 136 + role: item.role, 137 137 content: if content_text.is_empty() { 138 138 None 139 139 } else { ··· 164 164 .collect(); 165 165 166 166 messages.push(OpenRouterMessage { 167 - role: USER_ROLE.to_owned(), 167 + role: MessageRole::User, 168 168 content: Some(req.message.clone()), 169 169 tool_calls: None, 170 170 tool_call_id: None, ··· 212 212 #[derive(Clone, Serialize, Deserialize)] 213 213 struct OpenRouterResponseMessage { 214 214 /// Usually "assistant" 215 - role: String, 215 + role: MessageRole, 216 216 /// The model response 217 217 #[serde(skip_serializing_if = "Option::is_none")] 218 218 content: Option<String>, ··· 378 378 379 379 // Append the assistant's response to chat history 380 380 chat_history.push(OpenRouterMessage { 381 - role: choice.message.role.clone(), 381 + role: choice.message.role, 382 382 content: choice.message.content.clone(), 383 383 tool_calls: choice.message.tool_calls.clone(), 384 384 tool_call_id: None, ··· 420 420 421 421 // Append the new response to chat history 422 422 chat_history.push(OpenRouterMessage { 423 - role: choice.message.role.clone(), 423 + role: choice.message.role, 424 424 content: choice.message.content.clone(), 425 425 tool_calls: choice.message.tool_calls.clone(), 426 426 tool_call_id: None, ··· 503 503 }, 504 504 choices: vec![OpenRouterResponseChoices { 505 505 message: OpenRouterResponseMessage { 506 - role: String::from("assistant"), 506 + role: MessageRole::Assistant, 507 507 content: Some(String::from("Hi there! How can I help you today?")), 508 508 tool_calls: None, 509 509 refusal: Some(String::new()), ··· 589 589 }, 590 590 choices: vec![OpenRouterResponseChoices { 591 591 message: OpenRouterResponseMessage { 592 - role: "assistant".into(), 592 + role: MessageRole::Assistant, 593 593 content: None, 594 594 tool_calls: Some(vec![OpenRouterToolCall { 595 595 id: "call-1".into(), ··· 619 619 }, 620 620 choices: vec![OpenRouterResponseChoices { 621 621 message: OpenRouterResponseMessage { 622 - role: "assistant".into(), 622 + role: MessageRole::Assistant, 623 623 content: Some("Done!".into()), 624 624 tool_calls: None, 625 625 refusal: None, ··· 686 686 let second_message = ChatRequest { 687 687 chat_history: vec![ 688 688 ChatHistoryItem { 689 - role: USER_ROLE.into(), 689 + role: MessageRole::User, 690 690 content: vec![ChatHistoryContent::Text(first_message.message.clone())], 691 691 }, 692 692 chat_history_contents,
+3 -2
zqa-rag/src/llm/tools.rs
··· 14 14 15 15 use crate::llm::{ 16 16 base::{ 17 - ChatHistoryContent, ChatHistoryItem, ContentType, ToolCallResponse, ToolUseStats, USER_ROLE, 17 + ChatHistoryContent, ChatHistoryItem, ContentType, MessageRole, ToolCallResponse, 18 + ToolUseStats, 18 19 }, 19 20 errors::LLMError, 20 21 }; ··· 251 252 }; 252 253 253 254 let chat_history_item = ChatHistoryItem { 254 - role: USER_ROLE.into(), 255 + role: MessageRole::User, 255 256 content: vec![ChatHistoryContent::ToolCallResponse(ToolCallResponse { 256 257 id: tool_call_id, 257 258 tool_name: tool_call.tool_name.clone(),
+5 -5
zqa/src/cli/handlers/conversation.rs
··· 151 151 use chrono::Local; 152 152 use temp_env; 153 153 use zqa_macros::{test_contains, test_eq}; 154 - use zqa_rag::llm::base::{ASSISTANT_ROLE, ChatHistoryContent, ChatHistoryItem, USER_ROLE}; 154 + use zqa_rag::llm::base::{ChatHistoryContent, ChatHistoryItem, MessageRole}; 155 155 156 156 use super::resume_with_reader; 157 157 use crate::{ ··· 178 178 temp_env::with_var("ZQA_STATE_DIR", Some(temp_dir.path()), || { 179 179 let history_a = vec![ 180 180 ChatHistoryItem { 181 - role: USER_ROLE.into(), 181 + role: MessageRole::User, 182 182 content: vec![ChatHistoryContent::Text("What is attention?".into())], 183 183 }, 184 184 ChatHistoryItem { 185 - role: ASSISTANT_ROLE.into(), 185 + role: MessageRole::Assistant, 186 186 content: vec![ChatHistoryContent::Text( 187 187 "Attention is a mechanism...".into(), 188 188 )], 189 189 }, 190 190 ]; 191 191 let history_b = vec![ChatHistoryItem { 192 - role: USER_ROLE.into(), 192 + role: MessageRole::User, 193 193 content: vec![ChatHistoryContent::Text( 194 194 "Tell me about transformers.".into(), 195 195 )], ··· 232 232 temp_env::with_var("ZQA_STATE_DIR", Some(temp_dir.path()), || { 233 233 save_conversation(&SavedChatHistory { 234 234 history: vec![ChatHistoryItem { 235 - role: USER_ROLE.into(), 235 + role: MessageRole::User, 236 236 content: vec![ChatHistoryContent::Text("Hello".into())], 237 237 }], 238 238 date: Local::now(),
+3 -4
zqa/src/cli/handlers/query.rs
··· 10 10 use zqa_rag::{ 11 11 llm::{ 12 12 base::{ 13 - ASSISTANT_ROLE, ApiClient, ChatHistoryContent, ChatHistoryItem, ChatRequest, 14 - ToolUseStats, USER_ROLE, 13 + ApiClient, ChatHistoryContent, ChatHistoryItem, ChatRequest, MessageRole, ToolUseStats, 15 14 }, 16 15 factory::get_client_with_config, 17 16 tools::{CallbackFn, Tool}, ··· 345 344 .lock() 346 345 .expect("Could not obtain lock on chat history."); 347 346 history.push(ChatHistoryItem { 348 - role: USER_ROLE.into(), 347 + role: MessageRole::User, 349 348 content: vec![ChatHistoryContent::Text(query.clone())], 350 349 }); 351 350 history.push(ChatHistoryItem { 352 - role: ASSISTANT_ROLE.into(), 351 + role: MessageRole::Assistant, 353 352 content: vec![ChatHistoryContent::Text(model_response_text)], 354 353 }); 355 354 ctx.state.dirty.store(true, atomic::Ordering::Relaxed);
+4 -4
zqa/src/state.rs
··· 497 497 use chrono::Local; 498 498 use clap::builder::OsStr; 499 499 use zqa_macros::test_ok; 500 - use zqa_rag::llm::base::{ChatHistoryContent, ChatHistoryItem, USER_ROLE}; 500 + use zqa_rag::llm::base::{ChatHistoryContent, ChatHistoryItem, MessageRole}; 501 501 502 502 use crate::state::{ 503 503 SavedChatHistory, get_conversation_history, get_state_dir, oobe, save_conversation, ··· 537 537 date: Local::now(), 538 538 title: "foo".into(), 539 539 history: vec![ChatHistoryItem { 540 - role: USER_ROLE.into(), 540 + role: MessageRole::User, 541 541 content: vec![ChatHistoryContent::Text("Hello!".into())], 542 542 }], 543 543 }; ··· 557 557 date: Local::now(), 558 558 title: "foo".into(), 559 559 history: vec![ChatHistoryItem { 560 - role: USER_ROLE.into(), 560 + role: MessageRole::User, 561 561 content: vec![ChatHistoryContent::Text("Hello!".into())], 562 562 }], 563 563 }; ··· 575 575 assert_eq!(conversations.len(), 1); 576 576 assert_eq!(conversations[0].title, "foo"); 577 577 assert_eq!(conversations[0].history.len(), 1); 578 - assert_eq!(conversations[0].history[0].role, USER_ROLE); 578 + assert_eq!(conversations[0].history[0].role, MessageRole::User); 579 579 assert_eq!(conversations[0].history[0].content.len(), 1); 580 580 assert_eq!( 581 581 conversations[0].history[0].content[0],