zeitlich 0.1.1 → 0.2.0

package/src/lib/state-manager.ts

@@ -0,0 +1,268 @@
+import { defineQuery, setHandler } from "@temporalio/workflow";
+import {
+  type AgentStatus,
+  type BaseAgentState,
+  type WorkflowTask,
+  isTerminalStatus,
+} from "./types";
+import type { ToolDefinition } from "./tool-router";
+
+/**
+ * JSON primitive types that Temporal can serialize
+ */
+export type JsonPrimitive = string | number | boolean | null | undefined;
+
+/**
+ * JSON-serializable value (recursive type for Temporal compatibility)
+ */
+export type JsonValue =
+  | JsonPrimitive
+  | JsonValue[]
+  | { [key: string]: JsonValue };
+
+/**
+ * Type constraint ensuring T only contains JSON-serializable values.
+ * Use this for custom state to ensure Temporal workflow compatibility.
+ *
+ * Allows: primitives, arrays, plain objects, and JsonValue
+ * Rejects: functions, symbols, undefined, class instances with methods
+ */
+export type JsonSerializable<T> = {
+  [K in keyof T]: T[K] extends JsonValue
+    ? T[K]
+    : T[K] extends JsonPrimitive
+      ? T[K]
+      : T[K] extends (infer U)[]
+        ? U extends JsonValue
+          ? T[K]
+          : JsonSerializable<U>[]
+        : T[K] extends object
+          ? JsonSerializable<T[K]>
+          : never;
+};
+
+/**
+ * Full state type combining base state with custom state
+ */
+export type AgentState<TCustom extends JsonSerializable<TCustom>> =
+  BaseAgentState & TCustom;
+
+/**
+ * Agent state manager interface
+ * Note: Temporal handlers must be set up in the workflow file due to
+ * Temporal's workflow isolation requirements. This manager provides
+ * the state and helpers needed for those handlers.
+ */
+export interface AgentStateManager<TCustom extends JsonSerializable<TCustom>> {
+  /** Get current status */
+  getStatus(): AgentStatus;
+  /** Check if agent is running */
+  isRunning(): boolean;
+  /** Check if agent is in terminal state */
+  isTerminal(): boolean;
+  /** Get current state version */
+  getVersion(): number;
+
+  /** Set status to RUNNING */
+  run(): void;
+  /** Set status to WAITING_FOR_INPUT */
+  waitForInput(): void;
+  /** Set status to COMPLETED */
+  complete(): void;
+  /** Set status to FAILED */
+  fail(): void;
+  /** Set status to CANCELLED */
+  cancel(): void;
+
+  /** Increment state version (call after state changes) */
+  incrementVersion(): void;
+
+  /** Increment turns (call after each turn) */
+  incrementTurns(): void;
+
+  /** Get current turns */
+  getTurns(): number;
+
+  /** Get a custom state value by key */
+  get<K extends keyof TCustom>(key: K): TCustom[K];
+
+  /** Set a custom state value by key */
+  set<K extends keyof TCustom>(key: K, value: TCustom[K]): void;
+
+  /** Get full state for query handler */
+  getCurrentState(): AgentState<TCustom>;
+
+  /** Check if should return from waitForStateChange */
+  shouldReturnFromWait(lastKnownVersion: number): boolean;
+
+  // Task management methods
+  /** Get all tasks */
+  getTasks(): WorkflowTask[];
+  /** Get a task by ID */
+  getTask(id: string): WorkflowTask | undefined;
+  /** Add or update a task */
+  setTask(task: WorkflowTask): void;
+  /** Delete a task by ID */
+  deleteTask(id: string): boolean;
+
+  /** Set the tools */
+  setTools(newTools: ToolDefinition[]): void;
+}
+
+export const getStateQuery = defineQuery<BaseAgentState>("getState");
+
+/**
+ * Creates an agent state manager for tracking workflow state.
+ *
+ * @param initialState - Optional initial values for base and custom state
+ *   Base state defaults: status="RUNNING", version=0, turns=0, tasks=empty, fileTree=[]
+ *
+ * Note: Due to Temporal's workflow isolation, handlers must be set up
+ * in the workflow file using defineQuery/defineUpdate and setHandler.
+ * This manager provides the state and logic needed for those handlers.
+ */
+export function createAgentStateManager<
+  TCustom extends JsonSerializable<TCustom> = Record<string, never>,
+>(
+  initialState?: Partial<BaseAgentState> & TCustom
+): AgentStateManager<TCustom> {
+  // Default state (BaseAgentState fields)
+  let status: AgentStatus = initialState?.status ?? "RUNNING";
+  let version = initialState?.version ?? 0;
+  let turns = initialState?.turns ?? 0;
+  let tools = initialState?.tools ?? [];
+
+  // Tasks state
+  const tasks = new Map<string, WorkflowTask>(initialState?.tasks);
+
+  // Custom state - extract only custom fields (exclude base state keys)
+  const {
+    status: _,
+    version: __,
+    turns: ___,
+    tasks: ____,
+    tools: _____,
+    ...custom
+  } = initialState ?? {};
+  const customState = custom as TCustom;
+
+  function buildState(): AgentState<TCustom> {
+    return {
+      status,
+      version,
+      turns,
+      tools,
+      ...customState,
+    } as AgentState<TCustom>;
+  }
+
+  setHandler(getStateQuery, () => {
+    return buildState();
+  });
+
+  return {
+    getStatus(): AgentStatus {
+      return status;
+    },
+
+    isRunning(): boolean {
+      return status === "RUNNING";
+    },
+
+    isTerminal(): boolean {
+      return isTerminalStatus(status);
+    },
+
+    getTurns(): number {
+      return turns;
+    },
+
+    getVersion(): number {
+      return version;
+    },
+
+    run(): void {
+      status = "RUNNING";
+      version++;
+    },
+
+    waitForInput(): void {
+      status = "WAITING_FOR_INPUT";
+      version++;
+    },
+
+    complete(): void {
+      status = "COMPLETED";
+      version++;
+    },
+
+    fail(): void {
+      status = "FAILED";
+      version++;
+    },
+
+    cancel(): void {
+      status = "CANCELLED";
+      version++;
+    },
+
+    incrementVersion(): void {
+      version++;
+    },
+
+    incrementTurns(): void {
+      turns++;
+    },
+
+    get<K extends keyof TCustom>(key: K): TCustom[K] {
+      return customState[key];
+    },
+
+    set<K extends keyof TCustom>(key: K, value: TCustom[K]): void {
+      customState[key] = value;
+      version++;
+    },
+
+    getCurrentState(): AgentState<TCustom> {
+      return buildState();
+    },
+
+    shouldReturnFromWait(lastKnownVersion: number): boolean {
+      return version > lastKnownVersion || isTerminalStatus(status);
+    },
+
+    getTasks(): WorkflowTask[] {
+      return Array.from(tasks.values());
+    },
+
+    getTask(id: string): WorkflowTask | undefined {
+      return tasks.get(id);
+    },
+
+    setTask(task: WorkflowTask): void {
+      tasks.set(task.id, task);
+      version++;
+    },
+
+    setTools(newTools: ToolDefinition[]): void {
+      tools = newTools;
+    },
+
+    deleteTask(id: string): boolean {
+      const deleted = tasks.delete(id);
+      if (deleted) {
+        version++;
+      }
+      return deleted;
+    },
+  };
+}
+
+/**
+ * Handler names used across agents
+ */
+export const AGENT_HANDLER_NAMES = {
+  getAgentState: "getAgentState",
+  waitForStateChange: "waitForStateChange",
+  addMessage: "addMessage",
+} as const;

package/src/lib/thread-manager.ts

@@ -0,0 +1,169 @@
+import type Redis from "ioredis";
+
+import {
+  type $InferMessageContent,
+  AIMessage,
+  HumanMessage,
+  type MessageContent,
+  type MessageStructure,
+  type StoredMessage,
+  SystemMessage,
+  ToolMessage,
+} from "@langchain/core/messages";
+import { v4 as uuidv4 } from "uuid";
+
+const THREAD_TTL_SECONDS = 60 * 60 * 24 * 90; // 90 days
+
+function getThreadKey(threadId: string, key: string): string {
+  return `thread:${threadId}:${key}`;
+}
+
+/**
+ * Content for a tool message response.
+ * Can be a simple string or complex content parts (text, images, cache points, etc.)
+ */
+export type ToolMessageContent = $InferMessageContent<MessageStructure, "tool">;
+
+export interface ThreadManagerConfig {
+  redis: Redis;
+  threadId: string;
+  /** Thread key, defaults to 'messages' */
+  key?: string;
+}
+
+export interface ThreadManager {
+  /** Append a system message to the thread */
+  appendSystemMessage(content: string): Promise<void>;
+  /** Initialize an empty thread */
+  initialize(): Promise<void>;
+  /** Load all messages from the thread */
+  load(): Promise<StoredMessage[]>;
+  /** Append messages to the thread */
+  append(messages: StoredMessage[]): Promise<void>;
+  /** Delete the thread */
+  delete(): Promise<void>;
+
+  /** Create a SystemMessage (returns StoredMessage for storage) */
+  createSystemMessage(content: string): StoredMessage;
+
+  /** Create a HumanMessage (returns StoredMessage for storage) */
+  createHumanMessage(content: string | MessageContent): StoredMessage;
+
+  /** Create an AIMessage with optional additional kwargs */
+  createAIMessage(
+    content: string | MessageContent,
+    kwargs?: { header?: string; options?: string[]; multiSelect?: boolean }
+  ): StoredMessage;
+
+  /** Create a ToolMessage */
+  createToolMessage(
+    content: ToolMessageContent,
+    toolCallId: string
+  ): StoredMessage;
+
+  /** Create and append a HumanMessage */
+  appendHumanMessage(content: string | MessageContent): Promise<void>;
+  /** Create and append a ToolMessage */
+  appendToolMessage(
+    content: ToolMessageContent,
+    toolCallId: string
+  ): Promise<void>;
+  /** Create and append an AIMessage */
+  appendAIMessage(content: string | MessageContent): Promise<void>;
+}
+
+/**
+ * Creates a thread manager for handling conversation state in Redis.
+ */
+export function createThreadManager(
+  config: ThreadManagerConfig
+): ThreadManager {
+  const { redis, threadId, key = "messages" } = config;
+  const redisKey = getThreadKey(threadId, key);
+
+  return {
+    async initialize(): Promise<void> {
+      await redis.del(redisKey);
+    },
+
+    async load(): Promise<StoredMessage[]> {
+      const data = await redis.lrange(redisKey, 0, -1);
+      return data.map((item) => JSON.parse(item) as StoredMessage);
+    },
+
+    async append(messages: StoredMessage[]): Promise<void> {
+      if (messages.length > 0) {
+        await redis.rpush(redisKey, ...messages.map((m) => JSON.stringify(m)));
+        await redis.expire(redisKey, THREAD_TTL_SECONDS);
+      }
+    },
+
+    async delete(): Promise<void> {
+      await redis.del(redisKey);
+    },
+
+    createHumanMessage(content: string | MessageContent): StoredMessage {
+      return new HumanMessage({
+        id: uuidv4(),
+        content: content as string,
+      }).toDict();
+    },
+
+    createAIMessage(
+      content: string,
+      kwargs?: { header?: string; options?: string[]; multiSelect?: boolean }
+    ): StoredMessage {
+      return new AIMessage({
+        id: uuidv4(),
+        content,
+        additional_kwargs: kwargs
+          ? {
+              header: kwargs.header,
+              options: kwargs.options,
+              multiSelect: kwargs.multiSelect,
+            }
+          : undefined,
+      }).toDict();
+    },
+
+    createToolMessage(
+      content: ToolMessageContent,
+      toolCallId: string
+    ): StoredMessage {
+      return new ToolMessage({
+        // Cast needed due to langchain type compatibility
+        content: content as MessageContent,
+        tool_call_id: toolCallId,
+      }).toDict();
+    },
+
+    createSystemMessage(content: string): StoredMessage {
+      return new SystemMessage({
+        content,
+      }).toDict();
+    },
+
+    async appendSystemMessage(content: string): Promise<void> {
+      const message = this.createSystemMessage(content);
+      await this.append([message]);
+    },
+
+    async appendHumanMessage(content: string | MessageContent): Promise<void> {
+      const message = this.createHumanMessage(content);
+      await this.append([message]);
+    },
+
+    async appendToolMessage(
+      content: ToolMessageContent,
+      toolCallId: string
+    ): Promise<void> {
+      const message = this.createToolMessage(content, toolCallId);
+      await this.append([message]);
+    },
+
+    async appendAIMessage(content: string | MessageContent): Promise<void> {
+      const message = this.createAIMessage(content);
+      await this.append([message]);
+    },
+  };
+}