410 lines
16 KiB
Python
410 lines
16 KiB
Python
import json
|
|
import logging
|
|
from typing import Dict, Any, List
|
|
|
|
import httpx
|
|
|
|
from ai.debug_lmm import DebugRequest, DebugResponse, DebugTool
|
|
from ai.mcp_server import DummyMCPServer
|
|
|
|
logger = logging.getLogger("MCPClient")
|
|
|
|
|
|
class InProcessMCPClientNativeTools:
|
|
"""MCP Client that uses the MCP server directly in the same process"""
|
|
|
|
ID = "Native Tools"
|
|
|
|
def __init__(self, session, settings_manager, ollama_host: str):
|
|
self.mcp_server = DummyMCPServer(session, settings_manager)
|
|
self.ollama_host = ollama_host
|
|
self.available_tools = self.mcp_server.list_tools()
|
|
self.system_prompt_1 = (
|
|
"Tu es un assistant intelligent qui simplifie l'utilisation de l'application 'My Managing Tools'. "
|
|
"Tu va recevoir des rêquetes sur l'état actuel de l'application, comme la liste des 'repositories' ou la liste des 'tables'."
|
|
"Tu devras répondre avec une réponse simple et concise, si possible en une seule phrase. Tu ne dois pas indiquer comment obtenir la réponse. "
|
|
"Tu devras juste donner la réponse.")
|
|
self.system_prompt_2 = ("You are a helpful AI assistant that provides accurate and concise answers"
|
|
" based on the provided context. Do not respond in a long form. Do not hallucinate."
|
|
"If you don't know the answer, just say that you don't know.")
|
|
|
|
self.conversation_history = [
|
|
{"role": "system", "content": self.system_prompt_1}
|
|
]
|
|
|
|
async def generate_with_mcp_context(self, prompt: str, use_tools: bool = True) -> str:
|
|
"""Generate a response using native function calls"""
|
|
|
|
logger.debug(f"Calling LLM with prompt: {prompt}")
|
|
self.conversation_history.append({"role": "user", "content": prompt})
|
|
|
|
if use_tools and self.available_tools:
|
|
logger.debug(f" Using tools: {self.available_tools.keys()}")
|
|
tools = self._format_tools_for_mistral()
|
|
response = await self._call_ollama_with_tools(prompt, tools)
|
|
logger.debug(f" LLM response: {response}")
|
|
|
|
# Process tool calls if present
|
|
if "message" in response and "tool_calls" in response['message']:
|
|
tool_calls = response['message']["tool_calls"]
|
|
logger.debug(f" LLM requested tools: {[tool["function"]["name"] for tool in tool_calls]}")
|
|
|
|
tool_results = []
|
|
for tool_call in tool_calls:
|
|
result = await self.mcp_server.call_tool(tool_call["function"]["name"],
|
|
tool_call["function"]["arguments"])
|
|
tool_results.append(result)
|
|
|
|
logger.debug(f" MCP server tools calls: {tool_results}")
|
|
|
|
# Generate final response with tool results
|
|
final_response = await self._call_ollama_with_tool_results(prompt,
|
|
tools,
|
|
tool_calls,
|
|
tool_results)
|
|
logger.debug(f" LLM Final result: {final_response}")
|
|
final_response_text = final_response.get("message", {}).get("content", "")
|
|
|
|
else:
|
|
# no tool requested, just return the response
|
|
final_response_text = response.get("message", {}).get("content", "")
|
|
|
|
else:
|
|
logger.debug(f" Not using tools {use_tools=}, available_tools={self.available_tools.keys()}")
|
|
response = await self._call_ollama_simple(prompt)
|
|
logger.debug(f" LLM response: {response}")
|
|
final_response_text = response.get("message", {}).get("content", "")
|
|
|
|
self.conversation_history.append({"role": "assistant", "content": final_response_text})
|
|
return final_response_text
|
|
|
|
def _format_tools_for_mistral(self) -> List[Dict[str, Any]]:
|
|
"""Format MCP tools for Mistral's function calling API"""
|
|
tools = []
|
|
for tool_name, tool_info in self.available_tools.items():
|
|
tool_schema = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": tool_name,
|
|
"description": tool_info["description"],
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {}, # added later
|
|
"required": []
|
|
}
|
|
}
|
|
}
|
|
|
|
# Add parameters
|
|
schema_parameters = tool_schema["function"]["parameters"]
|
|
for param_name, param_info in tool_info["parameters"].items():
|
|
schema_parameters["properties"][param_name] = {
|
|
"type": param_info.get("type", "string"),
|
|
"description": param_info.get("description", f"Parameter {param_name}")
|
|
}
|
|
schema_parameters["required"].append(param_name)
|
|
|
|
tools.append(tool_schema)
|
|
|
|
return tools
|
|
|
|
async def _call_ollama_with_tools(self, prompt: str, tools: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
"""Call Ollama with function calls support"""
|
|
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
response = await client.post(
|
|
f"{self.ollama_host}/api/chat",
|
|
json={
|
|
"model": "mistral",
|
|
"messages": [
|
|
{"role": "system", "content": self.system_prompt_1},
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
"tools": tools,
|
|
"stream": False,
|
|
"options": {
|
|
"temperature": 0.0
|
|
}
|
|
|
|
}
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
return response.json()
|
|
else:
|
|
raise Exception(f"Ollama error: {response.status_code}")
|
|
|
|
async def _call_ollama_with_tool_results(self,
|
|
original_prompt: str,
|
|
tools: List[Dict[str, Any]],
|
|
tool_calls: List[Dict[str, Any]],
|
|
tool_results: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
"""Call Ollama with function call results"""
|
|
messages = [
|
|
{"role": "system", "content": self.system_prompt_2},
|
|
{"role": "user", "content": original_prompt},
|
|
{"role": "assistant", "tool_calls": tool_calls}
|
|
]
|
|
|
|
# Add tool results
|
|
for i, (tool_call, result) in enumerate(zip(tool_calls, tool_results)):
|
|
messages.append({
|
|
"role": "tool",
|
|
"tool_call_id": tool_call.get("id", f"call_{i}"),
|
|
"content": json.dumps(result)
|
|
})
|
|
|
|
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
response = await client.post(
|
|
f"{self.ollama_host}/api/chat",
|
|
json={
|
|
"model": "mistral",
|
|
"messages": messages,
|
|
"tools": tools,
|
|
"stream": False,
|
|
"options": {
|
|
"temperature": 0.0
|
|
}
|
|
}
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
return response.json()
|
|
else:
|
|
raise Exception(f"Ollama error: {response.status_code}")
|
|
|
|
async def _call_ollama_simple(self, prompt: str) -> Dict[str, Any]:
|
|
"""Call Ollama for simple generation without tools"""
|
|
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
response = await client.post(
|
|
f"{self.ollama_host}/api/chat",
|
|
json={
|
|
"model": "mistral",
|
|
"messages": [
|
|
{"role": "system", "content": self.system_prompt_2},
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
"stream": False
|
|
}
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
return response.json()
|
|
else:
|
|
raise Exception(f"Ollama error: {response.status_code}")
|
|
|
|
|
|
class InProcessMCPClientCustomTools:
|
|
"""MCP Client that uses the MCP server directly in the same process"""
|
|
|
|
ID = "Custom Tools"
|
|
|
|
def __init__(self, session, settings_manager, ollama_host: str, model: str = "mistral"):
|
|
self.mcp_server = DummyMCPServer(session, settings_manager)
|
|
self.ollama_host = ollama_host
|
|
self.model = model
|
|
self.available_tools = self.mcp_server.list_tools()
|
|
self.max_tool_iterations = 5 # Prevent infinite loops
|
|
|
|
async def generate_with_mcp_context(self, debug: DebugRequest, prompt: str, use_tools: bool = True) -> str:
|
|
"""Generate response with the LLM potentially using MCP tools"""
|
|
|
|
# Enhance prompt with available tools if requested
|
|
if use_tools and self.available_tools:
|
|
logger.debug(f"Using tools: {list(self.available_tools.keys())}")
|
|
debug.available_tools = list(self.available_tools.keys())
|
|
|
|
tools_description = self._format_tools_for_prompt()
|
|
enhanced_prompt = f"""
|
|
{prompt}
|
|
|
|
You have access to the following tools:
|
|
{tools_description}
|
|
|
|
If you need to use tools, you can use multiple tools in sequence. For each tool, respond with JSON format only:
|
|
{{"use_tool": true, "tool_name": "tool_name", "arguments": {{"param": "value"}}}}
|
|
|
|
You can chain multiple tool calls by putting each JSON on a separate line. For example:
|
|
{{"use_tool": true, "tool_name": "first_tool", "arguments": {{}}}}
|
|
{{"use_tool": true, "tool_name": "second_tool", "arguments": {{"param": "value_from_first_tool"}}}}
|
|
|
|
When you're done with tools, provide your final answer normally (not in JSON format).
|
|
"""
|
|
else:
|
|
logger.debug(f"Not using tools {use_tools=}, available_tools={list(self.available_tools.keys())}")
|
|
enhanced_prompt = prompt
|
|
|
|
debug.enhanced_prompt = enhanced_prompt
|
|
|
|
# Execute tool chain
|
|
if use_tools:
|
|
return await self._execute_tool_chain(debug, enhanced_prompt, prompt)
|
|
else:
|
|
debug_response = DebugResponse(enhanced_prompt)
|
|
debug.responses.append(debug_response)
|
|
return await self._call_ollama(debug_response, enhanced_prompt)
|
|
|
|
async def _execute_tool_chain(self, debug: DebugRequest, enhanced_prompt: str, original_prompt: str) -> str:
|
|
"""Execute a chain of tool calls until completion"""
|
|
current_prompt = enhanced_prompt
|
|
tool_results = []
|
|
iteration = 0
|
|
|
|
while iteration < self.max_tool_iterations:
|
|
logger.debug(f"Tool chain iteration {iteration}")
|
|
|
|
debug_response = DebugResponse(current_prompt)
|
|
debug.responses.append(debug_response)
|
|
|
|
# Call LLM
|
|
llm_response = await self._call_ollama(debug_response, current_prompt)
|
|
|
|
# Parse all tool requests from the response
|
|
tool_requests = self._parse_multiple_tool_requests(llm_response)
|
|
logger.debug(f"Tools requested : {[t["tool_name"] for t in tool_requests]}")
|
|
debug_response.tools_requested = [t["tool_name"] for t in tool_requests]
|
|
|
|
if not tool_requests:
|
|
# No more tools requested, this is the final response
|
|
return llm_response
|
|
|
|
# Execute all requested tools
|
|
current_results = []
|
|
debug_response.tools_called = []
|
|
for tool_request in tool_requests:
|
|
if tool_request["tool_name"] not in self.available_tools:
|
|
logger.warning(f"Requested tool not available: {tool_request['tool_name']}")
|
|
error_msg = f"Sorry, the requested tool '{tool_request['tool_name']}' is not available."
|
|
debug_response.error_response = error_msg
|
|
return error_msg
|
|
|
|
tool_name = tool_request["tool_name"]
|
|
tool_args = tool_request.get("arguments", {})
|
|
|
|
debug_tool = DebugTool(tool_name, tool_args)
|
|
debug_response.tools_called.append(debug_tool)
|
|
|
|
try:
|
|
tool_result = await self.mcp_server.call_tool(tool_name, tool_args)
|
|
|
|
current_results.append({"tool_name": tool_name, "arguments": tool_args, "result": tool_result})
|
|
debug_tool.result = tool_result
|
|
|
|
logger.debug(f"Tool {tool_request['tool_name']} executed successfully with result={tool_result}")
|
|
except Exception as e:
|
|
error_msg = f"Tool execution failed: {e}"
|
|
logger.error(error_msg)
|
|
debug_tool.error = error_msg
|
|
return f"Sorry, there was an error executing the tool {tool_request['tool_name']}: {str(e)}"
|
|
|
|
tool_results.extend(current_results)
|
|
|
|
# Prepare prompt for next iteration
|
|
current_prompt = self._build_continuation_prompt(original_prompt, tool_results)
|
|
|
|
iteration += 1
|
|
|
|
# Max iterations reached, generate final response
|
|
logger.warning(f"Max tool iterations ({self.max_tool_iterations}) reached")
|
|
final_prompt = self._build_final_prompt(original_prompt, tool_results)
|
|
return await self._call_ollama(final_prompt)
|
|
|
|
def _parse_multiple_tool_requests(self, response: str) -> List[Dict[str, Any]]:
|
|
"""Parse LLM response to detect multiple tool requests"""
|
|
tool_requests = []
|
|
|
|
try:
|
|
# Look for JSON in each line
|
|
for line in response.strip().split('\n'):
|
|
line = line.strip()
|
|
if line.startswith('{') and line.endswith('}'):
|
|
try:
|
|
parsed = json.loads(line)
|
|
if parsed.get("use_tool") and "tool_name" in parsed:
|
|
tool_requests.append(parsed)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
except Exception as e:
|
|
logger.debug(f"Error parsing tool requests: {e}")
|
|
|
|
return tool_requests
|
|
|
|
def _build_continuation_prompt(self, original_prompt: str, tool_results: List[Dict]) -> str:
|
|
"""Build prompt for continuing the tool chain"""
|
|
results_text = ""
|
|
for result in tool_results:
|
|
results_text += f"\nTool: {result['tool_name']} with arguments {result['arguments']}\nResult: {result['result']}\n"
|
|
|
|
tools_description = self._format_tools_for_prompt()
|
|
|
|
return f"""
|
|
Original question: {original_prompt}
|
|
|
|
Previous tool results:{results_text}
|
|
|
|
Give the final result. It must be concise, with no explanation. A simple sentence is preferred.
|
|
"""
|
|
|
|
def _build_final_prompt(self, original_prompt: str, tool_results: List[Dict]) -> str:
|
|
"""Build prompt for final response generation"""
|
|
results_text = ""
|
|
for result in tool_results:
|
|
results_text += f"\nTool: {result['tool_name']} - Result: {result['result']}"
|
|
|
|
return f"""
|
|
Original question: {original_prompt}
|
|
Tool results used:{results_text}
|
|
|
|
Generate a concise final response based on this information.
|
|
Make it as short as possible. Do not mention the tools used, only the final result. One or two sentences preferred.
|
|
"""
|
|
|
|
async def _call_ollama(self, debug_response: DebugResponse, prompt: str) -> str:
|
|
"""Call Ollama to generate a response"""
|
|
logger.debug(f"LLM Request {prompt=}")
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
response = await client.post(
|
|
f"{self.ollama_host}/api/generate",
|
|
json={
|
|
"model": self.model,
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
"options": {"temperature": 0.0}
|
|
}
|
|
)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
logger.debug(f"LLM Response {result=}")
|
|
|
|
llm_response = result.get("response", "")
|
|
debug_response.response = llm_response # Store the LLM response in the debug object
|
|
|
|
return llm_response
|
|
except httpx.TimeoutException:
|
|
error_msg = "Ollama request timeout"
|
|
logger.error(error_msg)
|
|
debug_response.error_response = error_msg
|
|
raise Exception(error_msg)
|
|
except httpx.HTTPStatusError as e:
|
|
error_msg = f"Ollama HTTP error: {e.response.status_code}"
|
|
logger.error(error_msg)
|
|
debug_response.error_response = error_msg
|
|
raise Exception(error_msg)
|
|
except Exception as e:
|
|
error_msg = f"Unexpected Ollama error: {e}"
|
|
logger.error(error_msg)
|
|
debug_response.error_response = error_msg
|
|
raise
|
|
|
|
|
|
def _format_tools_for_prompt(self) -> str:
|
|
"""Format available tools for the prompt"""
|
|
tools_list = []
|
|
for tool_name, tool_info in self.available_tools.items():
|
|
params_desc = ", ".join([f"{k}: {v.get('description', k)}" for k, v in tool_info['parameters'].items()])
|
|
tools_list.append(f"- {tool_name}: {tool_info['description']} (params: {params_desc})")
|
|
return "\n".join(tools_list)
|
|
|
|
|
|
MPC_CLIENTS_IDS = [InProcessMCPClientNativeTools.ID, InProcessMCPClientCustomTools.ID] |