GustavoNogueira/A-Team-Security-Infra-Agent-Engineer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add RAG pipeline with OCI GenAI embeddings and ADB vector search

Browse Source 
Integrate Retrieval-Augmented Generation into the chat flow using OCI
GenAI embed_text API and Oracle Autonomous Database vector storage.
The chat automatically queries ADB for relevant context when an active
ADB config with a linked GenAI config exists.

- Add EMBEDDING_MODELS catalog (Cohere Embed v3.0/light)
- Add _embed_text() using OCI GenAI SDK embed endpoint
- Add _vector_search() with VECTOR_DISTANCE cosine similarity
- Add _get_adb_connection(), _ensure_embeddings_table(), _build_rag_context()
- Add document ingestion endpoint (POST /api/adb/{vid}/ingest)
- Add table creation endpoint (POST /api/adb/{vid}/ensure-table)
- Modify chat endpoint with automatic RAG augmentation (non-fatal)
- Add GenAI config and embedding model selectors to ADB UI
- Add RAG status indicator in chat toolbar
- Add document ingestion section in ADB tab
This commit is contained in:
nogueiraguh
2026-03-02 19:48:53 -03:00
parent 11343af59a
commit 2c3fb724bf
2 changed files with 340 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

260
backend/app.py
View File

@@ -60,6 +60,13 @@ GENAI_MODELS = {
"xai.grok-3": {"provider":"xai","name":"xAI Grok 3","api_format":"GENERIC"},
}
EMBEDDING_MODELS = {
"cohere.embed-english-v3.0": {"name":"Cohere Embed English v3.0","dims":1024},
"cohere.embed-multilingual-v3.0": {"name":"Cohere Embed Multilingual v3.0","dims":1024},
"cohere.embed-english-light-v3.0": {"name":"Cohere Embed English Light v3.0","dims":384},
"cohere.embed-multilingual-light-v3.0": {"name":"Cohere Embed Multilingual Light v3.0","dims":384},
}
GENAI_REGIONS = [
"us-chicago-1","us-ashburn-1","us-phoenix-1","uk-london-1",
"eu-frankfurt-1","ap-tokyo-1","ap-osaka-1","sa-saopaulo-1",
@@ -84,7 +91,11 @@ def init_db():
with db() as c:
c.executescript("""
CREATE TABLE IF NOT EXISTS users (
id TEXT PRIMARY KEY, username TEXT UNIQUE NOT NULL, email TEXT,
id TEXT PRIMARY KEY,
first_name TEXT NOT NULL,
last_name TEXT NOT NULL,
username TEXT UNIQUE NOT NULL,
email TEXT,
password_hash TEXT NOT NULL, role TEXT NOT NULL DEFAULT 'viewer',
mfa_secret TEXT, mfa_enabled INTEGER DEFAULT 0,
is_active INTEGER DEFAULT 1,
@@ -171,11 +182,17 @@ def init_db():
ip TEXT, created_at TEXT DEFAULT (datetime('now'))
);
""")
# ── Migrations ──
for col in ["genai_config_id TEXT", "embedding_model_id TEXT DEFAULT 'cohere.embed-multilingual-v3.0'"]:
try:
c.execute(f"ALTER TABLE adb_vector_configs ADD COLUMN {col}")
except sqlite3.OperationalError:
pass
adm = c.execute("SELECT id FROM users WHERE username='admin'").fetchone()
if not adm:
c.execute(
"INSERT INTO users (id,username,email,password_hash,role) VALUES (?,?,?,?,?)",
(str(uuid.uuid4()), "admin", "admin@local", _hash_pw("admin123"), "admin")
"INSERT INTO users (id,first_name,last_name,username,email,password_hash,role) VALUES (?,?,?,?,?,?,?)",
(str(uuid.uuid4()), "Admin", "Sistema", "admin", "admin@local", _hash_pw("admin123"), "admin")
)
log.info("Default admin created: admin / admin123")
@@ -228,7 +245,7 @@ def _audit(uid, uname, action, resource=None, details=None, ip=None):
class LoginReq(BaseModel):
username: str; password: str; totp_code: Optional[str] = None
class RegisterReq(BaseModel):
username: str; email: str; password: str; role: str = "viewer"
first_name: str; last_name: str; username: str; email: str; password: str; role: str = "viewer"
class TOTPVerify(BaseModel):
totp_code: str
class ChangePwReq(BaseModel):
@@ -251,6 +268,9 @@ class GenAIConfigReq(BaseModel):
class ADBVectorReq(BaseModel):
config_name: str; dsn: str; username: str; password: str
wallet_password: Optional[str] = None; table_name: str = "CIS_EMBEDDINGS"; use_mtls: bool = True
genai_config_id: Optional[str] = None; embedding_model_id: str = "cohere.embed-multilingual-v3.0"
class IngestDocReq(BaseModel):
adb_config_id: str; documents: List[Dict[str, Any]]
class MCPServerReq(BaseModel):
name: str; description: Optional[str] = None; server_type: str = "stdio"
command: Optional[str] = None; args: Optional[List[str]] = None
@@ -274,7 +294,7 @@ async def login(req: LoginReq, request: Request):
c.execute("UPDATE users SET last_login=datetime('now') WHERE id=?", (u["id"],))
_audit(u["id"], u["username"], "login", ip=request.client.host if request.client else None)
return {"token":_make_token(u["id"],u["role"],sid),
"user":{"id":u["id"],"username":u["username"],"email":u["email"],"role":u["role"],"mfa_enabled":bool(u["mfa_enabled"])}}
"user":{"id":u["id"],"first_name":u["first_name"],"last_name":u["last_name"],"username":u["username"],"email":u["email"],"role":u["role"],"mfa_enabled":bool(u["mfa_enabled"])}}
@app.post("/api/auth/logout")
async def logout_ep(u=Depends(current_user)):
@@ -287,7 +307,8 @@ async def register(req: RegisterReq, adm=Depends(require("admin"))):
uid = str(uuid.uuid4())
with db() as c:
if c.execute("SELECT 1 FROM users WHERE username=?", (req.username,)).fetchone(): raise HTTPException(400, "Usuário já existe")
c.execute("INSERT INTO users (id,username,email,password_hash,role) VALUES (?,?,?,?,?)", (uid, req.username, req.email, _hash_pw(req.password), req.role))
c.execute("INSERT INTO users (id,first_name,last_name,username,email,password_hash,role) VALUES (?,?,?,?,?,?,?)",
(uid, req.first_name, req.last_name, req.username, req.email, _hash_pw(req.password), req.role))
_audit(adm["id"], adm["username"], "create_user", uid, f"user={req.username} role={req.role}")
return {"id": uid, "username": req.username, "role": req.role}
@@ -320,12 +341,12 @@ async def mfa_disable(user_id: str, adm=Depends(require("admin"))):
# ── Users ─────────────────────────────────────────────────────────────────────
@app.get("/api/users")
async def list_users(u=Depends(require("admin"))):
with db() as c: rows = c.execute("SELECT id,username,email,role,mfa_enabled,is_active,created_at,last_login FROM users").fetchall()
with db() as c: rows = c.execute("SELECT id,first_name,last_name,username,email,role,mfa_enabled,is_active,created_at,last_login FROM users").fetchall()
return [dict(r) for r in rows]
@app.get("/api/users/me")
async def me(u=Depends(current_user)):
return {k: u[k] for k in ("id","username","email","role","mfa_enabled")}
return {k: u[k] for k in ("id","first_name","last_name","username","email","role","mfa_enabled")}
@app.put("/api/users/{uid}")
async def update_user(uid: str, req: UserUpdateReq, adm=Depends(require("admin"))):
@@ -492,7 +513,7 @@ async def explore_buckets(cid: str, compartment_id: str = Query(None), u=Depends
# ── OCI GenAI Config & Chat ───────────────────────────────────────────────────
@app.get("/api/genai/models")
async def list_genai_models(u=Depends(current_user)):
return {"models": GENAI_MODELS, "regions": GENAI_REGIONS}
return {"models": GENAI_MODELS, "regions": GENAI_REGIONS, "embedding_models": EMBEDDING_MODELS}
@app.post("/api/genai/config")
async def save_genai(req: GenAIConfigReq, u=Depends(require("admin","user"))):
@@ -644,6 +665,130 @@ def _call_genai(gc: dict, message: str, history: list = None) -> str:
return contents[0].text
return str(resp)
# ── RAG Helpers ───────────────────────────────────────────────────────────────
RAG_SYSTEM_PROMPT = """You are the OCI CIS AI Agent, an expert assistant for Oracle Cloud Infrastructure security and compliance.
You have been provided with relevant context documents retrieved from a knowledge base. Use this context to provide accurate, specific answers. If the context does not contain relevant information for the question, say so and answer based on your general knowledge.
Always cite the source documents when using information from the context.
--- RETRIEVED CONTEXT ---
{context}
--- END CONTEXT ---
User question: {question}"""
def _get_adb_connection(cfg: dict):
"""Create an oracledb connection from an adb_vector_configs row."""
import oracledb
params = {"user": cfg["username"], "password": _dec(cfg["password_enc"]), "dsn": cfg["dsn"]}
if cfg["use_mtls"] and cfg.get("wallet_dir"):
params["config_dir"] = cfg["wallet_dir"]
params["wallet_location"] = cfg["wallet_dir"]
if cfg.get("wallet_password_enc"):
params["wallet_password"] = _dec(cfg["wallet_password_enc"])
return oracledb.connect(**params)
def _ensure_embeddings_table(cfg: dict):
"""Create the embeddings table in ADB if it doesn't exist."""
table_name = cfg.get("table_name", "CIS_EMBEDDINGS")
emb_model = cfg.get("embedding_model_id", "cohere.embed-multilingual-v3.0")
dims = EMBEDDING_MODELS.get(emb_model, {}).get("dims", 1024)
conn = _get_adb_connection(cfg)
try:
cur = conn.cursor()
cur.execute("SELECT COUNT(*) FROM user_tables WHERE table_name = :1", (table_name.upper(),))
if cur.fetchone()[0] == 0:
cur.execute(f"""
CREATE TABLE {table_name} (
ID VARCHAR2(100) PRIMARY KEY,
CONTENT CLOB,
EMBEDDING VECTOR({dims}, FLOAT64),
METADATA VARCHAR2(4000),
SOURCE VARCHAR2(500),
CREATED_AT TIMESTAMP DEFAULT SYSTIMESTAMP
)
""")
conn.commit()
log.info(f"Created embeddings table: {table_name} (dims={dims})")
cur.close()
finally:
conn.close()
def _embed_text(text: str, genai_cfg: dict, embedding_model_id: str) -> list:
"""Generate embedding using OCI GenAI embed endpoint."""
import oci
config_path = str(OCI_DIR / genai_cfg["oci_config_id"] / "config")
config = oci.config.from_file(config_path, "DEFAULT")
endpoint = genai_cfg["endpoint"]
client = oci.generative_ai_inference.GenerativeAiInferenceClient(
config=config, service_endpoint=endpoint,
retry_strategy=oci.retry.NoneRetryStrategy(), timeout=(10, 120)
)
embed_detail = oci.generative_ai_inference.models.EmbedTextDetails()
embed_detail.inputs = [text]
embed_detail.serving_mode = oci.generative_ai_inference.models.OnDemandServingMode(model_id=embedding_model_id)
embed_detail.compartment_id = genai_cfg["compartment_id"]
embed_detail.truncate = "NONE"
embed_detail.input_type = "SEARCH_QUERY"
response = client.embed_text(embed_detail)
return response.data.embeddings[0]
def _vector_search(cfg: dict, query_embedding: list, top_k: int = 5) -> list:
"""Search ADB vector store using cosine similarity. Returns top-K documents."""
import array
table_name = cfg.get("table_name", "CIS_EMBEDDINGS")
conn = _get_adb_connection(cfg)
try:
cur = conn.cursor()
vec = array.array('d', query_embedding)
cur.execute(f"""
SELECT ID, CONTENT, METADATA, SOURCE,
VECTOR_DISTANCE(EMBEDDING, :1, COSINE) AS distance
FROM {table_name}
ORDER BY distance ASC
FETCH FIRST :2 ROWS ONLY
""", [vec, top_k])
results = []
for row in cur:
content = row[1]
if hasattr(content, 'read'):
content = content.read()
results.append({
"id": row[0], "content": content or "",
"metadata": row[2], "source": row[3], "distance": float(row[4])
})
cur.close()
return results
finally:
conn.close()
def _build_rag_context(documents: list) -> str:
"""Format retrieved documents into a context string for the LLM prompt."""
if not documents:
return ""
parts = []
for i, doc in enumerate(documents, 1):
source = doc.get("source", "unknown")
content = doc.get("content", "")
if len(content) > 2000:
content = content[:2000] + "..."
parts.append(f"[Document {i} | Source: {source}]\n{content}")
return "\n\n---\n\n".join(parts)
def _get_active_adb_config(user_id: str) -> dict | None:
"""Get the first active ADB vector config with a linked GenAI config."""
with db() as c:
cfg = c.execute(
"SELECT * FROM adb_vector_configs WHERE user_id=? AND is_active=1 AND genai_config_id IS NOT NULL ORDER BY created_at DESC LIMIT 1",
(user_id,)
).fetchone()
if not cfg:
cfg = c.execute(
"SELECT * FROM adb_vector_configs WHERE is_active=1 AND genai_config_id IS NOT NULL ORDER BY created_at DESC LIMIT 1"
).fetchone()
return dict(cfg) if cfg else None
# ── MCP Servers ───────────────────────────────────────────────────────────────
@app.post("/api/mcp/servers")
async def register_mcp(req: MCPServerReq, u=Depends(require("admin","user"))):
@@ -705,9 +850,10 @@ async def save_adb(req: ADBVectorReq, u=Depends(require("admin","user"))):
vid = str(uuid.uuid4())
with db() as c:
c.execute(
"INSERT INTO adb_vector_configs (id,user_id,config_name,dsn,username,password_enc,wallet_password_enc,table_name,use_mtls) VALUES (?,?,?,?,?,?,?,?,?)",
"INSERT INTO adb_vector_configs (id,user_id,config_name,dsn,username,password_enc,wallet_password_enc,table_name,use_mtls,genai_config_id,embedding_model_id) VALUES (?,?,?,?,?,?,?,?,?,?,?)",
(vid, u["id"], req.config_name, req.dsn, req.username, _enc(req.password),
_enc(req.wallet_password) if req.wallet_password else None, req.table_name, int(req.use_mtls)))
_enc(req.wallet_password) if req.wallet_password else None, req.table_name, int(req.use_mtls),
req.genai_config_id, req.embedding_model_id))
_audit(u["id"], u["username"], "save_adb_config", vid, req.config_name)
return {"id": vid, "config_name": req.config_name}
@@ -724,8 +870,9 @@ async def upload_wallet(vid: str, wallet: UploadFile = File(...), u=Depends(requ
@app.get("/api/adb/configs")
async def list_adb(u=Depends(current_user)):
with db() as c:
if u["role"]=="admin": rows=c.execute("SELECT id,config_name,dsn,username,table_name,use_mtls,is_active,wallet_dir,created_at FROM adb_vector_configs").fetchall()
else: rows=c.execute("SELECT id,config_name,dsn,username,table_name,use_mtls,is_active,wallet_dir,created_at FROM adb_vector_configs WHERE user_id=?",(u["id"],)).fetchall()
cols = "id,config_name,dsn,username,table_name,use_mtls,is_active,wallet_dir,genai_config_id,embedding_model_id,created_at"
if u["role"]=="admin": rows=c.execute(f"SELECT {cols} FROM adb_vector_configs").fetchall()
else: rows=c.execute(f"SELECT {cols} FROM adb_vector_configs WHERE user_id=?",(u["id"],)).fetchall()
return [dict(r) for r in rows]
@app.post("/api/adb/test/{vid}")
@@ -734,14 +881,7 @@ async def test_adb(vid: str, u=Depends(require("admin","user"))):
cfg = c.execute("SELECT * FROM adb_vector_configs WHERE id=?",(vid,)).fetchone()
if not cfg: raise HTTPException(404)
try:
import oracledb
params = {"user": cfg["username"], "password": _dec(cfg["password_enc"]), "dsn": cfg["dsn"]}
if cfg["use_mtls"] and cfg.get("wallet_dir"):
params["config_dir"] = cfg["wallet_dir"]
params["wallet_location"] = cfg["wallet_dir"]
if cfg.get("wallet_password_enc"):
params["wallet_password"] = _dec(cfg["wallet_password_enc"])
conn = oracledb.connect(**params)
conn = _get_adb_connection(dict(cfg))
cur = conn.cursor(); cur.execute("SELECT 1 FROM DUAL"); cur.close(); conn.close()
return {"status":"success","message":"Conexão Autonomous DB OK!"}
except ImportError:
@@ -756,6 +896,63 @@ async def del_adb(vid: str, u=Depends(require("admin","user"))):
if d.exists(): shutil.rmtree(d)
return {"ok": True}
@app.post("/api/adb/{vid}/ensure-table")
async def ensure_table(vid: str, u=Depends(require("admin","user"))):
with db() as c:
cfg = c.execute("SELECT * FROM adb_vector_configs WHERE id=?", (vid,)).fetchone()
if not cfg: raise HTTPException(404)
try:
_ensure_embeddings_table(dict(cfg))
return {"ok": True, "table": cfg["table_name"]}
except Exception as e:
raise HTTPException(500, f"Falha ao criar tabela: {str(e)[:500]}")
@app.post("/api/adb/{vid}/ingest")
async def ingest_documents(vid: str, req: IngestDocReq, bg: BackgroundTasks, u=Depends(require("admin","user"))):
with db() as c:
cfg = c.execute("SELECT * FROM adb_vector_configs WHERE id=?", (vid,)).fetchone()
if not cfg: raise HTTPException(404, "ADB config not found")
cfg = dict(cfg)
if not cfg.get("genai_config_id"):
raise HTTPException(400, "GenAI config not linked to this ADB connection")
with db() as c:
gc = c.execute("SELECT * FROM genai_configs WHERE id=?", (cfg["genai_config_id"],)).fetchone()
if not gc: raise HTTPException(400, "Linked GenAI config not found")
bg.add_task(_ingest_documents_task, cfg, dict(gc), req.documents, u["id"], u["username"])
return {"ok": True, "message": f"Ingestão iniciada para {len(req.documents)} documentos", "config_id": vid}
def _ingest_documents_task(cfg: dict, genai_cfg: dict, documents: list, user_id: str, username: str):
"""Background task: embed and insert documents into ADB via OCI GenAI."""
import array
emb_model = cfg.get("embedding_model_id", "cohere.embed-multilingual-v3.0")
table_name = cfg.get("table_name", "CIS_EMBEDDINGS")
_ensure_embeddings_table(cfg)
conn = _get_adb_connection(cfg)
try:
cur = conn.cursor()
inserted = 0
for doc in documents:
try:
content = doc.get("content", "")
if not content: continue
embedding = _embed_text(content, genai_cfg, emb_model)
vec = array.array('d', embedding)
cur.execute(f"""
INSERT INTO {table_name} (ID, CONTENT, EMBEDDING, METADATA, SOURCE)
VALUES (:1, :2, :3, :4, :5)
""", [str(uuid.uuid4()), content, vec, doc.get("metadata", ""), doc.get("source", "manual_upload")])
inserted += 1
except Exception as e:
log.error(f"Failed to ingest document: {e}")
conn.commit()
cur.close()
log.info(f"Ingested {inserted}/{len(documents)} documents into {table_name}")
_audit(user_id, username, "ingest_documents", cfg["id"], f"{inserted} documents")
except Exception as e:
log.error(f"Ingestion task failed: {e}")
finally:
conn.close()
# ── Reports ───────────────────────────────────────────────────────────────────
@app.post("/api/reports/run")
async def run_report(req: RunReportReq, bg: BackgroundTasks, u=Depends(require("admin","user"))):
@@ -852,7 +1049,26 @@ async def chat(msg: ChatMsg, u=Depends(current_user)):
with db() as c:
prev = c.execute("SELECT role,content FROM chat_messages WHERE session_id=? AND role IN ('user','assistant') ORDER BY created_at ASC", (sid,)).fetchall()
history = [{"role":r["role"],"content":r["content"]} for r in prev]
resp = _call_genai(dict(genai_cfg), msg.message, history[:-1] if len(history) > 1 else None)
# ── RAG: augment with vector context if ADB config is active ──
rag_context = ""
adb_cfg = _get_active_adb_config(u["id"])
if adb_cfg:
try:
with db() as c:
emb_genai = c.execute("SELECT * FROM genai_configs WHERE id=?", (adb_cfg["genai_config_id"],)).fetchone()
if emb_genai:
emb_model = adb_cfg.get("embedding_model_id", "cohere.embed-multilingual-v3.0")
query_embedding = _embed_text(msg.message, dict(emb_genai), emb_model)
documents = _vector_search(adb_cfg, query_embedding, top_k=5)
if documents:
rag_context = _build_rag_context(documents)
log.info(f"RAG: Retrieved {len(documents)} documents for query")
except Exception as e:
log.warning(f"RAG retrieval failed (non-fatal): {e}")
augmented_message = RAG_SYSTEM_PROMPT.format(context=rag_context, question=msg.message) if rag_context else msg.message
resp = _call_genai(dict(genai_cfg), augmented_message, history[:-1] if len(history) > 1 else None)
except Exception as e:
resp = f"❌ Erro GenAI: {str(e)[:400]}"
else: