Merge pull request #6 from IBM/feature/mcp-support

Feature/mcp support

Author: MatthewGrigsby

.gitignore

@@ -5,5 +5,6 @@
 *.pyc
 *__pycache__
 sandbox.ipynb
+.webui_secret_key
 venv/
-.venv/
+.venv/

CHANGELOG.md

@@ -4,6 +4,30 @@ All notable changes to Flexo will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [v0.2.2] - 2025-03-25
+
+### Features & Improvements
+- Added `/models` endpoint and enabled CORS support.
+- Initiated MCP client and tool registry functionality (WIP).
+- Updated tool patterns for enhanced consistency.
+
+### Fixes & Updates
+- Fixed multi tool accumulation issue.
+- Fixed mid-response tool buffer leak.
+- Updated chat completions data models: updated `FunctionDetail` and removed `name` from `ToolCall`.
+
+### Notes
+- MCP configuration in `agent.yaml` is now commented out by default.
+
+[v0.2.2]: https://github.com/ibm/flexo/releases/tag/v0.2.2
+
+## [v0.2.1] - 2025-03-14
+
+### Fixes & Updates
+- Fixed issue with LLMFactory not recognizing openai-compat vendor names
+
+[v0.2.1]: https://github.com/ibm/flexo/releases/tag/v0.2.1
+
 ## [v0.2.0] - 2025-03-10
 
 ### Features & Improvements
@@ -38,6 +62,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Updated tool configuration and path structures
 - Added Elasticsearch SSL certificate documentation
 
+[v0.1.1]: https://github.com/ibm/flexo/releases/tag/v0.1.1
+
 ## [v0.1.0] - 2024-01-31
 
 ### Initial Release
@@ -52,5 +78,4 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Robust prompt building and parsing systems
 - Comprehensive LLM integration components
 
-[v0.1.1]: https://github.com/ibm/flexo/releases/tag/v0.1.1
-[v0.1.0]: https://github.com/ibm/flexo/releases/tag/v0.1.0
+[v0.1.0]: https://github.com/ibm/flexo/releases/tag/v0.1.0

src/agent/chat_agent_streaming.py

@@ -1,7 +1,7 @@
 # src/agent/chat_streaming_agent.py
 
 import os
-import json
+import json5
 import yaml
 import asyncio
 import logging
@@ -99,7 +99,7 @@ def __init__(self, config: Dict) -> None:
 
         # Load parser config if needed for manual detection
         self.logger.info("\n\n" + "=" * 60 + "\n" + "Agent Configuration Summary" + "\n" + "=" * 60 + "\n" +
-                         f"Main Chat Model Config:\n{json.dumps(self.main_chat_model_config, indent=4)}\n" +
+                         f"Main Chat Model Config:\n{json5.dumps(self.main_chat_model_config, indent=4)}\n" +
                          f"Tool Detection Mode: {self.detection_mode}\n" +
                          f"Vendor Chat API Mode: {self.use_vendor_chat_completions}\n" +
                          "\n" + "=" * 60 + "\n")
@@ -122,9 +122,12 @@ def __init__(self, config: Dict) -> None:
             vendor=self.main_chat_model_config.get('vendor')
         )
 
-        # Initialize ToolRegistry
-        self.tool_registry = ToolRegistry()
-        self.tool_registry.load_from_config(tool_configs=self.config.get("tools_config"))
+        # Initialize tool registry
+        self.tool_registry = ToolRegistry(
+            tools_config=self.config.get("tools_config"),
+            mcp_config=self.config.get("mcp_config")
+        )
+        asyncio.create_task(self.tool_registry.initialize_all_tools())
 
     @handle_streaming_errors
     async def stream_step(
@@ -160,7 +163,6 @@ async def stream_step(
 
                 case StreamState.STREAMING:
                     self.logger.info(f"--- Entering Streaming State ---")
-                    self.detection_strategy.reset()
                     async for item in self._handle_streaming(context):
                         yield item
 
@@ -205,6 +207,7 @@ async def _handle_streaming(self, context: StreamContext) -> AsyncGenerator[SSEC
         Raises:
             Exception: If maximum streaming iterations are exceeded
         """
+        self.detection_strategy.reset()
         context.streaming_entry_count += 1
         if context.streaming_entry_count > context.max_streaming_iterations:
             self.logger.error("Maximum streaming iterations reached. Aborting further streaming.")
@@ -218,7 +221,6 @@ async def _handle_streaming(self, context: StreamContext) -> AsyncGenerator[SSEC
             conversation_history=context.conversation_history,
             tool_definitions=context.tool_definitions if self.detection_mode == "manual" else None
         )
-        self.logger.debug(f"Prompt payload: {prompt_payload}")
 
         prompt_output: PromptBuilderOutput = (
             await self.prompt_builder.build_chat(prompt_payload) if self.use_vendor_chat_completions
@@ -253,6 +255,7 @@ async def _handle_streaming(self, context: StreamContext) -> AsyncGenerator[SSEC
                 return
 
         final_result = await self.detection_strategy.finalize_detection(context)
+        self.logger.debug(f"Final detection result: {final_result}")
 
         if final_result.state == DetectionState.COMPLETE_MATCH:
             async for chunk in self._handle_complete_match(context, final_result, accumulated_content):
@@ -328,7 +331,6 @@ async def _handle_tool_execution(
                 tool_results.append(result)
                 context.conversation_history.append(
                     ToolMessage(
-                        name=call.function.name,
                         content=result["result"],
                         tool_call_id=call.id
                     )
@@ -405,7 +407,7 @@ async def _initialize_context(
 
         return StreamContext(
             conversation_history=selected_history,
-            tool_definitions=self.tool_registry.get_tool_definitions(),
+            tool_definitions=await self.tool_registry.get_tool_definitions(),
             context=api_passed_context,
             llm_factory=self.llm_factory,
             current_state=StreamState.STREAMING,
@@ -458,13 +460,14 @@ async def _execute_tools_concurrently(self, context: StreamContext) -> List[Any]
         """
         async def run_tool(tool_call: ToolCall):
             try:
-                tool = self.tool_registry.get_tool(tool_call.function.name)
+                tool = await self.tool_registry.get_tool(tool_call.function.name)
                 if not tool:
                     raise RuntimeError(f"Tool {tool_call.function.name} not found")
-
+                tool_args = json5.loads(tool_call.function.arguments)
+                self.logger.info(f"Running tool {tool_call.function.name} with arguments: {tool_args}")
                 result = await tool.execute(
                     context=context,
-                    **tool_call.function.arguments
+                    **tool_args
                 )
                 return {"tool_name": tool_call.function.name, "result": result.result}
 

src/api/routes/chat_completions_api.py

@@ -1,5 +1,3 @@
-# src/api/routes/chat_completions_api.py
-
 """Chat Completions API Module.
 
 This module implements a FastAPI router for streaming chat completion functionality,
@@ -10,14 +8,11 @@
 - Message format conversion
 - Streaming chat completions
 - SSE (Server-Sent Events) response handling
-
-Dependencies:
-    - FastAPI for route handling
-    - StreamingChatAgent for chat processing
-    - Pydantic models for request/response validation
+- Agent info endpoint (compatible with OpenAI's /models format)
 """
 
 import os
+import time
 import yaml
 import logging
 from typing import List, Optional
@@ -37,6 +32,9 @@
 logger = logging.getLogger(__name__)
 router = APIRouter()
 
+# Capture the agent's start time once when the app starts.
+AGENT_START_TIME = int(time.time())
+
 # Security setup
 API_KEY_NAME = "X-API-KEY"
 api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
@@ -118,6 +116,22 @@ def convert_message_content(messages: List[TextChatMessage]) -> List[TextChatMes
     return converted
 
 
+def get_non_sensitive_config(config: dict) -> dict:
+    """Filter the agent configuration to include only non-sensitive parameters.
+
+    Args:
+        config (dict): The full agent configuration.
+
+    Returns:
+        dict: Filtered configuration excluding keys that may be sensitive.
+    """
+    # Exclude keys that include 'key', 'secret', or 'password' (case insensitive)
+    return {
+        k: v for k, v in config.items()
+        if not any(sensitive in k.lower() for sensitive in ["key", "secret", "password"])
+    }
+
+
 @router.post(
     "/chat/completions",
     summary="Generate streaming chat completions",
@@ -193,3 +207,37 @@ async def sse_generator():
     except Exception as e:
         logger.error("Error in /chat/completions: %s", str(e), exc_info=True)
         raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.get(
+    "/models",
+    summary="Agent information",
+    description="Returns information about the running agent including its default model and non-sensitive configuration parameters.",
+    tags=["Models"],
+    operation_id="getModels"
+)
+async def get_models(
+    agent: StreamingChatAgent = Depends(get_streaming_agent),
+    api_key: Optional[str] = Depends(get_api_key)
+):
+    """Retrieve agent information in a format similar to OpenAI's models endpoint.
+
+    This endpoint returns the agent's name, default model, creation time,
+    owner, and a selection of non-sensitive configuration parameters.
+    """
+    agent_name = agent.config.get("name", "default-agent")
+    default_model = agent.config.get("model", "default-model")
+    created = agent.config.get("created", AGENT_START_TIME)
+    owned_by = agent.config.get("owned_by", "user")
+    non_sensitive_params = get_non_sensitive_config(agent.config)
+
+    agent_info = {
+        "id": agent_name,
+        "object": "agent",
+        "default_model": default_model,
+        "created": created,
+        "owned_by": owned_by,
+        "params": non_sensitive_params
+    }
+
+    return {"data": [agent_info], "object": "list"}

src/configs/agent.yaml

@@ -10,6 +10,10 @@ max_streaming_iterations: 4 # Number of times the agent is allowed to return to
 timeouts:
   model_response_timeout: 60
 
+# CORS allowed origins (optional)
+allowed_origins:
+  - http://localhost:8080 # example for local Open WebUI
+
 # System prompt defining agent behavior and capabilities
 system_prompt: |
   I am a helpful AI assistant focused on clear, accurate, and direct communication. 
@@ -86,6 +90,20 @@ models_config:
 #     max_tokens: 2000
 #     temperature: 0.7
 
+# MCP configuration
+#mcp_config:
+#  # Example SSE
+#  transport: sse
+#  sse_url: "http://localhost:8001/sse"
+#  sampling_enabled: false
+
+  # Example stdio
+#  transport: stdio
+#  command: python
+#  args: ["weather/weather.py"]
+#  env: null
+#  sampling_enabled: false
+
 # Tool Configurations
 tools_config:
   # Weather API Integration

src/configs/prompt_builders.yaml

@@ -5,17 +5,25 @@ watsonx-llama:
     begin_text: "<|begin_of_text|>"
   system_prompt:
     header: "You have access to the following tools: {tools}. Current date: {date}"
-    tool_instructions: 'Each tool call must contain "name" and "parameters". To use a tool, emit a call in the following JSON format:\n<|tool_call|>[{\"name\": \"tool_name\", \"parameters\": {\"arg1\": \"value1\"}}]\n\n### Example:\n<|tool_call|>[{\"name\": \"tool_name_1\", \"parameters\": {\"param_1\": \"value_1\", \"param_2\": \"value_2\"}}]\n\nEnsure that all tool calls strictly follow this JSON format.'
+    tool_instructions: 'Each tool call must contain "name" and "parameters". To use a tool, emit a call in the following JSON format:
+<|tool_call|>[{"name": "tool_name", "parameters": {"arg1": "value1"}}]
+
+### Example:
+<|tool_call|>[{"name": "tool_name_1", "parameters": {"param_1": "value_1", "param_2": "value_2"}}]
+
+Ensure that all tool calls strictly follow this JSON format.'
 
 watsonx-granite:
   system_prompt:
     header: "You have access to the following tools: {tools}. Current date: {date}"
-    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:\n<|tool_call|>[{\"name\": \"tool_name\", \"arguments\": {\"arg1\": \"value1\"}}]'
+    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:
+<|tool_call|>[{"name": "tool_name", "arguments": {"arg1": "value1"}}]'
 
 watsonx-mistral:
   system_prompt:
     header: "You have access to the following tools: {tools}. Current date: {date}"
-    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:\n<|tool_call|>[{\"name\": \"tool_name\", \"arguments\": {\"arg1\": \"value1\"}}]'
+    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:
+<|tool_call|>[{"name": "tool_name", "arguments": {"arg1": "value1"}}]'
 
 openai:
   system_prompt:
@@ -25,26 +33,36 @@ openai:
 anthropic:
   system_prompt:
     header: "You can use the following tools: {tools}. Current date: {date}"
-    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:\n<|tool_call|>[{\"name\": \"tool_name\", \"arguments\": {\"arg1\": \"value1\"}}]'
+    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:
+<|tool_call|>[{"name": "tool_name", "arguments": {"arg1": "value1"}}]'
 
 mistralai:
   system_prompt:
     header: "You can use the following tools: {tools}. Current date: {date}"
-    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:\n<|tool_call|>[{\"name\": \"tool_name\", \"arguments\": {\"arg1\": \"value1\"}}]'
+    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:
+<|tool_call|>[{"name": "tool_name", "arguments": {"arg1": "value1"}}]'
 
 xai:
   system_prompt:
     header: "You can use the following tools: {tools}. Current date: {date}"
-    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:\n<|tool_call|>[{\"name\": \"tool_name\", \"arguments\": {\"arg1\": \"value1\"}}]'
+    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:
+<|tool_call|>[{"name": "tool_name", "arguments": {"arg1": "value1"}}]'
 
 openai-compat-llama:
   tokens:
     begin_text: "<|begin_of_text|>"
   system_prompt:
     header: "You have access to the following tools: {tools}. Current date: {date}"
-    tool_instructions: 'Each tool call must contain "name" and "parameters". To use a tool, emit a call in the following JSON format:\n<|tool_call|>[{\"name\": \"tool_name\", \"parameters\": {\"arg1\": \"value1\"}}]\n\n### Example:\n<|tool_call|>[{\"name\": \"tool_name_1\", \"parameters\": {\"param_1\": \"value_1\", \"param_2\": \"value_2\"}}]\n\nEnsure that all tool calls strictly follow this JSON format.'
+    tool_instructions: 'Each tool call must contain "name" and "parameters". To use a tool, emit a call in the following JSON format:
+<|tool_call|>[{"name": "tool_name", "parameters": {"arg1": "value1"}}]
+
+### Example:
+<|tool_call|>[{"name": "tool_name_1", "parameters": {"param_1": "value_1", "param_2": "value_2"}}]
+
+Ensure that all tool calls strictly follow this JSON format.'
 
 openai-compat-granite:
   system_prompt:
     header: "You have access to the following tools: {tools}. Current date: {date}"
-    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:\n<|tool_call|>[{\"name\": \"tool_name\", \"arguments\": {\"arg1\": \"value1\"}}]'
+    tool_instructions: 'Each tool call must contain "name" and "arguments". To use a tool, emit a function call in the following JSON format:
+<|tool_call|>[{"name": "tool_name", "arguments": {"arg1": "value1"}}]'