fix: Simplify Colab notebook, fix rank_bm25 import handling, and clean up prompts

ControlAI_Colab_Demo.ipynb

@@ -4,13 +4,13 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# ControlAI: Open-Source Safety-Critical AI Agent for Control Systems\n",
+    "# ControlAI: Open-Source AI Agent for Control Systems Engineering\n",
     "\n",
     "[![GitHub Repo](https://img.shields.io/badge/GitHub-atakankahya%2Fcontrolai--agent-blue?logo=github)](https://github.com/atakankahya/controlai-agent)\n",
     "[![Hugging Face Model](https://img.shields.io/badge/Hugging%20Face-ControlAI--Agent-blue)](https://huggingface.co/atakankahya/ControlAI-Agent)\n",
     "[![License](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)\n",
     "\n",
-    "**ControlAI** is an open-source AI agent specialized in Control Systems Engineering, Applied Mathematics, and Dynamical Systems Simulation.\n",
+    "Welcome to the official **ControlAI** interactive notebook! You can test control engineering problems, LQR synthesis, Bode plots, and ODE simulations directly in Google Colab.\n",
     "\n",
     "---"
    ]
@@ -19,7 +19,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Step 1: Clone Repository & Install PyTorch/Transformers Dependencies"
+    "### Step 1: Install Dependencies"
    ]
   },
   {
@@ -28,19 +28,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Clone the open-source repository\n",
     "!git clone https://github.com/atakankahya/controlai-agent.git\n",
     "%cd controlai-agent\n",
     "\n",
-    "# Install PyTorch, Transformers, and Scientific Control Libraries\n",
-    "!pip install -q torch transformers accelerate scipy control cvxpy rich pyngrok uvicorn fastapi sentencepiece protobuf python-multipart"
+    "!pip install -q torch transformers accelerate scipy control cvxpy rich rank-bm25 matplotlib"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Step 2: Quick CLI Test in Terminal Mode"
+    "### Step 2: Initialize the ControlAI Agent"
    ]
   },
   {
@@ -49,16 +47,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Run single control engineering query using PyTorch/Transformers\n",
-    "!python cli.py --model \"Qwen/Qwen2.5-3B-Instruct\" \"Design an LQR controller for A=[[0, 1], [-2, -3]], B=[[0], [1]], Q=diag([10, 1]), R=1 and simulate step response.\""
+    "import sys\n",
+    "from pathlib import Path\n",
+    "from IPython.display import display, Markdown, Image\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from controlai_agent.orchestrator import ControlAIAgent\n",
+    "\n",
+    "print(\"Loading ControlAI Engine with Qwen 3B...\")\n",
+    "agent = ControlAIAgent(model_path=\"Qwen/Qwen2.5-3B-Instruct\")\n",
+    "print(\"ControlAI Engine ready!\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Step 3: Launch Live Public Web Interface (Interactive Playground Link)\n",
-    "Run the cell below to launch the web console and get your public web URL!"
+    "### Step 3: Ask Any Control Engineering Question\n",
+    "You can edit the prompt below to test your own dynamical systems, transfer functions, or optimal control problems."
    ]
   },
   {
@@ -67,16 +73,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "import subprocess\n",
-    "import time\n",
+    "# Enter your engineering query here:\n",
+    "prompt = \"Design an LQR controller for A=[[0, 1], [-2, -3]], B=[[0], [1]], Q=diag([10, 1]), R=1 and simulate the step response.\"\n",
+    "\n",
+    "print(f\"User Query: {prompt}\\n\")\n",
+    "result = agent.run(prompt)\n",
     "\n",
-    "# Launch FastAPI server\n",
-    "server_process = subprocess.Popen([\"python\", \"-m\", \"uvicorn\", \"app:app\", \"--host\", \"0.0.0.0\", \"--port\", \"8000\"])\n",
-    "time.sleep(3)\n",
+    "# Display formatted response\n",
+    "display(Markdown(result.final_response))\n",
     "\n",
-    "# Expose public tunnel using localtunnel\n",
-    "!npx -y localtunnel --port 8000"
+    "# Display any generated simulation plots\n",
+    "for plot_path_rel in result.plots:\n",
+    "    clean_name = Path(plot_path_rel).name\n",
+    "    local_path = Path(\"outputs/plots\") / clean_name\n",
+    "    if local_path.exists():\n",
+    "        print(f\"\\nGenerated Simulation Plot: {clean_name}\")\n",
+    "        display(Image(filename=str(local_path)))"
    ]
   }
  ],

controlai_agent/prompts.py

@@ -1,4 +1,4 @@
-CONTROLAI_SYSTEM_PROMPT = """You are ControlAI, a premier AI research scientist and expert engineering agent specialized in control systems engineering, applied mathematics, robotics, and dynamical systems.
+CONTROLAI_SYSTEM_PROMPT = r"""You are ControlAI, a premier AI research scientist and expert engineering agent specialized in control systems engineering, applied mathematics, robotics, and dynamical systems.
 
 ### 4-Stage Mathematical Reasoning & Proof Standard:
 When answering theoretical principles, derivations, proofs, comparisons, or limitation questions:

controlai_rag/index.py

@@ -8,7 +8,10 @@ import re
 from pathlib import Path
 from typing import Any
 
-from rank_bm25 import BM25Okapi
+try:
+    from rank_bm25 import BM25Okapi
+except ImportError:
+    BM25Okapi = None
 
 from controlai_rag.chunker import Chunk
 
@@ -26,10 +29,12 @@ class ControlRAGIndex:
     def __init__(self, index_dir: Path = INDEX_DIR) -> None:
         self.index_dir = index_dir
         self.chunks: list[dict[str, Any]] = []
-        self.bm25: BM25Okapi | None = None
+        self.bm25: Any | None = None
         self._load_if_exists()
 
     def build_from_chunks(self, chunks: list[Chunk]) -> None:
+        if BM25Okapi is None:
+            return
         self.chunks = [c.to_dict() for c in chunks]
         corpus = [tokenize_corpus(c.text) for c in chunks]
         self.bm25 = BM25Okapi(corpus)

requirements.txt

@@ -0,0 +1,13 @@
+torch>=2.0.0
+transformers>=4.40.0
+accelerate>=0.28.0
+scipy>=1.11.0
+numpy>=1.24.0
+matplotlib>=3.7.0
+control>=0.9.4
+cvxpy>=1.4.0
+rich>=13.7.0
+rank-bm25>=0.2.2
+fastapi>=0.110.0
+uvicorn>=0.28.0
+python-multipart>=0.0.9