@falai/agent 0.8.1 → 0.9.0-alpha-2

package/src/adapters/MongoAdapter.ts

@@ -4,13 +4,15 @@
  */
 
 import type {
-  SessionRepository,
+  MessageData,
   MessageRepository,
+  PersistenceAdapter,
   SessionData,
-  MessageData,
+  SessionRepository,
   SessionStatus,
-  PersistenceAdapter,
-} from "../types/persistence";
+  CollectedStateData,
+  CreateSessionData,
+} from "../types";
 
 /**
  * MongoDB collection interface - matches mongodb driver
@@ -98,8 +100,10 @@ export interface MongoAdapterOptions {
  * });
  * ```
  */
-export class MongoAdapter implements PersistenceAdapter {
-  public readonly sessionRepository: SessionRepository;
+export class MongoAdapter<TData = Record<string, unknown>>
+  implements PersistenceAdapter<TData>
+{
+  public readonly sessionRepository: SessionRepository<TData>;
   public readonly messageRepository: MessageRepository;
   private client: MongoClient;
   private db: MongoDatabase;
@@ -111,7 +115,7 @@ export class MongoAdapter implements PersistenceAdapter {
     const sessionCollection = options.collections?.sessions || "agent_sessions";
     const messageCollection = options.collections?.messages || "agent_messages";
 
-    this.sessionRepository = new MongoSessionRepository(
+    this.sessionRepository = new MongoSessionRepository<TData>(
       this.db.collection(sessionCollection)
     );
 
@@ -128,16 +132,18 @@ export class MongoAdapter implements PersistenceAdapter {
 /**
  * MongoDB Session Repository
  */
-class MongoSessionRepository implements SessionRepository {
-  constructor(private collection: MongoCollection<SessionData>) {}
+class MongoSessionRepository<TData = Record<string, unknown>>
+  implements SessionRepository<TData>
+{
+  constructor(private collection: MongoCollection<SessionData<TData>>) {}
 
-  async create(
-    data: Omit<SessionData, "id" | "createdAt" | "updatedAt">
-  ): Promise<SessionData> {
+  async create(data: CreateSessionData<TData>): Promise<SessionData<TData>> {
     const now = new Date();
-    const session: SessionData = {
+    const session: SessionData<TData> = {
       ...data,
-      id: `session_${Date.now()}_${Math.random().toString(36).slice(2)}`,
+      id:
+        data.id ||
+        `session_${Date.now()}_${Math.random().toString(36).slice(2)}`,
       status: data.status || "active",
       messageCount: data.messageCount || 0,
       createdAt: now,
@@ -148,15 +154,18 @@ class MongoSessionRepository implements SessionRepository {
     return session;
   }
 
-  async findById(id: string): Promise<SessionData | null> {
+  async findById(id: string): Promise<SessionData<TData> | null> {
     return await this.collection.findOne({ id });
   }
 
-  async findActiveByUserId(userId: string): Promise<SessionData | null> {
+  async findActiveByUserId(userId: string): Promise<SessionData<TData> | null> {
     return await this.collection.findOne({ userId, status: "active" });
   }
 
-  async findByUserId(userId: string, limit = 100): Promise<SessionData[]> {
+  async findByUserId(
+    userId: string,
+    limit = 100
+  ): Promise<SessionData<TData>[]> {
     return await this.collection
       .find({ userId })
       .sort({ createdAt: -1 })
@@ -166,8 +175,8 @@ class MongoSessionRepository implements SessionRepository {
 
   async update(
     id: string,
-    data: Partial<Omit<SessionData, "id" | "createdAt">>
-  ): Promise<SessionData | null> {
+    data: Partial<Omit<SessionData<TData>, "id" | "createdAt">>
+  ): Promise<SessionData<TData> | null> {
     const result = await this.collection.updateOne(
       { id },
       { $set: { ...data, updatedAt: new Date() } }
@@ -181,7 +190,7 @@ class MongoSessionRepository implements SessionRepository {
     id: string,
     status: SessionStatus,
     completedAt?: Date
-  ): Promise<SessionData | null> {
+  ): Promise<SessionData<TData> | null> {
     const updateData: Record<string, unknown> = {
       status,
       updatedAt: new Date(),
@@ -201,8 +210,8 @@ class MongoSessionRepository implements SessionRepository {
 
   async updateCollectedData(
     id: string,
-    collectedData: Record<string, unknown>
-  ): Promise<SessionData | null> {
+    collectedData: CollectedStateData<TData>
+  ): Promise<SessionData<TData> | null> {
     return await this.update(id, { collectedData });
   }
 
@@ -210,14 +219,14 @@ class MongoSessionRepository implements SessionRepository {
     id: string,
     route?: string,
     step?: string
-  ): Promise<SessionData | null> {
+  ): Promise<SessionData<TData> | null> {
     return await this.update(id, {
       currentRoute: route,
       currentStep: step,
     });
   }
 
-  async incrementMessageCount(id: string): Promise<SessionData | null> {
+  async incrementMessageCount(id: string): Promise<SessionData<TData> | null> {
     const result = await this.collection.updateOne(
       { id },
       {

package/src/adapters/OpenSearchAdapter.ts

@@ -32,13 +32,16 @@
  * ```
  */
 
+import { cloneDeep } from "../utils/clone";
 import type {
   PersistenceAdapter,
   SessionRepository,
   MessageRepository,
   SessionData,
   MessageData,
-} from "../types/persistence.js";
+  CollectedStateData,
+  CreateSessionData,
+} from "../types";
 
 /**
  * OpenSearch Client interface (minimal typing for the official client)
@@ -137,8 +140,10 @@ export interface OpenSearchAdapterOptions {
  * Stores sessions and messages as documents in OpenSearch indices.
  * Compatible with OpenSearch 1.x, 2.x and Elasticsearch 7.x.
  */
-export class OpenSearchAdapter implements PersistenceAdapter {
-  readonly sessionRepository: SessionRepository;
+export class OpenSearchAdapter<TData = Record<string, unknown>>
+  implements PersistenceAdapter<TData>
+{
+  readonly sessionRepository: SessionRepository<TData>;
   readonly messageRepository: MessageRepository;
 
   private readonly client: OpenSearchClient;
@@ -157,7 +162,7 @@ export class OpenSearchAdapter implements PersistenceAdapter {
     this.autoCreateIndices = options.autoCreateIndices ?? true;
     this.refresh = options.refresh ?? false;
 
-    this.sessionRepository = new OpenSearchSessionRepository(
+    this.sessionRepository = new OpenSearchSessionRepository<TData>(
       this.client,
       this.sessionIndex,
       this.refresh
@@ -242,36 +247,41 @@ export class OpenSearchAdapter implements PersistenceAdapter {
 /**
  * OpenSearch-based session repository implementation
  */
-class OpenSearchSessionRepository implements SessionRepository {
+class OpenSearchSessionRepository<TData = Record<string, unknown>>
+  implements SessionRepository<TData>
+{
   constructor(
     private client: OpenSearchClient,
     private index: string,
     private refresh: boolean | "wait_for"
   ) {}
 
-  async create(
-    data: Omit<SessionData, "id" | "createdAt">
-  ): Promise<SessionData> {
-    const id = `sess_${Date.now()}_${Math.random().toString(36).slice(2)}`;
+  async create(data: CreateSessionData<TData>): Promise<SessionData<TData>> {
+    const id =
+      data.id || `session_${Date.now()}_${Math.random().toString(36).slice(2)}`;
     const now = new Date();
 
-    const session: SessionData = {
-      id,
+    const session: SessionData<TData> = {
       ...data,
+      id,
+      status: data.status || "active",
+      messageCount: data.messageCount || 0,
       createdAt: now,
+      updatedAt: now,
     };
+    const clonedSession = cloneDeep(session);
 
     await this.client.index({
       index: this.index,
       id,
-      body: this.serializeSession(session),
-      refresh: this.refresh,
+      body: clonedSession as unknown as Record<string, unknown>,
+      refresh: true, // Refresh to make the document immediately available for search
     });
 
     return session;
   }
 
-  async findById(id: string): Promise<SessionData | null> {
+  async findById(id: string): Promise<SessionData<TData> | null> {
     try {
       const response = await this.client.get({
         index: this.index,
@@ -287,7 +297,7 @@ class OpenSearchSessionRepository implements SessionRepository {
     }
   }
 
-  async findActiveByUserId(userId: string): Promise<SessionData | null> {
+  async findActiveByUserId(userId: string): Promise<SessionData<TData> | null> {
     const response = await this.client.search({
       index: this.index,
       body: {
@@ -309,7 +319,10 @@ class OpenSearchSessionRepository implements SessionRepository {
     return this.deserializeSession(hits[0]._source);
   }
 
-  async findByUserId(userId: string, limit = 100): Promise<SessionData[]> {
+  async findByUserId(
+    userId: string,
+    limit = 100
+  ): Promise<SessionData<TData>[]> {
     const response = await this.client.search({
       index: this.index,
       body: {
@@ -328,8 +341,8 @@ class OpenSearchSessionRepository implements SessionRepository {
 
   async update(
     id: string,
-    updates: Partial<Omit<SessionData, "id" | "createdAt">>
-  ): Promise<SessionData | null> {
+    updates: Partial<Omit<SessionData<TData>, "id" | "createdAt">>
+  ): Promise<SessionData<TData> | null> {
     const doc: Record<string, unknown> = {
       ...updates,
       updatedAt: new Date().toISOString(),
@@ -355,9 +368,9 @@ class OpenSearchSessionRepository implements SessionRepository {
 
   async updateStatus(
     id: string,
-    status: SessionData["status"],
+    status: SessionData<TData>["status"],
     completedAt?: Date
-  ): Promise<SessionData | null> {
+  ): Promise<SessionData<TData> | null> {
     const doc: Record<string, unknown> = {
       status,
       updatedAt: new Date().toISOString(),
@@ -379,8 +392,8 @@ class OpenSearchSessionRepository implements SessionRepository {
 
   async updateCollectedData(
     id: string,
-    collectedData: Record<string, unknown>
-  ): Promise<SessionData | null> {
+    collectedData: CollectedStateData<TData>
+  ): Promise<SessionData<TData> | null> {
     await this.client.update({
       index: this.index,
       id,
@@ -400,7 +413,7 @@ class OpenSearchSessionRepository implements SessionRepository {
     id: string,
     route?: string,
     step?: string
-  ): Promise<SessionData | null> {
+  ): Promise<SessionData<TData> | null> {
     const doc: Record<string, unknown> = {
       updatedAt: new Date().toISOString(),
     };
@@ -422,7 +435,7 @@ class OpenSearchSessionRepository implements SessionRepository {
     return await this.findById(id);
   }
 
-  async incrementMessageCount(id: string): Promise<SessionData | null> {
+  async incrementMessageCount(id: string): Promise<SessionData<TData> | null> {
     const session = await this.findById(id);
     if (!session) {
       return null;
@@ -462,25 +475,16 @@ class OpenSearchSessionRepository implements SessionRepository {
     }
   }
 
-  private serializeSession(session: SessionData): Record<string, unknown> {
-    return {
-      ...session,
-      createdAt: session.createdAt.toISOString(),
-      updatedAt: session.updatedAt?.toISOString(),
-      completedAt: session.completedAt?.toISOString(),
-    };
-  }
-
-  private deserializeSession(doc: Record<string, unknown>): SessionData {
+  private deserializeSession(doc: Record<string, unknown>): SessionData<TData> {
     return {
       id: doc.id as string,
       userId: doc.userId as string | undefined,
       agentName: doc.agentName as string | undefined,
-      status: doc.status as SessionData["status"],
+      status: doc.status as SessionData<TData>["status"],
       currentRoute: doc.currentRoute as string | undefined,
       currentStep: doc.currentStep as string | undefined,
-      collectedData: doc.collectedData as Record<string, unknown> | undefined,
-      messageCount: doc.messageCount as number | undefined,
+      collectedData: doc.collectedData as CollectedStateData<TData> | undefined,
+      messageCount: (doc.messageCount as number) || 0,
       createdAt: new Date(doc.createdAt as string),
       updatedAt: new Date(doc.updatedAt as string),
       lastMessageAt: doc.lastMessageAt

package/src/adapters/PostgreSQLAdapter.ts

@@ -4,13 +4,14 @@
  */
 
 import type {
+  CollectedStateData,
   SessionRepository,
   MessageRepository,
   SessionData,
   MessageData,
   SessionStatus,
   PersistenceAdapter,
-} from "../types/persistence";
+} from "../types";
 
 /**
  * PostgreSQL query result interface
@@ -76,8 +77,10 @@ export interface PostgreSQLAdapterOptions {
  * });
  * ```
  */
-export class PostgreSQLAdapter implements PersistenceAdapter {
-  public readonly sessionRepository: SessionRepository;
+export class PostgreSQLAdapter<TData = Record<string, unknown>>
+  implements PersistenceAdapter<TData>
+{
+  public readonly sessionRepository: SessionRepository<TData>;
   public readonly messageRepository: MessageRepository;
   private client: PgClient;
 
@@ -87,7 +90,7 @@ export class PostgreSQLAdapter implements PersistenceAdapter {
     const sessionTable = options.tables?.sessions || "agent_sessions";
     const messageTable = options.tables?.messages || "agent_messages";
 
-    this.sessionRepository = new PostgreSQLSessionRepository(
+    this.sessionRepository = new PostgreSQLSessionRepository<TData>(
       this.client,
       sessionTable
     );
@@ -155,16 +158,21 @@ export class PostgreSQLAdapter implements PersistenceAdapter {
 /**
  * PostgreSQL Session Repository
  */
-class PostgreSQLSessionRepository implements SessionRepository {
+class PostgreSQLSessionRepository<TData = Record<string, unknown>>
+  implements SessionRepository<TData>
+{
   constructor(private client: PgClient, private tableName: string) {}
 
   async create(
-    data: Omit<SessionData, "id" | "createdAt" | "updatedAt">
-  ): Promise<SessionData> {
-    const id = `session_${Date.now()}_${Math.random().toString(36).slice(2)}`;
+    data: Omit<SessionData<TData>, "createdAt" | "updatedAt"> & {
+      id?: string;
+    }
+  ): Promise<SessionData<TData>> {
+    const id =
+      data.id || `session_${Date.now()}_${Math.random().toString(36).slice(2)}`;
     const now = new Date();
 
-    const result = await this.client.query<SessionData>(
+    const result = await this.client.query<SessionData<TData>>(
       `INSERT INTO ${this.tableName} 
        (id, user_id, agent_name, status, collected_data, message_count, created_at, updated_at)
        VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
@@ -184,8 +192,8 @@ class PostgreSQLSessionRepository implements SessionRepository {
     return result.rows[0];
   }
 
-  async findById(id: string): Promise<SessionData | null> {
-    const result = await this.client.query<SessionData>(
+  async findById(id: string): Promise<SessionData<TData> | null> {
+    const result = await this.client.query<SessionData<TData>>(
       `SELECT * FROM ${this.tableName} WHERE id = $1`,
       [id]
     );
@@ -193,8 +201,8 @@ class PostgreSQLSessionRepository implements SessionRepository {
     return result.rows[0] || null;
   }
 
-  async findActiveByUserId(userId: string): Promise<SessionData | null> {
-    const result = await this.client.query<SessionData>(
+  async findActiveByUserId(userId: string): Promise<SessionData<TData> | null> {
+    const result = await this.client.query<SessionData<TData>>(
       `SELECT * FROM ${this.tableName} 
        WHERE user_id = $1 AND status = 'active'
        ORDER BY created_at DESC
@@ -205,8 +213,11 @@ class PostgreSQLSessionRepository implements SessionRepository {
     return result.rows[0] || null;
   }
 
-  async findByUserId(userId: string, limit = 100): Promise<SessionData[]> {
-    const result = await this.client.query<SessionData>(
+  async findByUserId(
+    userId: string,
+    limit = 100
+  ): Promise<SessionData<TData>[]> {
+    const result = await this.client.query<SessionData<TData>>(
       `SELECT * FROM ${this.tableName} 
        WHERE user_id = $1
        ORDER BY created_at DESC
@@ -219,8 +230,8 @@ class PostgreSQLSessionRepository implements SessionRepository {
 
   async update(
     id: string,
-    data: Partial<Omit<SessionData, "id" | "createdAt">>
-  ): Promise<SessionData | null> {
+    data: Partial<Omit<SessionData<TData>, "id" | "createdAt">>
+  ): Promise<SessionData<TData> | null> {
     const fields: string[] = [];
     const values: unknown[] = [];
     let paramIndex = 1;
@@ -259,7 +270,7 @@ class PostgreSQLSessionRepository implements SessionRepository {
 
     values.push(id);
 
-    const result = await this.client.query<SessionData>(
+    const result = await this.client.query<SessionData<TData>>(
       `UPDATE ${this.tableName} 
        SET ${fields.join(", ")}
        WHERE id = $${paramIndex}
@@ -274,14 +285,14 @@ class PostgreSQLSessionRepository implements SessionRepository {
     id: string,
     status: SessionStatus,
     completedAt?: Date
-  ): Promise<SessionData | null> {
+  ): Promise<SessionData<TData> | null> {
     return await this.update(id, { status, completedAt });
   }
 
   async updateCollectedData(
     id: string,
-    collectedData: Record<string, unknown>
-  ): Promise<SessionData | null> {
+    collectedData: CollectedStateData<TData>
+  ): Promise<SessionData<TData> | null> {
     return await this.update(id, { collectedData });
   }
 
@@ -289,15 +300,15 @@ class PostgreSQLSessionRepository implements SessionRepository {
     id: string,
     route?: string,
     step?: string
-  ): Promise<SessionData | null> {
+  ): Promise<SessionData<TData> | null> {
     return await this.update(id, {
       currentRoute: route,
       currentStep: step,
     });
   }
 
-  async incrementMessageCount(id: string): Promise<SessionData | null> {
-    const result = await this.client.query<SessionData>(
+  async incrementMessageCount(id: string): Promise<SessionData<TData> | null> {
+    const result = await this.client.query<SessionData<TData>>(
       `UPDATE ${this.tableName}
        SET message_count = message_count + 1,
            last_message_at = NOW(),