rusticai-litellm 0.0.1__tar.gz

rusticai_litellm-0.0.1/PKG-INFO

@@ -0,0 +1,38 @@
+Metadata-Version: 2.1
+Name: rusticai-litellm
+Version: 0.0.1
+Summary: Rustic AI module leveraging Litellm for LLMs
+Home-page: https://www.rustic.ai/
+License: Apache-2.0
+Author: Dragonscale Industries Inc.
+Author-email: dev@dragonscale.ai
+Requires-Python: >=3.12,<4.0
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: dotenv (>=0.9.9,<0.10.0)
+Requires-Dist: litellm (>=1.65.4.post1,<2.0.0)
+Requires-Dist: openai (>=1.71.0,<2.0.0)
+Requires-Dist: rusticai-core (>=0.0.4,<0.1.0)
+Project-URL: Repository, https://github.com/rustic-ai/python-framework
+Project-URL: Rustic AI Core, https://pypi.org/project/rusticai-core/
+Description-Content-Type: text/markdown
+
+# Rustic AI LiteLLM
+
+This module offers seamless integration with various LLMs through [LiteLLM](https://docs.litellm.ai/docs/), simplifying the use of language models within [Rustic AI](https://www.rustic.ai/) guilds.
+
+## Installing
+
+```shell
+pip install rusticai-litellm
+```
+**Note:** It depends on [rusticai-core](https://pypi.org/project/rusticai-core/)
+
+## Building from Source
+
+```shell
+poetry install --with dev
+poetry build
+```

rusticai_litellm-0.0.1/README.md

@@ -0,0 +1,17 @@
+# Rustic AI LiteLLM
+
+This module offers seamless integration with various LLMs through [LiteLLM](https://docs.litellm.ai/docs/), simplifying the use of language models within [Rustic AI](https://www.rustic.ai/) guilds.
+
+## Installing
+
+```shell
+pip install rusticai-litellm
+```
+**Note:** It depends on [rusticai-core](https://pypi.org/project/rusticai-core/)
+
+## Building from Source
+
+```shell
+poetry install --with dev
+poetry build
+```

rusticai_litellm-0.0.1/pyproject.toml

@@ -0,0 +1,66 @@
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "rusticai-litellm"
+version = "0.0.1"
+description = "Rustic AI module leveraging Litellm for LLMs"
+authors = ["Dragonscale Industries Inc. <dev@dragonscale.ai>"]
+license = "Apache-2.0"
+readme = "README.md"
+homepage = "https://www.rustic.ai/"
+repository = "https://github.com/rustic-ai/python-framework"
+packages = [{ include = "rustic_ai", from = "src" }]
+
+[tool.poetry.urls]
+"Rustic AI Core" = "https://pypi.org/project/rusticai-core/"
+
+[tool.poetry.dependencies]
+python = "^3.12"
+rusticai-core = {   version = "0.0.4"}
+litellm = "^1.65.4.post1"
+dotenv = "^0.9.9"
+openai = "^1.71.0"
+
+[tool.poetry-monorepo.deps]
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.3.4"
+black = { version = "^24.2.0", extras = ["jupyter"] }
+flake8 = "^7.1.2"
+tox = "^4.24.1"
+isort = "^6.0.0"
+mypy = "^1.15.0"
+
+
+[tool.black]
+line-length = 120
+target-version = ['py312']
+include = '\.pyi?$'
+exclude = '''
+/(
+    \.git
+  | \.mypy_cache
+  | \.tox
+  | \.venv
+  | _build
+  | buck-out
+  | build
+  | dist
+)/
+'''
+
+[tool.mypy]
+python_version = "3.12"
+ignore_missing_imports = true
+check_untyped_defs = true
+plugins = "pydantic.mypy"
+explicit_package_bases = true
+
+
+[tool.isort]
+profile = "black"
+known_third_party = ["pydantic"]
+known_first_party = ["rustic_ai.core", "rustic_ai.litellm"]
+

rusticai_litellm-0.0.1/src/rustic_ai/litellm/__init__.py

@@ -0,0 +1,4 @@
+from .agent import LiteLLMAgent
+from .conf import LiteLLMConf
+
+__all__ = ["LiteLLMAgent", "LiteLLMConf"]

rusticai_litellm-0.0.1/src/rustic_ai/litellm/agent.py

@@ -0,0 +1,127 @@
+import logging
+from collections import deque
+from typing import Deque, List, Union
+
+import openai
+from dotenv import load_dotenv
+
+import litellm
+from litellm import validate_environment
+from rustic_ai.core.guild.agent import (
+    Agent,
+    AgentMode,
+    AgentType,
+    ProcessContext,
+    processor,
+)
+from rustic_ai.core.guild.agent_ext.depends.llm.models import (
+    AssistantMessage,
+    ChatCompletionError,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionTool,
+    FunctionMessage,
+    ResponseCodes,
+    SystemMessage,
+    ToolMessage,
+    UserMessage,
+)
+from rustic_ai.core.guild.dsl import AgentSpec
+
+from .conf import LiteLLMConf
+
+load_dotenv()
+
+
+class LiteLLMAgent(Agent[LiteLLMConf]):
+    """
+    LiteLLM agent for invoking various LLM models using LiteLLM.
+    """
+
+    litellm.drop_params = True
+
+    def __init__(self, agent_spec: AgentSpec[LiteLLMConf]):
+        validation_result = validate_environment(agent_spec.props.model)
+        if not validation_result.get("keys_in_environment"):
+            missing_keys = validation_result.get("missing_keys", [])
+            raise RuntimeError(f"Required environment variable(s) {'{}'.join(missing_keys)} not set")
+
+        super().__init__(agent_spec, agent_type=AgentType.BOT, agent_mode=AgentMode.LOCAL)
+        self.pre_messages = agent_spec.props.messages
+        self.pre_tools = agent_spec.props.tools
+        self.model = agent_spec.props.model
+        self.client_props = agent_spec.props.model_dump(
+            mode="json", exclude_unset=True, exclude_none=True, exclude={"message_memory"}
+        )
+        self.message_memory_size = agent_spec.props.message_memory or 0
+        self.message_queue: Deque[SystemMessage | UserMessage | AssistantMessage | ToolMessage | FunctionMessage] = (
+            deque(maxlen=self.message_memory_size)
+        )
+
+    @processor(ChatCompletionRequest)
+    def llm_completion(self, ctx: ProcessContext[ChatCompletionRequest]):
+
+        prompt = ctx.payload
+
+        messages = self.pre_messages if self.pre_messages else []
+
+        all_messages = messages + list(self.message_queue) + prompt.messages
+
+        messages_dict = [m.model_dump(exclude_none=True) for m in all_messages]
+
+        tools: List[ChatCompletionTool] = self.pre_tools if self.pre_tools else []
+        if prompt.tools:
+            tools.extend(prompt.tools)
+
+        full_prompt = {
+            **self.client_props,
+            **prompt.model_dump(exclude_unset=True, exclude_none=True),
+            "messages": messages_dict,
+        }
+
+        if tools:
+            full_prompt["tools"] = tools
+
+        try:
+            completion = litellm.completion(**full_prompt)
+
+            response = ChatCompletionResponse.from_litellm_completion(completion)
+
+            self.message_queue.extend(prompt.messages)  # Append the prompt messages (from user) to the message queue
+            if response.choices and response.choices[0].message:
+                self.message_queue.append(
+                    response.choices[0].message
+                )  # Append the response message (from LLM) to the message queue
+
+            if response:
+                ctx.send(response)
+
+        except openai.APIStatusError as e:  # pragma: no cover
+            logging.error(f"Error in LLM completion: {e}")
+            # Publish the error message
+            self.process_api_status_error(
+                ctx=ctx,
+                status_code=ResponseCodes(e.status_code),
+                error=e,
+                messages=messages,
+            )
+
+    def process_api_status_error(
+        self,
+        ctx: ProcessContext[ChatCompletionRequest],
+        status_code: ResponseCodes,
+        error: openai.APIStatusError,
+        messages: List[Union[SystemMessage, UserMessage, AssistantMessage, ToolMessage, FunctionMessage]],
+    ):  # pragma: no cover
+        """
+        Process API status error and return a ChatCompletionError object.
+        """
+        ctx.send_error(
+            ChatCompletionError(
+                status_code=status_code,
+                message=error.message,
+                response=error.response.text if error.response else None,
+                model=self.model,
+                request_messages=messages,
+            )
+        )  # pragma: no cover

rusticai_litellm-0.0.1/src/rustic_ai/litellm/conf.py

@@ -0,0 +1,81 @@
+from typing import Annotated, List, Optional, Union
+
+from pydantic import ConfigDict, Field
+
+from rustic_ai.core.guild.agent_ext.depends.llm.models import (
+    AssistantMessage,
+    ChatCompletionTool,
+    FunctionMessage,
+    Models,
+    SystemMessage,
+    ToolMessage,
+    UserMessage,
+)
+from rustic_ai.core.guild.dsl import BaseAgentProps
+
+
+class LiteLLMConf(BaseAgentProps):
+    """
+    Properties to apply to all LiteLLM calls made by the agent.
+    """
+
+    model_config = ConfigDict(extra="ignore")
+
+    model: Annotated[Union[str, Models], Field(examples=["gpt-4-turbo"])]
+    """
+    ID of the model to use. See the [model endpoint compatibility](/docs/models/model-endpoint-compatibility) table
+    for details on which models work with the Chat API.
+    """
+    stream: Optional[bool] = False
+    """
+    If set, partial message deltas will be sent, like in ChatGPT. Tokens will be sent as data-only
+    [server-sent events](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#Event_stream_format)
+    as they become available, with the stream terminated by a `data: [DONE]` message.
+    [Example Python code](https://cookbook.openai.com/examples/how_to_stream_completions).
+
+    """
+    # set api_base, api_version, api_key
+    base_url: Optional[str] = None
+    api_version: Optional[str] = None
+
+    """
+    If provided, return a mock completion response for testing or debugging purposes (default is None).
+    """
+
+    custom_llm_provider: Optional[str] = None
+    """
+    Used for Non-OpenAI LLMs, Example usage for bedrock, set model="amazon.titan-tg1-large" and custom_llm_provider="bedrock"
+    """
+
+    timeout: Optional[float] = None
+    """
+    The maximum time in seconds to wait for the completion.
+    If the completion is not returned within this time, the request will time out.
+    """
+
+    messages: Optional[
+        List[
+            Union[
+                SystemMessage,
+                UserMessage,
+                AssistantMessage,
+                ToolMessage,
+                FunctionMessage,
+            ]
+        ]
+    ] = None
+    """
+    List of messages to send to the model before the prompt.
+    These messages will be prependended to the list of messages in the prompt.
+    """
+
+    tools: Optional[List[ChatCompletionTool]] = None
+    """
+    List of tools to enable for completion.
+    These tools will be appended to the list of tools in the prompt.
+    """
+
+    message_memory: Optional[int] = None
+    """
+    Number of messages to remember in the conversation history.
+    """

rusticai_litellm-0.0.1/src/rustic_ai/litellm/dependency.py

@@ -0,0 +1,78 @@
+from typing import List
+
+import litellm
+from rustic_ai.core.guild.agent_ext.depends.dependency_resolver import (
+    DependencyResolver,
+)
+from rustic_ai.core.guild.agent_ext.depends.llm import LLM
+from rustic_ai.core.guild.agent_ext.depends.llm.models import (
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionTool,
+)
+
+from .conf import LiteLLMConf
+
+
+class LiteLLM(LLM):
+    def __init__(self, props: LiteLLMConf):
+        self.preset_messages = props.messages
+        self.preset_tools = props.tools
+        self._model = props.model
+        self.client_props = props.model_dump(mode="json", exclude_unset=True, exclude_none=True)
+
+    def _prep_prompt(self, prompt: ChatCompletionRequest) -> dict:
+        messages = self.preset_messages if self.preset_messages else []
+
+        all_messages = messages + prompt.messages
+
+        messages_dict = [m.model_dump(exclude_none=True) for m in all_messages]
+
+        tools: List[ChatCompletionTool] = self.preset_tools if self.preset_tools else []
+        if prompt.tools:
+            tools.extend(prompt.tools)
+
+        full_prompt = {
+            **self.client_props,
+            **prompt.model_dump(exclude_unset=True, exclude_none=True),
+            "messages": messages_dict,
+        }
+
+        if tools:
+            full_prompt["tools"] = tools
+
+        return full_prompt
+
+    def completion(self, prompt: ChatCompletionRequest):
+        full_prompt = self._prep_prompt(prompt)
+
+        completion = litellm.completion(**full_prompt)
+        response = ChatCompletionResponse.from_litellm_completion(completion)
+        return response
+
+    async def async_completion(self, prompt: ChatCompletionRequest):
+        full_prompt = self._prep_prompt(prompt)
+
+        completion = await litellm.acompletion(**full_prompt)
+        response = ChatCompletionResponse.from_litellm_completion(completion)
+        return response
+
+    @property
+    def model(self) -> str:
+        return self._model
+
+    def get_config(self) -> dict:
+        return self.client_props
+
+
+class LiteLLMResolver(DependencyResolver[LLM]):
+    memoize_resolution: bool = False
+
+    def __init__(self, model: str, conf: dict = {}):
+        super().__init__()
+        conf["model"] = model
+        self.props = LiteLLMConf.model_validate(conf)
+        self.LiteLLM = LiteLLM(self.props)
+
+    def resolve(self, guild_id: str, agent_id: str) -> LLM:
+        return self.LiteLLM  # We can always return the same instance of LiteLLM