Folder-based functions for ChatGPT's function calling with Pydantic support 🚀
from typing import List

from vectordb import Memory

from sageai.types.abstract_vectordb import AbstractVectorDB


class CustomVectorDB(AbstractVectorDB):
    def __init__(self):
        super().__init__()
        self.memory = Memory()

    def index(self):
        self.memory.clear()
        for func_name, func in self.function_map.items():
            self.memory.save(func_name, func)

    def search(self, query: str, n: int) -> List[str]:
        return self.memory.search(query, top_n=n)

examples/advanced/src/customvectordb.py
0xnenlabs/SageAI
import random
from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field

from sageai.types.function import Function


class UnitTypes(str, Enum):
    CELSIUS = "celsius"
    FAHRENHEIT = "fahrenheit"


class FunctionInput(BaseModel):
    min: Optional[int] = Field(0, description="The minimum number.")
    max: Optional[int] = Field(100, description="The maximum number.")


class FunctionOutput(BaseModel):
    number: int

    # In this case, we don't have a way to determine if the output is correct, so we just return True to make
    # the tests pass.
    def __eq__(self, other):
        if not isinstance(other, FunctionOutput):
            return False
        return True


def get_random_number(params: FunctionInput) -> FunctionOutput:
    number = random.randint(params.min, params.max)
    return FunctionOutput(number=number)


function = Function(
    function=get_random_number,
    description="Get a random number.",
)

examples/advanced/src/functions/get_random_number/function.py
0xnenlabs/SageAI
from sageai.types.function import Function
from examples.advanced.src.functions.weather.get_current_weather.types import (
    FunctionInput,
    FunctionOutput,
)


def get_current_weather(params: FunctionInput) -> FunctionOutput:
    weather = (
        f"The weather in {params.location} is currently 22 degrees {params.unit.value}."
    )
    return FunctionOutput(weather=weather)


function = Function(
    function=get_current_weather,
    description="Get the current weather in a given location.",
)

examples/advanced/src/functions/weather/get_current_weather/function.py
0xnenlabs/SageAI
from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field


class UnitTypes(str, Enum):
    CELSIUS = "Celsius"
    FAHRENHEIT = "Fahrenheit"


class FunctionInput(BaseModel):
    location: str = Field(..., description="The city, e.g. San Francisco")
    unit: Optional[UnitTypes] = Field(
        UnitTypes.CELSIUS, description="The unit of temperature."
    )


class FunctionOutput(BaseModel):
    weather: str

    def __eq__(self, other):
        if not isinstance(other, FunctionOutput):
            return False
        return self.weather == other.weather

examples/advanced/src/functions/weather/get_current_weather/types.py
0xnenlabs/SageAI
from sageai.types.function import Function
from examples.advanced.src.functions.weather.get_forecast_weather.types import (
    FunctionInput,
    FunctionOutput,
)


def some_helper(params: FunctionInput) -> str:
    return (
        f"The weather {params.forecast} in {params.location} is going to be 22 degrees."
    )


def get_forecast_weather(params: FunctionInput) -> FunctionOutput:
    forecast = some_helper(params)
    return FunctionOutput(forecast=forecast)


function = Function(
    function=get_forecast_weather,
    description="Get the forecasted weather in a given location.",
)

examples/advanced/src/functions/weather/get_forecast_weather/function.py
0xnenlabs/SageAI
from pydantic import BaseModel, Field


class FunctionInput(BaseModel):
    location: str = Field(..., description="The city, e.g. San Francisco")
    forecast: str = Field(..., description="The forecast, e.g. today, tomorrow.")


class FunctionOutput(BaseModel):
    forecast: str

    def __eq__(self, other):
        if not isinstance(other, FunctionOutput):
            return False
        return self.forecast == other.forecast

examples/advanced/src/functions/weather/get_forecast_weather/types.py
0xnenlabs/SageAI
from sageai import SageAI

from customvectordb import CustomVectorDB

# Init on startup
sage = SageAI(
    openai_key="",
    vectordb=CustomVectorDB,
)

# In a CI/CD pipeline or in dev mode on startup/hot reload
sage.index()

# Anywhere in the codebase
message = "What's the weather like in Toronto right now?"
print(f"Message: {message}")
response = sage.chat(
    messages=[dict(role="user", content=message)],
    model="gpt-3.5-turbo-0613",
    top_n=5,
)
print(f"Response: {response}")

message = "Give me a number between 1 and 10."
print(f"\nMessage: {message}")
response = sage.chat(
    messages=[dict(role="user", content=message)],
    model="gpt-3.5-turbo-0613",
    top_n=5,
)
print(f"Response: {response}")

message = "What's the weather tomorrow in Toronto?"
print(f"\nMessage: {message}")
response = sage.chat(
    messages=[dict(role="user", content=message)],
    model="gpt-3.5-turbo-0613",
    top_n=5,
)
print(f"Response: {response}")

examples/advanced/src/main.py
0xnenlabs/SageAI
from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field

from sageai.types.function import Function


class UnitTypes(str, Enum):
    CELSIUS = "Celsius"
    FAHRENHEIT = "Fahrenheit"


class FunctionInput(BaseModel):
    location: str = Field(..., description="The city, e.g. San Francisco")
    unit: Optional[UnitTypes] = Field(
        UnitTypes.CELSIUS, description="The unit of temperature."
    )


class FunctionOutput(BaseModel):
    weather: str

    def __eq__(self, other):
        if not isinstance(other, FunctionOutput):
            return False
        return self.weather == other.weather


def get_current_weather(params: FunctionInput) -> FunctionOutput:
    weather = (
        f"The weather in {params.location} is currently 22 degrees {params.unit.value}."
    )

    return FunctionOutput(weather=weather)


function = Function(
    function=get_current_weather,
    description="Get the current weather in a given location.",
)

examples/basic/src/functions/get_current_weather/function.py
0xnenlabs/SageAI
from sageai import SageAI

# Init on startup
sage = SageAI(openai_key="")

# In a CI/CD pipeline or in dev mode on startup/hot reload
sage.index()

# Anywhere in the codebase
message = "What's the weather like in Toronto right now?"
response = sage.chat(
    messages=[dict(role="user", content=message)],
    model="gpt-3.5-turbo-0613",
    top_n=5,
)

print(f"Message: {message}")
print(f"Response: {response}")

examples/basic/src/main.py
0xnenlabs/SageAI
from sageai.sageai import SageAI

sageai/__init__.py
0xnenlabs/SageAI
from typing import Optional, Type

from pydantic import BaseModel, Field, ValidationError

from sageai.services.defaultvectordb_service import DefaultVectorDBService
from sageai.types.abstract_vectordb import AbstractVectorDB
from sageai.types.log_level import LogLevel
from sageai.utils.file_utilities import generate_functions_map
from sageai.utils.format_config_args import format_config_args


class Config(BaseModel):
    openai_key: str
    functions_directory: Optional[str] = Field(
        "functions", description="The directory of functions."
    )
    vectordb: Optional[Type[AbstractVectorDB]] = Field(
        DefaultVectorDBService, description="VectorDB class reference."
    )
    log_level: Optional[LogLevel] = Field(
        LogLevel.ERROR, description="The desired log level for output."
    )

    class Config:
        arbitrary_types_allowed = True


_config = Config(openai_key="")
function_map = {}


def set_config(**kwargs):
    """Set the configuration parameters."""
    global _config
    global function_map
    try:
        kwargs = format_config_args(**kwargs)
        _config = Config(**kwargs)
        function_map = generate_functions_map()
        return _config
    except ValidationError as e:
        raise ValidationError(f"Invalid configuration: {e}")


def get_config():
    """Retrieve the configuration."""
    return _config


def get_function_map():
    return function_map

sageai/config.py
0xnenlabs/SageAI
import json
from typing import Any, Dict, Optional, Tuple, Type

from sageai.config import LogLevel, get_config, get_function_map, set_config
from sageai.services.openai_service import OpenAIService
from sageai.types.abstract_vectordb import AbstractVectorDB
from sageai.utils.inspection_utilities import get_input_parameter_type
from sageai.utils.openai_utilities import get_latest_user_message

__all__ = ["SageAI"]


class SageAI:
    def __init__(
        self,
        *,
        openai_key: str,
        functions_directory: Optional[str] = None,
        vectordb: Optional[Type[AbstractVectorDB]] = None,
        log_level: Optional[LogLevel] = None,
    ):
        if openai_key is None:
            raise Exception("No OpenAI key provided.")

        config_args = {"openai_key": openai_key}

        if functions_directory is not None:
            config_args["functions_directory"] = functions_directory
        if vectordb is not None:
            config_args["vectordb"] = vectordb
        if log_level is not None:
            config_args["log_level"] = LogLevel(log_level)

        set_config(**config_args)
        self.config = get_config()

        self.openai = OpenAIService()
        self.vectordb = self.config.vectordb()

    def index(self):
        self.vectordb.index()

    def chat(self, *args, **kwargs) -> Dict[str, Any]:
        merged = {i: v for i, v in enumerate(args)}
        merged.update(kwargs)

        top_n = merged.pop("top_n") if "top_n" in merged else None

        if merged.get("model") is None:
            raise Exception("No model provided.")

        if merged.get("messages") is None:
            raise Exception("No messages provided.")

        if top_n is None:
            raise Exception("No top_n provided.")

        latest_user_message = get_latest_user_message(merged.get("messages"))
        if latest_user_message is None:
            raise Exception("No user message found.")

        top_functions = self.get_top_n_functions(
            query=latest_user_message["content"], top_n=top_n
        )

        function_name, function_args = self.call_openai(merged, top_functions)
        function_response = self.run_function(name=function_name, args=function_args)

        base_return = dict(name=function_name, args=function_args)

        if "error" in function_response:
            base_return["error"] = function_response["error"]
        else:
            base_return["result"] = function_response

        return base_return

    def get_top_n_functions(self, *, query: str, top_n: int):
        return self.vectordb.search_impl(query=query, top_n=top_n)

    def call_openai(
        self, openai_args: Dict[str, Any], top_functions: list[Dict[str, Any]]
    ) -> Tuple[str, Dict[str, Any]]:
        openai_result = self.openai.chat(**openai_args, functions=top_functions)

        if "function_call" not in openai_result:
            raise Exception("No function call found in OpenAI response.")

        function_name = openai_result["function_call"]["name"]
        function_args = json.loads(openai_result["function_call"]["arguments"])
        return function_name, function_args

    @staticmethod
    def run_function(*, name: str, args: Dict[str, Any]) -> Dict[str, Any]:
        try:
            function_map = get_function_map()
            func = function_map[name]
            function_input_type = get_input_parameter_type(func.function)
            func_args = function_input_type(**args)
            func_result = func(func_args)
            return func_result.dict()
        except Exception as e:
            return dict(error=str(e))

sageai/sageai.py
0xnenlabs/SageAI
from typing import List

from qdrant_client import QdrantClient
from qdrant_client.http.models import models

from sageai.services.openai_service import OpenAIService
from sageai.types.abstract_vectordb import AbstractVectorDB
from sageai.types.function import Function


class DefaultVectorDBService(AbstractVectorDB):
    def __init__(self):
        super().__init__()

        self.openai = OpenAIService()
        self.client = QdrantClient(":memory:")
        self.collection = "functions"

    def index(self):
        self.client.recreate_collection(
            collection_name=self.collection,
            vectors_config=models.VectorParams(
                size=self.openai.get_embeddings_size(),
                distance=models.Distance.COSINE,
            ),
        )

        def format_func_embedding(func: Function) -> str:
            return func.name.replace("_", " ") + " - " + func.description

        records = [
            models.Record(
                id=idx,
                vector=self.openai.create_embeddings(
                    model="text-embedding-ada-002", input=format_func_embedding(func)
                ),
                payload={"func_name": func_name},
            )
            for idx, (func_name, func) in enumerate(self.function_map.items())
        ]
        self.client.upload_records(
            collection_name=self.collection,
            records=records,
        )

    def search(self, *, query: str, top_n: int) -> List[str]:
        query_embedding = self.openai.create_embeddings(
            model="text-embedding-ada-002",
            input=query,
        )
        hits = self.client.search(
            collection_name=self.collection,
            query_vector=query_embedding,
            limit=top_n,
        )
        func_names = [hit.payload["func_name"] for hit in hits]
        return func_names

sageai/services/defaultvectordb_service.py
0xnenlabs/SageAI
from typing import Any, Dict, List

import openai


class OpenAIService:
    def __init__(self):
        from sageai.config import get_config

        config = get_config()
        openai.api_key = config.openai_key
        self.openai = openai

    def create_embeddings(self, *args, **kwargs) -> List[float]:
        response = self.openai.Embedding.create(*args, **kwargs)
        embeddings = response["data"][0]["embedding"]
        return embeddings

    def chat(self, *args, **kwargs) -> Dict[str, Any]:
        response = self.openai.ChatCompletion.create(*args, **kwargs)
        response_message = response["choices"][0]["message"]
        return response_message

    @staticmethod
    def get_embeddings_size():
        return 1536

sageai/services/openai_service.py
0xnenlabs/SageAI
import json
import os
from time import sleep
from typing import cast, get_type_hints
import pydantic

import pytest

from sageai.sageai import SageAI
from sageai.types.function import Function
from sageai.types.log_level import LogLevel
from sageai.utils.file_utilities import get_functions_directories, load_module_from_file
from sageai.utils.logger import get_logger

logger = get_logger("IntegrationTest", LogLevel.INFO)


@pytest.mark.parametrize(
    "dirpath",
    get_functions_directories(
        logger=logger, functions_directory_path=os.getenv("TEST_DIRECTORY")
    ),
)
def test_integration(dirpath):
    sleep(2)
    test_file = os.path.join(dirpath, "test.json")
    function_file = os.path.join(dirpath, "function.py")
    function_module = load_module_from_file("function", function_file)
    function_to_test = cast(Function, getattr(function_module, "function"))

    directory_full_path = os.path.abspath(os.getcwd())
    functions_directory = os.path.join(directory_full_path, dirpath)

    openai_key = os.environ["OPENAI_KEY"]
    sageai = SageAI(openai_key=openai_key, functions_directory=functions_directory)
    sageai.index()

    try:
        output_data_model_class = get_type_hints(function_to_test.function)["return"]

        with open(test_file) as f:
            test_cases = json.load(f)
            for i, test_case in enumerate(test_cases):
                logger.info(
                    f"{function_module.function.name} - Running test case {i+1}/{len(test_cases)}: {test_case['message']}",
                )
                try:
                    vector_db_result = sageai.vectordb.search(
                        query=test_case["message"], top_n=5
                    )
                    result = sageai.chat(
                        messages=[dict(role="user", content=test_case["message"])],
                        model="gpt-3.5-turbo-0613",
                        top_n=5,
                    )
                    if result is None:
                        logger.error(
                            "No result returned from Sennin for function: {}".format(
                                function_module.function.name
                            )
                        )
                        continue

                    test_output = None
                    try:
                        test_output = output_data_model_class(
                            **test_case["output"]
                        ).dict()
                    except pydantic.ValidationError as ve:
                        expected_fields = ", ".join(
                            get_type_hints(output_data_model_class).keys()
                        )
                        actual_fields = test_case["output"]
                        logger.error(
                            f"Validation error for function output {function_module.function.name}:\nExpected fields: {expected_fields}\nGot: {actual_fields}.\n{str(ve)}"
                        )
                        continue

                    if function_module.function.name not in vector_db_result:
                        logger.error(
                            f"VectorDB did not return expected function for {function_module.function.name}. Expected: {function_module.function.name} in {vector_db_result}."
                        )
                        continue
                    if result["name"] != function_module.function.name:
                        logger.error(
                            f"Wrong function chosen by model for {function_module.function.name}. Expected: {function_module.function.name}, Got: {result['name']}"
                        )
                        continue
                    if result["result"] != test_output:
                        logger.error(
                            f"Output mismatch for {function_module.function.name}. Expected: {test_output}, Got: {result['result']}"
                        )
                        continue

                except Exception as e:
                    logger.error(
                        f"Function {function_module.function.name} execution failed: {str(e)}"
                    )

    except Exception as e:
        logger.error(f"Entire test failed: {str(e)} for directory {dirpath}")

sageai/tests/integration.py
0xnenlabs/SageAI
import argparse
import os

import pytest


def main():
    parser = argparse.ArgumentParser(description="Run tests for functions.")
    parser.add_argument("--directory", type=str, help="Name of function to test")
    parser.add_argument(
        "--apikey", type=str, help="OpenAI API key for integration tests"
    )
    parser.add_argument("--unit", action="store_true", help="Run only the unit tests")
    parser.add_argument(
        "--integration", action="store_true", help="Run only the integration tests"
    )
    args = parser.parse_args()

    if args.directory:
        os.environ["TEST_DIRECTORY"] = args.directory
    if args.apikey:
        os.environ["OPENAI_KEY"] = args.apikey

    current_directory = os.path.dirname(os.path.abspath(__file__))

    if not args.integration or args.unit:
        unit_test_path = os.path.join(current_directory, "unit.py")
        pytest.main(["-x", "-s", unit_test_path, "--disable-warnings"])
    if not args.unit or args.integration:
        integration_test_path = os.path.join(current_directory, "integration.py")
        pytest.main(["-x", "-s", integration_test_path, "--disable-warnings"])

sageai/tests/main.py
0xnenlabs/SageAI
import json
import os
from typing import cast, get_type_hints
import pydantic

import pytest

from sageai.types.function import Function
from sageai.types.log_level import LogLevel
from sageai.utils.file_utilities import get_functions_directories, load_module_from_file
from sageai.utils.logger import get_logger

logger = get_logger("UnitTest", LogLevel.INFO)


def assert_function_module_correct(folder_name, function_module):
    if not hasattr(function_module, "function"):
        msg = f"{folder_name}/function.py does not export a variable called 'function'"
        raise ValueError(msg)


def assert_test_cases_correct(test_cases, folder_name):
    for idx, test_case in enumerate(test_cases):
        if not isinstance(test_case, dict):
            msg = f"Test case #{idx + 1} in {folder_name}/test.json is not an object"
            raise ValueError(msg)

        missing_properties = [
            prop for prop in ["input", "output", "message"] if prop not in test_case
        ]
        if missing_properties:
            msg = f"Test case #{idx + 1} in {folder_name}/tests.json missing properties: {', '.join(missing_properties)}"
            raise ValueError(msg)

        if not isinstance(test_case.get("message", ""), str):
            msg = f"Test case #{idx + 1} in {folder_name}/tests.json has a 'message' property which is not of type string"
            raise ValueError(msg)


@pytest.mark.parametrize(
    "dirpath",
    get_functions_directories(
        logger, functions_directory_path=os.getenv("TEST_DIRECTORY")
    ),
)
def test_unit(dirpath):
    folder_name = os.path.basename(dirpath)
    function_file = os.path.join(dirpath, "function.py")
    test_file = os.path.join(dirpath, "test.json")

    try:
        # Check necessary files
        if not os.path.exists(function_file):
            msg = f"Missing function.py in {folder_name}"
            raise ValueError(msg)

        if not os.path.exists(test_file):
            msg = f"Missing test.json in {folder_name}"
            raise ValueError(msg)

        function_module = load_module_from_file("function", function_file)
        function_to_test = cast(Function, getattr(function_module, "function"))
        assert_function_module_correct(folder_name, function_module)

        input_model_class = function_to_test.input_type
        output_data_model_class = get_type_hints(function_to_test.function)["return"]

        with open(test_file) as f:
            test_cases = json.load(f)
            assert_test_cases_correct(test_cases, folder_name)

            for test_case in test_cases:
                logger.info(f"Running message {test_case['message']}")
                func_output = function_to_test(input_model_class(**test_case["input"]))

                test_output = None
                try:
                    test_output = output_data_model_class(**test_case["output"]).dict()
                except pydantic.ValidationError as ve:
                    expected_fields = ", ".join(
                        get_type_hints(output_data_model_class).keys()
                    )
                    actual_fields = test_case["output"]
                    logger.error(
                        f"Validation error for function output {function_module.function.name}:\nExpected fields: {expected_fields}\nGot: {actual_fields}.\n{str(ve)}"
                    )
                    continue

                if func_output != test_output:
                    logger.error(
                        f"Function {folder_name} output mismatch for message '{test_case['message']}':"
                    )
                    logger.error(f"Expected: {test_output}")
                    logger.error(f"Got: {func_output}")
                    raise AssertionError("Expected output does not match actual output")

    except Exception as e:
        logger.error(f"Function {folder_name} failed. Error: {str(e)}")

sageai/tests/unit.py
0xnenlabs/SageAI
import re
from typing import List

from lasso_sennin.utils.address_mapping_utilities import test_address_mappings


def replace_with_index(string: str, pattern: str, replacement_list: List[str]):
    for match in re.finditer(pattern, string):
        index = int(match.group(1)) - 1
        string = string.replace(match.group(0), replacement_list[index])
    return string


def replace_without_index(string: str, pattern: str, replacement: str):
    return string.replace(pattern, replacement)


def replace_addresses_with_test_addresses(string: str):
    for key, value in test_address_mappings.items():
        pattern_base = re.escape(
            key[:-1],
        )
        string = replace_with_index(
            string,
            pattern_base + r"_(\d+)>",
            value,
        )
        string = replace_without_index(
            string,
            pattern_base + r">",
            value[0],
        )
    return string

sageai/tests/utils.py
0xnenlabs/SageAI
from abc import ABC, abstractmethod
from typing import Any, Dict, List


class AbstractVectorDB(ABC):
    def __init__(self):
        from sageai.config import get_function_map

        self.function_map = get_function_map()

    @abstractmethod
    def index(self) -> None:
        """Indexes the vector db based on the functions directory."""
        pass

    @abstractmethod
    def search(self, *, query: str, top_n: int) -> List[str]:
        """Actual search logic, which should be implemented in derived classes.
        It should return a list of function names
        """
        pass

    def format_search_result(
        self, *, function_names: List[str]
    ) -> List[Dict[str, Any]]:
        potential_functions = [
            self.function_map[func_name].parameters for func_name in function_names
        ]
        return potential_functions

    def search_impl(self, *, query: str, top_n: int) -> List[Dict[str, Any]]:
        """Search vector db based on a query and return top n function names."""
        results = self.search(query=query, top_n=top_n)
        return self.format_search_result(function_names=results)

sageai/types/abstract_vectordb.py
0xnenlabs/SageAI
from typing import Any, Callable, Dict, List, Type

from pydantic import BaseModel

from sageai.utils.inspection_utilities import get_input_parameter_type
from sageai.utils.model_utilities import (
    get_array_item_type,
    get_enum_array_values,
    get_possible_values,
    is_array,
    is_enum,
    is_enum_array,
    is_optional,
)


class Function(BaseModel):
    # required
    function: Callable
    description: str

    # generated from input
    name: str
    parameters: Dict[str, Any]
    input_type: Type[BaseModel]

    def __init__(
        self,
        function: Callable,
        description: str,
    ) -> None:
        name = function.__name__
        input_parameter_type = get_input_parameter_type(function)
        formatted_parameters = self._format_parameters(
            input_parameter_type=input_parameter_type,
            name=name,
            description=description,
        )

        super().__init__(
            function=function,
            description=description,
            name=name,
            parameters=formatted_parameters,
            input_type=input_parameter_type,
        )

    @staticmethod
    def _format_parameters(input_parameter_type, name: str, description: str):
        parameters = input_parameter_type.schema()
        new_schema = {
            "name": name,
            "description": description,
            "parameters": {
                "type": "object",
                "properties": {},
                "required": [],
            },
        }

        for key, value in parameters["properties"].items():
            if is_enum_array(input_parameter_type, key):
                enum_values = get_enum_array_values(input_parameter_type, key)
                enum_type = "integer" if isinstance(enum_values[0], int) else "string"
                property_schema = {
                    "type": "array",
                    "description": value.get("description", ""),
                    "items": {"type": enum_type, "enum": enum_values},
                }
            elif is_enum(input_parameter_type, key):
                property_schema = {
                    "type": "string",
                    "description": value["description"],
                    "enum": get_possible_values(input_parameter_type, key),
                }
            elif is_array(input_parameter_type, key):
                property_schema = {
                    "type": "array",
                    "description": value.get("description", ""),
                    "items": {
                        "type": get_array_item_type(input_parameter_type, key),
                    },
                }
            else:
                property_type = value.get("type")
                if property_type is None:
                    raise Exception("Property type is None")
                property_schema = {
                    "type": property_type,
                    "description": value.get("description", ""),
                }

            new_schema["parameters"]["properties"][key] = property_schema

            # Check if the property is required
            if not is_optional(input_parameter_type, key):
                new_schema["parameters"]["required"].append(key)

        return new_schema

    def __call__(self, *args, **kwargs):
        return self.function(*args, **kwargs)

    def __str__(self):
        return f"{self.name}: {self.description}"

sageai/types/function.py
0xnenlabs/SageAI
from enum import Enum


class LogLevel(str, Enum):
    DEBUG = "DEBUG"
    INFO = "INFO"
    WARNING = "WARNING"
    ERROR = "ERROR"
    CRITICAL = "CRITICAL"

sageai/types/log_level.py
0xnenlabs/SageAI
import os
from importlib import util
from types import ModuleType
from typing import List

from sageai.types.function import Function
from sageai.utils.logger import get_logger


def load_module_from_file(module_name: str, filepath: str) -> ModuleType:
    spec = util.spec_from_file_location(module_name, filepath)
    module = util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def get_functions_directories(
    logger,
    functions_directory_path: str = None,
) -> List[str]:
    logger.info(
        f"Getting functions directories from {functions_directory_path}",
    )

    function_directories = [
        dirpath
        for dirpath, dirnames, filenames in os.walk(functions_directory_path)
        if "function.py" in filenames
    ]
    logger.info(f"Found {len(function_directories)} function directories")
    return sorted(function_directories)


def generate_functions_map() -> dict[str, Function]:
    from sageai.config import get_config

    config = get_config()
    functions_directory_path = config.functions_directory
    log_level = config.log_level
    logger = get_logger("Utils", log_level)
    available_functions = {}

    logger.info("Generating function map")
    functions_directory = get_functions_directories(logger, functions_directory_path)

    for dirpath in functions_directory:
        folder_name = os.path.basename(dirpath)
        function_file = os.path.join(dirpath, "function.py")
        function_module = load_module_from_file(folder_name, function_file)

        if not hasattr(function_module, "function"):
            raise Exception(
                f"Function {folder_name} does not have a function attribute",
            )

        available_functions[function_module.function.name] = function_module.function

    if len(available_functions) == 0:
        raise Exception("No functions found")
    logger.info(f"Function map generated with {len(available_functions)} functions")
    available_functions = dict(sorted(available_functions.items()))

    return available_functions

sageai/utils/file_utilities.py
0xnenlabs/SageAI
import inspect
import json
import os


def get_user_package_path():
    caller_frame = inspect.stack()[4]
    caller_module = inspect.getmodule(caller_frame[0])

    if caller_module is None or not hasattr(caller_module, "__file__"):
        raise ValueError("Cannot determine the path of the caller's module")

    module_path = os.path.abspath(caller_module.__file__)
    return os.path.dirname(module_path)


def inject_env_vars(**kwargs):
    current_directory = os.path.dirname(os.path.abspath(__file__))
    env_file_path = os.path.join(current_directory, "../env.json")

    if not os.path.exists(env_file_path):
        return kwargs

    with open(env_file_path) as f:
        env_vars = json.load(f)

    if env_vars["OPENAI_KEY"]:
        kwargs["openai_key"] = env_vars["OPENAI_KEY"]

    return kwargs


def format_config_args(**kwargs):
    kwargs = inject_env_vars(**kwargs)

    if "functions_directory" not in kwargs:
        kwargs["functions_directory"] = "functions"
    if not os.path.isabs(kwargs["functions_directory"]) and not kwargs[
        "functions_directory"
    ].startswith("./"):
        base_path = get_user_package_path()
        kwargs["functions_directory"] = os.path.join(
            base_path, kwargs["functions_directory"]
        )
    elif kwargs["functions_directory"].startswith("./"):
        base_path = get_user_package_path()
        # Strip './' from the beginning and then join with the base path
        relative_path = kwargs["functions_directory"][2:]
        kwargs["functions_directory"] = os.path.join(base_path, relative_path)
    return kwargs

sageai/utils/format_config_args.py
0xnenlabs/SageAI
import inspect
from typing import Callable, Type

from pydantic import BaseModel


def get_input_parameter_type(function: Callable) -> Type[BaseModel]:
    signature = inspect.signature(function)
    for name, param in signature.parameters.items():
        if "Input" in str(param.annotation):
            return param.annotation
    raise Exception(f"Could not find Input parameter type from function arguments.")

sageai/utils/inspection_utilities.py
0xnenlabs/SageAI
import logging
from sageai.types.log_level import LogLevel


def get_logger(logger_name: str, log_level: LogLevel):
    logger = logging.getLogger(logger_name)

    log_level_value = log_level.value
    logger.setLevel(log_level_value)

    formatter = logging.Formatter("%(asctime)s | [%(levelname)s] %(message)s")

    if not logger.handlers:
        console_handler = logging.StreamHandler()
        console_handler.setLevel(log_level_value)
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)
    else:
        logger.handlers[0].setLevel(log_level_value)  # assumes only one handler

    return logger

sageai/utils/logger.py
0xnenlabs/SageAI
from enum import Enum
from types import ModuleType
from typing import Any, List, Optional, Type, Union, get_args, get_origin

from pydantic import BaseModel


def find_pydantic_model(
    module: ModuleType,
    keyword: str,
) -> Type[BaseModel] | None:
    models = [
        cls
        for name, cls in vars(module).items()
        if isinstance(cls, type)
        and issubclass(cls, BaseModel)
        and name.endswith(
            keyword,
        )
    ]
    return models[0] if models else None


def get_input_model(module: ModuleType) -> Type[BaseModel]:
    return find_pydantic_model(module, keyword="Input")


def get_output_data_model(module: ModuleType) -> Type[BaseModel]:
    return find_pydantic_model(module, keyword="Data")


def is_enum(model: Type[BaseModel], key: str) -> bool:
    if key in model.__annotations__:
        annotation = model.__annotations__[key]
        # Check if annotation is a class
        if isinstance(annotation, type):
            return issubclass(annotation, Enum)
        # Handle Optional[Enum] and List[Enum]
        origin = get_origin(annotation)
        if origin in {Optional, List, Union}:
            args = get_args(annotation)
            if args and isinstance(args[0], type):
                return issubclass(args[0], Enum)
    return False


def get_possible_values(model: Type[BaseModel], key: str) -> List[str] | List[int]:
    if is_enum(model, key) or is_enum_array(model, key):
        enum_type = model.__annotations__[key]
        if get_origin(enum_type) in {Optional, Union, List}:
            enum_type = get_args(enum_type)[0]
        return [e.value for e in enum_type]
    return []


def is_optional(model: Type[BaseModel], key: str) -> bool:
    annotation = model.__annotations__[key]
    return type(None) in get_args(annotation)


def is_array(model: Type[BaseModel], key: str) -> bool:
    if key in model.__annotations__:
        annotation = model.__annotations__[key]
        if get_origin(annotation) is list:
            return True

        if get_origin(annotation) in {Optional, Union}:
            args = get_args(annotation)
            return any(get_origin(arg) is list for arg in args)
    return False


def get_array_item_type(model: Type[BaseModel], key: str) -> str | None:
    type_map = {"str": "string", "int": "integer"}
    if not is_array(model, key):
        raise Exception(f"{key} is not an array")
    field = model.__annotations__[key]
    if get_origin(field) is list:
        return type_map[get_args(field)[0].__name__]
    elif get_origin(field) in {Optional, Union}:
        args = get_args(field)
        for arg in args:
            if get_origin(arg) is list:
                return type_map[get_args(arg)[0].__name__]
    return None


def is_enum_array(model: Type[BaseModel], key: str) -> bool:
    enum_type = get_array_item_enum_type(model, key)
    return issubclass(enum_type, Enum) if enum_type is not None else False


def get_enum_type_from_annotation(field) -> Type[Enum] | None:
    enum_type = None
    if get_origin(field) is list:
        enum_type = get_args(field)[0]
    elif get_origin(field) in {Optional, Union}:
        args = get_args(field)
        for arg in args:
            if get_origin(arg) is list:
                enum_type = get_args(arg)[0]
    if enum_type and issubclass(enum_type, Enum):
        return enum_type
    return None


def get_array_item_enum_type(model: Type[BaseModel], key: str) -> Type | None:
    if not is_array(model, key):
        return None
    field = model.__annotations__[key]
    return get_enum_type_from_annotation(field)


def get_enum_array_values(
    model: Type[BaseModel],
    key: str,
) -> List[str] | List[int]:
    if not is_enum_array(model, key):
        raise Exception(f"{key} is not an array of Enums")
    field = model.__annotations__[key]
    enum_type = get_enum_type_from_annotation(field)
    if enum_type:
        return [e.value for e in enum_type.__members__.values()]
    return []


def get_default_value(model: Type[BaseModel], key: str) -> Any:
    if is_enum(model, key):
        default_value = model.__fields__[key].default
        if is_enum(model, key) and isinstance(default_value, Enum):
            return default_value.value
        return default_value

    if key in model.__fields__:
        default_value = model.__fields__[key].default
        if default_value == "None" or default_value is None:
            return None
        return default_value

    return None

sageai/utils/model_utilities.py
0xnenlabs/SageAI
from typing import Dict, List


def get_latest_user_message(messages: List[Dict]) -> Dict | None:
    for message in reversed(messages):
        if message.get("role") == "user":
            return message
    return None

sageai/utils/openai_utilities.py
0xnenlabs/SageAI