program.py

# This is a re-implementation of LlamaIndex's FunctionCallingProgram. Thanks to the LlamaIndex team for sharing their implementation!
# https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/program/function_program.py

import logging
from typing import Any, Dict, Optional, Type, Union, List
from pydantic import BaseModel, ValidationError
from openai.types.chat import ChatCompletionMessage
from config import get_client, get_model_settings
import json

_logger = logging.getLogger(__name__)
model_settings = get_model_settings()


def num_valid_fields(obj: Union[BaseModel, List[BaseModel], Dict[str, BaseModel]]) -> int:
    if isinstance(obj, BaseModel):
        return sum(1 for value in obj.__dict__.values() if value is not None)
    elif isinstance(obj, (list, tuple)):
        return sum(num_valid_fields(item) for item in obj)
    elif isinstance(obj, dict):
        return sum(num_valid_fields(item) for item in obj.values())
    else:
        return 1 if obj is not None else 0

class FunctionCallingProgram:
    def __init__(
        self,
        output_cls: Type[BaseModel],
        prompt: str,
        allow_parallel_tool_calls: bool = False,
        verbose: bool = False,
    ):
        self._output_cls = output_cls
        self._prompt = prompt
        self._verbose = verbose
        self._allow_parallel_tool_calls = allow_parallel_tool_calls
        self._client = get_client()

    @classmethod
    def from_defaults(
        cls,
        output_cls: Type[BaseModel],
        prompt: str,
        verbose: bool = False,
        allow_parallel_tool_calls: bool = False,
    ) -> "FunctionCallingProgram":
        return cls(
            output_cls=output_cls,
            prompt=prompt,
            allow_parallel_tool_calls=allow_parallel_tool_calls,
            verbose=verbose,
        )

    def _get_function_schema(self) -> Dict[str, Any]:
        schema = self._output_cls.model_json_schema()
        return {
            "type": "function",
            "function": {
                "name": self._output_cls.__name__,
                "description": "",
                "parameters": schema
            }
        }

    def _process_objects(
        self,
        chat_response: ChatCompletionMessage,
        output_cls: Type[BaseModel],
        cur_objects: Optional[List[BaseModel]] = None,
    ) -> Union[BaseModel, List[BaseModel]]:
        tool_calls = chat_response.tool_calls or []
        
        if len(tool_calls) == 0:
            return output_cls()

        objects = []
        for tool_call in tool_calls:
            try:
                function_args = json.loads(tool_call.function.arguments)
                objects.append(output_cls.model_validate(function_args))
            except (json.JSONDecodeError, ValidationError) as e:
                _logger.warning(f"Failed to parse tool call: {e}")
                objects.append(output_cls())

        if cur_objects is None or num_valid_fields(objects) > num_valid_fields(cur_objects):
            cur_objects = objects

        new_cur_objects = []
        for obj in cur_objects:
            try:
                new_obj = self._output_cls.model_validate(obj.model_dump())
            except ValidationError as e:
                _logger.warning(f"Failed to parse object: {e}")
                new_obj = obj
            new_cur_objects.append(new_obj)

        if self._allow_parallel_tool_calls:
            return new_cur_objects
        else:
            if len(new_cur_objects) > 1:
                _logger.warning(
                    "Multiple outputs found, returning first one. "
                    "If you want to return all outputs, set allow_parallel_tool_calls=True."
                )
            return new_cur_objects[0]

    def __call__(
        self,
        *args: Any,
        llm_kwargs: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Union[BaseModel, List[BaseModel]]:
        llm_kwargs = llm_kwargs or {}
        messages = [
            {"role": "system", "content": self._prompt},
            {"role": "user", "content": json.dumps(kwargs)}
        ]
        function_schema = self._get_function_schema()

        try:
            response = self._client.chat.completions.create(
                model=model_settings["model"],
                temperature=model_settings["temperature"],
                messages=messages,
                tools=[function_schema],
                tool_choice="auto",
            )
            return self._process_objects(response.choices[0].message, self._output_cls)
        except Exception as e:
            _logger.error(f"An error occurred: {str(e)}")
            raise