From a5587ed7909603bb2939b9835057e1740f05e08c Mon Sep 17 00:00:00 2001 From: cozyty Date: Fri, 13 Mar 2026 01:41:27 -0400 Subject: [PATCH 01/21] feat: add api_base support to RAGService for custom LLM endpoints - Add api_base parameter to RAGService __init__ - Pass api_base to dspy.Embedder to support non-OpenAI endpoints - Pass api_base through mcp_svc.py when constructing RAGService - Enables compat with other embedding endpoints --- app/mcp_svc.py | 2 +- app/rag.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/app/mcp_svc.py b/app/mcp_svc.py index e32b796..63806a1 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -70,7 +70,7 @@ def _build_rag_service_from_files(self, filenames, api_key: str, embed_model: st with open(path, "r", encoding="utf-8") as f: bundles.append(json.load(f)) - rag = RAGService(api_key=api_key, log=self.log) + rag = RAGService(api_key=api_key, api_base=llm_config.get("api_base"), log=self.log) if topk: rag.topk_objects_to_retrieve = int(topk) rag.initialize_from_bundles(bundles, embed_model=embed_model or 'openai/text-embedding-3-small') diff --git a/app/rag.py b/app/rag.py index 8ead001..76159b9 100644 --- a/app/rag.py +++ b/app/rag.py @@ -6,13 +6,14 @@ class RAGService: """RAG service for CTI (Cyber Threat Intelligence) data retrieval using STIX bundles.""" - def __init__(self, stix_bundle_path: Optional[str] = None, api_key: Optional[str] = None, log: Optional[logging.Logger] = None): + def __init__(self, stix_bundle_path: Optional[str] = None, api_key: Optional[str] = None, api_base: Optional[str] = None, log: Optional[logging.Logger] = None): self.max_characters = 6000 self.topk_objects_to_retrieve = 5 self.corpus = [] self.adv_step = {} self.search = None self.api_key = api_key + self.api_base = api_base self.log = log or logging.getLogger("plugins.mcp") self.log.info(f"Loading STIX bundle from: {stix_bundle_path}") @@ -45,7 +46,7 @@ def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = ' self.adv_step = all_adv_step self.log.info("Initializing embeddings and retriever for STIX corpus") - embedder = dspy.Embedder(embed_model, api_key=self.api_key) + embedder = dspy.Embedder(embed_model, api_key=self.api_key, api_base=self.api_base) self.search = dspy.retrievers.Embeddings( corpus=self.corpus, embedder=embedder, From 4fcab48c3b07456f302cc4e9d935a5dee907459d Mon Sep 17 00:00:00 2001 From: cozyty Date: Fri, 13 Mar 2026 01:51:19 -0400 Subject: [PATCH 02/21] feat: use nvidia/llama-3.2-nv-embedqa-1b-v2 as default embedding model --- app/mcp_svc.py | 4 ++-- app/rag.py | 4 ++-- gui/views/mcp.vue | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/app/mcp_svc.py b/app/mcp_svc.py index 63806a1..9f71a71 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -73,7 +73,7 @@ def _build_rag_service_from_files(self, filenames, api_key: str, embed_model: st rag = RAGService(api_key=api_key, api_base=llm_config.get("api_base"), log=self.log) if topk: rag.topk_objects_to_retrieve = int(topk) - rag.initialize_from_bundles(bundles, embed_model=embed_model or 'openai/text-embedding-3-small') + rag.initialize_from_bundles(bundles, embed_model=embed_model or 'nvidia/llama-3.2-nv-embedqa-1b-v2') return rag async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: dict = None): @@ -99,7 +99,7 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d self.log.warning(f"[MCP] Failed to configure LM: {e}") rag_files = run_config.get("rag_files") or [] - rag_embed_model = run_config.get("rag_embed_model") or 'openai/text-embedding-3-small' + rag_embed_model = run_config.get("rag_embed_model") or 'nvidia/llama-3.2-nv-embedqa-1b-v2' rag_topk = run_config.get("rag_topk") # Use RAG if explicitly requested via focus or if files were selected diff --git a/app/rag.py b/app/rag.py index 76159b9..ccee42a 100644 --- a/app/rag.py +++ b/app/rag.py @@ -22,7 +22,7 @@ def __init__(self, stix_bundle_path: Optional[str] = None, api_key: Optional[str if stix_bundle_path: self.load_stix_bundle(stix_bundle_path) - def load_stix_bundle(self, stix_bundle_path: str, embed_model: str = 'openai/text-embedding-3-small'): + def load_stix_bundle(self, stix_bundle_path: str, embed_model: str = 'nvidia/llama-3.2-nv-embedqa-1b-v2'): """Load STIX bundle from file path and build embeddings.""" try: with open(stix_bundle_path, 'r') as f: @@ -33,7 +33,7 @@ def load_stix_bundle(self, stix_bundle_path: str, embed_model: str = 'openai/tex except json.JSONDecodeError: raise ValueError(f"Invalid JSON in STIX bundle: {stix_bundle_path}") - def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = 'openai/text-embedding-3-small'): + def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = 'nvidia/llama-3.2-nv-embedqa-1b-v2'): """Initialize the RAG service with multiple STIX bundles and create retriever.""" all_corpus = [] all_adv_step = {} diff --git a/gui/views/mcp.vue b/gui/views/mcp.vue index 1ac348a..624158d 100644 --- a/gui/views/mcp.vue +++ b/gui/views/mcp.vue @@ -162,7 +162,7 @@ class="input" type="text" v-model="globalConfig.ragEmbedModel" - placeholder="openai/text-embedding-3-small" + placeholder="nvidia/llama-3.2-nv-embedqa-1b-v2" /> @@ -500,7 +500,7 @@ const globalConfig = reactive({ apiKey: savedConfig?.apiKey || '', maxToolCalls: savedConfig?.maxToolCalls || 5, maxTokens: savedConfig?.maxTokens || 10000, - ragEmbedModel: savedConfig?.ragEmbedModel || 'openai/text-embedding-3-small', + ragEmbedModel: savedConfig?.ragEmbedModel || 'nvidia/llama-3.2-nv-embedqa-1b-v2', ragTopK: savedConfig?.ragTopK ?? 5 }) From 785e9393985db3f0aa9bc8e5ab27e89fb1bf39ad Mon Sep 17 00:00:00 2001 From: cozyty Date: Fri, 13 Mar 2026 13:50:25 -0400 Subject: [PATCH 03/21] feat: add api_base field to Global Model Config UI --- gui/views/mcp.vue | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/gui/views/mcp.vue b/gui/views/mcp.vue index 624158d..7f27be2 100644 --- a/gui/views/mcp.vue +++ b/gui/views/mcp.vue @@ -114,6 +114,17 @@ /> +
+ +
+ +
+
@@ -403,6 +414,7 @@ async function handleCustomSubmit() { config: { model: globalConfig.modelName, api_key: globalConfig.apiKey, + api_base: globalConfig.apiBase, temperature: globalConfig.temperature, max_tokens: globalConfig.maxTokens, max_tool_calls: globalConfig.maxToolCalls @@ -484,7 +496,8 @@ function saveConfig(config) { maxToolCalls: config.maxToolCalls, maxTokens: config.maxTokens, ragEmbedModel: config.ragEmbedModel, - ragTopK: config.ragTopK + ragTopK: config.ragTopK, + apiBase: config.apiBase, }) } catch (e) { console.warn('[MCP] Failed to save config:', e) @@ -498,6 +511,7 @@ const globalConfig = reactive({ modelName: savedConfig?.modelName || 'gpt-4o', temperature: savedConfig?.temperature ?? 0.5, apiKey: savedConfig?.apiKey || '', + apiBase: savedConfig?.apiBase || '', maxToolCalls: savedConfig?.maxToolCalls || 5, maxTokens: savedConfig?.maxTokens || 10000, ragEmbedModel: savedConfig?.ragEmbedModel || 'nvidia/llama-3.2-nv-embedqa-1b-v2', From f450a4fa0cf950ec29dfac80bfff4bee6a52afe2 Mon Sep 17 00:00:00 2001 From: cozyty Date: Fri, 13 Mar 2026 13:55:40 -0400 Subject: [PATCH 04/21] fix: pass api_base to dspy.LM for custom OpenAI-compatible endpoints --- app/mcp_svc.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/app/mcp_svc.py b/app/mcp_svc.py index 9f71a71..ee48125 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -92,7 +92,8 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d dspy.configure(lm=dspy.LM( model=lm_obj.get("model"), api_key=lm_obj.get("api_key"), - temperature=lm_obj.get("temperature"), + api_base=lm_obj.get("api_base"), + temperature=lm_obj.get("temperature"), max_tokens=lm_obj.get("max_tokens"), )) except Exception as e: @@ -163,4 +164,4 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d mlflow.log_param("error", str(e)) finally: - mlflow.end_run() \ No newline at end of file + mlflow.end_run() From 1d7d9325e17f9e770676011feac4957944ae6d1a Mon Sep 17 00:00:00 2001 From: cozyty Date: Fri, 13 Mar 2026 14:03:25 -0400 Subject: [PATCH 05/21] fix: use openai/ prefix for litellm and dspy.context for async safety --- app/mcp_svc.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/app/mcp_svc.py b/app/mcp_svc.py index ee48125..a45290b 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -60,7 +60,7 @@ async def execute(self, focus: str, prompt: str, model_config: dict, file: dict )) return {"run_id": run_id} - def _build_rag_service_from_files(self, filenames, api_key: str, embed_model: str, topk: int): + def _build_rag_service_from_files(self, filenames, api_key: str, embed_model: str, topk: int, api_base: str = None): base_dir = Path(__file__).resolve().parent.parent / "data" bundles = [] for name in filenames or []: @@ -70,7 +70,7 @@ def _build_rag_service_from_files(self, filenames, api_key: str, embed_model: st with open(path, "r", encoding="utf-8") as f: bundles.append(json.load(f)) - rag = RAGService(api_key=api_key, api_base=llm_config.get("api_base"), log=self.log) + rag = RAGService(api_key=api_key, api_base=api_base, log=self.log) if topk: rag.topk_objects_to_retrieve = int(topk) rag.initialize_from_bundles(bundles, embed_model=embed_model or 'nvidia/llama-3.2-nv-embedqa-1b-v2') @@ -89,13 +89,14 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d # Configure LM globally if provided if lm_obj and lm_obj.get("api_key"): try: - dspy.configure(lm=dspy.LM( - model=lm_obj.get("model"), + lm = dspy.LM( + model="openai/" + lm_obj.get("model"), api_key=lm_obj.get("api_key"), - api_base=lm_obj.get("api_base"), - temperature=lm_obj.get("temperature"), + api_base=lm_obj.get("api_base"), + temperature=lm_obj.get("temperature"), max_tokens=lm_obj.get("max_tokens"), - )) + ) + dspy.context(lm=lm) except Exception as e: self.log.warning(f"[MCP] Failed to configure LM: {e}") @@ -113,7 +114,8 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d rag = self._build_rag_service_from_files( filenames=rag_files, api_key=(lm_obj or {}).get("api_key"), - embed_model=rag_embed_model, + api_base=(lm_obj or {}).get("api_base"), + embed_model=rag_embed_model, topk=rag_topk or 5 ) rag_context = rag.get_context_for_task(prompt) From 22281bb6c965d28c5eda1ae9f43aabf05967e146 Mon Sep 17 00:00:00 2001 From: cozyty Date: Fri, 13 Mar 2026 14:09:33 -0400 Subject: [PATCH 06/21] fix: add openai/ prefix and DSPY_API_BASE support across all LM configs --- app/factory.py | 12 +++++++++--- app/mcp_factory_client.py | 29 ++++++++--------------------- app/mcp_planner_client.py | 2 +- 3 files changed, 18 insertions(+), 25 deletions(-) diff --git a/app/factory.py b/app/factory.py index 9008b32..71f2348 100644 --- a/app/factory.py +++ b/app/factory.py @@ -6,11 +6,17 @@ def configure_dspy_from_env(): model = os.environ.get('DSPY_MODEL', 'gpt-4o') api_key = os.environ.get('DSPY_API_KEY', '') + api_base = os.environ.get('DSPY_API_BASE', '') or None temperature = float(os.environ.get('DSPY_TEMPERATURE', '0.5')) max_tokens = int(os.environ.get('DSPY_MAX_TOKENS', '10000')) - - if api_key: # Only configure if we have an API key - lm = dspy.LM(model=model, api_key=api_key, temperature=temperature, max_tokens=max_tokens) + if api_key: + lm = dspy.LM( + model="openai/" + model, + api_key=api_key, + api_base=api_base, + temperature=temperature, + max_tokens=max_tokens + ) dspy.configure(lm=lm) configure_dspy_from_env() diff --git a/app/mcp_factory_client.py b/app/mcp_factory_client.py index 22851d7..00f6af5 100644 --- a/app/mcp_factory_client.py +++ b/app/mcp_factory_client.py @@ -46,19 +46,17 @@ def get_env(lm_settings=None): else: env['PYTHONPATH'] = venv_site_packages - # Pass LLM config to subprocess via environment variables if lm_settings: - # Use 'or' to handle None values and ensure we always get strings env['DSPY_MODEL'] = str(lm_settings.get('model') or 'gpt-4o') env['DSPY_API_KEY'] = str(lm_settings.get('api_key') or '') env['DSPY_TEMPERATURE'] = str(lm_settings.get('temperature') or 0.5) env['DSPY_MAX_TOKENS'] = str(lm_settings.get('max_tokens') or 10000) + env['DSPY_API_BASE'] = str(lm_settings.get('api_base') or '') # FIXED return env mlflow.set_tracking_uri("http://localhost:5000") mlflow.set_experiment("caldera-mcp-FACTORY-client-1") -# mlflow.dspy.autolog() current_dir = os.path.dirname(os.path.abspath(__file__)) @@ -93,7 +91,6 @@ class DSPyCalderaFactoryClientWithRAG(dspy.Signature): ) ) -# Factory function to create tool functions with proper closure def create_tool_function(session, tool_name, tool_description): async def tool_function(**kwargs): mlflow.set_tag("stage", f"Tool.{tool_name}") @@ -103,19 +100,16 @@ async def tool_function(**kwargs): return tool_function def format_rag_context(rag_context): - """Format RAG context into a string for the DSPy signature.""" if not rag_context: return "No CTI context available." formatted_parts = [] - # Add search results summary if "search_results" in rag_context: formatted_parts.append("Relevant CTI findings:") for i, result in enumerate(rag_context["search_results"][:3], 1): formatted_parts.append(f"{i}. {result}") - # Add detailed context if "detailed_context" in rag_context: formatted_parts.append("\nDetailed CTI Information:") for ctx in rag_context["detailed_context"]: @@ -125,14 +119,14 @@ def format_rag_context(rag_context): return "\n".join(formatted_parts) async def run(adversary_emulation_task: str, lm_obj = None, rag_context=None, run_id=None): - # Build LM settings safely (support defaults) lm_settings = {} - max_tool_calls = 5 # Default value + max_tool_calls = 5 if lm_obj: lm_obj_safe = copy.deepcopy(lm_obj) or {} lm_settings = { "model": lm_obj_safe.get("model") or "gpt-4o", "api_key": lm_obj_safe.get("api_key") or "", + "api_base": lm_obj_safe.get("api_base") or "", # FIXED "temperature": lm_obj_safe.get("temperature") or 0.5, "max_tokens": lm_obj_safe.get("max_tokens") or 10000, } @@ -142,12 +136,12 @@ async def run(adversary_emulation_task: str, lm_obj = None, rag_context=None, ru lm_settings = { "model": llm_config.get("model") or "gpt-4o", "api_key": llm_config.get("api_key") or "", + "api_base": llm_config.get("api_base") or "", # FIXED "temperature": llm_config.get("temperature") or 0.5, "max_tokens": llm_config.get("max_tokens") or 10000, } max_tool_calls = llm_config.get("max_tool_calls") or 5 - # Validate API key is provided if not lm_settings.get("api_key"): error_msg = "API key is required but not provided. Please set your API key in the Global Model Configuration." print(f"[MCP] ERROR: {error_msg}") @@ -161,7 +155,6 @@ async def run(adversary_emulation_task: str, lm_obj = None, rag_context=None, ru mlflow.end_run() raise ValueError(error_msg) - # Use the passed-in run_id to continue the MLflow run if provided created_local_run = False if not run_id: run = mlflow.start_run(run_name="MCP Ability Factory") @@ -172,7 +165,6 @@ async def run(adversary_emulation_task: str, lm_obj = None, rag_context=None, ru mlflow.set_tag("stage", "initializing") mlflow.log_param("prompt", adversary_emulation_task) - # Create server params with LLM settings passed via environment server_params = StdioServerParameters( command="python", args=[current_dir+"/mcp_server.py"], @@ -182,17 +174,17 @@ async def run(adversary_emulation_task: str, lm_obj = None, rag_context=None, ru try: async with stdio_client(server_params) as (read, write): async with ClientSession(read, write) as session: - # Initialize MCP session and list tools mlflow.set_tag("stage", "initializing MCP session") await session.initialize() mlflow.set_tag("stage", "listing tools") tools = await session.list_tools() - # Use context to set LM for this task/run + # FIXED: add openai/ prefix and api_base with dspy.context(lm=dspy.LM( - lm_settings['model'], + "openai/" + lm_settings['model'], api_key=lm_settings['api_key'], + api_base=lm_settings['api_base'], temperature=lm_settings['temperature'], max_tokens=lm_settings['max_tokens'] )): @@ -203,8 +195,7 @@ async def run(adversary_emulation_task: str, lm_obj = None, rag_context=None, ru signature = DSPyCalderaFactoryClientWithRAG formatted_context = format_rag_context(rag_context) - # Log CTI context being sent to LLM for verification - mlflow.log_param("cti_context_preview", formatted_context[:1000]) # First 1000 chars + mlflow.log_param("cti_context_preview", formatted_context[:1000]) mlflow.set_tag("cti_context_length", len(formatted_context)) mlflow.set_tag("cti_search_results_count", len(rag_context.get("search_results", []))) mlflow.set_tag("cti_detailed_context_count", len(rag_context.get("detailed_context", []))) @@ -219,13 +210,10 @@ async def run(adversary_emulation_task: str, lm_obj = None, rag_context=None, ru mlflow.set_tag("stage", "executing DSPy ReAct") result = await react.acall(adversary_emulation_task=adversary_emulation_task) - # Log outputs and trajectory mlflow.set_tag("stage", "completed") mlflow.set_tag("status", "complete") mlflow.set_tag("reasoning", result.reasoning) - # Prefer param for process_result to match status API mlflow.log_param("process_result", result.process_result) - # Keep tag for backward compatibility (optional) mlflow.set_tag("process_result", result.process_result) for k, v in result.trajectory.items(): @@ -235,7 +223,6 @@ async def run(adversary_emulation_task: str, lm_obj = None, rag_context=None, ru print(json.dumps(result.toDict(), indent=4)) - # End the run only if we created it locally if created_local_run: mlflow.end_run() diff --git a/app/mcp_planner_client.py b/app/mcp_planner_client.py index 0098d5c..6881161 100644 --- a/app/mcp_planner_client.py +++ b/app/mcp_planner_client.py @@ -31,7 +31,7 @@ def build_lm_from_dict(settings: dict) -> dspy.LM: raise ValueError("API key is required but not provided. Please set your API key in the Global Model Configuration.") lm_kwargs = { - "model": settings.get("model") or "gpt-4o", + "model": "openai/" + (settings.get("model") or "gpt-4o"), "api_key": api_key, "api_base": settings.get("api_base"), } From 9c022c5badc2265265e52b3b5a012e1e26dda61c Mon Sep 17 00:00:00 2001 From: cozyty Date: Fri, 13 Mar 2026 14:17:29 -0400 Subject: [PATCH 07/21] fix: pass api_base from globalConfig in factory and planner payloads --- gui/views/local_mcp_ability_factory.vue | 3 ++- gui/views/public_mcp_ability_factory.vue | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/gui/views/local_mcp_ability_factory.vue b/gui/views/local_mcp_ability_factory.vue index a2296f0..7a45617 100644 --- a/gui/views/local_mcp_ability_factory.vue +++ b/gui/views/local_mcp_ability_factory.vue @@ -310,7 +310,8 @@ async function handleSubmit() { max_tokens: globalConfig.maxTokens, rag_files: selectedRag.value, rag_embed_model: globalConfig.ragEmbedModel, - rag_topk: globalConfig.ragTopK + rag_topk: globalConfig.ragTopK, + api_base: globalConfig.apiBase, } } diff --git a/gui/views/public_mcp_ability_factory.vue b/gui/views/public_mcp_ability_factory.vue index bc5e5b5..2105ffb 100644 --- a/gui/views/public_mcp_ability_factory.vue +++ b/gui/views/public_mcp_ability_factory.vue @@ -322,7 +322,8 @@ async function handleSubmit() { max_tokens: globalConfig.maxTokens, rag_files: selectedRag.value, rag_embed_model: globalConfig.ragEmbedModel, - rag_topk: globalConfig.ragTopK + rag_topk: globalConfig.ragTopK, + api_base: globalConfig.apiBase, } } From cc45e2b33bbd04bc0870c9d55ca8ec7061d917b4 Mon Sep 17 00:00:00 2001 From: cozyty Date: Fri, 13 Mar 2026 14:41:40 -0400 Subject: [PATCH 08/21] update --- app/mcp_svc.py | 28 ++++++++++++++++++++++------ conf/default.yml | 14 +++++++++----- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/app/mcp_svc.py b/app/mcp_svc.py index a45290b..b6e1e9d 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -36,6 +36,7 @@ def _create_dspy_client(self, model_config: dict): "temperature": model_config.get("temperature"), "max_tokens": model_config.get("max_tokens"), "max_tool_calls": model_config.get("max_tool_calls"), + "api_base": model_config.get("api_base") } return lm @@ -96,7 +97,6 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d temperature=lm_obj.get("temperature"), max_tokens=lm_obj.get("max_tokens"), ) - dspy.context(lm=lm) except Exception as e: self.log.warning(f"[MCP] Failed to configure LM: {e}") @@ -140,18 +140,34 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d if use_rag: if focus in [ExecuteStyle.LLMplanner.value, ExecuteStyle.RAGplanner.value]: self.log.info(f"[MCP] Executing RAG-enhanced planner with prompt: {prompt}") - result = await planner_run(prompt, lm_obj, rag_context=rag_context, run_id=run_id) + if lm: + with dspy.context(lm=lm): # ✅ Context entered here! + result = await planner_run(prompt, lm_obj, rag_context=rag_context, run_id=run_id) + else: + result = await planner_run(prompt, lm_obj, rag_context=rag_context, run_id=run_id) else: self.log.info(f"[MCP] Executing RAG-enhanced factory with prompt: {prompt}") - result = await factory_run(prompt, lm_obj, rag_context=rag_context, run_id=run_id) + if lm: + with dspy.context(lm=lm): # ✅ Context entered here! + result = await factory_run(prompt, lm_obj, rag_context=rag_context, run_id=run_id) + else: + result = await factory_run(prompt, lm_obj, rag_context=rag_context, run_id=run_id) else: if focus == ExecuteStyle.LLMplanner.value: self.log.info(f"[MCP] Executing planner with prompt: {prompt}") - result = await planner_run(prompt, lm_obj, run_id=run_id) + if lm: + with dspy.context(lm=lm): # ✅ Context entered here! + result = await planner_run(prompt, lm_obj, run_id=run_id) + else: + result = await planner_run(prompt, lm_obj, run_id=run_id) else: self.log.info(f"[MCP] Executing factory with prompt: {prompt}") - result = await factory_run(prompt, lm_obj, run_id=run_id) - + if lm: + with dspy.context(lm=lm): # ✅ Context entered here! + result = await factory_run(prompt, lm_obj, run_id=run_id) + else: + result = await factory_run(prompt, lm_obj, run_id=run_id) + mlflow.set_tag("stage", "complete") mlflow.set_tag("status", "success") # Store process_result as a tag instead of param to avoid conflicts diff --git a/conf/default.yml b/conf/default.yml index 21de9db..e69f9ea 100644 --- a/conf/default.yml +++ b/conf/default.yml @@ -1,10 +1,14 @@ --- llm: - model: gpt-4o - api_key: - offline: true - use_mock: false + model: moonshotai/kimi-k2-instruct-0905 + api_key: + api_base: https://integrate.api.nvidia.com/v1 + temperature: 0.5 + max_tokens: 10000 + max_tool_calls: 10 factory: - model: gpt-4o + model: moonshotai/kimi-k2-instruct-0905 + api_key: + api_base: https://integrate.api.nvidia.com/v1 api_key: temperature: 0.4 From d377ce8bd36798e86da3474a96966cc359d6777d Mon Sep 17 00:00:00 2001 From: cozyty Date: Fri, 13 Mar 2026 14:45:16 -0400 Subject: [PATCH 09/21] config: remove api_key from default.yml for UI-provided keys --- conf/default.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/conf/default.yml b/conf/default.yml index e69f9ea..9030437 100644 --- a/conf/default.yml +++ b/conf/default.yml @@ -1,7 +1,7 @@ --- llm: model: moonshotai/kimi-k2-instruct-0905 - api_key: + api_key: api_base: https://integrate.api.nvidia.com/v1 temperature: 0.5 max_tokens: 10000 From feb15724c64cabecbae19cafd97bb38edc07e2c7 Mon Sep 17 00:00:00 2001 From: hiCozyty <101854573+hiCozyty@users.noreply.github.com> Date: Fri, 13 Mar 2026 14:48:38 -0400 Subject: [PATCH 10/21] Remove redundant api_key entry from default.yml --- conf/default.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/conf/default.yml b/conf/default.yml index 9030437..d4a37a4 100644 --- a/conf/default.yml +++ b/conf/default.yml @@ -10,5 +10,4 @@ factory: model: moonshotai/kimi-k2-instruct-0905 api_key: api_base: https://integrate.api.nvidia.com/v1 - api_key: temperature: 0.4 From eff414e40c7e6e7669c16a0c5a7ea4bf6c90848b Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 16:28:36 -0400 Subject: [PATCH 11/21] fix --- app/factory.py | 2 +- app/mcp_factory_client.py | 3 +-- app/mcp_planner_client.py | 2 +- app/mcp_svc.py | 2 +- 4 files changed, 4 insertions(+), 5 deletions(-) diff --git a/app/factory.py b/app/factory.py index 71f2348..b61dc5c 100644 --- a/app/factory.py +++ b/app/factory.py @@ -11,7 +11,7 @@ def configure_dspy_from_env(): max_tokens = int(os.environ.get('DSPY_MAX_TOKENS', '10000')) if api_key: lm = dspy.LM( - model="openai/" + model, + model=f"nvidia_nim/" + model, api_key=api_key, api_base=api_base, temperature=temperature, diff --git a/app/mcp_factory_client.py b/app/mcp_factory_client.py index 00f6af5..5549509 100644 --- a/app/mcp_factory_client.py +++ b/app/mcp_factory_client.py @@ -180,9 +180,8 @@ async def run(adversary_emulation_task: str, lm_obj = None, rag_context=None, ru mlflow.set_tag("stage", "listing tools") tools = await session.list_tools() - # FIXED: add openai/ prefix and api_base with dspy.context(lm=dspy.LM( - "openai/" + lm_settings['model'], + f"nvidia_nim/" + lm_settings['model'], api_key=lm_settings['api_key'], api_base=lm_settings['api_base'], temperature=lm_settings['temperature'], diff --git a/app/mcp_planner_client.py b/app/mcp_planner_client.py index 6881161..be14b09 100644 --- a/app/mcp_planner_client.py +++ b/app/mcp_planner_client.py @@ -31,7 +31,7 @@ def build_lm_from_dict(settings: dict) -> dspy.LM: raise ValueError("API key is required but not provided. Please set your API key in the Global Model Configuration.") lm_kwargs = { - "model": "openai/" + (settings.get("model") or "gpt-4o"), + "model": f"nvidia_nim/" + (settings.get("model") or "gpt-4o"), "api_key": api_key, "api_base": settings.get("api_base"), } diff --git a/app/mcp_svc.py b/app/mcp_svc.py index b6e1e9d..b0d99d5 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -91,7 +91,7 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d if lm_obj and lm_obj.get("api_key"): try: lm = dspy.LM( - model="openai/" + lm_obj.get("model"), + model=f"nvidia_nim/" + lm_obj.get("model"), api_key=lm_obj.get("api_key"), api_base=lm_obj.get("api_base"), temperature=lm_obj.get("temperature"), From 97e540e9499cb9d4f1034f0b6aed735dbb4d0ba1 Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 16:51:33 -0400 Subject: [PATCH 12/21] nvidia_nim change --- app/mcp_svc.py | 4 ++-- app/rag.py | 8 ++++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/app/mcp_svc.py b/app/mcp_svc.py index b0d99d5..062ab95 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -74,7 +74,7 @@ def _build_rag_service_from_files(self, filenames, api_key: str, embed_model: st rag = RAGService(api_key=api_key, api_base=api_base, log=self.log) if topk: rag.topk_objects_to_retrieve = int(topk) - rag.initialize_from_bundles(bundles, embed_model=embed_model or 'nvidia/llama-3.2-nv-embedqa-1b-v2') + rag.initialize_from_bundles(bundles, embed_model=embed_model or 'nvidia_nim/llama-3.2-nv-embedqa-1b-v2') return rag async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: dict = None): @@ -101,7 +101,7 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d self.log.warning(f"[MCP] Failed to configure LM: {e}") rag_files = run_config.get("rag_files") or [] - rag_embed_model = run_config.get("rag_embed_model") or 'nvidia/llama-3.2-nv-embedqa-1b-v2' + rag_embed_model = run_config.get("rag_embed_model") or 'nvidia_nim/llama-3.2-nv-embedqa-1b-v2' rag_topk = run_config.get("rag_topk") # Use RAG if explicitly requested via focus or if files were selected diff --git a/app/rag.py b/app/rag.py index ccee42a..0e544d7 100644 --- a/app/rag.py +++ b/app/rag.py @@ -33,7 +33,7 @@ def load_stix_bundle(self, stix_bundle_path: str, embed_model: str = 'nvidia/lla except json.JSONDecodeError: raise ValueError(f"Invalid JSON in STIX bundle: {stix_bundle_path}") - def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = 'nvidia/llama-3.2-nv-embedqa-1b-v2'): + def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = 'nvidia_nim/llama-3.2-nv-embedqa-1b-v2'): """Initialize the RAG service with multiple STIX bundles and create retriever.""" all_corpus = [] all_adv_step = {} @@ -46,7 +46,11 @@ def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = ' self.adv_step = all_adv_step self.log.info("Initializing embeddings and retriever for STIX corpus") - embedder = dspy.Embedder(embed_model, api_key=self.api_key, api_base=self.api_base) + embedder = dspy.Embedder( + embed_model, + api_key=self.api_key, + api_base=self.api_base + ) self.search = dspy.retrievers.Embeddings( corpus=self.corpus, embedder=embedder, From 325c0d7020ac34eb2c02f0d8e489bf5f8f8ebb71 Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 17:03:47 -0400 Subject: [PATCH 13/21] infinite loop fix --- app/mcp_planner_client.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/app/mcp_planner_client.py b/app/mcp_planner_client.py index be14b09..ef94afc 100644 --- a/app/mcp_planner_client.py +++ b/app/mcp_planner_client.py @@ -226,21 +226,21 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ mlflow.end_run() raise - # Optional streaming updates (if desired for parity) - client = MlflowClient() - latest_thought = None - latest_observation = None + # # Optional streaming updates (if desired for parity) + # client = MlflowClient() + # latest_thought = None + # latest_observation = None - while True: - run = client.get_run(run_id) - tags = run.data.tags + # while True: + # run = client.get_run(run_id) + # tags = run.data.tags - if tags.get("latest_thought") != latest_thought: - latest_thought = tags["latest_thought"] - client.set_tag(run_id, "frontend_thought", latest_thought) + # if tags.get("latest_thought") != latest_thought: + # latest_thought = tags["latest_thought"] + # client.set_tag(run_id, "frontend_thought", latest_thought) - if tags.get("latest_observation") != latest_observation: - latest_observation = tags["latest_observation"] - client.set_tag(run_id, "frontend_observation", latest_observation) + # if tags.get("latest_observation") != latest_observation: + # latest_observation = tags["latest_observation"] + # client.set_tag(run_id, "frontend_observation", latest_observation) - await asyncio.sleep(2) + # await asyncio.sleep(2) From 63188e6fd85810dc44a3343dff36692559611fc5 Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 17:07:39 -0400 Subject: [PATCH 14/21] revert --- app/mcp_planner_client.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/app/mcp_planner_client.py b/app/mcp_planner_client.py index ef94afc..be14b09 100644 --- a/app/mcp_planner_client.py +++ b/app/mcp_planner_client.py @@ -226,21 +226,21 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ mlflow.end_run() raise - # # Optional streaming updates (if desired for parity) - # client = MlflowClient() - # latest_thought = None - # latest_observation = None + # Optional streaming updates (if desired for parity) + client = MlflowClient() + latest_thought = None + latest_observation = None - # while True: - # run = client.get_run(run_id) - # tags = run.data.tags + while True: + run = client.get_run(run_id) + tags = run.data.tags - # if tags.get("latest_thought") != latest_thought: - # latest_thought = tags["latest_thought"] - # client.set_tag(run_id, "frontend_thought", latest_thought) + if tags.get("latest_thought") != latest_thought: + latest_thought = tags["latest_thought"] + client.set_tag(run_id, "frontend_thought", latest_thought) - # if tags.get("latest_observation") != latest_observation: - # latest_observation = tags["latest_observation"] - # client.set_tag(run_id, "frontend_observation", latest_observation) + if tags.get("latest_observation") != latest_observation: + latest_observation = tags["latest_observation"] + client.set_tag(run_id, "frontend_observation", latest_observation) - # await asyncio.sleep(2) + await asyncio.sleep(2) From 83a5a584f09088ae95dded829a2bc36a309a84b3 Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 17:13:05 -0400 Subject: [PATCH 15/21] nvidia_nim change --- gui/views/mcp.vue | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gui/views/mcp.vue b/gui/views/mcp.vue index 7f27be2..afab95c 100644 --- a/gui/views/mcp.vue +++ b/gui/views/mcp.vue @@ -173,7 +173,7 @@ class="input" type="text" v-model="globalConfig.ragEmbedModel" - placeholder="nvidia/llama-3.2-nv-embedqa-1b-v2" + placeholder="nvidia_nim/llama-3.2-nv-embedqa-1b-v2" />
@@ -514,7 +514,7 @@ const globalConfig = reactive({ apiBase: savedConfig?.apiBase || '', maxToolCalls: savedConfig?.maxToolCalls || 5, maxTokens: savedConfig?.maxTokens || 10000, - ragEmbedModel: savedConfig?.ragEmbedModel || 'nvidia/llama-3.2-nv-embedqa-1b-v2', + ragEmbedModel: savedConfig?.ragEmbedModel || 'nvidia_nim/llama-3.2-nv-embedqa-1b-v2', ragTopK: savedConfig?.ragTopK ?? 5 }) From a2b2feca1d6cf36a11e36b324cb7f058c17a61ed Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 17:17:58 -0400 Subject: [PATCH 16/21] debugging logs --- app/mcp_svc.py | 2 ++ app/rag.py | 3 ++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/app/mcp_svc.py b/app/mcp_svc.py index 062ab95..6a08e13 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -72,6 +72,7 @@ def _build_rag_service_from_files(self, filenames, api_key: str, embed_model: st bundles.append(json.load(f)) rag = RAGService(api_key=api_key, api_base=api_base, log=self.log) + print(f"[RAG SVC DEBUG] embed_model passed to RAG: '{embed_model}', api_base: '{api_base}'") if topk: rag.topk_objects_to_retrieve = int(topk) rag.initialize_from_bundles(bundles, embed_model=embed_model or 'nvidia_nim/llama-3.2-nv-embedqa-1b-v2') @@ -97,6 +98,7 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d temperature=lm_obj.get("temperature"), max_tokens=lm_obj.get("max_tokens"), ) + print(f"[LM DEBUG] model='{lm_obj.get('model')}', api_base='{lm_obj.get('api_base')}'") except Exception as e: self.log.warning(f"[MCP] Failed to configure LM: {e}") diff --git a/app/rag.py b/app/rag.py index 0e544d7..8f801a1 100644 --- a/app/rag.py +++ b/app/rag.py @@ -46,10 +46,11 @@ def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = ' self.adv_step = all_adv_step self.log.info("Initializing embeddings and retriever for STIX corpus") + print(f"[RAG DEBUG] embed_model='{embed_model}', api_key_set={bool(self.api_key)}, api_base='{self.api_base}'") embedder = dspy.Embedder( embed_model, api_key=self.api_key, - api_base=self.api_base + api_base=self.api_base or "https://integrate.api.nvidia.com/v1" ) self.search = dspy.retrievers.Embeddings( corpus=self.corpus, From 11b02307ad8751a175680437e12f594f0e8630d4 Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 17:24:45 -0400 Subject: [PATCH 17/21] apibase fix --- gui/views/mcp.vue | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gui/views/mcp.vue b/gui/views/mcp.vue index afab95c..209eb75 100644 --- a/gui/views/mcp.vue +++ b/gui/views/mcp.vue @@ -511,7 +511,7 @@ const globalConfig = reactive({ modelName: savedConfig?.modelName || 'gpt-4o', temperature: savedConfig?.temperature ?? 0.5, apiKey: savedConfig?.apiKey || '', - apiBase: savedConfig?.apiBase || '', + apiBase: savedConfig?.apiBase || 'https://integrate.api.nvidia.com/v1', maxToolCalls: savedConfig?.maxToolCalls || 5, maxTokens: savedConfig?.maxTokens || 10000, ragEmbedModel: savedConfig?.ragEmbedModel || 'nvidia_nim/llama-3.2-nv-embedqa-1b-v2', From ed0f314ddff2afbc0bb51d57cb2625756564e34d Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 17:49:15 -0400 Subject: [PATCH 18/21] embed model fix --- app/mcp_svc.py | 4 ++-- app/rag.py | 2 +- gui/views/mcp.vue | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/app/mcp_svc.py b/app/mcp_svc.py index 6a08e13..5ca255c 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -75,7 +75,7 @@ def _build_rag_service_from_files(self, filenames, api_key: str, embed_model: st print(f"[RAG SVC DEBUG] embed_model passed to RAG: '{embed_model}', api_base: '{api_base}'") if topk: rag.topk_objects_to_retrieve = int(topk) - rag.initialize_from_bundles(bundles, embed_model=embed_model or 'nvidia_nim/llama-3.2-nv-embedqa-1b-v2') + rag.initialize_from_bundles(bundles, embed_model=embed_model or 'nvidia/llama-3.2-nv-embedqa-1b-v2') return rag async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: dict = None): @@ -103,7 +103,7 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d self.log.warning(f"[MCP] Failed to configure LM: {e}") rag_files = run_config.get("rag_files") or [] - rag_embed_model = run_config.get("rag_embed_model") or 'nvidia_nim/llama-3.2-nv-embedqa-1b-v2' + rag_embed_model = run_config.get("rag_embed_model") or 'nvidia/llama-3.2-nv-embedqa-1b-v2' rag_topk = run_config.get("rag_topk") # Use RAG if explicitly requested via focus or if files were selected diff --git a/app/rag.py b/app/rag.py index 8f801a1..329c1e3 100644 --- a/app/rag.py +++ b/app/rag.py @@ -33,7 +33,7 @@ def load_stix_bundle(self, stix_bundle_path: str, embed_model: str = 'nvidia/lla except json.JSONDecodeError: raise ValueError(f"Invalid JSON in STIX bundle: {stix_bundle_path}") - def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = 'nvidia_nim/llama-3.2-nv-embedqa-1b-v2'): + def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = 'nvidia/llama-3.2-nv-embedqa-1b-v2'): """Initialize the RAG service with multiple STIX bundles and create retriever.""" all_corpus = [] all_adv_step = {} diff --git a/gui/views/mcp.vue b/gui/views/mcp.vue index 209eb75..2a48643 100644 --- a/gui/views/mcp.vue +++ b/gui/views/mcp.vue @@ -173,7 +173,7 @@ class="input" type="text" v-model="globalConfig.ragEmbedModel" - placeholder="nvidia_nim/llama-3.2-nv-embedqa-1b-v2" + placeholder="nvidia/llama-3.2-nv-embedqa-1b-v2" /> @@ -514,7 +514,7 @@ const globalConfig = reactive({ apiBase: savedConfig?.apiBase || 'https://integrate.api.nvidia.com/v1', maxToolCalls: savedConfig?.maxToolCalls || 5, maxTokens: savedConfig?.maxTokens || 10000, - ragEmbedModel: savedConfig?.ragEmbedModel || 'nvidia_nim/llama-3.2-nv-embedqa-1b-v2', + ragEmbedModel: savedConfig?.ragEmbedModel || 'nvidia/llama-3.2-nv-embedqa-1b-v2', ragTopK: savedConfig?.ragTopK ?? 5 }) From da4cb707d4293043ae3ea2a1e5747dd760d39188 Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 18:20:34 -0400 Subject: [PATCH 19/21] rag fix --- app/rag.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/app/rag.py b/app/rag.py index 329c1e3..5e5bd44 100644 --- a/app/rag.py +++ b/app/rag.py @@ -46,6 +46,10 @@ def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = ' self.adv_step = all_adv_step self.log.info("Initializing embeddings and retriever for STIX corpus") + + if embed_model.startswith("nvidia_nim/nvidia/"): + embed_model = embed_model[len("nvidia_nim/"):] + print(f"[RAG DEBUG] embed_model='{embed_model}', api_key_set={bool(self.api_key)}, api_base='{self.api_base}'") embedder = dspy.Embedder( embed_model, From fc08537c6bbccac263d7291e5a166f6521102b46 Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 18:26:57 -0400 Subject: [PATCH 20/21] rag fix --- app/mcp_svc.py | 1 + app/rag.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/app/mcp_svc.py b/app/mcp_svc.py index 5ca255c..6a12c03 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -89,6 +89,7 @@ async def _run_execution(self, focus, prompt, run_id, lm_obj=None, run_config: d mlflow.log_param("prompt", prompt) # Configure LM globally if provided + lm = None if lm_obj and lm_obj.get("api_key"): try: lm = dspy.LM( diff --git a/app/rag.py b/app/rag.py index 5e5bd44..6a640d4 100644 --- a/app/rag.py +++ b/app/rag.py @@ -54,7 +54,8 @@ def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = ' embedder = dspy.Embedder( embed_model, api_key=self.api_key, - api_base=self.api_base or "https://integrate.api.nvidia.com/v1" + api_base=self.api_base or "https://integrate.api.nvidia.com/v1", + encoding_format="float", ) self.search = dspy.retrievers.Embeddings( corpus=self.corpus, From 41f87eddc3cf28f43558f6ed2aceff1477c8523b Mon Sep 17 00:00:00 2001 From: cozyty Date: Wed, 15 Apr 2026 18:31:19 -0400 Subject: [PATCH 21/21] rag fix --- app/rag.py | 1 + 1 file changed, 1 insertion(+) diff --git a/app/rag.py b/app/rag.py index 6a640d4..65f897e 100644 --- a/app/rag.py +++ b/app/rag.py @@ -56,6 +56,7 @@ def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str = ' api_key=self.api_key, api_base=self.api_base or "https://integrate.api.nvidia.com/v1", encoding_format="float", + input_type="passage" ) self.search = dspy.retrievers.Embeddings( corpus=self.corpus,