fast-agent-mcp 0.1.12__py3-none-any.whl → 0.2.0__py3-none-any.whl

mcp_agent/{workflows/llm → llm/providers}/augmented_llm_openai.py

@@ -1,39 +1,37 @@
 import os
-from typing import List, Type, TYPE_CHECKING
+from typing import List, Type
 
-from pydantic_core import from_json
-
-from mcp_agent.workflows.llm.providers.multipart_converter_openai import OpenAIConverter
-from mcp_agent.workflows.llm.providers.sampling_converter_openai import (
-    OpenAISamplingConverter,
+from mcp.types import (
+    CallToolRequest,
+    CallToolRequestParams,
+    CallToolResult,
 )
-
-if TYPE_CHECKING:
-    from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
-from openai import OpenAI, AuthenticationError
+from openai import AuthenticationError, OpenAI
 
 # from openai.types.beta.chat import
 from openai.types.chat import (
-    ChatCompletionMessageParam,
     ChatCompletionMessage,
+    ChatCompletionMessageParam,
     ChatCompletionSystemMessageParam,
     ChatCompletionToolParam,
     ChatCompletionUserMessageParam,
 )
-from mcp.types import (
-    CallToolRequestParams,
-    CallToolRequest,
-    CallToolResult,
-)
+from pydantic_core import from_json
+from rich.text import Text
 
-from mcp_agent.workflows.llm.augmented_llm import (
+from mcp_agent.core.exceptions import ProviderKeyError
+from mcp_agent.core.prompt import Prompt
+from mcp_agent.llm.augmented_llm import (
     AugmentedLLM,
     ModelT,
     RequestParams,
 )
-from mcp_agent.core.exceptions import ProviderKeyError
+from mcp_agent.llm.providers.multipart_converter_openai import OpenAIConverter
+from mcp_agent.llm.providers.sampling_converter_openai import (
+    OpenAISamplingConverter,
+)
 from mcp_agent.logging.logger import get_logger
-from rich.text import Text
+from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
 
 _logger = get_logger(__name__)
 
@@ -41,23 +39,21 @@ DEFAULT_OPENAI_MODEL = "gpt-4o"
 DEFAULT_REASONING_EFFORT = "medium"
 
 
-class OpenAIAugmentedLLM(
-    AugmentedLLM[ChatCompletionMessageParam, ChatCompletionMessage]
-):
+class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletionMessage]):
     """
     The basic building block of agentic systems is an LLM enhanced with augmentations
     such as retrieval, tools, and memory provided from a collection of MCP servers.
     This implementation uses OpenAI's ChatCompletion as the LLM.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, provider_name: str = "OpenAI", *args, **kwargs) -> None:
         # Set type_converter before calling super().__init__
         if "type_converter" not in kwargs:
             kwargs["type_converter"] = OpenAISamplingConverter
 
         super().__init__(*args, **kwargs)
 
-        self.provider = "OpenAI"
+        self.provider = provider_name
         # Initialize logger with name if available
         self.logger = get_logger(f"{__name__}.{self.name}" if self.name else __name__)
 
@@ -70,9 +66,7 @@ class OpenAIAugmentedLLM(
                 self._reasoning_effort = self.context.config.openai.reasoning_effort
 
         # Determine if we're using a reasoning model
-        chosen_model = (
-            self.default_request_params.model if self.default_request_params else None
-        )
+        chosen_model = self.default_request_params.model if self.default_request_params else None
         self._reasoning = chosen_model and (
             chosen_model.startswith("o3") or chosen_model.startswith("o1")
         )
@@ -85,11 +79,6 @@ class OpenAIAugmentedLLM(
         """Initialize OpenAI-specific default parameters"""
         chosen_model = kwargs.get("model", DEFAULT_OPENAI_MODEL)
 
-        # Get default model from config if available
-        if self.context and self.context.config and self.context.config.openai:
-            if hasattr(self.context.config.openai, "default_model"):
-                chosen_model = self.context.config.openai.default_model
-
         return RequestParams(
             model=chosen_model,
             systemPrompt=self.instruction,
@@ -120,11 +109,9 @@ class OpenAIAugmentedLLM(
         return api_key
 
     def _base_url(self) -> str:
-        return (
-            self.context.config.openai.base_url if self.context.config.openai else None
-        )
+        return self.context.config.openai.base_url if self.context.config.openai else None
 
-    async def generate(
+    async def generate_internal(
         self,
         message,
         request_params: RequestParams | None = None,
@@ -149,25 +136,21 @@ class OpenAIAugmentedLLM(
 
         system_prompt = self.instruction or params.systemPrompt
         if system_prompt:
-            messages.append(
-                ChatCompletionSystemMessageParam(role="system", content=system_prompt)
-            )
+            messages.append(ChatCompletionSystemMessageParam(role="system", content=system_prompt))
 
         # Always include prompt messages, but only include conversation history
         # if use_history is True
         messages.extend(self.history.get(include_history=params.use_history))
 
         if isinstance(message, str):
-            messages.append(
-                ChatCompletionUserMessageParam(role="user", content=message)
-            )
+            messages.append(ChatCompletionUserMessageParam(role="user", content=message))
         elif isinstance(message, list):
             messages.extend(message)
         else:
             messages.append(message)
 
         response = await self.aggregator.list_tools()
-        available_tools: List[ChatCompletionToolParam] = [
+        available_tools: List[ChatCompletionToolParam] | None = [
             ChatCompletionToolParam(
                 type="function",
                 function={
@@ -180,17 +163,10 @@ class OpenAIAugmentedLLM(
             for tool in response.tools
         ]
         if not available_tools:
-            available_tools = []
+            available_tools = None  # deepseek does not allow empty array
 
         responses: List[ChatCompletionMessage] = []
-        model = await self.select_model(params)
-        chat_turn = len(messages) // 2
-        if self._reasoning:
-            self.show_user_message(
-                str(message), f"{model} ({self._reasoning_effort})", chat_turn
-            )
-        else:
-            self.show_user_message(str(message), model, chat_turn)
+        model = self.default_request_params.model
 
         # we do NOT send stop sequences as this causes errors with mutlimodal processing
         for i in range(params.max_iterations):
@@ -214,7 +190,7 @@ class OpenAIAugmentedLLM(
                 arguments = {**arguments, **params.metadata}
 
             self.logger.debug(f"{arguments}")
-            self._log_chat_progress(chat_turn, model=model)
+            self._log_chat_progress(self.chat_turn(), model=model)
 
             if response_model is None:
                 executor_result = await self.executor.execute(
@@ -248,21 +224,14 @@ class OpenAIAugmentedLLM(
                 # No response from the model, we're done
                 break
 
-            # TODO: saqadri - handle multiple choices for more complex interactions.
-            # Keeping it simple for now because multiple choices will also complicate memory management
             choice = response.choices[0]
             message = choice.message
             responses.append(message)
 
-            converted_message = self.convert_message_to_message_param(
-                message, name=self.name
-            )
+            converted_message = self.convert_message_to_message_param(message, name=self.name)
             messages.append(converted_message)
             message_text = converted_message.content
-            if (
-                choice.finish_reason in ["tool_calls", "function_call"]
-                and message.tool_calls
-            ):
+            if choice.finish_reason in ["tool_calls", "function_call"] and message.tool_calls:
                 if message_text:
                     await self.show_assistant_message(
                         message_text,
@@ -290,9 +259,7 @@ class OpenAIAugmentedLLM(
                         method="tools/call",
                         params=CallToolRequestParams(
                             name=tool_call.function.name,
-                            arguments=from_json(
-                                tool_call.function.arguments, allow_partial=True
-                            ),
+                            arguments=from_json(tool_call.function.arguments, allow_partial=True),
                         ),
                     )
                     result = await self.call_tool(tool_call_request, tool_call.id)
@@ -300,18 +267,14 @@ class OpenAIAugmentedLLM(
 
                     tool_results.append((tool_call.id, result))
 
-                messages.extend(
-                    OpenAIConverter.convert_function_results_to_openai(tool_results)
-                )
+                messages.extend(OpenAIConverter.convert_function_results_to_openai(tool_results))
 
                 self.logger.debug(
                     f"Iteration {i}: Tool call results: {str(tool_results) if tool_results else 'None'}"
                 )
             elif choice.finish_reason == "length":
                 # We have reached the max tokens limit
-                self.logger.debug(
-                    f"Iteration {i}: Stopping because finish_reason is 'length'"
-                )
+                self.logger.debug(f"Iteration {i}: Stopping because finish_reason is 'length'")
                 if request_params and request_params.maxTokens is not None:
                     message_text = Text(
                         f"the assistant has reached the maximum token limit ({request_params.maxTokens})",
@@ -334,9 +297,7 @@ class OpenAIAugmentedLLM(
                 # TODO: saqadri - would be useful to return the reason for stopping to the caller
                 break
             elif choice.finish_reason == "stop":
-                self.logger.debug(
-                    f"Iteration {i}: Stopping because finish_reason is 'stop'"
-                )
+                self.logger.debug(f"Iteration {i}: Stopping because finish_reason is 'stop'")
                 if message_text:
                     await self.show_assistant_message(message_text, "")
                 break
@@ -361,7 +322,7 @@ class OpenAIAugmentedLLM(
         self,
         message,
         request_params: RequestParams | None = None,
-    ):
+    ) -> str:
         """
         Process a query using an LLM and available tools.
         The default implementation uses OpenAI's ChatCompletion as the LLM.
@@ -371,11 +332,8 @@ class OpenAIAugmentedLLM(
         - "***SAVE_HISTORY <filename.md>" - Saves the conversation history to the specified file
           in MCP prompt format with user/assistant delimiters.
         """
-        # Check if this is a special command to save history
-        if isinstance(message, str) and message.startswith("***SAVE_HISTORY "):
-            return await self._save_history_to_file(message)
 
-        responses = await self.generate(
+        responses = await self.generate_internal(
             message=message,
             request_params=request_params,
         )
@@ -393,23 +351,11 @@ class OpenAIAugmentedLLM(
 
         return "\n".join(final_text)
 
-    async def _apply_prompt_template_provider_specific(
+    async def _apply_prompt_provider_specific(
         self,
         multipart_messages: List["PromptMessageMultipart"],
         request_params: RequestParams | None = None,
-    ) -> str:
-        """
-        OpenAI-specific implementation of apply_prompt_template that handles
-        multimodal content natively.
-
-        Args:
-            multipart_messages: List of PromptMessageMultipart objects parsed from the prompt template
-
-        Returns:
-            String representation of the assistant's response if generated,
-            or the last assistant message in the prompt
-        """
-
+    ) -> PromptMessageMultipart:
         # TODO -- this is very similar to Anthropic (just the converter class changes).
         # TODO -- potential refactor to base class, standardize Converter interface
         # Check the last message role
@@ -417,9 +363,7 @@ class OpenAIAugmentedLLM(
 
         # Add all previous messages to history (or all messages if last is from assistant)
         messages_to_add = (
-            multipart_messages[:-1]
-            if last_message.role == "user"
-            else multipart_messages
+            multipart_messages[:-1] if last_message.role == "user" else multipart_messages
         )
         converted = []
         for msg in messages_to_add:
@@ -428,95 +372,87 @@ class OpenAIAugmentedLLM(
 
         if last_message.role == "user":
             # For user messages: Generate response to the last one
-            self.logger.debug(
-                "Last message in prompt is from user, generating assistant response"
-            )
+            self.logger.debug("Last message in prompt is from user, generating assistant response")
             message_param = OpenAIConverter.convert_to_openai(last_message)
-            return await self.generate_str(message_param, request_params)
+            return Prompt.assistant(await self.generate_str(message_param, request_params))
         else:
             # For assistant messages: Return the last message content as text
-            self.logger.debug(
-                "Last message in prompt is from assistant, returning it directly"
-            )
-            return str(last_message)
+            self.logger.debug("Last message in prompt is from assistant, returning it directly")
+            return last_message
 
-    async def _save_history_to_file(self, command: str) -> str:
+    async def structured(
+        self,
+        prompt: List[PromptMessageMultipart],
+        model: Type[ModelT],
+        request_params: RequestParams | None = None,
+    ) -> ModelT | None:
         """
-        Save the conversation history to a file in MCP prompt format.
+        Apply the prompt and return the result as a Pydantic model.
+
+        Uses OpenAI's beta parse feature when compatible, falling back to standard
+        JSON parsing if the beta feature fails or is unavailable.
 
         Args:
-            command: The command string, expected format: "***SAVE_HISTORY <filename.md>"
+            prompt: List of messages to process
+            model: Pydantic model to parse the response into
+            request_params: Optional request parameters
 
         Returns:
-            Success or error message
+            The parsed response as a Pydantic model, or None if parsing fails
         """
-        try:
-            # Extract the filename from the command
-            parts = command.split(" ", 1)
-            if len(parts) != 2 or not parts[1].strip():
-                return "Error: Invalid format. Expected '***SAVE_HISTORY <filename.md>'"
 
-            filename = parts[1].strip()
+        if not "OpenAI" == self.provider:
+            return await super().structured(prompt, model, request_params)
 
-            # Get all messages from history
-            messages = self.history.get(include_history=True)
+        logger = get_logger(__name__)
 
-            # Import required utilities
-            from mcp_agent.workflows.llm.openai_utils import (
-                openai_message_param_to_prompt_message_multipart,
-            )
-            from mcp_agent.mcp.prompt_serialization import (
-                multipart_messages_to_delimited_format,
-            )
+        # First try to use OpenAI's beta.chat.completions.parse feature
+        try:
+            # Convert the multipart messages to OpenAI format
+            messages = []
+            for msg in prompt:
+                messages.append(OpenAIConverter.convert_to_openai(msg))
+
+            # Add system prompt if available and not already present
+            if self.instruction and not any(m.get("role") == "system" for m in messages):
+                system_msg = ChatCompletionSystemMessageParam(
+                    role="system", content=self.instruction
+                )
+                messages.insert(0, system_msg)
 
-            # Convert message params to PromptMessageMultipart objects
-            multipart_messages = []
-            for msg in messages:
-                # Skip system messages - PromptMessageMultipart only supports user and assistant roles
-                if isinstance(msg, dict) and msg.get("role") == "system":
-                    continue
+            # Use the beta parse feature
+            try:
+                openai_client = OpenAI(api_key=self._api_key(), base_url=self._base_url())
+                model_name = self.default_request_params.model
 
-                # Convert the message to a multipart message
-                multipart_messages.append(
-                    openai_message_param_to_prompt_message_multipart(msg)
+                logger.debug(
+                    f"Using OpenAI beta parse with model {model_name} for structured output"
+                )
+                response = await self.executor.execute(
+                    openai_client.beta.chat.completions.parse,
+                    model=model_name,
+                    messages=messages,
+                    response_format=model,
                 )
 
-            # Convert to delimited format
-            delimited_content = multipart_messages_to_delimited_format(
-                multipart_messages,
-                user_delimiter="---USER",
-                assistant_delimiter="---ASSISTANT",
-            )
+                if response and isinstance(response[0], BaseException):
+                    raise response[0]
 
-            # Write to file
-            with open(filename, "w", encoding="utf-8") as f:
-                f.write("\n\n".join(delimited_content))
+                parsed_result = response[0].choices[0].message
+                logger.debug("Successfully used OpenAI beta parse feature for structured output")
+                return parsed_result.parsed
 
-            self.logger.info(f"Saved conversation history to {filename}")
-            return f"Done. Saved conversation history to {filename}"
+            except (ImportError, AttributeError, NotImplementedError) as e:
+                # Beta feature not available, log and continue to fallback
+                logger.debug(f"OpenAI beta parse feature not available: {str(e)}")
+                # Continue to fallback
 
         except Exception as e:
-            self.logger.error(f"Error saving history: {str(e)}")
-            return f"Error saving history: {str(e)}"
-
-    async def generate_structured(
-        self,
-        message,
-        response_model: Type[ModelT],
-        request_params: RequestParams | None = None,
-    ) -> ModelT:
-        responses = await self.generate(
-            message=message,
-            request_params=request_params,
-            response_model=response_model,
-        )
-        return responses[0].parsed
+            logger.debug(f"OpenAI beta parse failed: {str(e)}, falling back to standard method")
+            # Continue to standard method as fallback
 
-    async def generate_prompt(
-        self, prompt: "PromptMessageMultipart", request_params: RequestParams | None
-    ) -> str:
-        converted_prompt = OpenAIConverter.convert_to_openai(prompt)
-        return await self.generate_str(converted_prompt, request_params)
+        # Fallback to standard method (inheriting from base class)
+        return await super().structured(prompt, model, request_params)
 
     async def pre_tool_call(self, tool_call_id: str | None, request: CallToolRequest):
         return request
@@ -525,30 +461,3 @@ class OpenAIAugmentedLLM(
         self, tool_call_id: str | None, request: CallToolRequest, result: CallToolResult
     ):
         return result
-
-    def message_param_str(self, message: ChatCompletionMessageParam) -> str:
-        """Convert an input message to a string representation."""
-        if message.get("content"):
-            content = message["content"]
-            if isinstance(content, str):
-                return content
-            else:  # content is a list
-                final_text: List[str] = []
-                for part in content:
-                    text_part = part.get("text")
-                    if text_part:
-                        final_text.append(str(text_part))
-                    else:
-                        final_text.append(str(part))
-
-                return "\n".join(final_text)
-
-        return str(message)
-
-    def message_str(self, message: ChatCompletionMessage) -> str:
-        """Convert an output message to a string representation."""
-        content = message.content
-        if content:
-            return content
-
-        return str(message)