Chat Flow

// Frontend request (POST /memories/chat)
ChatRequest {
  question: string (required)
  chat_id: string (optional — omitted for new conversations)
  verbose: boolean (optional — returns trace when true)
}

// Dispatcher response (synchronous, returned within ~1-2s)
DispatcherResponse {
  chat_id: string
  message_id: string   // placeholder message ID for polling
  status: "thinking"
}

// ChatMessage — returned by ChatMessagesFunction and pollForMessage
ChatMessage {
  message_id: string (ULID)
  chat_id: string
  role: "user" | "assistant"
  content: string
  ready: boolean        // false while GPU is processing; true when answer is saved
  gpu_pending: boolean  // true when message is queued for GPU retry
  created_at: string (ISO 8601)
}

GetChatMessagesRequest {
  chat_id: string (required)
  message_id: string (optional — filter to a single message for polling)
  cursor: string | null (optional — DynamoDB ExclusiveStartKey)
  limit: number (optional, default 20)
}

GetChatMessagesResponse {
  messages: ChatMessage[]
  next_cursor: string | null
}

// === FRONTEND (chat.tsx) ===

askQuestion(question):
  append user message bubble
  setIsLoading(true)
  { chat_id, message_id, status, answer, trace } = await chatWithMemory(question, chatId, isDebug)
  setChatId(chat_id)

  if status === "thinking" and message_id:
    // Async path: dispatcher returned handle; poll for ready=true
    append placeholder assistant bubble (id=message_id, content="")
    readyMsg = await pollForMessage({ chat_id, message_id })  // 500ms interval, 60s max
    replace placeholder bubble content with readyMsg.content
    invalidate ["chats", "feed"]

  else if answer is defined:
    // Legacy/verbose path: answer returned directly
    append assistant bubble with answer (+ trace if present)
    invalidate ["chats", "feed"]

  catch GpuUnavailableError:
    // GPU timed out and retry was enqueued backend-side
    append assistant bubble: "Starting up the GPU, we will notify you when we have an answer"
    // push notification will arrive when ChatGpuRetryFunction resolves
  catch PollTimeoutError:
    append assistant bubble: "Sorry, I couldn't process that. Please try again."

// chatWithMemory.ts
POST /memories/chat { question, chat_id?, verbose? }
  timeout: 60_000 ms (client-side axios timeout)
  return { chat_id, message_id, status?, answer?, trace? }

// pollForMessage.ts — polls POST /chats/messages?message_id=<id>
pollForMessage({ chat_id, message_id, intervalMs=500, maxWaitMs=60000 }):
  deadline = now + maxWaitMs
  loop:
    { messages } = getChatMessages({ chat_id, message_id })
    msg = messages[0]
    if msg.ready === true: return msg
    if now >= deadline: throw PollTimeoutError
    sleep(intervalMs)
  // NOTE: if the GPU retry Lambda resolves the message after PollTimeoutError,
  // the push notification will alert the user and they can re-open the chat
  // to see the answer loaded via the history fetch.

// === BACKEND DISPATCHER (lambda_handler → _dispatch_async) ===

lambda_handler(event):
  if event._async_worker:
    → run worker path (see below)
  else:
    → run _dispatch_async() (dispatcher path)

_dispatch_async(payload, auth_context):
  validate question != "" else raise ValueError
  validate user_id in auth_context else raise ValueError
  repo = ChatRepository(MESSAGES_TABLE_NAME, CHATS_TABLE_NAME)
  if chat_id is None:
    chat_id = repo.create_chat(user_id, title=question[:120])
  else:
    verify repo.chat_owned_by(user_id, chat_id) else raise Forbidden
  async_payload = { ...payload, _async_worker: True, _user_id: user_id, chat_id }
  lambda.invoke(FunctionName=self, InvocationType="Event", Payload=async_payload)
  repo.save_message(chat_id, "user", question)  // saved AFTER invoke succeeds
  return { chat_id, status: "thinking" }

// === ASYNC WORKER PATH ===

lambda_handler(event) when event._async_worker == True:
  auth = { user_id: event._user_id }
  result = implementation(payload, auth, _skip_user_message_save=True)
  return { statusCode: 200, body: result }

implementation(payload, auth, _skip_user_message_save=False):
  // 1. DB init
  configure PostgreSQL from SSM secrets
  orm.initDb()

  // 2. DynamoDB setup
  repo = ChatRepository(...)
  if not _skip_user_message_save:
    repo.save_message(chat_id, "user", question)

  // 3. GPU resolution
  gpu_instance_id = normalize(GPU_INSTANCE_ID env var)
  if not gpu_instance_id:
    → find running by tag
    → find stopped by tag and start_instances
    → launch from template with type fallback
  else:
    gpu_worker_url = _resolve_gpu_url(instance_id, port)

  // 4. GPU health polling
  if gpu_worker_url:
    is_healthy, elapsed = wait_for_gpu_health(url, GPU_MAX_WAIT_MS=15000, initial=2000)
    if not is_healthy: gpu_worker_url = None

  // 5. Fallback
  if not gpu_worker_url:
    if message_id:  // async dispatcher path
      repo.mark_gpu_pending(chat_id, message_id)  // sets gpu_pending=True, ready=False, content=""
      _enqueue_gpu_retry(chat_id, message_id, user_id, question, retry_count=0, delay_seconds=120)
      return { status: "gpu_retry_enqueued", chat_id, message_id }
      // chatWithMemory.ts raises GpuUnavailableError on this response
    else:  // legacy synchronous path (no message_id)
      repo.save_message(chat_id, "assistant", "GPU starting up…")
      repo.touch_chat(user_id, chat_id)
      return fallback

  // 6. Answer
  result = handle_chat(user_id, question, chat_id, gpu_worker_url)
  return result

handle_chat(user_id, question, chat_id, gpu_worker_url):
  vlm2vec = VLM2VecClient(gpu_worker_url)
  reasoning_llm = GPULLMClient(vlm2vec)
  embedder = TextEmbedder(vlm2vec)
  episodic_graphs = load_episodic_graphs(user_id)
  semantic_graph = load_semantic_graph(user_id)
  agent = ReasoningAgent(reasoning_llm, response_llm, { episodic, semantic, visual })
  answer, trace = agent.answer(question)
  repo.save_message(chat_id, "assistant", answer)
  repo.touch_chat(user_id, chat_id)
  return { answer, chat_id, trace? }

// SQS message body (ChatGpuRetryQueue)
GpuRetryRecord {
  chat_id: string
  message_id: string
  user_id: string
  question: string
  retry_count: number   // 0-based; compared against GPU_RETRY_MAX_ATTEMPTS (default 15)
}

retry_handler(event):
  for record in event.Records:
    body = parse(record.body)
    { chat_id, message_id, user_id, question, retry_count } = body

    if retry_count >= GPU_RETRY_MAX_ATTEMPTS (default 15):
      repo.update_message_ready(chat_id, message_id, failure_message)
      push_token = get_push_token(user_id)  // reads Sessions table
      if push_token: send_push(push_token, "Couldn't answer", failure_body)
      return

    // Resolve running GPU by tag
    running = _find_running_gpu_instance_by_tag(region)
    if not running:
      _re_enqueue(chat_id, message_id, user_id, question, retry_count)
      return

    gpu_worker_url = http://{running.PrivateIpAddress}:{GPU_WORKER_PORT}
    is_healthy, _ = wait_for_gpu_health(gpu_worker_url, GPU_RETRY_WAIT_MS=240000, initial_delay_ms=5000)

    if not is_healthy:
      _re_enqueue(chat_id, message_id, user_id, question, retry_count)
      return

    // GPU healthy: run full reasoning pipeline
    handle_chat(user_id, question, chat_id, gpu_worker_url, message_id=message_id)
    // handle_chat calls update_message_ready internally (ready=True, content=answer)

    push_token = get_push_token(user_id)
    if push_token:
      send_push(push_token, "Your answer is ready", question[:80])

// PUT /users/push-token request body
PushTokenRequest {
  push_token: string (required, non-empty)
}

// PUT /users/push-token response
PushTokenResponse {
  success: true
}

// _layout.tsx — runs after successful auth
useEffect([user]):
  if user:
    token = await registerForPushNotifications()
    if token: savePushToken(token)  // PUT /users/push-token

// push-notifications.ts
registerForPushNotifications():
  if not Device.isDevice: return null
  { status } = getPermissionsAsync()
  if status != "granted": { status } = requestPermissionsAsync()
  if status != "granted": return null
  projectId = EXPO_PUBLIC_EAS_PROJECT_ID
  if not projectId: return null
  token = Notifications.getExpoPushTokenAsync({ projectId })
  return token.data  // "ExponentPushToken[...]"

// UsersPushTokenFunction (PUT /users/push-token)
handle_push_token(payload, auth_context):
  push_token = payload.push_token.strip()
  if not push_token: raise InvalidInputError
  user_id = auth_context.user_id
  Sessions.update_item(Key={sessionId: user_id}, SET push_token=:t, updated_at=:u)
  return { success: True }

// GPU resolution when instance_id is set:
_resolve_gpu_url(instance_id, port, launch_template_id, region):
  state = ec2.describe_instances([instance_id]).State.Name
  if state in (stopped, stopping):
    running = _find_running_gpu_instance_by_tag()
    if running: adopt running instance, update SSM, return URL
    if state == stopped:
      try: ec2.start_instances([instance_id])
      except InsufficientInstanceCapacity:
        _update_ssm_gpu_instance_id("none")
        if launch_template_id: _launch_gpu_with_type_fallback(...)
    return None
  if state == terminated:
    running = _find_running_gpu_instance_by_tag()
    if running: adopt running instance, update SSM, return URL
    if launch_template_id: _launch_gpu_with_type_fallback(...)
    return None
  ip = instance.PrivateIpAddress or instance.PublicIpAddress
  return f"http://{ip}:{port}"

// GPU resolution when instance_id is None (stale SSM / first deploy):
  1. _find_gpu_instance_by_tag(state=running) → adopt if found
  2. _find_gpu_instance_by_tag(state=stopped) → start_instances if found
  3. _launch_gpu_with_type_fallback(launch_template_id) if still None

_launch_gpu_with_type_fallback(launch_template_id):
  for type in GPU_FALLBACK_INSTANCE_TYPES:
    try:
      ami = _gpu_ami_for_type(type)  // per-type AMI override from GPU_INSTANCE_TYPE_AMIS
      resp = ec2.run_instances(LaunchTemplate, InstanceType=type, Tags=[Name=encache-gpu-worker])
      _update_ssm_gpu_instance_id(resp.InstanceId)
      return resp.InstanceId
    except InsufficientInstanceCapacity:
      continue  // try next type
    except other ClientError:
      return None  // surface in logs, abort
  return None  // all types exhausted

wait_for_gpu_health(gpu_url, max_wait_ms=15000, initial_delay_ms=2000, lambda_context):
  start = time.time()
  retry_delay_ms = initial_delay_ms
  retry_count = 0
  lambda_timeout_guard_ms = 10000  // reserve 10s for response assembly

  loop:
    is_healthy = GET {gpu_url}/health (timeout=3s) == 200
    if is_healthy: return (True, elapsed_ms)
    elapsed_ms = (time.time() - start) * 1000
    if elapsed_ms >= max_wait_ms: return (False, elapsed_ms)  // max-wait exhausted
    if lambda_remaining_ms - elapsed_ms <= 10000: return (False, elapsed_ms)  // Lambda timeout guard
    sleep(retry_delay_ms / 1000)
    retry_delay_ms = min(retry_delay_ms * 2, 60000)  // exponential backoff, cap 60s

// GPU worker /health endpoint (server.py):
GET /health:
  if _embed_model is None or _caption_model is None:
    return HTTP 503 (models still loading)
  return { status: "ready", model: "VLM2Vec-2B+Qwen2-VL-7B", device: "cuda:0" }

// Backend: ChatMessagesFunction
implementation(payload, auth_context):
  validate user_id, chat_id
  repo.get_messages(user_id, chat_id, cursor, limit=20)

// ChatRepository.get_messages:
  verify chat_owned_by(user_id, chat_id) else Forbidden
  DynamoDB.query(
    KeyConditionExpression: "chat_id = :cid",
    ScanIndexForward: False,   // newest first
    Limit: limit,
    ExclusiveStartKey: cursor  // if paginating
  )
  return (items, LastEvaluatedKey)

// Frontend: chat.tsx history load
useEffect (on mount, if params.chatId):
  { messages } = await getChatMessages({ chat_id: params.chatId })
  setMessages([...messages].reverse())  // reverse to chronological order

ReasoningAgent.answer(question):
  rounds = [], accumulated_context = [], trace = []
  for round in range(MAX_ROUNDS=5):
    messages = _build_reasoning_messages(question, rounds)
    raw = reasoning_llm.call(model=gpu, json_object)
    action = parse_agent_action(raw)  // normalizes {decision/action} field
    if action.action == "ANSWER": break
    if action.action == "SEARCH":
      results = retrievers[action.memory_type](action.query)
      // retrievers: "episodic", "semantic", "visual"
      rounds.append({ memory_type, query, results })
      accumulated_context.extend(results)
      trace.append({ round, decision: SEARCH, memory_type, agent_query, results, result_count })
  response = _generate_response(question, accumulated_context)
  return response, trace