1 files changed, 24 insertions, 12 deletions

diff --git a/src/config/session.rs b/src/config/session.rs
index 1aebd64..644255f 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -1,12 +1,14 @@
+use super::input::resolve_data_url;
 use super::role::Role;
-use super::Model;
+use super::{Input, Model};
 
-use crate::client::{Message, MessageRole};
+use crate::client::{Message, MessageContent, MessageRole};
 use crate::render::MarkdownRender;
 
 use anyhow::{bail, Context, Result};
 use serde::{Deserialize, Serialize};
 use serde_json::json;
+use std::collections::HashMap;
 use std::fs::{self, read_to_string};
 use std::path::Path;
 
@@ -18,6 +20,7 @@ pub struct Session {
     model_id: String,
     temperature: Option<f64>,
     messages: Vec<Message>,
+    data_urls: HashMap<String, String>,
     #[serde(skip)]
     pub name: String,
     #[serde(skip)]
@@ -37,6 +40,7 @@ impl Session {
             model_id: model.id(),
             temperature,
             messages: vec![],
+            data_urls: Default::default(),
             name: name.to_string(),
             path: None,
             dirty: false,
@@ -121,6 +125,7 @@ impl Session {
 
         if !self.is_empty() {
             lines.push("".into());
+            let resolve_url_fn = |url: &str| resolve_data_url(&self.data_urls, url.to_string());
 
             for message in &self.messages {
                 match message.role {
@@ -128,11 +133,17 @@ impl Session {
                         continue;
                     }
                     MessageRole::Assistant => {
-                        lines.push(render.render(&message.content));
+                        if let MessageContent::Text(text) = &message.content {
+                            lines.push(render.render(text));
+                        }
                         lines.push("".into());
                     }
                     MessageRole::User => {
-                        lines.push(format!("{}）{}", self.name, message.content));
+                        lines.push(format!(
+                            "{}）{}",
+                            self.name,
+                            message.content.render_input(resolve_url_fn)
+                        ));
                     }
                 }
             }
@@ -218,7 +229,7 @@ impl Session {
         self.messages.is_empty()
     }
 
-    pub fn add_message(&mut self, input: &str, output: &str) -> Result<()> {
+    pub fn add_message(&mut self, input: &Input, output: &str) -> Result<()> {
         let mut need_add_msg = true;
         if self.messages.is_empty() {
             if let Some(role) = self.role.as_ref() {
@@ -229,35 +240,36 @@ impl Session {
         if need_add_msg {
             self.messages.push(Message {
                 role: MessageRole::User,
-                content: input.to_string(),
+                content: input.to_message_content(),
             });
         }
+        self.data_urls.extend(input.data_urls());
         self.messages.push(Message {
             role: MessageRole::Assistant,
-            content: output.to_string(),
+            content: MessageContent::Text(output.to_string()),
         });
         self.dirty = true;
         Ok(())
     }
 
-    pub fn echo_messages(&self, content: &str) -> String {
-        let messages = self.build_emssages(content);
+    pub fn echo_messages(&self, input: &Input) -> String {
+        let messages = self.build_emssages(input);
         serde_yaml::to_string(&messages).unwrap_or_else(|_| "Unable to echo message".into())
     }
 
-    pub fn build_emssages(&self, content: &str) -> Vec<Message> {
+    pub fn build_emssages(&self, input: &Input) -> Vec<Message> {
         let mut messages = self.messages.clone();
         let mut need_add_msg = true;
         if messages.is_empty() {
             if let Some(role) = self.role.as_ref() {
-                messages = role.build_messages(content);
+                messages = role.build_messages(input);
                 need_add_msg = false;
             }
         };
         if need_add_msg {
             messages.push(Message {
                 role: MessageRole::User,
-                content: content.into(),
+                content: input.to_message_content(),
             });
         }
         messages