update schema to properly support model server and identifier

serverlessnext · Jul 20, 2024 · ef51d52 · ef51d52
1 parent 1fe66df
commit ef51d52
Show file tree

Hide file tree

Showing 9 changed files with 164 additions and 120 deletions.
diff --git a/lumni/src/apps/builtin/llm/prompt/src/chat/db/mod.rs b/lumni/src/apps/builtin/llm/prompt/src/chat/db/mod.rs
@@ -5,7 +5,7 @@ mod schema;
 mod store;
 
 pub use reader::ConversationReader;
-pub use schema::{ConversationCache, ConversationId, Message, MessageId};
+pub use schema::{ConversationCache, ConversationId, Model, ModelIdentifier, ModelServerName, Message, MessageId};
 pub use store::ConversationDatabaseStore;
 
 pub use super::PromptRole;
diff --git a/lumni/src/apps/builtin/llm/prompt/src/chat/db/schema.rs b/lumni/src/apps/builtin/llm/prompt/src/chat/db/schema.rs
@@ -4,8 +4,17 @@ use serde::{Deserialize, Serialize};
 
 use super::PromptRole;
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
-pub struct ModelId(pub i64);
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub struct ModelIdentifier(pub String);
+
+impl ModelIdentifier {
+    pub fn new(provider: &str, name: &str) -> Self {
+        ModelIdentifier(format!("{}::{}", provider, name))
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub struct ModelServerName(pub String);
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct ConversationId(pub i64);
@@ -16,23 +25,38 @@ pub struct MessageId(pub i64);
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct AttachmentId(pub i64);
 
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Model {
-    pub model_id: ModelId,
-    pub model_name: String,
-    pub model_service: String,
+    pub identifier: ModelIdentifier,
+    pub info: Option<serde_json::Value>,
+    pub config: Option<serde_json::Value>,
+    pub context_window_size: Option<i64>,
+    pub input_token_limit: Option<i64>,
+}
+
+impl Model {
+    pub fn new(identifier: ModelIdentifier) -> Self {
+        Model {
+            identifier,
+            info: None,
+            config: None,
+            context_window_size: None,
+            input_token_limit: None,
+        }
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Conversation {
     pub id: ConversationId,
     pub name: String,
-    pub metadata: serde_json::Value,
-    pub model_id: ModelId,
+    pub info: serde_json::Value,
+    pub model_identifier: ModelIdentifier,
+    pub model_server: ModelServerName,
     pub parent_conversation_id: Option<ConversationId>,
     pub fork_message_id: Option<MessageId>, // New field
     pub completion_options: Option<serde_json::Value>,
-    pub schema_version: i64,
     pub created_at: i64,
     pub updated_at: i64,
     pub is_deleted: bool,
@@ -73,7 +97,7 @@ pub struct Attachment {
 #[derive(Debug)]
 pub struct ConversationCache {
     conversation_id: ConversationId,
-    models: HashMap<ModelId, Model>,
+    models: HashMap<ModelIdentifier, Model>,
     messages: Vec<Message>, // messages have to be ordered
     attachments: HashMap<AttachmentId, Attachment>,
     message_attachments: HashMap<MessageId, Vec<AttachmentId>>,
@@ -107,7 +131,7 @@ impl ConversationCache {
     }
 
     pub fn add_model(&mut self, model: Model) {
-        self.models.insert(model.model_id, model);
+        self.models.insert(model.identifier.clone(), model);
     }
 
     pub fn add_message(&mut self, message: Message) {

diff --git a/lumni/src/apps/builtin/llm/prompt/src/chat/db/schema.sql b/lumni/src/apps/builtin/llm/prompt/src/chat/db/schema.sql
@@ -7,27 +7,29 @@ CREATE TABLE metadata (
 );
 
 CREATE TABLE models (
-    model_id INTEGER PRIMARY KEY AUTOINCREMENT,
-    model_name TEXT NOT NULL,
-    model_service TEXT NOT NULL UNIQUE
+    identifier TEXT PRIMARY KEY,
+    info TEXT, -- JSON string
+    config TEXT, -- JSON string
+    context_window_size INTEGER,
+    input_token_limit INTEGER
 );
 
 CREATE TABLE conversations (
     id INTEGER PRIMARY KEY AUTOINCREMENT,
     name TEXT,
-    metadata TEXT, -- JSON string including description and other metadata
-    model_id INTEGER NOT NULL,
+    info TEXT, -- JSON string including description and other metadata
+    completion_options TEXT, -- JSON string
+    model_identifier TEXT NOT NULL,
+    model_server TEXT NOT NULL,
     parent_conversation_id INTEGER,
     fork_message_id INTEGER,
-    completion_options TEXT, -- JSON string
-    schema_version INTEGER NOT NULL,
     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
     updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
     message_count INTEGER DEFAULT 0,
     total_tokens INTEGER DEFAULT 0,
     is_deleted BOOLEAN DEFAULT FALSE,
     FOREIGN KEY (parent_conversation_id) REFERENCES conversations(id),
-    FOREIGN KEY (model_id) REFERENCES models(model_id),
+    FOREIGN KEY (model_identifier) REFERENCES models(identifier),
     FOREIGN KEY (fork_message_id) REFERENCES messages(id)
 );
 
@@ -61,9 +63,8 @@ CREATE TABLE attachments (
     CHECK ((file_uri IS NULL) != (file_data IS NULL))
 );
 
-CREATE INDEX idx_model_service ON models(model_service);
 CREATE INDEX idx_parent_conversation ON conversations(parent_conversation_id);
-CREATE INDEX idx_conversation_model_id ON conversations(model_id);  
+CREATE INDEX idx_conversation_model_identifier ON conversations(model_identifier);  
 CREATE INDEX idx_attachment_message ON attachments(message_id);
 CREATE INDEX idx_conversation_is_deleted_updated ON conversations(is_deleted, updated_at);
 CREATE INDEX idx_message_conversation_created ON messages(conversation_id, created_at);