458c2d9cab
AI-generated trees can have circular next_node_id references (e.g., node A → B → A). The parent mapping now checks for cycles before assigning parent_node_id, preventing FK deadlocks during insert. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
583 lines
20 KiB
Python
583 lines
20 KiB
Python
"""KB Accelerator AI conversion service.
|
|
|
|
Converts extracted KB article text into ResolutionFlow tree structures
|
|
using the Anthropic API (via the shared AI provider layer).
|
|
"""
|
|
import json
|
|
import logging
|
|
import re
|
|
import time
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.ai_provider import get_ai_provider
|
|
from app.core.ai_quota_service import record_ai_usage, get_user_plan
|
|
from app.core.config import settings
|
|
from app.models.kb_import import KBImport, KBImportNode
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Cost estimation (Sonnet pricing)
|
|
COST_PER_INPUT_TOKEN = 3.0 / 1_000_000
|
|
COST_PER_OUTPUT_TOKEN = 15.0 / 1_000_000
|
|
|
|
|
|
def _strip_markdown_fences(text: str) -> str:
|
|
"""Strip markdown code fences if the model wrapped its JSON response."""
|
|
text = text.strip()
|
|
match = re.match(r"^```(?:json)?\s*([\s\S]*?)```$", text)
|
|
if match:
|
|
return match.group(1).strip()
|
|
return text
|
|
|
|
|
|
def _try_repair_json(text: str) -> dict | None:
|
|
"""Attempt to repair common JSON issues from AI responses.
|
|
|
|
Handles: trailing commas, unclosed brackets/braces, truncated responses.
|
|
Returns parsed dict on success, None on failure.
|
|
"""
|
|
# Strip trailing commas before closing brackets/braces
|
|
repaired = re.sub(r",\s*([}\]])", r"\1", text)
|
|
|
|
# Try parsing after comma cleanup
|
|
try:
|
|
return json.loads(repaired)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# Try closing unclosed brackets/braces (truncated response)
|
|
# Count open vs close brackets
|
|
open_braces = repaired.count("{") - repaired.count("}")
|
|
open_brackets = repaired.count("[") - repaired.count("]")
|
|
|
|
if open_braces > 0 or open_brackets > 0:
|
|
# Remove any trailing partial key-value pair or string
|
|
# Find the last complete value (ends with }, ], ", number, true, false, null)
|
|
truncated = repaired.rstrip()
|
|
# Strip trailing partial string or key
|
|
truncated = re.sub(r',\s*"[^"]*$', "", truncated) # trailing "partial_key
|
|
truncated = re.sub(r',\s*$', "", truncated) # trailing comma
|
|
|
|
# Close remaining brackets/braces
|
|
truncated += "]" * max(0, open_brackets)
|
|
truncated += "}" * max(0, open_braces)
|
|
|
|
# Re-strip trailing commas that may have appeared
|
|
truncated = re.sub(r",\s*([}\]])", r"\1", truncated)
|
|
|
|
try:
|
|
return json.loads(truncated)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
return None
|
|
|
|
|
|
def _estimate_cost(input_tokens: int, output_tokens: int) -> float:
|
|
return (input_tokens * COST_PER_INPUT_TOKEN) + (output_tokens * COST_PER_OUTPUT_TOKEN)
|
|
|
|
|
|
# ── System Prompts ──
|
|
|
|
TROUBLESHOOTING_SYSTEM_PROMPT = """You are an MSP documentation specialist for ResolutionFlow. Your task is to convert a knowledge base article into an interactive troubleshooting decision tree.
|
|
|
|
Analyze the article and produce a JSON array of nodes that form a troubleshooting flow. Each node represents either a diagnostic question (decision point) or a resolution (solution).
|
|
|
|
## Node Types
|
|
|
|
- **question**: A diagnostic question with multiple answer options. Each option leads to another node.
|
|
- **resolution**: A terminal node with the solution/fix text.
|
|
- **action**: An instruction step that leads to the next node via next_node_id.
|
|
- **warning**: A caution or important note.
|
|
|
|
## Output Format
|
|
|
|
Return a JSON object with this structure:
|
|
```json
|
|
{
|
|
"title": "Flow title derived from the article",
|
|
"description": "Brief description of what this flow troubleshoots",
|
|
"nodes": [
|
|
{
|
|
"id": "unique-node-id",
|
|
"type": "question",
|
|
"question": "What symptom is the user experiencing?",
|
|
"options": [
|
|
{"label": "Cannot connect", "next_node_id": "check-network"},
|
|
{"label": "Slow performance", "next_node_id": "check-resources"}
|
|
],
|
|
"confidence": 0.95,
|
|
"source_excerpt": "The exact text from the article this node was derived from"
|
|
},
|
|
{
|
|
"id": "check-network",
|
|
"type": "action",
|
|
"question": "Check the network connection and ping the server",
|
|
"next_node_id": "network-result",
|
|
"confidence": 0.88,
|
|
"source_excerpt": "Step 1: Verify network connectivity..."
|
|
},
|
|
{
|
|
"id": "solution-restart",
|
|
"type": "resolution",
|
|
"question": "Restart the service. The issue should now be resolved.",
|
|
"confidence": 0.92,
|
|
"source_excerpt": "Restarting the service resolves the connectivity issue."
|
|
}
|
|
]
|
|
}
|
|
```
|
|
|
|
## Rules
|
|
|
|
1. Every node MUST have a unique `id` (descriptive kebab-case).
|
|
2. Every node MUST have a `confidence` score between 0.0 and 1.0.
|
|
3. Every node MUST have a `source_excerpt` — the exact text from the source article it was derived from.
|
|
4. The first node is the root of the decision tree.
|
|
5. All `next_node_id` and option `next_node_id` references must point to existing node IDs.
|
|
6. Detect implicit branching logic (e.g., "If X, do Y; otherwise Z") and create decision nodes.
|
|
7. Produce at least 3 nodes. Maximum 50 nodes.
|
|
8. Use high confidence (0.9+) for directly stated steps, medium (0.7-0.89) for reasonable inferences, low (<0.7) for significant interpretation.
|
|
9. Return ONLY valid JSON — no markdown fences, no explanation text."""
|
|
|
|
PROCEDURAL_SYSTEM_PROMPT = """You are an MSP documentation specialist for ResolutionFlow. Your task is to convert a knowledge base article into a procedural (step-by-step) flow.
|
|
|
|
Analyze the article and produce a JSON object with sequential steps and detected variables.
|
|
|
|
## Step Types
|
|
|
|
- **step**: A regular instruction step.
|
|
- **section_header**: A section divider/title (no action, just organizational).
|
|
- **warning**: A caution or important note that should be highlighted.
|
|
|
|
## Variable Detection
|
|
|
|
Identify values that would change between executions (server names, IPs, usernames, domains, etc.) and replace them with `[VAR:variable_name]` tokens. Also produce an intake_form that captures these variables before execution.
|
|
|
|
## Output Format
|
|
|
|
Return a JSON object:
|
|
```json
|
|
{
|
|
"title": "Procedure title derived from the article",
|
|
"description": "Brief description of what this procedure accomplishes",
|
|
"steps": [
|
|
{
|
|
"id": "unique-step-id",
|
|
"type": "step",
|
|
"content": "Open Server Manager and navigate to Add Roles on [VAR:server_name]",
|
|
"confidence": 0.95,
|
|
"source_excerpt": "Step 1: Open Server Manager on DC01..."
|
|
},
|
|
{
|
|
"id": "warning-dns",
|
|
"type": "warning",
|
|
"content": "WARNING: This will restart DNS and cause brief connectivity loss",
|
|
"confidence": 0.90,
|
|
"source_excerpt": "Note: Restarting DNS will cause a brief outage"
|
|
},
|
|
{
|
|
"id": "section-verification",
|
|
"type": "section_header",
|
|
"content": "Verification Steps",
|
|
"confidence": 1.0,
|
|
"source_excerpt": "Verification"
|
|
}
|
|
],
|
|
"intake_form": [
|
|
{
|
|
"variable_name": "server_name",
|
|
"label": "Server Name",
|
|
"field_type": "text",
|
|
"required": true,
|
|
"display_order": 1
|
|
},
|
|
{
|
|
"variable_name": "ip_address",
|
|
"label": "IP Address",
|
|
"field_type": "text",
|
|
"required": true,
|
|
"display_order": 2
|
|
}
|
|
]
|
|
}
|
|
```
|
|
|
|
## Variable Type Mapping
|
|
|
|
- IP addresses → field_type: "text", variable like `ip_address`
|
|
- Server/computer names → field_type: "text", variable like `server_name`
|
|
- Domain names → field_type: "text", variable like `domain_name`
|
|
- Usernames/email → field_type: "text", variable like `username`
|
|
- Port numbers → field_type: "number", variable like `port`
|
|
|
|
## Rules
|
|
|
|
1. Every step MUST have a unique `id` (descriptive kebab-case).
|
|
2. Every step MUST have a `confidence` score between 0.0 and 1.0.
|
|
3. Every step MUST have a `source_excerpt` — the exact text from the source article.
|
|
4. Preserve the original step ordering from the article.
|
|
5. Detect ALL instance-specific values and replace with `[VAR:name]` tokens.
|
|
6. Generate an intake_form entry for each unique variable detected.
|
|
7. Produce at least 2 steps. Maximum 100 steps.
|
|
8. Use high confidence (0.9+) for directly stated steps, medium (0.7-0.89) for inferences, low (<0.7) for significant interpretation.
|
|
9. Return ONLY valid JSON — no markdown fences, no explanation text."""
|
|
|
|
|
|
def _build_user_message(
|
|
source_text: str,
|
|
source_metadata: dict[str, Any] | None,
|
|
source_filename: str | None,
|
|
) -> str:
|
|
"""Build the user message containing the extracted text and metadata."""
|
|
parts = []
|
|
|
|
if source_filename:
|
|
parts.append(f"Source file: {source_filename}")
|
|
|
|
if source_metadata:
|
|
headings = source_metadata.get("headings", [])
|
|
if headings:
|
|
heading_text = ", ".join(
|
|
f"H{h['level']}: {h['text']}" for h in headings[:20]
|
|
)
|
|
parts.append(f"Detected headings: {heading_text}")
|
|
|
|
lists = source_metadata.get("lists", [])
|
|
if lists:
|
|
parts.append(f"Detected {len(lists)} list(s) in the document.")
|
|
|
|
tables = source_metadata.get("tables", [])
|
|
if tables:
|
|
parts.append(f"Detected {len(tables)} table(s) in the document.")
|
|
|
|
parts.append(f"\n--- ARTICLE CONTENT ---\n\n{source_text}")
|
|
|
|
return "\n".join(parts)
|
|
|
|
|
|
def _parse_troubleshooting_response(
|
|
data: dict[str, Any],
|
|
kb_import_id: UUID,
|
|
) -> tuple[list[KBImportNode], str, str | None]:
|
|
"""Parse AI response into KBImportNode records for troubleshooting flows.
|
|
|
|
Returns (nodes, title, description).
|
|
"""
|
|
title = data.get("title", "Imported Troubleshooting Flow")
|
|
description = data.get("description")
|
|
raw_nodes = data.get("nodes", [])
|
|
|
|
if not raw_nodes:
|
|
raise ValueError("AI returned no nodes")
|
|
|
|
# Build parent mapping from the tree structure
|
|
# First node is root (no parent). For others, trace via options/next_node_id.
|
|
node_id_to_parent: dict[str, str | None] = {}
|
|
node_id_to_data: dict[str, dict[str, Any]] = {}
|
|
for node in raw_nodes:
|
|
nid = node.get("id", "")
|
|
node_id_to_data[nid] = node
|
|
if nid not in node_id_to_parent:
|
|
node_id_to_parent[nid] = None # default: no parent
|
|
|
|
# Trace parent relationships (only set if it won't create a cycle)
|
|
def _would_cycle(child: str, parent: str) -> bool:
|
|
"""Check if setting child's parent to parent creates a cycle."""
|
|
visited: set[str] = set()
|
|
cur: str | None = parent
|
|
while cur:
|
|
if cur == child:
|
|
return True
|
|
if cur in visited:
|
|
break
|
|
visited.add(cur)
|
|
cur = node_id_to_parent.get(cur)
|
|
return False
|
|
|
|
for node in raw_nodes:
|
|
nid = node.get("id", "")
|
|
# Options point to children
|
|
for opt in node.get("options", []):
|
|
child_id = opt.get("next_node_id")
|
|
if child_id and child_id in node_id_to_data and not _would_cycle(nid, child_id):
|
|
node_id_to_parent[child_id] = nid
|
|
# next_node_id points to child
|
|
next_id = node.get("next_node_id")
|
|
if next_id and next_id in node_id_to_data and not _would_cycle(nid, next_id):
|
|
node_id_to_parent[next_id] = nid
|
|
|
|
# Create import node records preserving order
|
|
import uuid as uuid_mod
|
|
node_id_map: dict[str, uuid_mod.UUID] = {}
|
|
nodes: list[KBImportNode] = []
|
|
|
|
for order, raw_node in enumerate(raw_nodes):
|
|
node_uuid = uuid_mod.uuid4()
|
|
nid = raw_node.get("id", f"node-{order}")
|
|
node_id_map[nid] = node_uuid
|
|
|
|
for order, raw_node in enumerate(raw_nodes):
|
|
nid = raw_node.get("id", f"node-{order}")
|
|
node_type = raw_node.get("type", "question")
|
|
if node_type == "decision":
|
|
node_type = "question"
|
|
|
|
parent_str_id = node_id_to_parent.get(nid)
|
|
parent_uuid = node_id_map.get(parent_str_id) if parent_str_id else None
|
|
|
|
# Build content JSONB
|
|
content: dict[str, Any] = {
|
|
"original_id": nid,
|
|
"question": raw_node.get("question", ""),
|
|
}
|
|
if raw_node.get("options"):
|
|
content["options"] = raw_node["options"]
|
|
if raw_node.get("next_node_id"):
|
|
content["next_node_id"] = raw_node["next_node_id"]
|
|
|
|
import_node = KBImportNode(
|
|
id=node_id_map[nid],
|
|
kb_import_id=kb_import_id,
|
|
node_order=order,
|
|
node_type=node_type,
|
|
content=content,
|
|
parent_node_id=parent_uuid,
|
|
source_excerpt=raw_node.get("source_excerpt"),
|
|
confidence_score=float(raw_node.get("confidence", 0.5)),
|
|
user_edited=False,
|
|
user_approved=False,
|
|
)
|
|
nodes.append(import_node)
|
|
|
|
return nodes, title, description
|
|
|
|
|
|
def _parse_procedural_response(
|
|
data: dict[str, Any],
|
|
kb_import_id: UUID,
|
|
) -> tuple[list[KBImportNode], str, str | None, list[dict[str, Any]] | None]:
|
|
"""Parse AI response into KBImportNode records for procedural flows.
|
|
|
|
Returns (nodes, title, description, intake_form).
|
|
"""
|
|
title = data.get("title", "Imported Procedure")
|
|
description = data.get("description")
|
|
raw_steps = data.get("steps", [])
|
|
intake_form = data.get("intake_form")
|
|
|
|
if not raw_steps:
|
|
raise ValueError("AI returned no steps")
|
|
|
|
import uuid as uuid_mod
|
|
nodes: list[KBImportNode] = []
|
|
|
|
for order, raw_step in enumerate(raw_steps):
|
|
content: dict[str, Any] = {
|
|
"original_id": raw_step.get("id", f"step-{order}"),
|
|
"content": raw_step.get("content", ""),
|
|
}
|
|
|
|
node_type = raw_step.get("type", "step")
|
|
if node_type not in ("step", "section_header", "warning"):
|
|
node_type = "step"
|
|
|
|
import_node = KBImportNode(
|
|
id=uuid_mod.uuid4(),
|
|
kb_import_id=kb_import_id,
|
|
node_order=order,
|
|
node_type=node_type,
|
|
content=content,
|
|
parent_node_id=None, # Procedural flows are linear
|
|
source_excerpt=raw_step.get("source_excerpt"),
|
|
confidence_score=float(raw_step.get("confidence", 0.5)),
|
|
user_edited=False,
|
|
user_approved=False,
|
|
)
|
|
nodes.append(import_node)
|
|
|
|
return nodes, title, description, intake_form
|
|
|
|
|
|
async def convert_document(
|
|
kb_import: KBImport,
|
|
db: AsyncSession,
|
|
) -> list[KBImportNode]:
|
|
"""Run AI conversion on an extracted KB article.
|
|
|
|
Creates KBImportNode records and updates the kb_import status.
|
|
Returns the created nodes.
|
|
"""
|
|
start_time = time.monotonic()
|
|
|
|
# Select system prompt based on target type
|
|
if kb_import.target_type == "troubleshooting":
|
|
system_prompt = TROUBLESHOOTING_SYSTEM_PROMPT
|
|
else:
|
|
system_prompt = PROCEDURAL_SYSTEM_PROMPT
|
|
|
|
user_message = _build_user_message(
|
|
source_text=kb_import.source_text,
|
|
source_metadata=kb_import.source_metadata,
|
|
source_filename=kb_import.source_filename,
|
|
)
|
|
|
|
# Get AI provider with model routing
|
|
model = settings.get_model_for_action("kb_convert")
|
|
provider = get_ai_provider(model=model)
|
|
|
|
try:
|
|
raw_text, input_tokens, output_tokens = await provider.generate_json(
|
|
system_prompt=system_prompt,
|
|
messages=[{"role": "user", "content": user_message}],
|
|
max_tokens=16384,
|
|
)
|
|
except Exception as e:
|
|
logger.error("AI conversion failed for kb_import=%s: %s", kb_import.id, e)
|
|
kb_import.status = "failed"
|
|
kb_import.error_message = f"AI processing error: {str(e)}"
|
|
kb_import.processing_time_ms = int((time.monotonic() - start_time) * 1000)
|
|
await db.flush()
|
|
|
|
# Record failed usage
|
|
plan = await get_user_plan(kb_import.account_id, db)
|
|
await record_ai_usage(
|
|
user_id=kb_import.created_by,
|
|
account_id=kb_import.account_id,
|
|
conversation_id=None,
|
|
generation_type="kb_convert",
|
|
tier=plan,
|
|
input_tokens=0,
|
|
output_tokens=0,
|
|
estimated_cost=0.0,
|
|
succeeded=False,
|
|
counts_toward_quota=False,
|
|
error_code="ai_error",
|
|
extra_data={"kb_import_id": str(kb_import.id)},
|
|
db=db,
|
|
)
|
|
return []
|
|
|
|
# Parse JSON response
|
|
raw_text = _strip_markdown_fences(raw_text)
|
|
try:
|
|
data = json.loads(raw_text)
|
|
except json.JSONDecodeError as e:
|
|
# Attempt JSON repair before giving up
|
|
data = _try_repair_json(raw_text)
|
|
if data is None:
|
|
logger.error(
|
|
"KB conversion JSON parse failed for kb_import=%s (%d chars). "
|
|
"Parse error: %s. Raw response (first 2000 chars): %s",
|
|
kb_import.id, len(raw_text), e, raw_text[:2000],
|
|
)
|
|
kb_import.status = "failed"
|
|
kb_import.error_message = (
|
|
"AI response could not be parsed as valid JSON. "
|
|
"This can happen with very long articles — try again or simplify the article."
|
|
)
|
|
kb_import.processing_time_ms = int((time.monotonic() - start_time) * 1000)
|
|
kb_import.ai_tokens_input = input_tokens
|
|
kb_import.ai_tokens_output = output_tokens
|
|
await db.flush()
|
|
return []
|
|
else:
|
|
logger.info(
|
|
"KB conversion JSON repaired for kb_import=%s (%d chars)",
|
|
kb_import.id, len(raw_text),
|
|
)
|
|
|
|
# Parse into nodes based on target type
|
|
try:
|
|
intake_form = None
|
|
if kb_import.target_type == "troubleshooting":
|
|
nodes, title, description = _parse_troubleshooting_response(
|
|
data, kb_import.id
|
|
)
|
|
else:
|
|
nodes, title, description, intake_form = _parse_procedural_response(
|
|
data, kb_import.id
|
|
)
|
|
except (ValueError, KeyError, TypeError) as e:
|
|
logger.error("KB node parsing failed for kb_import=%s: %s", kb_import.id, e)
|
|
kb_import.status = "failed"
|
|
kb_import.error_message = f"Failed to parse AI response: {e}"
|
|
kb_import.processing_time_ms = int((time.monotonic() - start_time) * 1000)
|
|
kb_import.ai_tokens_input = input_tokens
|
|
kb_import.ai_tokens_output = output_tokens
|
|
await db.flush()
|
|
return []
|
|
|
|
# Persist nodes — insert roots first to satisfy parent_node_id FK,
|
|
# then children in subsequent passes until all are inserted.
|
|
remaining = list(nodes)
|
|
inserted_ids: set[Any] = set()
|
|
while remaining:
|
|
batch = [
|
|
n for n in remaining
|
|
if n.parent_node_id is None or n.parent_node_id in inserted_ids
|
|
]
|
|
if not batch:
|
|
# Circular reference or orphan — force insert remaining to surface the real error
|
|
for n in remaining:
|
|
db.add(n)
|
|
break
|
|
for n in batch:
|
|
db.add(n)
|
|
inserted_ids.add(n.id)
|
|
await db.flush()
|
|
remaining = [n for n in remaining if n.id not in inserted_ids]
|
|
|
|
# Update import record
|
|
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
|
confidence_scores = [n.confidence_score for n in nodes]
|
|
avg_confidence = sum(confidence_scores) / len(confidence_scores) if confidence_scores else 0.0
|
|
|
|
kb_import.status = "ready"
|
|
kb_import.confidence_avg = avg_confidence
|
|
kb_import.processing_time_ms = elapsed_ms
|
|
kb_import.ai_tokens_input = input_tokens
|
|
kb_import.ai_tokens_output = output_tokens
|
|
|
|
# Store parsed metadata for commit phase
|
|
if not kb_import.source_metadata:
|
|
kb_import.source_metadata = {}
|
|
kb_import.source_metadata["_conversion"] = {
|
|
"title": title,
|
|
"description": description,
|
|
"node_count": len(nodes),
|
|
}
|
|
if intake_form:
|
|
kb_import.source_metadata["_intake_form"] = intake_form
|
|
|
|
await db.flush()
|
|
|
|
# Record successful usage
|
|
plan = await get_user_plan(kb_import.account_id, db)
|
|
cost = _estimate_cost(input_tokens, output_tokens)
|
|
await record_ai_usage(
|
|
user_id=kb_import.created_by,
|
|
account_id=kb_import.account_id,
|
|
conversation_id=None,
|
|
generation_type="kb_convert",
|
|
tier=plan,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
estimated_cost=cost,
|
|
succeeded=True,
|
|
counts_toward_quota=True,
|
|
error_code=None,
|
|
extra_data={"kb_import_id": str(kb_import.id), "node_count": len(nodes)},
|
|
db=db,
|
|
)
|
|
|
|
logger.info(
|
|
"KB conversion complete: import=%s, nodes=%d, confidence=%.2f, time=%dms, tokens=%d/%d",
|
|
kb_import.id, len(nodes), avg_confidence, elapsed_ms, input_tokens, output_tokens,
|
|
)
|
|
|
|
return nodes
|