This commit is contained in:
duanfuxiang
2025-01-05 11:51:39 +08:00
commit 0c7ee142cb
215 changed files with 20611 additions and 0 deletions

View File

@@ -0,0 +1,162 @@
import { SerializedEditorState } from 'lexical'
import { App } from 'obsidian'
import { editorStateToPlainText } from '../../../components/chat-view/chat-input/utils/editor-state-to-plain-text'
import { ChatAssistantMessage, ChatConversationMeta, ChatMessage, ChatUserMessage } from '../../../types/chat'
import { ContentPart } from '../../../types/llm/request'
import { Mentionable, SerializedMentionable } from '../../../types/mentionable'
import { deserializeMentionable, serializeMentionable } from '../../../utils/mentionable'
import { DBManager } from '../../database-manager'
import { InsertMessage } from '../../schema'
import { ConversationRepository } from './conversation-repository'
export class ConversationManager {
private app: App
private repository: ConversationRepository
private dbManager: DBManager
constructor(app: App, dbManager: DBManager) {
this.app = app
this.dbManager = dbManager
const db = dbManager.getPgClient()
if (!db) throw new Error('Database not initialized')
this.repository = new ConversationRepository(app, db)
}
async createConversation(id: string, title = 'New chat'): Promise<void> {
const conversation = {
id,
title,
createdAt: new Date(),
updatedAt: new Date(),
}
await this.repository.create(conversation)
await this.dbManager.save()
}
async saveConversation(id: string, messages: ChatMessage[]): Promise<void> {
const conversation = await this.repository.findById(id)
if (!conversation) {
let title = 'New chat'
if (messages.length > 0 && messages[0].role === 'user') {
const query = editorStateToPlainText(messages[0].content)
if (query.length > 20) {
title = `${query.slice(0, 20)}...`
} else {
title = query
}
}
await this.createConversation(id, title)
}
// Delete existing messages
await this.repository.deleteAllMessagesFromConversation(id)
// Insert new messages
for (const message of messages) {
const insertMessage = this.serializeMessage(message, id)
await this.repository.createMessage(insertMessage)
}
// Update conversation timestamp
await this.repository.update(id, { updatedAt: new Date() })
await this.dbManager.save()
}
async findConversation(id: string): Promise<ChatMessage[] | null> {
const conversation = await this.repository.findById(id)
if (!conversation) {
return null
}
const messages = await this.repository.findMessagesByConversationId(id)
return messages.map(msg => this.deserializeMessage(msg))
}
async deleteConversation(id: string): Promise<void> {
await this.repository.delete(id)
await this.dbManager.save()
}
getAllConversations(callback: (conversations: ChatConversationMeta[]) => void): void {
const db = this.dbManager.getPgClient()
db?.live.query('SELECT * FROM conversations ORDER BY updated_at', [], (results) => {
callback(results.rows.map(conv => ({
id: conv.id,
title: conv.title,
schemaVersion: 2,
createdAt: conv.createdAt instanceof Date ? conv.createdAt.getTime() : conv.createdAt,
updatedAt: conv.updatedAt instanceof Date ? conv.updatedAt.getTime() : conv.updatedAt,
})))
})
}
async updateConversationTitle(id: string, title: string): Promise<void> {
await this.repository.update(id, { title })
await this.dbManager.save()
}
private serializeMessage(message: ChatMessage, conversationId: string): InsertMessage {
const base = {
id: message.id,
conversationId,
role: message.role,
createdAt: new Date(),
}
if (message.role === 'user') {
const userMessage: ChatUserMessage = message
return {
...base,
content: userMessage.content ? JSON.stringify(userMessage.content) : null,
promptContent: userMessage.promptContent
? typeof userMessage.promptContent === 'string'
? userMessage.promptContent
: JSON.stringify(userMessage.promptContent)
: null,
mentionables: JSON.stringify(userMessage.mentionables.map(serializeMentionable)),
similaritySearchResults: userMessage.similaritySearchResults
? JSON.stringify(userMessage.similaritySearchResults)
: null,
}
} else {
const assistantMessage: ChatAssistantMessage = message
return {
...base,
content: assistantMessage.content,
metadata: assistantMessage.metadata ? JSON.stringify(assistantMessage.metadata) : null,
}
}
}
private deserializeMessage(message: InsertMessage): ChatMessage {
if (message.role === 'user') {
return {
id: message.id,
role: 'user',
content: message.content ? JSON.parse(message.content) as SerializedEditorState : null,
promptContent: message.promptContent
? message.promptContent.startsWith('{')
? JSON.parse(message.promptContent) as ContentPart[]
: message.promptContent
: null,
mentionables: message.mentionables
? (JSON.parse(message.mentionables) as SerializedMentionable[])
.map(m => deserializeMentionable(m, this.app))
.filter((m: Mentionable | null): m is Mentionable => m !== null)
: [],
similaritySearchResults: message.similaritySearchResults
? JSON.parse(message.similaritySearchResults)
: undefined,
}
} else {
return {
id: message.id,
role: 'assistant',
content: message.content || '',
metadata: message.metadata ? JSON.parse(message.metadata) : undefined,
}
}
}
}

View File

@@ -0,0 +1,131 @@
import { PGliteInterface } from '@electric-sql/pglite'
import { App } from 'obsidian'
import {
InsertConversation,
InsertMessage,
SelectConversation,
SelectMessage,
} from '../../schema'
type QueryResult<T> = {
rows: T[]
}
export class ConversationRepository {
private app: App
private db: PGliteInterface
constructor(app: App, db: PGliteInterface) {
this.app = app
this.db = db
}
async create(conversation: InsertConversation): Promise<SelectConversation> {
const result = await this.db.query<SelectConversation>(
`INSERT INTO conversations (id, title, created_at, updated_at)
VALUES ($1, $2, $3, $4)
RETURNING *`,
[
conversation.id,
conversation.title,
conversation.createdAt || new Date(),
conversation.updatedAt || new Date()
]
) as QueryResult<SelectConversation>
return result.rows[0]
}
async createMessage(message: InsertMessage): Promise<SelectMessage> {
const result = await this.db.query<SelectMessage>(
`INSERT INTO messages (
id, conversation_id, role, content,
prompt_content, metadata, mentionables,
similarity_search_results, created_at
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
RETURNING *`,
[
message.id,
message.conversationId,
message.role,
message.content,
message.promptContent,
message.metadata,
message.mentionables,
message.similaritySearchResults,
message.createdAt || new Date()
]
) as QueryResult<SelectMessage>
return result.rows[0]
}
async findById(id: string): Promise<SelectConversation | undefined> {
const result = await this.db.query<SelectConversation>(
`SELECT * FROM conversations WHERE id = $1 LIMIT 1`,
[id]
) as QueryResult<SelectConversation>
return result.rows[0]
}
async findMessagesByConversationId(conversationId: string): Promise<SelectMessage[]> {
const result = await this.db.query<SelectMessage>(
`SELECT * FROM messages
WHERE conversation_id = $1
ORDER BY created_at`,
[conversationId]
) as QueryResult<SelectMessage>
return result.rows
}
async findAll(): Promise<SelectConversation[]> {
const result = await this.db.query<SelectConversation>(
`SELECT * FROM conversations ORDER BY updated_at DESC`
) as QueryResult<SelectConversation>
return result.rows
}
async update(id: string, data: Partial<InsertConversation>): Promise<SelectConversation> {
const setClauses: string[] = []
const values: any[] = []
let paramIndex = 1
if (data.title !== undefined) {
setClauses.push(`title = $${paramIndex}`)
values.push(data.title)
paramIndex++
}
// Always update updated_at
setClauses.push(`updated_at = $${paramIndex}`)
values.push(new Date())
paramIndex++
// Add id as the last parameter
values.push(id)
const result = await this.db.query<SelectConversation>(
`UPDATE conversations
SET ${setClauses.join(', ')}
WHERE id = $${paramIndex}
RETURNING *`,
values
) as QueryResult<SelectConversation>
return result.rows[0]
}
async delete(id: string): Promise<boolean> {
const result = await this.db.query<SelectConversation>(
`DELETE FROM conversations WHERE id = $1 RETURNING *`,
[id]
) as QueryResult<SelectConversation>
return result.rows.length > 0
}
async deleteAllMessagesFromConversation(conversationId: string): Promise<void> {
await this.db.query(
`DELETE FROM messages WHERE conversation_id = $1`,
[conversationId]
)
}
}