diff --git a/app/factory.py b/app/factory.py
index 9008b32..b61dc5c 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=f"nvidia_nim/" + 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..5549509 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,16 @@ 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
with dspy.context(lm=dspy.LM(
- 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'],
max_tokens=lm_settings['max_tokens']
)):
@@ -203,8 +194,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 +209,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 +222,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..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": 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 e32b796..6a12c03 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
@@ -60,7 +61,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,10 +71,11 @@ 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=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 '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):
@@ -87,19 +89,22 @@ 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:
- dspy.configure(lm=dspy.LM(
- model=lm_obj.get("model"),
+ lm = dspy.LM(
+ 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"),
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}")
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
@@ -112,7 +117,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)
@@ -137,18 +143,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
@@ -163,4 +185,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()
diff --git a/app/rag.py b/app/rag.py
index 8ead001..65f897e 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}")
@@ -21,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:
@@ -32,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 = {}
@@ -45,7 +46,18 @@ 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)
+
+ 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,
+ 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,
embedder=embedder,
diff --git a/conf/default.yml b/conf/default.yml
index 21de9db..d4a37a4 100644
--- a/conf/default.yml
+++ b/conf/default.yml
@@ -1,10 +1,13 @@
---
llm:
- model: gpt-4o
+ model: moonshotai/kimi-k2-instruct-0905
api_key:
- offline: true
- use_mock: false
+ api_base: https://integrate.api.nvidia.com/v1
+ temperature: 0.5
+ max_tokens: 10000
+ max_tool_calls: 10
factory:
- model: gpt-4o
- api_key:
+ model: moonshotai/kimi-k2-instruct-0905
+ api_key:
+ api_base: https://integrate.api.nvidia.com/v1
temperature: 0.4
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/mcp.vue b/gui/views/mcp.vue
index 1ac348a..2a48643 100644
--- a/gui/views/mcp.vue
+++ b/gui/views/mcp.vue
@@ -114,6 +114,17 @@
/>
+
@@ -162,7 +173,7 @@
class="input"
type="text"
v-model="globalConfig.ragEmbedModel"
- placeholder="openai/text-embedding-3-small"
+ placeholder="nvidia/llama-3.2-nv-embedqa-1b-v2"
/>
@@ -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,9 +511,10 @@ const globalConfig = reactive({
modelName: savedConfig?.modelName || 'gpt-4o',
temperature: savedConfig?.temperature ?? 0.5,
apiKey: savedConfig?.apiKey || '',
+ apiBase: savedConfig?.apiBase || 'https://integrate.api.nvidia.com/v1',
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
})
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,
}
}