openai.pyā¢3.96 kB
import json
from typing import Any, Optional
from openai import OpenAI
from tests.clients.grafana import GrafanaMCPClient
class OpenAIMCPClient:
def __init__(
self,
test_client: Any,
openai_api_key: Optional[str],
mcp_api_key: str = "test-key",
):
"""
Initialize the OpenAI-MCP Client for testing.
"""
if not openai_api_key:
raise ValueError("OpenAI API key must be provided for testing.")
self.openai_client = OpenAI(api_key=openai_api_key)
self.mcp_client = GrafanaMCPClient(test_client=test_client, api_key=mcp_api_key)
self.mcp_tools = self._get_mcp_tools()
def _get_mcp_tools(self):
"""Fetch and format tools from the MCP server."""
try:
mcp_tools_raw = self.mcp_client.list_tools()
formatted_tools = []
for tool in mcp_tools_raw:
tool_name = tool.get("name")
if isinstance(tool_name, str):
formatted_tools.append(
{
"type": "function",
"function": {
"name": tool_name,
"description": tool.get("description"),
"parameters": tool.get("parameters", {"type": "object", "properties": {}}),
},
}
)
return formatted_tools
except Exception as e:
print(f"Failed to fetch or format MCP tools: {e}")
return []
def chat(self, messages: list, model: str = "gpt-4o", **kwargs):
"""
Send a chat request to OpenAI, handling MCP tool calls. Only non-streaming mode is supported: returns the full response as a string.
"""
completion = self.openai_client.chat.completions.create(
model=model,
messages=messages,
tools=self.mcp_tools,
tool_choice="auto",
stream=False,
temperature=0,
**kwargs,
)
content = ""
tool_calls = None
for choice in completion.choices:
if hasattr(choice, "message") and choice.message:
if choice.message.content:
content += choice.message.content
if hasattr(choice.message, "tool_calls") and choice.message.tool_calls:
tool_calls = choice.message.tool_calls
if not tool_calls:
return content
# If there are tool calls, execute them and get the final response
messages.append({"role": "assistant", "tool_calls": [tc.to_dict() for tc in tool_calls]})
for tool_call in tool_calls:
function_name = tool_call.function.name
function_args = json.loads(tool_call.function.arguments)
try:
tool_result = self.mcp_client.execute_tool(tool_name=function_name, parameters=function_args)
except Exception as e:
tool_result = {"error": str(e)}
messages.append(
{
"tool_call_id": tool_call.id,
"role": "tool",
"name": function_name,
"content": json.dumps(tool_result),
}
)
# Get the final response after tool call(s)
completion2 = self.openai_client.chat.completions.create(model=model, messages=messages, stream=False, temperature=0, **kwargs)
final_content = ""
for choice in completion2.choices:
if hasattr(choice, "message") and choice.message and choice.message.content:
final_content += choice.message.content
print("final_content", final_content)
return final_content
def close(self):
"""Close the client session."""
self.mcp_client.close_session()