chibi-bot 1.6.0b0__py3-none-any.whl

chibi/services/providers/tools/tool.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Coroutine, ParamSpec, TypeVar, cast
+
+from loguru import logger
+from openai.types.chat import ChatCompletionToolParam
+from telegram import Update
+from telegram.ext import ContextTypes
+
+from chibi.config import gpt_settings
+from chibi.services.providers.tools.schemas import ToolResponse
+from chibi.services.providers.tools.utils import AdditionalOptions
+from chibi.services.providers.utils import escape_and_truncate
+from chibi.services.task_manager import task_manager
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+ToolFunction = Callable[P, Coroutine[Any, Any, ToolResponse]]
+RegisteredFunctionsMap = dict[str, ToolFunction]
+ToolsDefinitionMap = dict[str, ChatCompletionToolParam]
+
+
+class ChibiTool:
+    register: bool
+    definition: ChatCompletionToolParam
+    name: str
+    run_in_background_by_default: bool = False
+    allow_model_to_change_background_mode: bool = True
+
+    @classmethod
+    def add_global_params(cls) -> dict[str, Any]:
+        return {
+            "run_in_background": {
+                "type": "boolean",
+                "description": "Execute task in background. You'll receive a result when it done.",
+                "default": cls.run_in_background_by_default,
+            }
+        }
+
+    @classmethod
+    async def _get_tool_call_result(cls, *args, **kwargs) -> ToolResponse:
+        try:
+            result = await cls.function(**kwargs)
+            logger.log(
+                "CALL",
+                (
+                    f"[{kwargs.get('model', 'Unknown model')}] Function '{cls.name}' called, "
+                    f"result retrieved: {escape_and_truncate(result)}"
+                ),
+            )
+            return ToolResponse(tool_name=cls.name, status="ok", result=result)
+        except Exception as e:
+            logger.warning(f"[{kwargs.get('model', 'Unknown model')}] Tool {cls.name} raised an exception: {e}")
+            return ToolResponse(tool_name=cls.name, status="error", result=str(e))
+
+    @classmethod
+    async def _get_and_send_tool_call_result(
+        cls, *args, update: Update, context: ContextTypes.DEFAULT_TYPE, **kwargs
+    ) -> None:
+        from chibi.services.bot import handle_tool_response
+
+        tool_call_result = await cls._get_tool_call_result(*args, **kwargs)
+        await handle_tool_response(tool_response=tool_call_result, update=update, context=context)
+
+    @classmethod
+    async def tool(cls, *args, run_in_background: bool | None = None, **kwargs: Any) -> ToolResponse:
+        non_printable_kwargs = list(AdditionalOptions.__annotations__.keys())
+        printable_kwargs = {k: v for k, v in kwargs.items() if k not in non_printable_kwargs}
+        telegram_context = kwargs.get("telegram_context")
+        telegram_update = kwargs.get("telegram_update")
+        if run_in_background is None:
+            run_in_background = cls.run_in_background_by_default
+        background_run_ready = run_in_background and telegram_update and telegram_context
+
+        logger.log(
+            "CALL",
+            (
+                f"[{kwargs.get('model', 'Unknown model')}] Calling a function '{cls.name}' "
+                f"in {'background' if background_run_ready else 'foreground'} mode. "
+                f"Args: {escape_and_truncate(printable_kwargs)}"
+            ),
+        )
+        if not background_run_ready:
+            return await cls._get_tool_call_result(*args, **kwargs)
+
+        assert telegram_update
+        assert telegram_context
+
+        coro = cls._get_and_send_tool_call_result(*args, update=telegram_update, context=telegram_context, **kwargs)
+        task_manager.run_task(coro)
+        return ToolResponse(
+            tool_name=cls.name,
+            status="tool is running in background",
+            result="in progress",
+            additional_details="you'll receive the result when it ready",
+        )
+
+    @classmethod
+    async def function(cls, *args: Any, **kwargs: Any) -> dict[str, Any]:
+        raise NotImplementedError
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+
+        if not cls.register:
+            return None
+
+        if gpt_settings.tools_whitelist and cls.name not in gpt_settings.tools_whitelist:
+            return None
+
+        if cls.allow_model_to_change_background_mode:
+            cast(dict, cls.definition["function"]["parameters"]["properties"]).update(cls.add_global_params())
+
+        RegisteredChibiTools.register(cls)
+
+
+class RegisteredChibiTools:
+    tools_map: dict[str, type[ChibiTool]] = {}
+
+    @classmethod
+    def get_tool_definitions(cls) -> list[ChatCompletionToolParam]:
+        return [tool.definition for tool in cls.tools_map.values()]
+
+    @classmethod
+    def get_registered_functions(cls) -> RegisteredFunctionsMap:
+        registered_functions = {name: tool.tool for name, tool in cls.tools_map.items()}
+        registered_functions["stub_function"] = cls._stub_function
+        return registered_functions
+
+    @classmethod
+    def register(cls, tool: type[ChibiTool]) -> None:
+        cls.tools_map[tool.name] = tool
+
+    @classmethod
+    def deregister_tools(cls, tool_names: list[str]) -> None:
+        for tool_name in tool_names:
+            if tool_name not in cls.tools_map:
+                continue
+            cls.tools_map.pop(tool_name)
+            logger.info(f"The tool {tool_name} had been deregistered.")
+
+    @classmethod
+    def get(cls, tool_name: str) -> ToolFunction:
+        if chibi_tool_class := cls.tools_map.get(tool_name):
+            return chibi_tool_class.tool
+        logger.error(f"Function {tool_name} called but it's not registered.")
+        return cls._stub_function
+
+    @classmethod
+    async def call(cls, tool_name: str, tools_args: dict[str, Any]) -> ToolResponse:
+        tool = cls.get(tool_name)
+        return await tool(**tools_args)
+
+    @classmethod
+    async def _stub_function(cls, *args: Any, **kwargs: Any) -> ToolResponse:
+        """A stub function that is executed when the LLM calls a non-existent function.
+
+        Returns:
+            A ToolResponse object describing the error.
+        """
+        logger.log("TOOL", f"Running stub function. Args: {args}, kwargs: {kwargs}")
+        return ToolResponse(tool_name="stub", status="error", result="A non-existent function called")

chibi/services/providers/tools/utils.py

@@ -0,0 +1,146 @@
+import datetime
+import json
+import urllib.parse
+from typing import TYPE_CHECKING, ParamSpec, TypedDict, TypeVar
+
+import httpx
+from fake_useragent import UserAgent
+from httpx import Response
+from loguru import logger
+from telegram import Update
+from telegram.ext import ContextTypes
+
+from chibi.config import gpt_settings
+from chibi.constants import SUB_EXECUTOR_PROMPT
+from chibi.models import Message
+from chibi.schemas.app import ChatResponseSchema
+from chibi.storage.abstract import Database
+from chibi.storage.database import inject_database
+
+if TYPE_CHECKING:
+    from chibi.services.providers.provider import Provider
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+ua_generator = UserAgent()
+
+
+class AdditionalOptions(TypedDict, total=False):
+    user_id: int | None
+    model: str | None
+    telegram_context: ContextTypes.DEFAULT_TYPE | None
+    telegram_update: Update | None
+
+
+def _generate_google_search_referrer(target_url: str) -> str:
+    """Generates a fake Google search referrer URL for a given target URL.
+
+    This helps simulate traffic coming from a Google search result link,
+    which can sometimes affect how websites serve content.
+
+    Args:
+        target_url: The URL that the fake referrer should point to.
+
+    Returns:
+        A string representing the generated Google referrer URL.
+    """
+    encoded_target_url = urllib.parse.quote(target_url, safe="")
+
+    fake_ved = "2ahUKEwj_0sL5yPaFAxW_FRAIHeYxBpUQwgF6BAgGEAA"
+    fake_opi = "89974493"
+
+    referrer = (
+        f"https://www.google.com/url?sa=t&rct=j&q={encoded_target_url}&esrc=s&source=web&"
+        f"cd=1&cad=rja&uact=8&ved={fake_ved}&url={encoded_target_url}&opi={fake_opi}"
+    )
+
+    return referrer
+
+
+async def _get_url(url: str) -> Response:
+    """Fetch content from a given URL.
+
+    It uses configured proxy, retries, and timeout settings from gpt_settings,
+    and includes various headers including a generated Google referrer and a
+    random User-Agent to mimic a real browser request.
+
+    Args:
+        url: The URL to fetch content from.
+
+    Returns:
+        An httpx.Response object containing the response from the URL.
+
+    Raises:
+        Httpx exceptions if the request fails (e.g., network errors).
+    """
+    transport = httpx.AsyncHTTPTransport(retries=gpt_settings.retries, proxy=gpt_settings.proxy)
+    headers: dict[str, str] = {
+        "User-Agent": ua_generator.random,
+        "Referer": _generate_google_search_referrer(target_url=url),
+        "Accept": (
+            "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,"
+            "image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7"
+        ),
+        "Accept-Encoding": "gzip, deflate, br",
+        "Accept-Language": "en-US,en;q=0.8",
+        "Connection": "keep-alive",
+        "Upgrade-Insecure-Requests": "1",
+        "Sec-Fetch-Site": "none",
+        "Sec-Fetch-Dest": "document",
+        "Sec-Fetch-Mode": "navigate",
+        "Sec-Fetch-User": "?1",
+    }
+    async with httpx.AsyncClient(transport=transport, timeout=gpt_settings.timeout, proxy=gpt_settings.proxy) as client:
+        return await client.get(url=url, headers=headers)
+
+
+@inject_database
+async def get_sub_agent_response(
+    db: Database,
+    user_id: int,
+    prompt: str,
+    model_name: str | None = None,
+    provider_name: str | None = None,
+) -> ChatResponseSchema:
+    user = await db.get_or_create_user(user_id=user_id)
+    provider: Provider | None
+    if not model_name or not provider_name:
+        provider = user.active_gpt_provider
+        model = user.selected_gpt_model_name
+    else:
+        provider = user.providers.get(provider_name=provider_name)
+        model = model_name
+
+    if not provider:
+        raise ValueError(f"No provider found. Provided provider name: {provider_name}")
+
+    user_prompt = {
+        "user_type": "llm",
+        "current_working_dir": user.working_dir,
+        "prompt": prompt,
+        "datetime_now": datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S %Z%z"),
+    }
+
+    user_message = Message(role="user", content=json.dumps(user_prompt))
+    conversation_messages = [
+        user_message,
+    ]
+
+    chat_response, _ = await provider.get_chat_response(
+        messages=conversation_messages, user=user, model=model, system_prompt=SUB_EXECUTOR_PROMPT
+    )
+    return chat_response
+
+
+async def download(url: str) -> bytes | None:
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url, timeout=90.0)  # TODO: move timeout to settings or use one of existent
+            response.raise_for_status()
+            data = response.content
+            logger.log("TOOL", f"Downloaded data from URL {url}: {len(data)} bytes")
+            return data
+    except Exception as e:
+        logger.error(f"Failed to download file from {url}: {e}")
+    return None

chibi/services/providers/tools/web.py

@@ -0,0 +1,261 @@
+from typing import Any, Unpack
+
+import httpx
+from ddgs import DDGS
+from httpx import Response
+from loguru import logger
+from openai.types.chat import ChatCompletionToolParam
+from openai.types.shared_params import FunctionDefinition
+from trafilatura import extract
+
+from chibi.config import gpt_settings
+from chibi.services.providers.tools.exceptions import ToolException
+from chibi.services.providers.tools.tool import ChibiTool
+from chibi.services.providers.tools.utils import AdditionalOptions, _get_url
+
+
+class SearchNewsTool(ChibiTool):
+    register = True
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="search_news",
+            description="Searches for current news articles based on the given search query at duckduckgo.com",
+            parameters={
+                "type": "object",
+                "properties": {
+                    "search_phrase": {
+                        "type": "string",
+                        "description": "The text of the search query for news searching.",
+                    },
+                    "max_results": {
+                        "type": "integer",
+                        "description": "The maximum number of news articles to return (default is 10).",
+                    },
+                },
+                "required": ["search_phrase"],
+            },
+        ),
+    )
+    name = "search_news"
+
+    @classmethod
+    async def function(
+        cls, search_phrase: str, max_results: int = 10, **kwargs: Unpack[AdditionalOptions]
+    ) -> dict[str, Any]:
+        """Search for news articles using DuckDuckGo News.
+
+        Args:
+            search_phrase: The keywords or phrase to search for in news.
+            max_results: The maximum number of news results to return (default is 10).
+
+        Returns:
+            A JSON formatted string containing the list of news articles found,
+            or an error message string if the search fails.
+        """
+        logger.log(
+            "TOOL",
+            f"[{kwargs.get('model', 'Unknown model')}] Searching news for '{search_phrase}', max_results={max_results}",
+        )
+        try:
+            result = DDGS(proxy=gpt_settings.proxy).news(query=search_phrase, max_results=max_results, region="wt-wt")
+        except Exception as e:
+            raise ToolException(f"Couldn't find news for '{search_phrase}', max_results={max_results}. Error: {e}")
+        return {
+            "news": result,
+        }
+
+
+class DDGSWebSearchTool(ChibiTool):
+    register = True
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="ddgs_web_search",
+            description=(
+                "Search for information on the internet using the DDGS python library. "
+                "Use this function if other web search functions are unavailable or not working."
+            ),
+            parameters={
+                "type": "object",
+                "properties": {
+                    "search_phrase": {
+                        "type": "string",
+                        "description": "The text of the search query for web searching.",
+                    },
+                    "max_results": {
+                        "type": "integer",
+                        "description": "The maximum number of web search results to return (default is 10).",
+                    },
+                },
+                "required": ["search_phrase"],
+            },
+        ),
+    )
+    name = "ddgs_web_search"
+
+    @classmethod
+    async def function(
+        cls, search_phrase: str, max_results: int = 10, **kwargs: Unpack[AdditionalOptions]
+    ) -> dict[str, Any]:
+        """Perform a general web search using DDGS python library.
+
+        Args:
+            search_phrase: The keywords or phrase to search for on the web.
+            max_results: The maximum number of search results to return (default is 10).
+
+        Returns:
+            A JSON formatted string containing the list of search results found,
+            or an error message string if the search fails.
+        """
+        logger.log(
+            "TOOL",
+            (
+                f"[{kwargs.get('model', 'Unknown model')}] Using web-search for '{search_phrase}', "
+                f"max_results={max_results}"
+            ),
+        )
+        try:
+            result = DDGS(proxy=gpt_settings.proxy).text(query=search_phrase, max_results=max_results, region="wt-wt")
+        except Exception as e:
+            raise ToolException(
+                f"Couldn't get search result for '{search_phrase}', max_results={max_results}. Error: {e}"
+            )
+
+        return {
+            "search_results": result,
+        }
+
+
+class GoogleSearchTool(ChibiTool):
+    register = gpt_settings.google_search_client_set
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="google_web_search",
+            description=("Search for information on the internet via Google Web Search API."),
+            parameters={
+                "type": "object",
+                "properties": {
+                    "search_phrase": {
+                        "type": "string",
+                        "description": "The text of the search query for web searching.",
+                    },
+                },
+                "required": ["search_phrase"],
+            },
+        ),
+    )
+    name = "google_web_search"
+
+    @classmethod
+    async def function(cls, search_phrase: str, **kwargs: Unpack[AdditionalOptions]) -> dict[str, Any]:
+        """Perform a general web search using Google Web Search.
+
+        TODO: upgrade to using `max_results` arg.
+
+        Args:
+            search_phrase: The keywords or phrase to search for on the web.
+
+        Returns:
+            A JSON formatted string containing the list of search results found,
+            or an error message string if the search fails.
+        """
+        logger.log("TOOL", f"[{kwargs.get('model', 'Unknown model')}] Using Google web-search for '{search_phrase}'")
+        transport = httpx.AsyncHTTPTransport(retries=gpt_settings.retries, proxy=gpt_settings.proxy)
+        params = {
+            "key": gpt_settings.google_search_api_key,
+            "cx": gpt_settings.google_search_cx,
+            "q": search_phrase,
+        }
+        url = "https://www.googleapis.com/customsearch/v1"
+        try:
+            async with httpx.AsyncClient(
+                transport=transport,
+                timeout=gpt_settings.timeout,
+                proxy=gpt_settings.proxy,
+            ) as client:
+                response = await client.get(
+                    url=url,
+                    params=params,
+                )
+                response.raise_for_status()
+        except Exception as e:
+            raise ToolException(f"An error occurred while calling the Google Search API: {e}")
+
+        data = response.json()
+        items = data.get("items")
+        if not items:
+            logger.warning(f"{cls.name} tool returned an empty list of results. Search phrase: {search_phrase}.")
+            return {"search_results": "Ooops, the search returned an empty list of results."}
+
+        target_keys = ["title", "link", "snippet"]
+        search_results = [{key: item.get(key) for key in target_keys} for item in items]
+        return {
+            "search_results": search_results,
+        }
+
+
+class ReadWebPageTool(ChibiTool):
+    register = True
+    definition = ChatCompletionToolParam(
+        type="function",
+        function=FunctionDefinition(
+            name="read_web_page",
+            description=(
+                "Read the content of the web page. Be prepared that trafilatura may not cope and "
+                "will not be able to retrieve information either due to captcha or because of js."
+            ),
+            parameters={
+                "type": "object",
+                "properties": {
+                    "url": {"type": "string", "description": "Web page URL to fetch."},
+                },
+                "required": ["url"],
+            },
+        ),
+    )
+    name = "read_web_page"
+
+    @classmethod
+    async def function(cls, url: str, **kwargs: Unpack[AdditionalOptions]) -> dict[str, Any]:
+        """Fetch and extract the main content from a given web page URL.
+
+        Args:
+            url: The URL of the web page to read.
+
+        Returns:
+            A string containing the extracted main text of the page, the raw HTML
+            content if extraction fails, or an error message string if fetching
+            fails or status code is not 200.
+        """
+        logger.log("TOOL", f"[{kwargs.get('model', 'Unknown model')}] Reading URL: {url}")
+        try:
+            response: Response = await _get_url(url)
+        except Exception as e:
+            raise ToolException(f"Couldn't read URL: {url}. Error: {e}")
+
+        if response.status_code != 200:
+            raise ToolException(f"Failed to get URL: {url}. Status code: {response.status_code}")
+
+        data = response.text
+        if not data:
+            raise ToolException(f"Failed to extract data from URL: {url}. Empty response received.")
+
+        content = extract(filecontent=data, include_links=True)
+        if not content:
+            msg = f"Failed to extract URL: {url}. Empty extracted data. Trying to send raw HTML to model"
+            logger.warning(f"[{kwargs.get('model', 'Unknown model')}] {msg}")
+            return {
+                "data": data,
+                "warning": msg,
+            }
+
+        logger.log(
+            "TOOL",
+            f"[{kwargs.get('model', 'Unknown model')}] The data from the URL {url} seems to be successfully extracted",
+        )
+
+        return {
+            "content": content,
+        }