feat: add api_base support to RAGService for custom LLM endpoints

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()

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()
 

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"),
     }

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()
+            mlflow.end_run()

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, 

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

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,
       }
     }