Function calling vllm

src/happy_vllm/core/resources.py

@@ -31,6 +31,7 @@
 
 from fastapi import FastAPI
 from ..model.model_base import Model
+from ..function_tools.functions import update_tools, clean_tools
 
 logger = logging.getLogger(__name__)
 
@@ -46,8 +47,13 @@ async def lifespan(app: FastAPI):
         logger.info("Model loaded")
 
         RESOURCES[RESOURCE_MODEL] = model
+
+        # Load the tools
+        update_tools(args=args)
+
         yield
 
         # Clean up the ML models and release the resources
         RESOURCES.clear()
+        clean_tools()
     return lifespan

src/happy_vllm/function_tools/functions.py

@@ -0,0 +1,205 @@
+import logging
+from copy import copy
+from typing import Union
+from argparse import Namespace
+
+
+class ToolFunctions:
+    """
+    Represents a tool function with specific attributes.
+
+    Attributes:
+    description (str): Description of the tool function.
+    parameters (dict): Parameters required for the tool function.
+    name (str): Name of the tool function.
+    tool_type (str): Type of the tool function.
+
+    Methods:
+    __init__(description: Union[str, None], parameters: Union[dict, None], name: Union[str, None], tool_type: Union[str, None]):
+        Initializes a ToolFunctions instance with the provided attributes. Raises NotImplementedError if any attribute is None.
+
+    _check_attributes():
+        Checks if the required attributes (description, parameters, name, tool_type) are not None.
+        Raises NotImplementedError if any attribute is None.
+
+    generate_dict() -> dict:
+        Generates and returns a dictionary representation of the tool function, including its type, name, description, and parameters.
+    """
+    def __init__(self, description:Union[str, None], parameters:Union[dict, None], name:Union[str, None], tool_type:Union[str, None]):
+        self.description:str = description
+        self.parameters:dict = parameters
+        self.name:str = name
+        self.tool_type:str = tool_type
+        self._check_attributes()
+
+    def _check_attributes(self):
+        if not self.description:
+            raise NotImplementedError("This attributes must be different to None")
+        if not self.parameters:
+            raise NotImplementedError("This attributes must be different to None")
+        if not self.name:
+            raise NotImplementedError("This attributes must be different to None")
+        if not self.tool_type:
+            raise NotImplementedError("This attributes must be different to None")
+
+    def generate_dict(self):
+        return {
+            "type": self.tool_type,
+            "function": {
+                "name": self.name,         
+                "description": self.description,
+                "parameters": self.parameters,
+
+            }
+        } 
+
+
+class Weather(ToolFunctions):
+    """
+    Represents a example tool function about the weather.
+    """
+    def __init__(self):
+        tool_type = "function"
+        name = "get_current_weather"
+        description = "Get current weather"
+        parameters = {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "The city and state, e.g. San Francisco, CA",
+                    },
+                    "format": {
+                        "type": "string",
+                        "enum": ["celsius", "fahrenheit"],
+                        "description": "The temperature unit to use. Infer this from the users location.",
+                    },
+                },
+                "required": ["location", "format"]
+            }
+        super().__init__(description=description, parameters=parameters, name=name, tool_type=tool_type)
+
+
+class Music(ToolFunctions):
+    """
+    Represents a example tool function about the music.
+    """
+    def __init__(self):
+        tool_type = "function"
+        name = "ask_database"
+        description = "Use this function to answer user questions about music. Input should be a fully formed SQL query."
+        parameters = {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": f"""
+                                SQL query extracting info to answer the user's question.
+                                SQL should be written using this database schema:
+                                <schema>
+                                The query should be returned in plain text, not in JSON.
+                                """,
+                    }
+                },
+                "required": ["query"]
+            }
+        super().__init__(description=description, parameters=parameters, name=name, tool_type=tool_type)
+
+
+TOOLS_DICT = {
+    'weather': Weather,
+    'music': Music
+}
+TOOLS = []
+
+def get_tools():
+    return TOOLS_DICT, TOOLS
+
+def reset_tools_dict_and_tools():
+    """
+    Resets the global variables TOOLS_DICT and TOOLS with new default values.
+
+    Returns:
+    tuple: A tuple containing the updated values of TOOLS_DICT and TOOLS. 
+           TOOLS_DICT is a dictionary with keys for different tools and corresponding values, 
+           and TOOLS is an empty list.
+    """
+    global TOOLS_DICT
+    global TOOLS
+    TOOLS_DICT = {
+        'weather': Weather,
+        'music': Music
+    }
+    TOOLS = []
+
+
+def update_tools(args: Namespace):
+    """
+    Updates the global variables TOOLS_DICT and TOOLS based on the provided arguments.
+
+    Args:
+    args (Namespace): A Namespace object containing parsed command-line arguments.
+
+    Returns:
+    tuple: A tuple containing the updated values of TOOLS_DICT and TOOLS. 
+           TOOLS_DICT is updated with instances of selected tools or set to None if no tools are selected.
+           TOOLS is updated with names of selected tools or set to None if no tools are selected.
+    """
+    global TOOLS_DICT
+    global TOOLS
+    if args.tools and args.tool_choice and 'none' not in args.tool_choice:
+        tools = {}
+        for t in args.tool_choice:
+            if TOOLS_DICT.get(t.lower(), None):
+                tools[t.lower()] = TOOLS_DICT[t.lower()]()
+            else:
+                raise KeyError(f"The tool '{t.lower()}' is not available in TOOLS_DICT")
+        TOOLS_DICT = tools
+        TOOLS = [t.lower() for t in args.tool_choice]
+    else:
+        if args.tools and args.tool_choice is None:
+            raise ValueError("The argument '--tool-choice' is required when '--tools' is specified")
+        elif args.tools is None and args.tool_choice:
+            raise ValueError("The argument '--tools' is required when '--tool-choice' is specified")
+        TOOLS_DICT = None
+        TOOLS = None
+
+
+def clean_tools():
+    """
+    Clears the global variable TOOLS_DICT, removing all entries.
+
+    Returns:
+    dict: An empty dictionary representing the cleaned TOOLS_DICT after removal of all entries.
+    """
+    global TOOLS_DICT
+    if TOOLS_DICT:
+        TOOLS_DICT.clear()
+
+
+def get_tools_prompt() -> dict:
+    """
+    Returns a dictionary containing information about selected tools.
+
+    Returns:
+    dict or None: A dictionary containing information about selected tools, structured as follows:
+                  - "tools": A list of dictionaries, each representing a tool's generated dictionary.
+                  - "tool_choice": A dictionary containing type and function details of the first tool in the list,
+                                   or None if TOOLS is empty.
+                  Returns None if TOOLS is empty.
+    """
+    tools_dict = copy(TOOLS_DICT)
+    tools = copy(TOOLS)
+    if tools:
+        return {
+            "tools": [tools_dict[t].generate_dict() for t in tools],
+            "tool_choice": [
+                {
+                    "type": tools_dict[t].tool_type, 
+                    "function": {"name":tools_dict[t].name}
+                }
+                for t in tools
+            ][0]
+        }
+    else:
+        return None

src/happy_vllm/routers/functional.py

@@ -17,9 +17,9 @@
 import os
 import json
 from vllm.utils import random_uuid
-from fastapi import APIRouter, Body
 from pydantic import BaseModel, Field
 from starlette.requests import Request
+from fastapi import APIRouter, Body, Depends
 from vllm.sampling_params import SamplingParams
 from vllm.engine.async_llm_engine import AsyncLLMEngine
 from lmformatenforcer import TokenEnforcerTokenizerData
@@ -327,7 +327,7 @@ async def metadata_text(request: Request,
 
 
 @router.post("/v1/chat/completions", response_model=functional_schema.HappyvllmChatCompletionResponse)
-async def create_chat_completion(request: Annotated[vllm_protocol.ChatCompletionRequest, Body(openapi_examples=request_openapi_examples["chat_completions"])],
+async def create_chat_completion(request: Annotated[vllm_protocol.ChatCompletionRequest, Depends(functional_schema.update_chat_completion_request)],
                                  raw_request: Request):
     """Open AI compatible chat completion. See https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html for more details
     """

src/happy_vllm/routers/schemas/functional.py

@@ -19,12 +19,16 @@
 
 import os
 import json
+from fastapi import Depends
+from starlette.requests import Request
 from starlette.responses import JSONResponse
 from pydantic import BaseModel, Field, conint
 from typing import Any, List, Union, Optional
-from vllm.entrypoints.openai.protocol import ResponseFormat, CompletionResponse, ChatCompletionResponse
+from vllm.entrypoints.openai.protocol import ResponseFormat, CompletionResponse, ChatCompletionResponse, ChatCompletionRequest
 
 from .utils import NumpyArrayEncoder
+from ...function_tools.functions import get_tools_prompt
+
 
 # Load the response examples
 directory = os.path.dirname(os.path.abspath(__file__))
@@ -136,4 +140,30 @@ class HappyvllmCompletionResponse(CompletionResponse):
 
 
 class HappyvllmChatCompletionResponse(ChatCompletionResponse):
-    model_config = {"json_schema_extra": {"examples": [response_examples["chat_completion_response"]]}}
+    model_config = {"json_schema_extra": {"examples": [response_examples["chat_completion_response"]]}}
+
+
+async def update_chat_completion_request(request: Request, data: ChatCompletionRequest):
+    """
+    Updates a ChatCompletionRequest object with additional tools and settings if available.
+
+    Args:
+        request (Request): The incoming request object.
+        data (ChatCompletionRequest): The original ChatCompletionRequest object to be updated.
+
+    Returns:
+        ChatCompletionRequest: The updated ChatCompletionRequest object if tools are present,
+                               otherwise returns the original ChatCompletionRequest object.
+    """
+    tools : Union[dict, None] = get_tools_prompt()
+    if tools:
+        data_dict = data.dict()
+        if data_dict['tools']:
+            data_dict['tools'].extend(tools["tools"])
+        else:
+            data_dict['tools'] = tools["tools"]
+        if data_dict['top_logprobs'] == 0:
+            data_dict['top_logprobs'] = None
+        data_dict['tool_choice'] = tools["tool_choice"]
+        return ChatCompletionRequest(**data_dict)
+    return data

src/happy_vllm/utils_args.py

@@ -47,6 +47,8 @@
 DEFAULT_CHAT_TEMPLATE = None
 DEFAULT_RESPONSE_ROLE = "assistant"
 DEFAULT_WITH_LAUNCH_ARGUMENTS = False
+DEFAULT_TOOLS = None
+DEFAULT_TOOLS_CHOICE = None
 
 
 class ApplicationSettings(BaseSettings):
@@ -79,6 +81,8 @@ class ApplicationSettings(BaseSettings):
     chat_template : Optional[str] = DEFAULT_CHAT_TEMPLATE
     response_role: str = DEFAULT_RESPONSE_ROLE
     with_launch_arguments: bool = DEFAULT_WITH_LAUNCH_ARGUMENTS
+    tools: Optional[str] = DEFAULT_TOOLS
+    tools_choice: Optional[str] = DEFAULT_TOOLS_CHOICE
 
     model_config = SettingsConfigDict(env_file=".env", extra='ignore', protected_namespaces=('settings', ))
 
@@ -280,6 +284,18 @@ def get_parser() -> ArgumentParser:
                         type=bool,
                         default=application_settings.with_launch_arguments,
                         help="Whether the route launch_arguments should display the launch arguments")
+    parser.add_argument("--tools",
+                        nargs='+',
+                        type=str,
+                        default=application_settings.tools,
+                        help="List of function tools LLM model could use")
+    parser.add_argument("--tool-choice",
+                        nargs='+', 
+                        type=str,
+                        default=application_settings.tools_choice,
+                        help="Controls which (if any) tool is called by the model.")
+
+
     parser = AsyncEngineArgs.add_cli_args(parser)
     return parser
 

tests/conftest.py

@@ -87,7 +87,9 @@ def test_base_client() -> TestClient:
                                         allowed_methods=["*"],
                                         allowed_headers=["*"],
                                         root_path=None,
-                                        with_launch_arguments=True))
+                                        with_launch_arguments=True,
+                                        tools=None,
+                                        tool_choice=None))
     return TestClient(app)
 
 
@@ -114,7 +116,9 @@ def test_complete_client(monkeypatch) -> TestClient:
                                         allowed_methods=["*"],
                                         allowed_headers=["*"],
                                         root_path=None,
-                                        with_launch_arguments=True))
+                                        with_launch_arguments=True,
+                                        tools=None,
+                                        tool_choice=None))
 
     with TestClient(app) as client:
         yield client