rusticai-llm-agent 1.0.0__tar.gz

rusticai_llm_agent-1.0.0/PKG-INFO

@@ -0,0 +1,36 @@
+Metadata-Version: 2.4
+Name: rusticai-llm-agent
+Version: 1.0.0
+Summary: A complete set of LLM agents
+License-Expression: Apache-2.0
+Author: Dragonscale Industries Inc.
+Author-email: dev@dragonscale.ai
+Requires-Python: >=3.12,<3.13
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: test
+Requires-Dist: openai (>=1.107.0,<2.0.0)
+Requires-Dist: rusticai-core (>=1.0.0,<1.1.0)
+Project-URL: Homepage, https://www.rustic.ai/
+Project-URL: Repository, https://github.com/rustic-ai/rustic-ai
+Project-URL: Rustic AI Core, https://pypi.org/project/rusticai-core/
+Description-Content-Type: text/markdown
+
+# Rustic AI LLMAgent
+
+This module introduces a pluggable LLMAgent based on common patterns observed while using of language models within [Rustic AI](https://www.rustic.ai/) guilds.
+The implementation leverages the Rustic AI's LLM dependency injection to allow for custom implementations of the LLM.
+
+## Installing
+
+```shell
+pip install rusticai-llm-agent 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_llm_agent-1.0.0/README.md

@@ -0,0 +1,18 @@
+# Rustic AI LLMAgent
+
+This module introduces a pluggable LLMAgent based on common patterns observed while using of language models within [Rustic AI](https://www.rustic.ai/) guilds.
+The implementation leverages the Rustic AI's LLM dependency injection to allow for custom implementations of the LLM.
+
+## Installing
+
+```shell
+pip install rusticai-llm-agent 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_llm_agent-1.0.0/pyproject.toml

@@ -0,0 +1,77 @@
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+
+[project]
+name = "rusticai-llm-agent"
+version = "1.0.0"
+description = "A complete set of LLM agents"
+authors = [{name = "Dragonscale Industries Inc.", email = "dev@dragonscale.ai"}]
+license = "Apache-2.0"
+readme = "README.md"
+requires-python = ">=3.12,<3.13"
+dynamic = ["dependencies"]
+urls = { Homepage = "https://www.rustic.ai/", Repository = "https://github.com/rustic-ai/rustic-ai", "Rustic AI Core" = "https://pypi.org/project/rusticai-core/" }
+
+
+[tool.poetry]
+packages = [{ include = "rustic_ai", from = "src" }]
+
+[tool.poetry.dependencies]
+rusticai-core = {   version = "1.0.0"}
+openai = "^1.107.0"
+
+[tool.poetry-monorepo.deps]
+
+
+[tool.poetry.group.dev]
+optional = true
+
+[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"
+flaky = "^3.8.1"
+pytest-asyncio = "^0.26.0"
+rusticai-testing = { version = "1.0.0"}
+rusticai-litellm = { version = "1.0.0"}
+
+[tool.poetry.extras]
+test = ["pytest", "pytest-asyncio", "rusticai-testing","rusticai-litellm"]
+
+[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"
+force_sort_within_sections = true # For the sorting order issues
+lines_between_sections = 1 # For the newline issue
+known_third_party = ["pydantic"]
+known_first_party = ["rustic_ai.core", "rustic_ai.llm_agent"]
+known_local_folder = ["rustic_ai.testing"]
+

rusticai_llm_agent-1.0.0/src/rustic_ai/llm_agent/__init__.py

@@ -0,0 +1,32 @@
+from .llm_agent import LLMAgent
+from .llm_agent_conf import LLMAgentConfig
+from .llm_agent_utils import LLMAgentUtils
+from .memories import (
+    HistoryBasedMemoriesStore,
+    MemoriesStore,
+    QueueBasedMemoriesStore,
+    StateBackedMemoriesStore,
+)
+from .plugins import (
+    LLMCallWrapper,
+    PromptGenerator,
+    RequestPreprocessor,
+    ResponsePostprocessor,
+    TemplatedPromptGenerator,
+)
+
+__all__ = [
+    "LLMAgent",
+    "LLMAgentConfig",
+    "LLMAgentUtils",
+    "PromptGenerator",
+    "TemplatedPromptGenerator",
+    "MemoriesStore",
+    "QueueBasedMemoriesStore",
+    "HistoryBasedMemoriesStore",
+    "StateBackedMemoriesStore",
+    "LLMCallWrapper",
+    "RequestPreprocessor",
+    "ResponsePostprocessor",
+    "TemplatedPromptGenerator",
+]

rusticai_llm_agent-1.0.0/src/rustic_ai/llm_agent/llm_agent.py

@@ -0,0 +1,71 @@
+from typing import Optional
+
+from rustic_ai.core.guild.agent import Agent, ProcessContext, processor
+from rustic_ai.core.guild.agent_ext.depends.llm.llm import LLM
+from rustic_ai.core.guild.agent_ext.depends.llm.models import (
+    ChatCompletionRequest,
+    SystemMessage,
+)
+from rustic_ai.core.messaging.core.message import Message
+from rustic_ai.core.utils.json_utils import JsonDict
+from rustic_ai.llm_agent.llm_agent_conf import LLMAgentConfig
+from rustic_ai.llm_agent.llm_agent_helper import LLMAgentHelper
+from rustic_ai.llm_agent.llm_agent_utils import LLMAgentUtils
+
+
+class LLMAgent(Agent[LLMAgentConfig]):
+    """
+    A simple LLM Agent that simply invokes an LLM.
+    """
+
+    _system_prompt: Optional[str] = None
+    _updated: bool = False
+
+    @processor(
+        ChatCompletionRequest,
+        predicate=lambda self, msg: LLMAgentUtils.has_no_attachments(msg),
+        depends_on=["llm"],
+    )
+    def invoke_llm(self, ctx: ProcessContext[ChatCompletionRequest], llm: LLM):
+        """
+        Invoke the LLM with the given context. All the LLM configuration parameters are passed along.
+        """
+        prompt = ctx.payload
+
+        if self._system_prompt:
+            # If the system prompt was updated, add it to the messages.
+            messages = [SystemMessage(content=self._system_prompt)] + prompt.messages
+            prompt = prompt.model_copy(update={"messages": messages})
+        elif self.config.default_system_prompt:
+            # If there is a default system prompt, use it.
+            messages = [SystemMessage(content=self.config.default_system_prompt)] + prompt.messages
+            prompt = prompt.model_copy(update={"messages": messages})
+
+        LLMAgentHelper.invoke_llm_and_handle_response(
+            self,
+            self.config,
+            llm,
+            ctx,
+            prompt,
+        )
+
+    def _is_system_prompt_update(self, msg: Message) -> bool:
+        config = self.config
+        if config.system_prompt_generator and msg.format == config.system_prompt_generator.update_on_message_format:
+            return True
+        return False
+
+    @processor(JsonDict, predicate=lambda self, msg: self._is_system_prompt_update(msg))
+    def update_system_prompt(self, ctx: ProcessContext[JsonDict]):
+        """
+        Update the system prompt using the system prompt generator.
+        """
+        if not self.config.system_prompt_generator:
+            return
+
+        self._system_prompt = (
+            self.config.system_prompt_generator.generate_prompt(ctx.payload)
+            if self.config.system_prompt_generator
+            else self.config.default_system_prompt
+        )
+        self._updated = True

rusticai_llm_agent-1.0.0/src/rustic_ai/llm_agent/llm_agent_conf.py

@@ -0,0 +1,231 @@
+from enum import Enum
+from typing import Any, List, Optional, Type, TypeVar, Union, cast
+
+from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
+from typing_extensions import Annotated
+
+from rustic_ai.core.guild.dsl import BaseAgentProps
+from rustic_ai.core.utils.basic_class_utils import get_class_from_name
+from rustic_ai.llm_agent.plugins.llm_call_wrapper import LLMCallWrapper
+from rustic_ai.llm_agent.plugins.prompt_generators import PromptGenerator
+from rustic_ai.llm_agent.plugins.request_preprocessor import RequestPreprocessor
+from rustic_ai.llm_agent.plugins.response_postprocessor import ResponsePostprocessor
+
+
+class Models(str, Enum):
+    gpt_5 = "gpt-5"
+    gpt_5_chat = "gpt-5-chat"
+    gpt_5_mini = "gpt-5-mini"
+    gpt_5_nano = "gpt-5-nano"
+
+    gpt_4o = "gpt-4o"
+    gpt_4o_mini = "gpt-4o-mini"
+
+    gpt_4_1 = "gpt-4.1"
+    gpt_4_1_mini = "gpt-4.1-mini"
+    gpt_4_1_nano = "gpt-4.1-nano"
+
+    o3 = "o3"
+    o3_pro = "o3-pro"
+    o3_deep_research = "o3-deep-research"
+    o3_mini = "o3-mini"
+
+    o4_mini = "o4-mini"
+
+    gpt_4 = "gpt-4"
+    gpt_4_turbo = "gpt-4-turbo"
+
+    gemini_pro = "gemini-pro"
+    gemini_2_5_pro = "gemini-2.5-pro"
+    gemini_flash = "gemini-flash"
+    gemini_2_5_flash = "gemini-2.5-flash"
+    gemini_2_5_flash_lite = "gemini-2.5-flash-lite"
+
+    gemini_2_0_flash = "gemini-2.0-flash"
+    gemini_2_0_flash_lite = "gemini-2.0-flash-lite"
+
+    gemini_1_5_pro = "gemini-1.5-pro"
+    gemini_1_5_flash = "gemini-1.5-flash"
+
+    claude_opus_4_1 = "claude-opus-4-1"
+    claude_opus_4_0 = "claude-opus-4-0"
+    claude_sonnet_4_0 = "claude-sonnet-4-0"
+    claude_sonnet_3_5 = "claude-3-5-sonnet-latest"
+    claude_haiku_3_5 = "claude-3-5-haiku-latest"
+
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def _build_plugins(value: Any, base_type: type[T]) -> List[T]:
+    """
+    Accepts:
+      - list of dicts/instances (mixed allowed)
+      - single dict
+      - single instance
+
+    Dict shape: {"kind": "pkg.mod.Class", ...kwargs}
+    Instances are kept as-is.
+    """
+    if value is None:
+        return []
+
+    items = value if isinstance(value, list) else [value]
+    out: List[T] = []
+
+    for i, item in enumerate(items):
+        # Already constructed instance (programmatic usage)
+        if isinstance(item, base_type):
+            # mypy: item is a runtime instance of base_type (a subclass of BaseModel)
+            out.append(cast(T, item))
+            continue
+
+        # Dict spec with FQCN
+        if isinstance(item, dict):
+            cls_path = item.get("kind")
+            if not cls_path or not isinstance(cls_path, str):
+                raise TypeError(f"Plugin spec at index {i} must include string 'kind' (FQCN)")
+            cls: Type[Any] = get_class_from_name(cls_path)
+            if not issubclass(cls, base_type):
+                raise TypeError(f"{cls_path!r} is not a subclass of {base_type.__name__}")
+            kwargs = {k: v for k, v in item.items() if k != "kind"}
+            instance = cls(**kwargs)
+            out.append(cast(T, instance))
+            continue
+
+        raise TypeError(
+            f"Unsupported plugin spec at index {i}: expected dict or {base_type.__name__} instance; "
+            f"got {type(item).__name__}"
+        )
+
+    return out
+
+
+class LLMAgentConfig(BaseAgentProps):
+    """
+    Base configuration for LLM agents.
+    """
+
+    model_config = ConfigDict(extra="ignore")
+
+    model: Annotated[Union[str, Models], Field(examples=["gpt-5"])]
+    """
+    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.
+    """
+
+    base_url: Optional[str] = None
+    """Base URL for the LLM API."""
+
+    api_version: Optional[str] = None
+    """Version for the LLM API."""
+
+    custom_llm_provider: Optional[str] = None
+    """Custom LLM provider to use, e.g., 'bedrock' for Amazon Bedrock or 'vertex' for Google Vertex AI."""
+
+    timeout: Optional[float] = None
+    """
+    Timeout for the LLM API requests.
+    If the completion takes longer than this time, the request will be aborted.
+    """
+
+    max_retries: int = 0
+    """
+    Maximum number of retries for LLM requests if post processing fails.
+    This is useful if the LLM returns a response that can not be processed successfully
+    e.g. due to a parsing error in tools response.
+    """
+
+    default_system_prompt: Optional[str] = None
+    """
+    Default system prompt to use if no other mechanism updates the prompt.
+    """
+
+    system_prompt_generator: Optional[PromptGenerator] = Field(discriminator="type", default=None)
+    """
+    Mechanism to update the system prompt based on messages from other agents messages.
+    """
+
+    request_preprocessors: List[RequestPreprocessor] = Field(default_factory=list)
+    """
+    0 or more request preprocessors to apply before sending prompts to the LLM.
+    The order of preprocessors matters: they are applied in the order they are listed.
+    """
+
+    llm_request_wrappers: List[LLMCallWrapper] = Field(default_factory=list)
+    """
+    0 or more request wrap processors to apply before sending prompts to the LLM
+    and after receiving the response from the LLM. The plugin may modify the prompts before sending them to the LLM
+    and also act on the response after receiving it from the LLM.
+    The order of wrap processors matters: they are applied in the order they are listed.
+    """
+
+    response_postprocessors: List[ResponsePostprocessor] = Field(default_factory=list)
+    """
+    0 or more response postprocessors to apply after receiving the response from the LLM.
+    All the post processors get the response as is and can act on it.
+    Any modifications to the response don't have any affect on other post processors or the final response.
+    """
+
+    send_response: bool = True
+    """
+    Whether to send the original LLM response back to the conversation.
+    If False, only the post-processed messages (if any) will be sent back.
+    This is useful if the LLM response is only used for tool calls and not for direct responses.
+    """
+
+    vertex_location: Optional[str] = None
+    """
+    Place where vertex model is deployed (us-central1, asia-southeast1, etc.). If None, attempts to use the VERTEXAI_LOCATION environment variable.
+    """
+
+    vertex_project: Optional[str] = None
+    """
+    The Google Cloud project ID. If None, attempts to use the VERTEXAI_PROJECT environment variable.
+    """
+
+    @field_validator("request_preprocessors", mode="before")
+    @classmethod
+    def _coerce_req(cls, v):
+        return _build_plugins(v, RequestPreprocessor)
+
+    @field_validator("llm_request_wrappers", mode="before")
+    @classmethod
+    def _coerce_wrap(cls, v):
+        return _build_plugins(v, LLMCallWrapper)
+
+    @field_validator("response_postprocessors", mode="before")
+    @classmethod
+    def _coerce_resp(cls, v):
+        return _build_plugins(v, ResponsePostprocessor)
+
+    @field_serializer("request_preprocessors", mode="plain")
+    def _serialize_preprocessors(self, preprocessors):
+        """Ensure subclass fields are preserved during serialization."""
+        return [p.model_dump(exclude_none=True) for p in preprocessors]
+
+    @field_serializer("llm_request_wrappers", mode="plain")
+    def _serialize_wrappers(self, wrappers):
+        """Ensure subclass fields are preserved during serialization."""
+        return [w.model_dump(exclude_none=True) for w in wrappers]
+
+    @field_serializer("response_postprocessors", mode="plain")
+    def _serialize_postprocessors(self, postprocessors):
+        """Ensure subclass fields are preserved during serialization."""
+        return [p.model_dump(exclude_none=True) for p in postprocessors]
+
+    _non_llm_fields = {
+        "max_retries",
+        "request_preprocessors",
+        "llm_request_wrappers",
+        "response_postprocessors",
+        "default_system_prompt",
+        "system_prompt_generator",
+        "send_response",
+    }
+
+    def get_llm_params(self) -> dict:
+        """
+        Get the LLM parameters from the config, excluding non-LLM fields.
+        """
+        return self.model_dump(exclude={*self._non_llm_fields})

rusticai_llm_agent-1.0.0/src/rustic_ai/llm_agent/llm_agent_helper.py

@@ -0,0 +1,213 @@
+import logging
+from typing import List, Optional, Union
+
+import openai
+
+from rustic_ai.core.guild.agent import Agent, ProcessContext
+from rustic_ai.core.guild.agent_ext.depends.llm.llm import LLM
+from rustic_ai.core.guild.agent_ext.depends.llm.models import (
+    ChatCompletionError,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    LLMMessage,
+    ResponseCodes,
+)
+from rustic_ai.llm_agent.llm_agent_conf import LLMAgentConfig
+from rustic_ai.llm_agent.plugins.llm_call_wrapper import LLMCallWrapper
+from rustic_ai.llm_agent.plugins.response_postprocessor import ResponsePostprocessor
+
+
+class LLMAgentHelper:
+
+    @staticmethod
+    def invoke_llm_completion(
+        agent: Agent,
+        config: LLMAgentConfig,
+        llm: LLM,
+        ctx: ProcessContext[ChatCompletionRequest],
+        prompt: ChatCompletionRequest,
+    ) -> Union[ChatCompletionResponse, ChatCompletionError]:
+        """
+        Invoke the LLM completion with the given context.
+        The fields from the chat completion request, Agent Config, and the LLM are combined.
+        The LLM Configuration is used as the base, overwritten by Agent config and then
+        the Chat Completion Request.
+        """
+
+        pre_processors = config.request_preprocessors or []
+        wrap_processors = config.llm_request_wrappers or []
+        post_processors = config.response_postprocessors or []
+
+        for pre_processor in pre_processors:
+            prompt = pre_processor.preprocess(agent=agent, ctx=ctx, request=prompt, llm=llm)
+
+        for wrap_processor in wrap_processors:
+            prompt = wrap_processor.preprocess(agent=agent, ctx=ctx, request=prompt, llm=llm)
+
+        config_dict = config.get_llm_params()
+        prompt_dict = prompt.model_dump(exclude_none=True)
+
+        final_prompt = {
+            **config_dict,
+            **prompt_dict,
+        }
+
+        cc_request: ChatCompletionRequest = ChatCompletionRequest.model_validate(final_prompt)
+
+        max_retries = config.max_retries
+
+        return LLMAgentHelper._invoke_llm_completion(
+            agent=agent,
+            ctx=ctx,
+            llm=llm,
+            cc_request=cc_request,
+            wrap_processors=wrap_processors,
+            post_processors=post_processors,
+            max_retries=max_retries,
+            model=config.model,
+        )
+
+    @staticmethod
+    def _invoke_llm_completion(
+        agent: Agent,
+        ctx: ProcessContext[ChatCompletionRequest],
+        llm: LLM,
+        cc_request: ChatCompletionRequest,
+        wrap_processors: List[LLMCallWrapper],
+        post_processors: List[ResponsePostprocessor],
+        max_retries: int,
+        model: Optional[str] = None,
+    ) -> Union[ChatCompletionResponse, ChatCompletionError]:
+
+        response = llm.completion(cc_request, model)
+
+        new_messages = []
+
+        try:
+            for wrap_processor in reversed(wrap_processors):
+                post_msg = wrap_processor.postprocess(
+                    agent=agent,
+                    ctx=ctx,
+                    final_prompt=cc_request,
+                    llm_response=response,
+                    llm=llm,
+                )
+                if post_msg:
+                    new_messages.extend(post_msg)
+
+            for post_processor in post_processors:
+                post_msg = post_processor.postprocess(
+                    agent=agent,
+                    ctx=ctx,
+                    final_prompt=cc_request,
+                    llm_response=response,
+                    llm=llm,
+                )
+                if post_msg:
+                    new_messages.extend(post_msg)
+
+            reasoning = response.choices[0].message.reasoning_content if response.choices else ""
+
+            for msg in new_messages:
+                if reasoning:
+                    ctx.send(msg, reason=reasoning)
+                else:
+                    ctx.send(msg)
+
+        except Exception as e:  # pragma: no cover
+            logging.error(f"Error in post processing: {e}")
+            if max_retries > 0:
+                logging.info(f"Retrying LLM call, remaining retries: {max_retries}")
+                return LLMAgentHelper._invoke_llm_completion(
+                    agent=agent,
+                    ctx=ctx,
+                    llm=llm,
+                    cc_request=cc_request,
+                    wrap_processors=wrap_processors,
+                    post_processors=post_processors,
+                    max_retries=max_retries - 1,
+                )
+            else:
+                return ChatCompletionError(
+                    status_code=ResponseCodes.RESPONSE_PROCESSING_ERROR,
+                    message=str(e),
+                    model=agent.name,
+                    request_messages=ctx.payload.messages,
+                )
+
+        return response
+
+    @staticmethod
+    def process_api_status_error(
+        model_name: str,
+        ctx: ProcessContext[ChatCompletionRequest],
+        status_code: ResponseCodes,
+        error: openai.APIStatusError,
+        messages: List[Union[LLMMessage]],
+    ) -> ChatCompletionError:  # pragma: no cover
+        """
+        Process API status error and return a ChatCompletionError object.
+        """
+        cc_error = ChatCompletionError(
+            status_code=status_code,
+            message=error.message,
+            response=error.response.text if error.response else None,
+            model=model_name,
+            request_messages=messages,
+            body=error.body if hasattr(error, "body") else None,
+        )
+
+        return cc_error
+
+    @staticmethod
+    def invoke_llm_and_handle_response(
+        agent: Agent,
+        config: LLMAgentConfig,
+        llm: LLM,
+        ctx: ProcessContext[ChatCompletionRequest],
+        prompt: ChatCompletionRequest,
+    ) -> ChatCompletionResponse | ChatCompletionError:
+        """
+        Invoke the LLM and handle the response.
+        The fields from the chat completion request, Agent Config, and the LLM are combined.
+        The LLM Configuration is used as the base, overwritten by Agent config and then
+        the Chat Completion Request.
+        """
+        try:
+            chat_response = LLMAgentHelper.invoke_llm_completion(agent, config, llm, ctx, prompt)
+
+            if isinstance(chat_response, ChatCompletionError):
+                ctx.send_error(chat_response)
+                return chat_response
+
+            reasoning = chat_response.choices[0].message.reasoning_content if chat_response.choices else ""
+
+            if reasoning:
+                ctx.send(chat_response, reason=reasoning)
+            else:
+                ctx.send(chat_response)
+            return chat_response
+        except openai.APIStatusError as e:  # pragma: no cover
+            logging.error(f"Error in LLM completion: {e}")
+            # Publish the error message
+            error = LLMAgentHelper.process_api_status_error(
+                model_name=config.model,
+                ctx=ctx,
+                status_code=ResponseCodes(e.status_code),
+                error=e,
+                messages=ctx.payload.messages,
+            )
+
+            ctx.send_error(error)
+            return error
+        except Exception as e:  # pragma: no cover
+            logging.error(f"Unexpected error in LLM completion: {e}")
+            error = ChatCompletionError(
+                status_code=ResponseCodes.INTERNAL_SERVER_ERROR,
+                message=str(e),
+                model=agent.name,
+                request_messages=ctx.payload.messages,
+            )
+
+            ctx.send_error(error)
+            return error