refactor: integrate SessionManager into daemon and agent
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -1,79 +1,69 @@
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { NativeAgent } from './agent.js';
|
||||
import type { ModelClient, ChatResponse } from '../../models/types.js';
|
||||
|
||||
describe('NativeAgent', () => {
|
||||
it('processes message and returns response', async () => {
|
||||
const mockClient: ModelClient = {
|
||||
chat: vi.fn().mockResolvedValue({
|
||||
content: 'Hello! How can I help you?',
|
||||
stopReason: 'end_turn',
|
||||
usage: { inputTokens: 10, outputTokens: 8 },
|
||||
} satisfies ChatResponse),
|
||||
};
|
||||
|
||||
const agent = new NativeAgent({
|
||||
modelClient: mockClient,
|
||||
systemPrompt: 'You are Flynn, a helpful assistant.',
|
||||
});
|
||||
|
||||
const response = await agent.process('Hello');
|
||||
|
||||
expect(response).toBe('Hello! How can I help you?');
|
||||
expect(mockClient.chat).toHaveBeenCalledWith({
|
||||
messages: [{ role: 'user', content: 'Hello' }],
|
||||
system: 'You are Flynn, a helpful assistant.',
|
||||
});
|
||||
const createMockClient = (): ModelClient => ({
|
||||
chat: vi.fn().mockResolvedValue({
|
||||
content: 'Hello!',
|
||||
stopReason: 'end_turn',
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
} satisfies ChatResponse),
|
||||
});
|
||||
|
||||
it('maintains conversation history', async () => {
|
||||
const mockClient: ModelClient = {
|
||||
chat: vi.fn().mockResolvedValue({
|
||||
content: 'Response',
|
||||
stopReason: 'end_turn',
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
} satisfies ChatResponse),
|
||||
};
|
||||
|
||||
it('processes messages and maintains history', async () => {
|
||||
const mockClient = createMockClient();
|
||||
const agent = new NativeAgent({
|
||||
modelClient: mockClient,
|
||||
systemPrompt: 'System',
|
||||
systemPrompt: 'You are helpful.',
|
||||
});
|
||||
|
||||
await agent.process('First message');
|
||||
await agent.process('Second message');
|
||||
const response = await agent.process('Hi');
|
||||
|
||||
expect(mockClient.chat).toHaveBeenLastCalledWith({
|
||||
messages: [
|
||||
{ role: 'user', content: 'First message' },
|
||||
{ role: 'assistant', content: 'Response' },
|
||||
{ role: 'user', content: 'Second message' },
|
||||
],
|
||||
system: 'System',
|
||||
expect(response).toBe('Hello!');
|
||||
expect(mockClient.chat).toHaveBeenCalledWith({
|
||||
messages: [{ role: 'user', content: 'Hi' }],
|
||||
system: 'You are helpful.',
|
||||
});
|
||||
|
||||
const history = agent.getHistory();
|
||||
expect(history).toHaveLength(2);
|
||||
expect(history[0]).toEqual({ role: 'user', content: 'Hi' });
|
||||
expect(history[1]).toEqual({ role: 'assistant', content: 'Hello!' });
|
||||
});
|
||||
|
||||
it('resets conversation history', async () => {
|
||||
const mockClient: ModelClient = {
|
||||
chat: vi.fn().mockResolvedValue({
|
||||
content: 'Response',
|
||||
stopReason: 'end_turn',
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
} satisfies ChatResponse),
|
||||
const mockClient = createMockClient();
|
||||
const agent = new NativeAgent({
|
||||
modelClient: mockClient,
|
||||
systemPrompt: 'You are helpful.',
|
||||
});
|
||||
|
||||
await agent.process('Hi');
|
||||
agent.reset();
|
||||
|
||||
expect(agent.getHistory()).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('uses session when provided', async () => {
|
||||
const mockClient = createMockClient();
|
||||
const mockSession = {
|
||||
id: 'test-session',
|
||||
getHistory: vi.fn().mockReturnValue([]),
|
||||
addMessage: vi.fn(),
|
||||
clear: vi.fn(),
|
||||
};
|
||||
|
||||
const agent = new NativeAgent({
|
||||
modelClient: mockClient,
|
||||
systemPrompt: 'System',
|
||||
systemPrompt: 'You are helpful.',
|
||||
session: mockSession,
|
||||
});
|
||||
|
||||
await agent.process('Message 1');
|
||||
agent.reset();
|
||||
await agent.process('Message 2');
|
||||
await agent.process('Hi');
|
||||
|
||||
expect(mockClient.chat).toHaveBeenLastCalledWith({
|
||||
messages: [{ role: 'user', content: 'Message 2' }],
|
||||
system: 'System',
|
||||
});
|
||||
expect(mockSession.addMessage).toHaveBeenCalledTimes(2);
|
||||
expect(mockSession.addMessage).toHaveBeenNthCalledWith(1, { role: 'user', content: 'Hi' });
|
||||
expect(mockSession.addMessage).toHaveBeenNthCalledWith(2, { role: 'assistant', content: 'Hello!' });
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,59 +1,58 @@
|
||||
import type { ModelClient, Message } from '../../models/types.js';
|
||||
import type { SessionStore } from '../../session/index.js';
|
||||
import type { Session } from '../../session/index.js';
|
||||
|
||||
export interface NativeAgentConfig {
|
||||
modelClient: ModelClient;
|
||||
systemPrompt: string;
|
||||
sessionStore?: SessionStore;
|
||||
sessionId?: string;
|
||||
session?: Session;
|
||||
}
|
||||
|
||||
export class NativeAgent {
|
||||
private modelClient: ModelClient;
|
||||
private systemPrompt: string;
|
||||
private sessionStore?: SessionStore;
|
||||
private sessionId: string;
|
||||
private history: Message[] = [];
|
||||
private session?: Session;
|
||||
private inMemoryHistory: Message[] = [];
|
||||
|
||||
constructor(config: NativeAgentConfig) {
|
||||
this.modelClient = config.modelClient;
|
||||
this.systemPrompt = config.systemPrompt;
|
||||
this.sessionStore = config.sessionStore;
|
||||
this.sessionId = config.sessionId ?? 'default';
|
||||
this.session = config.session;
|
||||
}
|
||||
|
||||
// Load existing history from store
|
||||
if (this.sessionStore) {
|
||||
this.history = this.sessionStore.getMessages(this.sessionId);
|
||||
}
|
||||
private get history(): Message[] {
|
||||
return this.session?.getHistory() ?? [...this.inMemoryHistory];
|
||||
}
|
||||
|
||||
async process(userMessage: string): Promise<string> {
|
||||
const userMsg: Message = { role: 'user', content: userMessage };
|
||||
this.history.push(userMsg);
|
||||
|
||||
if (this.sessionStore) {
|
||||
this.sessionStore.addMessage(this.sessionId, userMsg);
|
||||
if (this.session) {
|
||||
this.session.addMessage(userMsg);
|
||||
} else {
|
||||
this.inMemoryHistory.push(userMsg);
|
||||
}
|
||||
|
||||
const response = await this.modelClient.chat({
|
||||
messages: [...this.history],
|
||||
messages: this.history,
|
||||
system: this.systemPrompt,
|
||||
});
|
||||
|
||||
const assistantMsg: Message = { role: 'assistant', content: response.content };
|
||||
this.history.push(assistantMsg);
|
||||
|
||||
if (this.sessionStore) {
|
||||
this.sessionStore.addMessage(this.sessionId, assistantMsg);
|
||||
if (this.session) {
|
||||
this.session.addMessage(assistantMsg);
|
||||
} else {
|
||||
this.inMemoryHistory.push(assistantMsg);
|
||||
}
|
||||
|
||||
return response.content;
|
||||
}
|
||||
|
||||
reset(): void {
|
||||
this.history = [];
|
||||
if (this.sessionStore) {
|
||||
this.sessionStore.clearSession(this.sessionId);
|
||||
if (this.session) {
|
||||
this.session.clear();
|
||||
} else {
|
||||
this.inMemoryHistory = [];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+12
-9
@@ -4,7 +4,7 @@ import type { Config } from '../config/index.js';
|
||||
import { AnthropicClient, OpenAIClient, OllamaClient, ModelRouter } from '../models/index.js';
|
||||
import { NativeAgent } from '../backends/index.js';
|
||||
import { createTelegramBot } from '../frontends/telegram/index.js';
|
||||
import { SessionStore } from '../session/index.js';
|
||||
import { SessionStore, SessionManager } from '../session/index.js';
|
||||
import { HookEngine } from '../hooks/index.js';
|
||||
import { resolve } from 'path';
|
||||
import { homedir } from 'os';
|
||||
@@ -16,6 +16,7 @@ export interface DaemonContext {
|
||||
bot: Bot;
|
||||
agent: NativeAgent;
|
||||
sessionStore: SessionStore;
|
||||
sessionManager: SessionManager;
|
||||
hookEngine: HookEngine;
|
||||
modelRouter: ModelRouter;
|
||||
}
|
||||
@@ -27,12 +28,10 @@ Keep responses focused and avoid unnecessary verbosity. Use markdown formatting
|
||||
function createModelRouter(config: Config): ModelRouter {
|
||||
const models = config.models;
|
||||
|
||||
// Create default client (required)
|
||||
const defaultClient = new AnthropicClient({
|
||||
model: models.default.model,
|
||||
});
|
||||
|
||||
// Create optional tier clients
|
||||
let fastClient;
|
||||
let complexClient;
|
||||
let localClient;
|
||||
@@ -54,7 +53,6 @@ function createModelRouter(config: Config): ModelRouter {
|
||||
}
|
||||
}
|
||||
|
||||
// Build fallback chain
|
||||
const fallbackChain = [];
|
||||
for (const providerName of models.fallback_chain) {
|
||||
if (providerName === 'openai') {
|
||||
@@ -80,8 +78,10 @@ export async function startDaemon(config: Config): Promise<DaemonContext> {
|
||||
const dataDir = resolve(homedir(), '.local/share/flynn');
|
||||
mkdirSync(dataDir, { recursive: true });
|
||||
|
||||
// Initialize session store
|
||||
// Initialize session store and manager
|
||||
const sessionStore = new SessionStore(resolve(dataDir, 'sessions.db'));
|
||||
const sessionManager = new SessionManager(sessionStore);
|
||||
|
||||
lifecycle.onShutdown(async () => {
|
||||
sessionStore.close();
|
||||
console.log('Session store closed');
|
||||
@@ -93,12 +93,15 @@ export async function startDaemon(config: Config): Promise<DaemonContext> {
|
||||
// Initialize model router
|
||||
const modelRouter = createModelRouter(config);
|
||||
|
||||
// Initialize native agent with session persistence
|
||||
// Get Telegram session
|
||||
const telegramUserId = String(config.telegram.allowed_chat_ids[0]);
|
||||
const session = sessionManager.getSession('telegram', telegramUserId);
|
||||
|
||||
// Initialize native agent with session
|
||||
const agent = new NativeAgent({
|
||||
modelClient: modelRouter,
|
||||
systemPrompt: SYSTEM_PROMPT,
|
||||
sessionStore,
|
||||
sessionId: `telegram-${config.telegram.allowed_chat_ids[0]}`,
|
||||
session,
|
||||
});
|
||||
|
||||
// Initialize Telegram bot with hook engine
|
||||
@@ -136,7 +139,7 @@ export async function startDaemon(config: Config): Promise<DaemonContext> {
|
||||
|
||||
console.log('Flynn daemon started');
|
||||
|
||||
return { config, lifecycle, bot, agent, sessionStore, hookEngine, modelRouter };
|
||||
return { config, lifecycle, bot, agent, sessionStore, sessionManager, hookEngine, modelRouter };
|
||||
}
|
||||
|
||||
export { Lifecycle } from './lifecycle.js';
|
||||
|
||||
Reference in New Issue
Block a user