yaicli 0.6.3__py3-none-any.whl → 0.7.0__py3-none-any.whl

yaicli/llms/providers/gemini_provider.py

@@ -1,25 +1,18 @@
 import json
-from functools import wraps
 from typing import Any, Callable, Dict, Generator, List
 
 import google.genai as genai
 from google.genai import types
 
+from yaicli.tools.mcp import get_mcp_manager
+
 from ...config import cfg
 from ...console import get_console
 from ...schemas import ChatMessage, LLMResponse
-from ...tools import get_func_name_map
+from ...tools.function import get_functions_gemini_format
 from ..provider import Provider
 
 
-def wrap_function(func):
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        return func(*args, **kwargs)
-
-    return wrapper
-
-
 class GeminiProvider(Provider):
     """Gemini provider implementation based on google-genai library"""
 
@@ -28,6 +21,7 @@ class GeminiProvider(Provider):
     def __init__(self, config: dict = cfg, verbose: bool = False, **kwargs):
         self.config = config
         self.enable_function = self.config["ENABLE_FUNCTIONS"]
+        self.enable_mcp = self.config["ENABLE_MCP"]
         self.verbose = verbose
 
         # Initialize client
@@ -67,16 +61,17 @@ class GeminiProvider(Provider):
             config_map["frequency_penalty"] = self.config["FREQUENCY_PENALTY"]
         if self.config.get("SEED"):
             config_map["seed"] = self.config["SEED"]
-        # Indicates whether to include thoughts in the response. If true, thoughts are returned only if the model supports thought and thoughts are available.
+        # Indicates whether to include thoughts in the response.
+        # If true, thoughts are returned only if the model supports thought and thoughts are available.
         thinking_config_map = {"include_thoughts": self.config.get("INCLUDE_THOUGHTS", True)}
         if self.config.get("THINKING_BUDGET"):
             thinking_config_map["thinking_budget"] = int(self.config["THINKING_BUDGET"])
         config_map["thinking_config"] = types.ThinkingConfig(**thinking_config_map)
-        config = types.GenerateContentConfig(**config_map)
-        if self.enable_function:
+        if self.enable_function or self.enable_mcp:
             # TODO: support disable automatic function calling
             # config.automatic_function_calling = types.AutomaticFunctionCallingConfig(disable=False)
-            config.tools = self.gen_gemini_functions()
+            config_map["tools"] = self.gen_gemini_functions()
+        config = types.GenerateContentConfig(**config_map)
         return config
 
     def _convert_messages(self, messages: List[ChatMessage]) -> List[types.Content]:
@@ -88,7 +83,9 @@ class GeminiProvider(Provider):
             content = types.Content(role=self._map_role(msg.role), parts=[types.Part(text=msg.content)])
             if msg.role == "tool":
                 content.role = "user"
-                content.parts = [types.Part.from_function_response(name=msg.name, response={"result": msg.content})]
+                content.parts = [
+                    types.Part.from_function_response(name=msg.name or "", response={"result": msg.content})
+                ]
             converted_messages.append(content)
         return converted_messages
 
@@ -101,15 +98,19 @@ class GeminiProvider(Provider):
 
     def gen_gemini_functions(self) -> List[Callable[..., Any]]:
         """Wrap Gemini functions from OpenAI functions for automatic function calling"""
-        func_name_map = get_func_name_map()
-        if not func_name_map:
-            return []
         funcs = []
-        for func_name, func in func_name_map.items():
-            wrapped_func = wrap_function(func.execute)
-            wrapped_func.__name__ = func_name
-            wrapped_func.__doc__ = func.__doc__
-            funcs.append(wrapped_func)
+
+        # Add regular functions
+        if self.enable_function:
+            funcs.extend(get_functions_gemini_format())
+
+        # Add MCP functions if enabled
+        if self.enable_mcp:
+            try:
+                mcp_tools = get_mcp_manager().to_gemini_tools()
+                funcs.extend(mcp_tools)
+            except (ImportError, Exception) as e:
+                self.console.print(f"Failed to load MCP tools for Gemini: {e}", style="red")
         return funcs
 
     def completion(
@@ -137,14 +138,14 @@ class GeminiProvider(Provider):
             self.console.print(gemini_messages)
         chat_config = self.get_chat_config()
         chat_config.system_instruction = messages[0].content
-        chat = self.client.chats.create(model=self.config["MODEL"], history=gemini_messages, config=chat_config)
+        chat = self.client.chats.create(model=self.config["MODEL"], history=gemini_messages, config=chat_config)  # type: ignore
         message = messages[-1].content
 
         if stream:
-            response = chat.send_message_stream(message=message)
+            response = chat.send_message_stream(message=message)  # type: ignore
             yield from self._handle_stream_response(response)
         else:
-            response = chat.send_message(message=message)
+            response = chat.send_message(message=message)  # type: ignore
             yield from self._handle_normal_response(response)
 
     def _handle_normal_response(self, response) -> Generator[LLMResponse, None, None]:
@@ -158,7 +159,7 @@ class GeminiProvider(Provider):
             return
         for part in response.candidates[0].content.parts:
             if part.thought:
-                yield LLMResponse(reasoning=part.text, content=None, finish_reason="stop")
+                yield LLMResponse(reasoning=part.text, finish_reason="stop")
             else:
                 yield LLMResponse(reasoning=None, content=part.text, finish_reason="stop")
 
@@ -181,7 +182,7 @@ class GeminiProvider(Provider):
                     reasoning = None
                 yield LLMResponse(
                     reasoning=reasoning,
-                    content=content,
+                    content=content or "",
                     tool_call=tool_call if finish_reason == "tool_calls" else None,
                     finish_reason=finish_reason or None,
                 )

yaicli/llms/providers/huggingface_provider.py

@@ -0,0 +1,40 @@
+from typing import Any, Dict
+
+from huggingface_hub import InferenceClient
+
+from .chatglm_provider import ChatglmProvider
+
+
+class HuggingFaceProvider(ChatglmProvider):
+    """
+    HuggingFaceProvider is a provider for the HuggingFace API.
+    """
+
+    CLIENT_CLS = InferenceClient
+    DEFAULT_PROVIDER = "auto"
+
+    COMPLETION_PARAMS_KEYS = {
+        "model": "MODEL",
+        "temperature": "TEMPERATURE",
+        "top_p": "TOP_P",
+        "max_tokens": "MAX_TOKENS",
+        "extra_body": "EXTRA_BODY",
+    }
+
+    def get_client_params(self) -> Dict[str, Any]:
+        client_params = {
+            "api_key": self.config["API_KEY"],
+            "timeout": self.config["TIMEOUT"],
+            "provider": self.config.get("HF_PROVIDER") or self.DEFAULT_PROVIDER,
+        }
+        if self.config["BASE_URL"]:
+            client_params["base_url"] = self.config["BASE_URL"]
+        if self.config["EXTRA_HEADERS"]:
+            client_params["headers"] = {
+                **self.config["EXTRA_HEADERS"],
+                "X-Title": self.APP_NAME,
+                "HTTP-Referer": self.APP_REFERER,
+            }
+        if self.config.get("BILL_TO"):
+            client_params["bill_to"] = self.config["BILL_TO"]
+        return client_params

yaicli/llms/providers/infiniai_provider.py

@@ -1,5 +1,6 @@
 from typing import Any, Dict
 
+from ...config import cfg
 from .openai_provider import OpenAIProvider
 
 
@@ -8,7 +9,7 @@ class InfiniAIProvider(OpenAIProvider):
 
     DEFAULT_BASE_URL = "https://cloud.infini-ai.com/maas/v1"
 
-    def __init__(self, config: dict = ..., **kwargs):
+    def __init__(self, config: dict = cfg, **kwargs):
         super().__init__(config, **kwargs)
         if self.enable_function:
             self.console.print("InfiniAI does not support functions, disabled", style="yellow")
@@ -16,5 +17,6 @@ class InfiniAIProvider(OpenAIProvider):
 
     def get_completion_params(self) -> Dict[str, Any]:
         params = super().get_completion_params()
-        params["max_tokens"] = params.pop("max_completion_tokens")
+        if "max_completion_tokens" in params:
+            params["max_tokens"] = params.pop("max_completion_tokens")
         return params

yaicli/llms/providers/modelscope_provider.py

@@ -10,5 +10,6 @@ class ModelScopeProvider(OpenAIProvider):
 
     def get_completion_params(self) -> Dict[str, Any]:
         params = super().get_completion_params()
-        params["max_tokens"] = params.pop("max_completion_tokens")
+        if "max_completion_tokens" in params:
+            params["max_tokens"] = params.pop("max_completion_tokens")
         return params

yaicli/llms/providers/openai_provider.py

@@ -10,6 +10,7 @@ from ...config import cfg
 from ...console import get_console
 from ...schemas import ChatMessage, LLMResponse, ToolCall
 from ...tools import get_openai_schemas
+from ...tools.mcp import get_mcp_manager
 from ..provider import Provider
 
 
@@ -19,7 +20,7 @@ class OpenAIProvider(Provider):
     DEFAULT_BASE_URL = "https://api.openai.com/v1"
     CLIENT_CLS = openai.OpenAI
     # Base mapping between config keys and API parameter names
-    _BASE_COMPLETION_PARAMS_KEYS = {
+    COMPLETION_PARAMS_KEYS = {
         "model": "MODEL",
         "temperature": "TEMPERATURE",
         "top_p": "TOP_P",
@@ -34,6 +35,7 @@ class OpenAIProvider(Provider):
         if not self.config.get("API_KEY"):
             raise ValueError("API_KEY is required")
         self.enable_function = self.config["ENABLE_FUNCTIONS"]
+        self.enable_mcp = self.config["ENABLE_MCP"]
         self.verbose = verbose
 
         # Initialize client
@@ -50,15 +52,12 @@ class OpenAIProvider(Provider):
         client_params = {
             "api_key": self.config["API_KEY"],
             "base_url": self.config["BASE_URL"] or self.DEFAULT_BASE_URL,
+            "default_headers": {"X-Title": self.APP_NAME, "HTTP_Referer": self.APP_REFERER},
         }
 
         # Add extra headers if set
         if self.config["EXTRA_HEADERS"]:
-            client_params["default_headers"] = {
-                **self.config["EXTRA_HEADERS"],
-                "X-Title": self.APP_NAME,
-                "HTTP-Referer": self.APP_REFERER,
-            }
+            client_params["default_headers"] = {**self.config["EXTRA_HEADERS"], **client_params["default_headers"]}
         return client_params
 
     def get_completion_params_keys(self) -> Dict[str, str]:
@@ -69,7 +68,7 @@ class OpenAIProvider(Provider):
         Returns:
             Dict[str, str]: Mapping from API parameter names to config keys
         """
-        return self._BASE_COMPLETION_PARAMS_KEYS.copy()
+        return self.COMPLETION_PARAMS_KEYS.copy()
 
     def get_completion_params(self) -> Dict[str, Any]:
         """
@@ -89,7 +88,7 @@ class OpenAIProvider(Provider):
         """Convert a list of ChatMessage objects to a list of OpenAI message dicts."""
         converted_messages = []
         for msg in messages:
-            message = {"role": msg.role, "content": msg.content or ""}
+            message: Dict[str, Any] = {"role": msg.role, "content": msg.content or ""}
 
             if msg.name:
                 message["name"] = msg.name
@@ -134,11 +133,21 @@ class OpenAIProvider(Provider):
         params = self.completion_params.copy()
         params["messages"] = openai_messages
         params["stream"] = stream
+        tools = []
 
         if self.enable_function:
-            tools = get_openai_schemas()
-            if tools:
-                params["tools"] = tools
+            tools.extend(get_openai_schemas())
+
+        # Add MCP tools if enabled
+        if self.enable_mcp:
+            try:
+                mcp_tools = get_mcp_manager().to_openai_tools()
+            except (ValueError, FileNotFoundError) as e:
+                self.console.print(f"Failed to load MCP tools: {e}", style="red")
+                mcp_tools = []
+            tools.extend(mcp_tools)
+        if tools:
+            params["tools"] = tools
 
         try:
             if stream:

yaicli/llms/providers/siliconflow_provider.py

@@ -10,5 +10,6 @@ class SiliconFlowProvider(OpenAIProvider):
 
     def get_completion_params(self) -> Dict[str, Any]:
         params = super().get_completion_params()
-        params["max_tokens"] = params.pop("max_completion_tokens")
+        if "max_completion_tokens" in params:
+            params["max_tokens"] = params.pop("max_completion_tokens")
         return params

yaicli/tools/__init__.py

@@ -0,0 +1,127 @@
+from typing import Any, Dict, List, Tuple, cast
+
+from json_repair import repair_json
+from mcp import types
+from rich.panel import Panel
+
+from ..config import cfg
+from ..console import get_console
+from ..schemas import ToolCall
+from .function import get_function, list_functions
+from .mcp import MCP_TOOL_NAME_PREFIX, get_mcp, get_mcp_manager, parse_mcp_tool_name
+
+console = get_console()
+
+
+def get_openai_schemas() -> List[Dict[str, Any]]:
+    """Get OpenAI-compatible function schemas
+
+    Returns:
+        List of function schemas in OpenAI format
+    """
+    transformed_schemas = []
+    for function in list_functions():
+        schema = {
+            "type": "function",
+            "function": {
+                "name": function.name,
+                "description": function.description,
+                "parameters": function.parameters,
+            },
+        }
+        transformed_schemas.append(schema)
+    return transformed_schemas
+
+
+def get_openai_mcp_tools() -> list[dict[str, Any]]:
+    """Get OpenAI-compatible function schemas
+
+    Returns:
+        List of function schemas in OpenAI format
+    """
+    return get_mcp_manager().to_openai_tools()
+
+
+def execute_mcp_tool(tool_name: str, tool_kwargs: dict) -> str:
+    """Execute an MCP tool
+
+    Args:
+        tool_name: The name of the tool to execute
+        tool_kwargs: The arguments to pass to the tool
+    """
+    manager = get_mcp_manager()
+    tool = manager.get_tool(tool_name)
+    try:
+        result = tool.execute(**tool_kwargs)
+        if isinstance(result, list) and len(result) > 0:
+            result = result[0]
+        if isinstance(result, types.TextContent):
+            return result.text
+        else:
+            return str(result)
+    except Exception as e:
+        error_msg = f"Call MCP tool error:\nTool name: {tool_name!r}\nArguments: {tool_kwargs!r}\nError: {e}"
+        console.print(error_msg, style="red")
+        return error_msg
+
+
+def execute_tool_call(tool_call: ToolCall) -> Tuple[str, bool]:
+    """Execute a tool call and return the result
+
+    Args:
+        tool_call: The tool call to execute
+
+    Returns:
+        Tuple[str, bool]: (result text, success flag)
+    """
+    is_function_call = not tool_call.name.startswith(MCP_TOOL_NAME_PREFIX)
+    if is_function_call:
+        get_tool_func = get_function
+        show_output = cfg["SHOW_FUNCTION_OUTPUT"]
+        _type = "function"
+    else:
+        tool_call.name = parse_mcp_tool_name(tool_call.name)
+        get_tool_func = get_mcp
+        show_output = cfg["SHOW_MCP_OUTPUT"]
+        _type = "mcp"
+
+    console.print(f"@{_type.title()} call: {tool_call.name}({tool_call.arguments})", style="blue")
+    # 1. Get the tool
+    try:
+        tool = get_tool_func(tool_call.name)
+    except ValueError as e:
+        error_msg = f"{_type.title()} '{tool_call.name!r}' not exists: {e}"
+        console.print(error_msg, style="red")
+        return error_msg, False
+
+    # 2. Parse tool arguments
+    try:
+        arguments = repair_json(tool_call.arguments, return_objects=True)
+        if not isinstance(arguments, dict):
+            error_msg = f"Invalid arguments type: {arguments!r}, should be JSON object"
+            console.print(error_msg, style="red")
+            return error_msg, False
+        arguments = cast(dict, arguments)
+    except Exception as e:
+        error_msg = f"Invalid arguments from llm: {e}\nRaw arguments: {tool_call.arguments!r}"
+        console.print(error_msg, style="red")
+        return error_msg, False
+
+    # 3. Execute the tool
+    try:
+        result = tool.execute(**arguments)
+        if show_output:
+            panel = Panel(
+                result,
+                title=f"{_type.title()} output",
+                title_align="left",
+                expand=False,
+                border_style="blue",
+                style="dim",
+            )
+            console.print(panel)
+        return result, True
+    except Exception as e:
+        error_msg = f"Call {_type} error: {e}\n{_type} name: {tool_call.name!r}\nArguments: {arguments!r}"
+        console.print(error_msg, style="red")
+        return error_msg, False

yaicli/tools/function.py

@@ -0,0 +1,90 @@
+import importlib.util
+import sys
+from typing import Callable, List, Optional
+
+from instructor import OpenAISchema
+
+from ..const import FUNCTIONS_DIR
+from ..utils import wrap_function
+
+
+class Function:
+    """Function description class"""
+
+    def __init__(self, function: type[OpenAISchema]):
+        self.name = function.openai_schema["name"]
+        self.description = function.openai_schema.get("description", "")
+        self.parameters = function.openai_schema.get("parameters", {})
+        self.execute = function.execute  # type: ignore
+
+
+_func_name_map: Optional[dict[str, Function]] = None
+
+
+def get_func_name_map() -> dict[str, Function]:
+    """Get function name map"""
+    global _func_name_map
+    if _func_name_map:
+        return _func_name_map
+    if not FUNCTIONS_DIR.exists():
+        FUNCTIONS_DIR.mkdir(parents=True, exist_ok=True)
+        return {}
+    functions = []
+    for file in FUNCTIONS_DIR.glob("*.py"):
+        if file.name.startswith("_"):
+            continue
+        module_name = str(file).replace("/", ".").rstrip(".py")
+        spec = importlib.util.spec_from_file_location(module_name, str(file))
+        module = importlib.util.module_from_spec(spec)  # type: ignore
+        sys.modules[module_name] = module
+        spec.loader.exec_module(module)  # type: ignore
+
+        if not issubclass(module.Function, OpenAISchema):
+            raise TypeError(f"Function {module_name} must be a subclass of instructor.OpenAISchema")
+        if not hasattr(module.Function, "execute"):
+            raise TypeError(f"Function {module_name} must have an 'execute' classmethod")
+
+        # Add to function list
+        functions.append(Function(function=module.Function))
+
+    # Cache the function list
+    _func_name_map = {func.name: func for func in functions}
+    return _func_name_map
+
+
+def list_functions() -> list[Function]:
+    """List all available buildin functions"""
+    global _func_name_map
+    if not _func_name_map:
+        _func_name_map = get_func_name_map()
+
+    return list(_func_name_map.values())
+
+
+def get_function(name: str) -> Function:
+    """Get a function by name
+
+    Args:
+        name: Function name
+
+    Returns:
+        Function execute method
+
+    Raises:
+        ValueError: If function not found
+    """
+    func_map = get_func_name_map()
+    if name in func_map:
+        return func_map[name]
+    raise ValueError(f"Function {name!r} not found")
+
+
+def get_functions_gemini_format() -> List[Callable]:
+    """Get functions in gemini format"""
+    gemini_functions = []
+    for func_name, func in get_func_name_map().items():
+        wrapped_func = wrap_function(func.execute)
+        wrapped_func.__name__ = func_name
+        wrapped_func.__doc__ = func.description
+        gemini_functions.append(wrapped_func)
+    return gemini_functions