From 1fcb4c84e52718c27a2daccc6ba21bf8c4d1457f Mon Sep 17 00:00:00 2001
From: Rassl <rasul.sahibjanli@gmail.com>
Date: Sun, 26 Jul 2026 07:31:03 +0700
Subject: [PATCH] feat: ontology page agent

---
 .../admin/ontology/ontology-agent-panel.tsx   | 264 ++++++++++++++++++
 src/app/admin/ontology/page.tsx               |  23 +-
 src/lib/graph-api.ts                          | 102 ++++++-
 3 files changed, 385 insertions(+), 4 deletions(-)
 create mode 100644 src/app/admin/ontology/ontology-agent-panel.tsx

diff --git a/src/app/admin/ontology/ontology-agent-panel.tsx b/src/app/admin/ontology/ontology-agent-panel.tsx
new file mode 100644
index 0000000..b0e41c2
--- /dev/null
+++ b/src/app/admin/ontology/ontology-agent-panel.tsx
@@ -0,0 +1,264 @@
+"use client"
+
+import { useCallback, useEffect, useRef, useState } from "react"
+import { Sparkles, X, ArrowUp, Loader2 } from "lucide-react"
+import {
+  triggerOntologyAgent,
+  getStakworkRun,
+  listReviews,
+  type OntologyAgentMessage,
+  type Review,
+  type StakworkRun,
+} from "@/lib/graph-api"
+import { ReviewRow } from "@/components/admin/review-row"
+import { useSchemaStore } from "@/stores/schema-store"
+import type { SchemaNode } from "@/lib/schema-types"
+
+type TurnStatus = "running" | "completed" | "failed"
+
+interface AgentTurn {
+  id: string
+  runRef?: string
+  status: TurnStatus
+  reviews: Review[]
+  error?: string
+}
+
+interface UserTurn {
+  id: string
+  text: string
+}
+
+type Turn = ({ role: "user" } & UserTurn) | ({ role: "agent" } & AgentTurn)
+
+const POLL_INTERVAL_MS = 5000
+
+// Normalize any backend status casing to the three states the panel tracks.
+function mapRunStatus(status: StakworkRun["status"]): TurnStatus {
+  const s = String(status).toUpperCase()
+  if (s === "COMPLETED") return "completed"
+  if (s === "FAILED" || s === "ERROR" || s === "HALTED") return "failed"
+  return "running"
+}
+
+export function OntologyAgentPanel({ onClose }: { onClose: () => void }) {
+  const schemas = useSchemaStore((s) => s.schemas)
+  const [turns, setTurns] = useState<Turn[]>([])
+  const [input, setInput] = useState("")
+  const [busy, setBusy] = useState(false)
+  const pollRef = useRef<ReturnType<typeof setInterval> | null>(null)
+  const scrollRef = useRef<HTMLDivElement>(null)
+
+  // Clear any in-flight poll on unmount.
+  useEffect(() => () => { if (pollRef.current) clearInterval(pollRef.current) }, [])
+
+  // Keep the transcript scrolled to the latest turn.
+  useEffect(() => {
+    scrollRef.current?.scrollTo({ top: scrollRef.current.scrollHeight, behavior: "smooth" })
+  }, [turns])
+
+  const patchAgentTurn = useCallback((id: string, patch: Partial<AgentTurn>) => {
+    setTurns((prev) =>
+      prev.map((t) => (t.role === "agent" && t.id === id ? { ...t, ...patch } : t))
+    )
+  }, [])
+
+  // Re-fetch this run's pending proposals (called after an approve/dismiss) and
+  // refresh the schema store so the graph + type list reflect any applied change.
+  const refreshTurn = useCallback(
+    async (turnId: string, runRef: string) => {
+      try {
+        const res = await listReviews({ run_ref_id: runRef, status: "pending" })
+        patchAgentTurn(turnId, { reviews: res.reviews })
+      } catch {
+        /* keep existing cards on transient failure */
+      }
+      useSchemaStore.getState().fetchAll()
+    },
+    [patchAgentTurn]
+  )
+
+  const startPoll = useCallback(
+    (turnId: string, runRef: string) => {
+      if (pollRef.current) clearInterval(pollRef.current)
+      pollRef.current = setInterval(async () => {
+        try {
+          const run = await getStakworkRun(runRef)
+          const status = run ? mapRunStatus(run.status) : "running"
+          if (status === "running") return
+          if (pollRef.current) clearInterval(pollRef.current)
+          pollRef.current = null
+          setBusy(false)
+          if (status === "failed") {
+            patchAgentTurn(turnId, { status: "failed", error: run?.error_message ?? "The agent run failed." })
+            return
+          }
+          const res = await listReviews({ run_ref_id: runRef, status: "pending" })
+          patchAgentTurn(turnId, { status: "completed", reviews: res.reviews })
+        } catch {
+          /* swallow — keep polling until a terminal state or unmount */
+        }
+      }, POLL_INTERVAL_MS)
+    },
+    [patchAgentTurn]
+  )
+
+  const handleSubmit = useCallback(async () => {
+    const instruction = input.trim()
+    if (!instruction || busy) return
+    setBusy(true)
+    setInput("")
+
+    const history: OntologyAgentMessage[] = turns.map((t) =>
+      t.role === "user"
+        ? { role: "user", content: t.text }
+        : { role: "agent", content: agentSummary(t) }
+    )
+
+    const userId = `u-${Date.now()}`
+    const agentId = `a-${Date.now()}`
+    setTurns((prev) => [
+      ...prev,
+      { role: "user", id: userId, text: instruction },
+      { role: "agent", id: agentId, status: "running", reviews: [] },
+    ])
+
+    try {
+      const { stakwork_run_ref_id } = await triggerOntologyAgent(instruction, history)
+      patchAgentTurn(agentId, { runRef: stakwork_run_ref_id })
+      startPoll(agentId, stakwork_run_ref_id)
+    } catch {
+      setBusy(false)
+      patchAgentTurn(agentId, { status: "failed", error: "Could not start the ontology agent." })
+    }
+  }, [input, busy, turns, patchAgentTurn, startPoll])
+
+  return (
+    <div className="w-[440px] shrink-0 border-l border-border flex flex-col bg-card">
+      {/* Header */}
+      <div className="flex items-center justify-between p-4 border-b border-border">
+        <div className="flex items-center gap-2">
+          <Sparkles className="h-4 w-4 shrink-0 text-primary" />
+          <div>
+            <p className="text-sm font-medium">Edit ontology with AI</p>
+            <p className="text-[11px] text-muted-foreground">Describe a change; approve the proposals.</p>
+          </div>
+        </div>
+        <button
+          onClick={onClose}
+          className="text-muted-foreground hover:text-foreground transition-colors shrink-0"
+          title="Close"
+        >
+          <X className="h-4 w-4" />
+        </button>
+      </div>
+
+      {/* Transcript */}
+      <div ref={scrollRef} className="flex-1 overflow-y-auto p-4 space-y-4">
+        {turns.length === 0 && (
+          <div className="flex h-full flex-col items-center justify-center gap-2 text-center">
+            <Sparkles className="h-6 w-6 text-muted-foreground/40" />
+            <p className="text-xs text-muted-foreground max-w-[240px]">
+              e.g. &ldquo;Add a Podcast type that extends Media with an episode_count number, and a
+              HOSTED_BY relationship to Person.&rdquo;
+            </p>
+          </div>
+        )}
+
+        {turns.map((t) =>
+          t.role === "user" ? (
+            <div key={t.id} className="flex justify-end">
+              <div className="max-w-[85%] rounded-2xl rounded-br-sm bg-primary/10 px-3 py-2 text-sm text-foreground">
+                {t.text}
+              </div>
+            </div>
+          ) : (
+            <AgentTurnView
+              key={t.id}
+              turn={t}
+              schemas={schemas}
+              onRefresh={() => t.runRef && refreshTurn(t.id, t.runRef)}
+            />
+          )
+        )}
+      </div>
+
+      {/* Composer */}
+      <div className="border-t border-border p-3">
+        <div className="flex items-end gap-2 rounded-lg border border-border bg-background p-2">
+          <textarea
+            value={input}
+            onChange={(e) => setInput(e.target.value)}
+            onKeyDown={(e) => {
+              if (e.key === "Enter" && !e.shiftKey) {
+                e.preventDefault()
+                handleSubmit()
+              }
+            }}
+            placeholder="Describe an ontology change…"
+            rows={2}
+            className="flex-1 resize-none bg-transparent text-sm outline-none placeholder:text-muted-foreground/60"
+          />
+          <button
+            onClick={handleSubmit}
+            disabled={busy || !input.trim()}
+            className="flex h-8 w-8 shrink-0 items-center justify-center rounded-md bg-primary text-primary-foreground transition-opacity disabled:opacity-40"
+            title="Send"
+          >
+            {busy ? <Loader2 className="h-4 w-4 animate-spin" /> : <ArrowUp className="h-4 w-4" />}
+          </button>
+        </div>
+        <p className="mt-1.5 px-1 text-[10px] text-muted-foreground/70">
+          Proposed changes appear as reviews here — approve to apply them to the schema.
+        </p>
+      </div>
+    </div>
+  )
+}
+
+function agentSummary(t: AgentTurn): string {
+  if (t.status === "running") return "Working on the proposal…"
+  if (t.status === "failed") return t.error ?? "The run failed."
+  const n = t.reviews.length
+  return n === 0 ? "No changes proposed." : `Proposed ${n} change${n === 1 ? "" : "s"}.`
+}
+
+function AgentTurnView({
+  turn,
+  schemas,
+  onRefresh,
+}: {
+  turn: AgentTurn
+  schemas: SchemaNode[]
+  onRefresh: () => void
+}) {
+  return (
+    <div className="space-y-2">
+      <div className="flex items-center gap-2 text-xs text-muted-foreground">
+        <Sparkles className="h-3.5 w-3.5 text-primary/70" />
+        {turn.status === "running" ? (
+          <span className="flex items-center gap-1.5">
+            <Loader2 className="h-3 w-3 animate-spin" /> Reasoning over the schema…
+          </span>
+        ) : turn.status === "failed" ? (
+          <span className="text-red-400">{turn.error ?? "The run failed."}</span>
+        ) : (
+          <span>{agentSummary(turn)}</span>
+        )}
+      </div>
+
+      {turn.status === "completed" && turn.reviews.length > 0 && (
+        <div className="overflow-hidden rounded-lg border border-border/60">
+          {turn.reviews.map((review) => (
+            <ReviewRow
+              key={review.ref_id}
+              review={review}
+              schemas={schemas}
+              onRefresh={onRefresh}
+            />
+          ))}
+        </div>
+      )}
+    </div>
+  )
+}
diff --git a/src/app/admin/ontology/page.tsx b/src/app/admin/ontology/page.tsx
index 5bcec8b..05697f4 100644
--- a/src/app/admin/ontology/page.tsx
+++ b/src/app/admin/ontology/page.tsx
@@ -7,7 +7,8 @@ import { OntologyGraph } from "./ontology-graph"
 import { TypeEditor } from "./type-editor"
 import { EdgeTypePanel } from "./edge-type-panel"
 import { EdgeCreatePanel, type NewEdgeParams } from "./edge-create-panel"
-import { Plus, ArrowLeft, Box, Grid2x2, Search, ArrowRight, HelpCircle } from "lucide-react"
+import { OntologyAgentPanel } from "./ontology-agent-panel"
+import { Plus, ArrowLeft, Box, Grid2x2, Search, ArrowRight, HelpCircle, Sparkles } from "lucide-react"
 import { useUserStore } from "@/stores/user-store"
 
 const OntologyGraph3D = dynamic(
@@ -39,6 +40,8 @@ export default function OntologyPage() {
   // Non-null while a draft (unsaved) new node type is being authored.
   const [draftType, setDraftType] = useState<SchemaNode | null>(null)
   const [showHelp, setShowHelp] = useState(false)
+  // When true, the AI ontology-editor panel takes over the right-panel slot.
+  const [showAgent, setShowAgent] = useState(false)
 
   useEffect(() => {
     if (isMocksEnabled()) {
@@ -276,6 +279,18 @@ export default function OntologyPage() {
             </button>
           </div>
 
+          {isAdmin && (
+            <Button
+              size="sm"
+              variant="ghost"
+              onClick={() => setShowAgent((v) => !v)}
+              className={`h-7 w-7 p-0 hover:text-foreground ${showAgent ? "text-primary" : "text-muted-foreground"}`}
+              title="Edit ontology with AI"
+            >
+              <Sparkles className="h-4 w-4" />
+            </Button>
+          )}
+
           <Button
             size="sm"
             variant="ghost"
@@ -462,8 +477,10 @@ export default function OntologyPage() {
         )}
       </div>
 
-      {/* Right panel — create flow takes precedence over the inspectors */}
-      {edgeCreate ? (
+      {/* Right panel — AI editor takes precedence, then create flow, then inspectors */}
+      {showAgent ? (
+        <OntologyAgentPanel onClose={() => setShowAgent(false)} />
+      ) : edgeCreate ? (
         <EdgeCreatePanel
           allSchemas={store.schemas}
           initialSource={edgeCreate.source}
diff --git a/src/lib/graph-api.ts b/src/lib/graph-api.ts
index dd613b1..4d34cb0 100644
--- a/src/lib/graph-api.ts
+++ b/src/lib/graph-api.ts
@@ -728,6 +728,104 @@ export async function getLatestStakworkRun(
   }
 }
 
+// -------------------------------------------------------------------------
+// Ontology Agent (plain-language schema editing)
+// -------------------------------------------------------------------------
+
+export interface OntologyAgentMessage {
+  role: "user" | "agent"
+  content: string
+}
+
+// Per-run mock poll counters (mirrors the deep-research mock pattern)
+const _mockOntologyPollCounts: Record<string, number> = {}
+let _mockOntologyRunCounter = 0
+
+// Admin-only, graph-scoped. Triggers the ontology_agent workflow; proposals
+// come back as reviews (polled via getStakworkRun → listReviews(run_ref_id)).
+export async function triggerOntologyAgent(
+  instruction: string,
+  history?: OntologyAgentMessage[],
+  signal?: AbortSignal
+): Promise<{ stakwork_run_ref_id: string }> {
+  if (isMocksEnabled()) {
+    const runRef = `mock-ontology-run-${++_mockOntologyRunCounter}`
+    _mockOntologyPollCounts[runRef] = 0
+    // Seed a couple of pending schema-change reviews tied to this run so the
+    // conversational panel can render + approve them once the run "completes".
+    getMockReviewsStore().unshift(...buildMockOntologyReviews(runRef, instruction))
+    return { stakwork_run_ref_id: runRef }
+  }
+  return api.post<{ stakwork_run_ref_id: string }>(
+    "/v2/schema/ontology-agent",
+    { instruction, history: history ?? [] },
+    undefined,
+    signal
+  )
+}
+
+// Poll a single run by its run ref_id (graph-scoped runs have no node ref_id).
+export async function getStakworkRun(
+  refId: string,
+  signal?: AbortSignal
+): Promise<StakworkRun | null> {
+  if (isMocksEnabled()) {
+    const count = (_mockOntologyPollCounts[refId] ?? 0) + 1
+    _mockOntologyPollCounts[refId] = count
+    const running = count <= 2
+    return {
+      ref_id: refId,
+      job_type: "ontology_agent",
+      status: running ? "RUNNING" : "COMPLETED",
+      created_at: Math.floor(Date.now() / 1000),
+      ...(running ? {} : { finished_at: Math.floor(Date.now() / 1000) }),
+    }
+  }
+  try {
+    return await api.get<StakworkRun>(`/v2/stakwork-runs/${refId}`, undefined, signal)
+  } catch (err) {
+    if (typeof err === "object" && err !== null && (err as { status?: number }).status === 404) return null
+    throw err
+  }
+}
+
+// Mock-only: fabricate schema-change reviews so the panel is demoable end-to-end.
+function buildMockOntologyReviews(runRef: string, instruction: string): Review[] {
+  const now = new Date().toISOString()
+  const base = { status: "pending" as ReviewStatus, priority: 5, subject_ids: [], subject_nodes: [], created_at: now, run_ref_id: runRef }
+  return [
+    {
+      ...base,
+      ref_id: `${runRef}-r1`,
+      type: "schema_change",
+      rationale: `Proposed from: "${instruction}"`,
+      action_name: "add_schema_node_type",
+      action_payload: {
+        type: "Podcast",
+        parent: "Media",
+        color: "#8b5cf6",
+        node_key: "name",
+        attributes: [
+          { key: "name", type: "string", required: true },
+          { key: "episode_count", type: "number", required: false },
+        ],
+      },
+      fingerprint: `${runRef}-r1`,
+      icon: "layers",
+    },
+    {
+      ...base,
+      ref_id: `${runRef}-r2`,
+      type: "schema_change",
+      rationale: `Proposed from: "${instruction}"`,
+      action_name: "add_schema_edge_type",
+      action_payload: { edge_type: "HOSTED_BY", source: "Podcast", target: "Person" },
+      fingerprint: `${runRef}-r2`,
+      icon: "network",
+    },
+  ]
+}
+
 // -------------------------------------------------------------------------
 // Reviews
 // -------------------------------------------------------------------------
@@ -778,7 +876,7 @@ function getMockReviewsStore(): Review[] {
 }
 
 export async function listReviews(
-  params?: { status?: ReviewStatus; type?: string; action_name?: string; sort?: string; skip?: number; limit?: number; search?: string; node_type?: string },
+  params?: { status?: ReviewStatus; type?: string; action_name?: string; run_ref_id?: string; sort?: string; skip?: number; limit?: number; search?: string; node_type?: string },
   signal?: AbortSignal
 ): Promise<ReviewsListResponse> {
   if (isMocksEnabled()) {
@@ -787,6 +885,7 @@ export async function listReviews(
     if (params?.status) filtered = filtered.filter((r) => r.status === params.status)
     if (params?.type) filtered = filtered.filter((r) => r.type === params.type)
     if (params?.action_name) filtered = filtered.filter((r) => r.action_name === params.action_name)
+    if (params?.run_ref_id) filtered = filtered.filter((r) => r.run_ref_id === params.run_ref_id)
     if (params?.search) {
       const q = params.search.toLowerCase()
       filtered = filtered.filter(
@@ -822,6 +921,7 @@ export async function listReviews(
   if (params?.status) qs.set("status", params.status)
   if (params?.type) qs.set("type", params.type)
   if (params?.action_name) qs.set("action_name", params.action_name)
+  if (params?.run_ref_id) qs.set("run_ref_id", params.run_ref_id)
   if (params?.sort) qs.set("sort", params.sort)
   if (params?.skip !== undefined) qs.set("skip", String(params.skip))
   if (params?.limit !== undefined) qs.set("limit", String(params.limit))