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
  status: "thinking"
}

// Full answer response (from async worker, accessed via polling)
ChatMessage {
  message_id: string (ULID)
  chat_id: string
  role: "user" | "assistant"
  content: string
  created_at: string (ISO 8601)
}

GetChatMessagesRequest {
  chat_id: string (required)
  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)  // shows "Thinking..." footer
  { answer, chat_id } = await chatWithMemory(question, chatId, isDebug)
  setChatId(chat_id)
  append assistant message bubble
  invalidate ["chats", "feed"] (refreshes side menu)
  setIsLoading(false)

// chatWithMemory.ts
POST /memories/chat { question, chat_id?, verbose? }
  timeout: 60_000 ms (client-side axios timeout)
  // Note: API Gateway hard limit is 29s — dispatcher returns well within that
  return { answer, chat_id, trace? }

// NOTE: the chat screen does NOT poll for the answer.
// chatWithMemory.ts awaits the full response from the dispatcher,
// which in production returns { chat_id, status: "thinking" }.
// The current UI flow BLOCKS on chatWithMemory and renders the answer
// from the synchronous response — the dispatcher path sends { chat_id, status }
// but the handler wraps it so the response to the client is the dispatcher result.
// When the async worker later saves the answer, it is accessible via
// POST /chats/messages but the current chat.tsx does NOT poll — it just awaits
// the initial chatWithMemory call which returns once the dispatcher responds.

// === 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:
    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? }

// 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