feat: granular CIS MCP server, chat memory compaction, and UX improvements
- Rewrite mcp_cis_server.py with per-section scan tools (IAM, Networking, Compute, Logging/Monitoring, Storage, Asset Management) instead of monolithic full-tenancy scan — 12 granular tools - Add cis_reports.py (Oracle CIS Benchmark checker) to backend - Add chat memory compaction: auto-summarize old messages when history exceeds ~8000 tokens, keeping 6 recent messages intact - Fix GenAI tool use loop: accumulate assistant+tool messages across iterations for proper conversation flow (Generic format) - Remove tenancy confirmation section from system prompt - Add thinking indicator and button disable in chat UI while waiting - Add requests dependency for cis_reports.py - Fix NoneType error in cis_reports.py region filtering
This commit is contained in:
nogueiraguh
197
backend/app.py
197
backend/app.py
@@ -38,6 +38,12 @@ for d in [DATA, OCI_DIR, REPORTS, MCP_DIR, WALLET_DIR]:
|
||||
|
||||
_running_reports: dict[str, asyncio.subprocess.Process] = {} # rid → subprocess
|
||||
|
||||
# ── Chat Memory Compaction Settings ──
|
||||
COMPACT_TOKEN_THRESHOLD = 8000 # estimated tokens before triggering compaction
|
||||
COMPACT_KEEP_RECENT = 6 # recent messages to keep uncompacted
|
||||
COMPACT_SUMMARY_MAX_TOKENS = 1000
|
||||
COMPACT_MIN_MESSAGES = 6
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
log = logging.getLogger("agent")
|
||||
|
||||
@@ -304,6 +310,15 @@ def init_db():
|
||||
model_id TEXT,
|
||||
created_at TEXT DEFAULT (datetime('now'))
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS chat_summaries (
|
||||
id TEXT PRIMARY KEY,
|
||||
session_id TEXT NOT NULL,
|
||||
user_id TEXT NOT NULL,
|
||||
summary TEXT NOT NULL,
|
||||
messages_compacted INTEGER NOT NULL,
|
||||
up_to_message_id TEXT NOT NULL,
|
||||
created_at TEXT DEFAULT (datetime('now'))
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS audit_log (
|
||||
id TEXT PRIMARY KEY, user_id TEXT, username TEXT,
|
||||
action TEXT NOT NULL, resource TEXT, details TEXT,
|
||||
@@ -871,7 +886,7 @@ async def test_genai(gid: str, u=Depends(require("admin","user"))):
|
||||
if not gc: raise HTTPException(404)
|
||||
gname = gc["name"]
|
||||
try:
|
||||
resp, _ = _call_genai(dict(gc), "Say 'connection successful' in one short sentence.")
|
||||
resp, _, _ = _call_genai(dict(gc), "Say 'connection successful' in one short sentence.")
|
||||
_config_log("genai", gid, gname, "success", "test", f"GenAI OK: {resp[:200]}", u["id"], u["username"])
|
||||
return {"status":"success","message":"GenAI OK","response":resp[:300]}
|
||||
except Exception as e:
|
||||
@@ -879,11 +894,118 @@ async def test_genai(gid: str, u=Depends(require("admin","user"))):
|
||||
_config_log("genai", gid, gname, "error", "test", msg, u["id"], u["username"])
|
||||
return {"status":"error","message":msg}
|
||||
|
||||
# ── Chat Memory Compaction ─────────────────────────────────────────────────────
|
||||
|
||||
def _estimate_tokens(text: str) -> int:
|
||||
return len(text) // 4
|
||||
|
||||
def _estimate_history_tokens(history: list) -> int:
|
||||
return sum(_estimate_tokens(h["content"]) for h in history)
|
||||
|
||||
def _should_compact(history: list) -> bool:
|
||||
if len(history) < COMPACT_MIN_MESSAGES:
|
||||
return False
|
||||
return _estimate_history_tokens(history) > COMPACT_TOKEN_THRESHOLD
|
||||
|
||||
def _generate_summary(genai_cfg: dict, messages_to_summarize: list) -> str:
|
||||
"""Use the same GenAI model to summarize older conversation messages."""
|
||||
conversation_text = ""
|
||||
for m in messages_to_summarize:
|
||||
role_label = "Usuário" if m["role"] == "user" else "Assistente"
|
||||
# Truncate very long messages in the summary input
|
||||
content = m["content"][:3000] if len(m["content"]) > 3000 else m["content"]
|
||||
conversation_text += f"{role_label}: {content}\n\n"
|
||||
|
||||
summary_prompt = (
|
||||
"Resuma a conversa abaixo em um parágrafo conciso que capture:\n"
|
||||
"- Tópicos principais discutidos\n"
|
||||
"- Decisões tomadas e conclusões\n"
|
||||
"- Resultados de ferramentas/tools executadas\n"
|
||||
"- Contexto OCI/CIS relevante para perguntas futuras\n\n"
|
||||
"Seja conciso mas preserve toda informação acionável.\n\n"
|
||||
"CONVERSA:\n" + conversation_text
|
||||
)
|
||||
|
||||
summary_cfg = dict(genai_cfg)
|
||||
summary_cfg["system_prompt"] = "Você é um sumarizador. Responda apenas com o resumo."
|
||||
summary_cfg["max_tokens"] = COMPACT_SUMMARY_MAX_TOKENS
|
||||
|
||||
try:
|
||||
text, _, _ = _call_genai(summary_cfg, summary_prompt, history=None, tools=None)
|
||||
return text.strip()
|
||||
except Exception as e:
|
||||
log.warning(f"Summary generation failed: {e}")
|
||||
return ""
|
||||
|
||||
def _compact_history(session_id: str, user_id: str, genai_cfg: dict, history: list) -> list:
|
||||
"""Compact conversation history by summarizing older messages."""
|
||||
if len(history) <= COMPACT_KEEP_RECENT:
|
||||
return history
|
||||
|
||||
messages_to_summarize = history[:-COMPACT_KEEP_RECENT]
|
||||
recent_messages = history[-COMPACT_KEEP_RECENT:]
|
||||
|
||||
# Check for existing summary
|
||||
with db() as c:
|
||||
row = c.execute(
|
||||
"SELECT summary, messages_compacted FROM chat_summaries "
|
||||
"WHERE session_id=? ORDER BY created_at DESC LIMIT 1",
|
||||
(session_id,)
|
||||
).fetchone()
|
||||
if row:
|
||||
existing_summary = row["summary"]
|
||||
prev_compacted = row["messages_compacted"]
|
||||
# Reuse if no new messages to summarize
|
||||
if len(messages_to_summarize) <= prev_compacted:
|
||||
return [{"role": "assistant", "content": f"[Resumo da conversa anterior: {existing_summary}]"}] + recent_messages
|
||||
else:
|
||||
existing_summary = None
|
||||
|
||||
# Include existing summary as prefix for incremental compaction
|
||||
if existing_summary:
|
||||
messages_to_summarize = [{"role": "assistant", "content": f"[Resumo anterior: {existing_summary}]"}] + messages_to_summarize
|
||||
|
||||
summary_text = _generate_summary(genai_cfg, messages_to_summarize)
|
||||
|
||||
if not summary_text:
|
||||
# Fallback: truncate keeping recent messages that fit in budget
|
||||
truncated = []
|
||||
budget = 6000
|
||||
for m in reversed(history):
|
||||
t = _estimate_tokens(m["content"])
|
||||
if budget - t < 0:
|
||||
break
|
||||
truncated.insert(0, m)
|
||||
budget -= t
|
||||
return truncated
|
||||
|
||||
# Save summary to DB
|
||||
actual_count = len(messages_to_summarize) - (1 if existing_summary else 0)
|
||||
with db() as c:
|
||||
last_msg = c.execute(
|
||||
"SELECT id FROM chat_messages WHERE session_id=? AND role IN ('user','assistant') "
|
||||
"ORDER BY created_at ASC LIMIT 1 OFFSET ?",
|
||||
(session_id, max(0, actual_count - 1))
|
||||
).fetchone()
|
||||
last_id = last_msg["id"] if last_msg else "unknown"
|
||||
c.execute(
|
||||
"INSERT INTO chat_summaries (id,session_id,user_id,summary,messages_compacted,up_to_message_id) VALUES (?,?,?,?,?,?)",
|
||||
(str(uuid.uuid4()), session_id, user_id, summary_text, actual_count, last_id)
|
||||
)
|
||||
|
||||
log.info(f"Generated summary for session {session_id}: {len(summary_text)} chars, {actual_count} msgs compacted")
|
||||
return [{"role": "assistant", "content": f"[Resumo da conversa anterior: {summary_text}]"}] + recent_messages
|
||||
|
||||
# ── GenAI Call ─────────────────────────────────────────────────────────────────
|
||||
|
||||
def _call_genai(gc: dict, message: str, history: list = None, tools: list = None,
|
||||
tool_results_cohere: list = None, tool_results_generic: list = None) -> tuple:
|
||||
tool_results_cohere: list = None,
|
||||
extra_messages: list = None) -> tuple:
|
||||
"""
|
||||
Call OCI Generative AI with optional tool use support.
|
||||
Returns (text, tool_calls) tuple. tool_calls is a list of dicts or None.
|
||||
Returns (text, tool_calls, tool_calls_raw) tuple.
|
||||
tool_calls is a list of dicts or None. tool_calls_raw is the raw OCI SDK objects (for Generic format continuations).
|
||||
extra_messages: list of raw OCI SDK message objects to append after the user message (for tool use loop accumulation).
|
||||
"""
|
||||
import oci
|
||||
|
||||
@@ -992,16 +1114,6 @@ def _call_genai(gc: dict, message: str, history: list = None, tools: list = None
|
||||
msg.content = [content]
|
||||
messages.append(msg)
|
||||
|
||||
# Tool results from previous iteration (Generic format)
|
||||
if tool_results_generic:
|
||||
for tr in tool_results_generic:
|
||||
tool_msg = oci.generative_ai_inference.models.ToolMessage()
|
||||
tool_msg.tool_call_id = tr["tool_call_id"]
|
||||
tool_content = oci.generative_ai_inference.models.TextContent()
|
||||
tool_content.text = tr["content"]
|
||||
tool_msg.content = [tool_content]
|
||||
messages.append(tool_msg)
|
||||
|
||||
# Current user message
|
||||
content = oci.generative_ai_inference.models.TextContent()
|
||||
content.text = message
|
||||
@@ -1010,6 +1122,10 @@ def _call_genai(gc: dict, message: str, history: list = None, tools: list = None
|
||||
user_message.content = [content]
|
||||
messages.append(user_message)
|
||||
|
||||
# Append accumulated tool use loop messages (assistant+tool_calls → tool_results → ...)
|
||||
if extra_messages:
|
||||
messages.extend(extra_messages)
|
||||
|
||||
chat_request.messages = messages
|
||||
|
||||
# Tool definitions for Generic format
|
||||
@@ -1058,10 +1174,11 @@ def _call_genai(gc: dict, message: str, history: list = None, tools: list = None
|
||||
"name": tc.name,
|
||||
"arguments": tc.parameters if isinstance(tc.parameters, dict) else {}
|
||||
})
|
||||
return (text or "", tool_calls)
|
||||
return (text or "", tool_calls, None)
|
||||
else:
|
||||
text = ""
|
||||
tool_calls = None
|
||||
tool_calls_raw = None # raw content list from assistant message for continuation
|
||||
if hasattr(resp, 'choices') and resp.choices:
|
||||
choice = resp.choices[0]
|
||||
if hasattr(choice, 'message') and choice.message:
|
||||
@@ -1078,12 +1195,14 @@ def _call_genai(gc: dict, message: str, history: list = None, tools: list = None
|
||||
"name": tc.name if hasattr(tc, 'name') else "",
|
||||
"arguments": args
|
||||
})
|
||||
# Preserve raw tool_calls for building assistant message in next iteration
|
||||
tool_calls_raw = choice.message.tool_calls
|
||||
# Extract text
|
||||
if hasattr(choice.message, 'content') and choice.message.content:
|
||||
contents = choice.message.content
|
||||
if contents and len(contents) > 0:
|
||||
text = contents[0].text if hasattr(contents[0], 'text') else ""
|
||||
return (text or "", tool_calls)
|
||||
return (text or "", tool_calls, tool_calls_raw)
|
||||
|
||||
# ── RAG Helpers ───────────────────────────────────────────────────────────────
|
||||
RAG_CONTEXT_TEMPLATE = """--- CONTEXTO RECUPERADO (Vector Search) ---
|
||||
@@ -1098,11 +1217,6 @@ RAG_DEFAULT_SYSTEM_PROMPT = """Você é um assistente RAG especializado em Oracl
|
||||
- Responda **SOMENTE** perguntas relacionadas a Oracle Cloud (OCI), CIS Benchmark, CIS Report e itens de inventário OCI (IAM/AssetManagement/Networking/StorageBlock/FileStorageService/Objectstore/Compute/LoggingandMonitoring).
|
||||
- Se a pergunta não for desse escopo, recuse educadamente e peça uma pergunta dentro do tema.
|
||||
|
||||
### Tenancy (obrigatório)
|
||||
1) Confirme a **TENANCY** alvo antes de responder, a menos que já esteja definida na conversa.
|
||||
2) Se a tenancy ainda não estiver clara, pergunte uma única vez: "Qual tenancy devo considerar? (ex.: MinhaTenancy)" e aguarde a resposta.
|
||||
3) Se o usuário trocar a tenancy, considere a nova tenancy como a atual.
|
||||
|
||||
### Bases vetorizadas disponíveis (contexto recuperado automaticamente)
|
||||
- **CIS_REPORT**: report coletado na tenancy (compliance, findings, critérios/auditoria, evidências).
|
||||
- **CIS_RECOMMENDATIONS**: recomendações práticas de correção/remediação para itens não compliant.
|
||||
@@ -2101,11 +2215,19 @@ async def chat(msg: ChatMsg, u=Depends(current_user)):
|
||||
tool_defs = [t["tool"] for t in mcp_tools]
|
||||
log.info(f"Chat with {len(tool_defs)} MCP tools available")
|
||||
|
||||
# ── Memory compaction: summarize old messages if history is too long ──
|
||||
if history and _should_compact(history):
|
||||
log.info(f"Compaction triggered: {len(history)} msgs, ~{_estimate_history_tokens(history)} est tokens")
|
||||
history = _compact_history(sid, u["id"], cfg_dict, history)
|
||||
log.info(f"Post-compaction: {len(history)} msgs, ~{_estimate_history_tokens(history)} est tokens")
|
||||
|
||||
hist = history[:-1] if len(history) > 1 else None
|
||||
resp_text, tool_calls = _call_genai(cfg_dict, augmented_message, hist, tools=tool_defs)
|
||||
resp_text, tool_calls, tool_calls_raw = _call_genai(cfg_dict, augmented_message, hist, tools=tool_defs)
|
||||
|
||||
# Tool use loop (max 5 iterations)
|
||||
# Accumulate extra_messages so the model sees the full tool use conversation
|
||||
all_tool_results = []
|
||||
accumulated_msgs = [] # raw OCI SDK message objects for Generic format
|
||||
iterations = 0
|
||||
api_format = GENAI_MODELS.get(genai_cfg.get("model_id", ""), {}).get("api_format", "GENERIC")
|
||||
while tool_calls and iterations < 5:
|
||||
@@ -2138,14 +2260,26 @@ async def chat(msg: ChatMsg, u=Depends(current_user)):
|
||||
tr_obj.call = tc_obj
|
||||
tr_obj.outputs = [{"result": tr["content"]}]
|
||||
cohere_results.append(tr_obj)
|
||||
resp_text, tool_calls = _call_genai(
|
||||
resp_text, tool_calls, tool_calls_raw = _call_genai(
|
||||
cfg_dict, augmented_message, hist, tools=tool_defs,
|
||||
tool_results_cohere=cohere_results)
|
||||
else:
|
||||
generic_results = [{"tool_call_id": tr["tool_call_id"], "content": tr["content"]} for tr in iteration_results]
|
||||
resp_text, tool_calls = _call_genai(
|
||||
import oci
|
||||
# Accumulate: add assistant message with tool_calls + tool result messages
|
||||
assistant_msg = oci.generative_ai_inference.models.AssistantMessage()
|
||||
assistant_msg.tool_calls = tool_calls_raw
|
||||
accumulated_msgs.append(assistant_msg)
|
||||
for tr in iteration_results:
|
||||
tool_msg = oci.generative_ai_inference.models.ToolMessage()
|
||||
tool_msg.tool_call_id = tr["tool_call_id"]
|
||||
tool_content = oci.generative_ai_inference.models.TextContent()
|
||||
tool_content.text = tr["content"]
|
||||
tool_msg.content = [tool_content]
|
||||
accumulated_msgs.append(tool_msg)
|
||||
|
||||
resp_text, tool_calls, tool_calls_raw = _call_genai(
|
||||
cfg_dict, augmented_message, hist, tools=tool_defs,
|
||||
tool_results_generic=generic_results)
|
||||
extra_messages=accumulated_msgs)
|
||||
|
||||
resp = resp_text
|
||||
except Exception as e:
|
||||
@@ -2324,6 +2458,19 @@ async def startup():
|
||||
orphaned = c.execute("UPDATE reports SET status='failed', error_msg='Interrompido: container reiniciado' WHERE status='running'").rowcount
|
||||
if orphaned:
|
||||
log.warning(f"Marked {orphaned} orphaned running report(s) as failed")
|
||||
# Auto-register CIS Compliance Scanner MCP server if not present
|
||||
with db() as c:
|
||||
existing = c.execute("SELECT id FROM mcp_servers WHERE name='CIS Compliance Scanner'").fetchone()
|
||||
if not existing:
|
||||
mid = str(uuid.uuid4())
|
||||
c.execute(
|
||||
"INSERT INTO mcp_servers (id, user_id, name, description, server_type, command, args) VALUES (?,?,?,?,?,?,?)",
|
||||
(mid, "system", "CIS Compliance Scanner",
|
||||
"OCI CIS Foundations Benchmark 3.0 - 48 CIS checks + 11 OBP checks via MCP",
|
||||
"stdio", "python3", json.dumps(["/app/mcp_cis_server.py"]))
|
||||
)
|
||||
log.info(f"Auto-registered CIS Compliance Scanner MCP server (id={mid})")
|
||||
|
||||
log.info(f"OCI CIS AI Agent v{VERSION} API started")
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user