Skip to content

πŸ”§ AI Engine

This describes the design, invocation flow, contracts, and deployment notes for the ai-engine.

Purpose

The ai-engine is a small Python Flask service that orchestrates a set of AI "agents" to generate structured mindmap (a.k.a. Visualli) JSON from long-form text. It is intended to be invoked by the backend (backend/server) after a user uploads any content from the frontend (PDF, text, ..).

Design

Components Overview

  • HTTP Ingress (Node server) β€” backend/server/server.js accepts uploads and text, extracts plain text, and calls the ai-engine.
  • Python Flask Service β€” ai-engine exposes /generate_mindmap that accepts { text: string } and returns validated mindmap JSON.
  • Agent Layer β€” agents.py initializes AI agents (MindMapAgent, ComplianceAgent) via the autogen library to generate and gate outputs.
  • Position Optimizer β€” position_optimizer.py lays out nodes in radial coordinates and writes position.x and position.y.
  • Persistence β€” The Node server receives the returned mindmap JSON and stores it in MongoDB (collection visuallis).

Sequence Diagram

sequenceDiagram
    autonumber

    participant User
    participant Node as Node Backend
    participant Flask as Flask ai-engine
    participant Agents
    participant Mongo as MongoDB

    User->>Node: Upload PDF/text\nPOST /api/upload or /api/upload/text
    Node->>Node: Extract text\npdf-parse (PDF) or direct blob
    Node->>Flask: POST /generate_mindmap
    Note over Node,Flask: Alternatively via proxy\nNode /api/generate_mindmap
    Flask->>Agents: agents.main(task)
    Agents->>Agents: initialize_agents()\nMindmapAgent + ComplianceAgent
    Agents->>Agents: MindmapAgent.run(task)\nproduce initial JSON
    Agents->>Agents: ComplianceAgent.run(task)\nvalidate to MindmapData
    Agents->>Agents: position_optimizer.tidy_radial(...)\ncompute x,y positions
    Agents-->>Flask: Return final JSON
    Flask-->>Node: JSON response
    Node->>Mongo: Persist JSON to collection "visuallis"
    Node-->>User: Success payload

I/O Contracts

Endpoint: POST /generate_mindmap - Request JSON body: - { "text": "" }

  • Success response: JSON matching MindmapData (see src/schemas/mindmap.py). Key fields:
  • metadata (title, description, version, created, lastModified)
  • nodes (list of nodes with at least: id, title, level, position {x,y}, color, optional parent, optional relationshipLabel, optional nested children)
  • topLevelConnections (list of { from, to, label })

Notes on shape: The code expects the returned JSON to be convertible to Python dicts and ultimately to be inserted into MongoDB. The Node server adds some top-level fields (created, lastModified, version, title) if missing.

Key Files

  • src/main.py β€” Flask app and the /generate_mindmap endpoint. Runs the async agent pipeline in a background thread.
  • src/agents/agents.py β€” AI Agent definitions and the async main(task) orchestration.
  • src/agents/position_optimizer.py β€” Radial layout algorithm used to assign x,y positions to nodes (tidy_radial function). Algorithm to ensure nodes in the visual don't overlap and diverge outwards radially.
  • src/schemas/mindmap.py β€” pydantic schema for the expected mindmap JSON shape (MindmapData).
  • src/wsgi/handler.py β€” WSGI Lambda handler (wraps the Flask app with serverless_wsgi.handle_request) for Serverless/AWS Lambda deployments.
  • src/wsgi/app.py β€” tiny module that imports and runs the Flask app for local execution under a module.

Error Handling & Edge Cases

  • Timeouts: Node sets a 5-minute timeout when calling Flask. The agent pipeline may take time depending on model responses.
  • Invalid/empty text: Node and Flask both check for empty extracted text and return 400.
  • Invalid response from Python: Node checks that response.data is an object; otherwise returns 500.
  • PDF parsing failures: Node has multiple fallbacks and returns 400 with helpful messages if parsing fails.
  • Model keys missing: agents/agents.py warns when OPENAI_API_KEY is not set and logs via Logtail; agent clients will fail if keys are missing.

Environment Vars Used

  • OPENAI_API_KEY β€” required for the OpenAI client used by agents.

Set these vars locally (example .env) in backend/ai-engine/.env or export them in your shell before starting.

Deployment Notes

  • Local development:
  • Start the ai-engine Flask app (it listens on port 5000 by default).
cd backend/ai-engine
python3 -m src.wsgi.app
  • Start the Node backend:
cd backend/server
node server.js
  • The Node backend calls http://127.0.0.1:5000/generate_mindmap by default, or uses the proxy route /api/generate_mindmap.

  • Serverless / Lambda:

  • src/wsgi/handler.py is prepared for Lambda via serverless_wsgi.handle_request(app, event, context).
  • serverless.yml is present in the ai-engine directory and can be configured to deploy the Flask app behind API Gateway.

Observability & Logging

  • The agents module configures logtail (BetterStack) as a handler. Make sure BETTERSTACK_SOURCE_TOKEN and BETTERSTACK_INGESTING_HOST are set to collect logs.
  • stdout is redirected to the logger via StdOutToLogger so agent library prints are captured.

Performance / cost considerations

  • Model usage: Agents instantiate multiple clients (gpt-4o, gpt-4o-mini, etc.) β€” these can be expensive. Consider rate limits and batch sizes.
  • Concurrency: Flask endpoint runs the full pipeline inside a background thread using a fresh event loop. If the server receives many concurrent requests, the ThreadPoolExecutor default size may need tuning.

Security considerations

  • Validate inputs: The Node server and Flask endpoint expect plain text; be careful to limit upload sizes (Node uses a 5MB limit) and sanitize any filename inputs.
  • Secrets: Keep OPENAI_API_KEY and BetterStack tokens out of source control (use .env and secrets manager in production).

Next steps / Improvements

  • Add unit tests around agents.main (mock model clients) and position_optimizer.tidy_radial.
  • Provide a health endpoint in the Flask service for readiness checks (used by load balancers).
  • Add request tracing (e.g., a request ID passed from Node to Flask) for easier debugging.
  • Limit the ThreadPoolExecutor size and implement an in-queue to avoid overload during spikes.