beeai-framework 0.1.0__py3-none-any.whl

beeai_framework/__init__.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+from beeai_framework.agents import BaseAgent
+from beeai_framework.agents.bee.agent import BeeAgent
+from beeai_framework.backend import (
+    AssistantMessage,
+    CustomMessage,
+    Message,
+    Role,
+    SystemMessage,
+    ToolMessage,
+    UserMessage,
+)
+from beeai_framework.llms import LLM, AgentInput, BaseLLM
+from beeai_framework.memory import BaseMemory, ReadOnlyMemory, TokenMemory, UnconstrainedMemory
+from beeai_framework.memory.serializable import Serializable
+from beeai_framework.tools import Tool, tool
+from beeai_framework.tools.weather.openmeteo import OpenMeteoTool
+from beeai_framework.utils.templates import Prompt
+
+__all__ = [
+    "LLM",
+    "AgentInput",
+    "AssistantMessage",
+    "BaseAgent",
+    "BaseLLM",
+    "BaseMemory",
+    "BeeAgent",
+    "CustomMessage",
+    "Message",
+    "OpenMeteoTool",
+    "Prompt",
+    "ReadOnlyMemory",
+    "Role",
+    "Serializable",
+    "SystemMessage",
+    "TokenMemory",
+    "Tool",
+    "ToolMessage",
+    "UnconstrainedMemory",
+    "UserMessage",
+    "tool",
+]

beeai_framework/adapters/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

beeai_framework/adapters/litellm/chat.py

@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+from collections.abc import AsyncGenerator
+from typing import Any
+
+import litellm
+from litellm import (
+    ModelResponse,
+    ModelResponseStream,
+    acompletion,
+    get_supported_openai_params,
+)
+from pydantic import BaseModel, ConfigDict
+
+from beeai_framework.backend.chat import (
+    ChatModel,
+    ChatModelInput,
+    ChatModelOutput,
+    ChatModelStructureInput,
+    ChatModelStructureOutput,
+)
+from beeai_framework.backend.errors import ChatModelError
+from beeai_framework.backend.message import AssistantMessage, Message, Role, ToolMessage
+from beeai_framework.backend.utils import parse_broken_json
+from beeai_framework.context import RunContext
+from beeai_framework.tools.tool import Tool
+from beeai_framework.utils.custom_logger import BeeLogger
+
+logger = BeeLogger(__name__)
+
+
+class LiteLLMParameters(BaseModel):
+    model: str
+    messages: list[dict[str, Any]]
+    tools: list[dict[str, Any]] | None = None
+    response_format: dict[str, Any] | type[BaseModel] | None = None
+
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
+
+
+class LiteLLMChatModel(ChatModel):
+    @property
+    def model_id(self) -> str:
+        return self._model_id
+
+    @property
+    def provider_id(self) -> str:
+        return self._provider_id
+
+    def __init__(self, model_id: str | None = None, **settings: Any) -> None:
+        llm_provider = "ollama_chat" if self.provider_id == "ollama" else self.provider_id
+        self.supported_params = get_supported_openai_params(model=self.model_id, custom_llm_provider=llm_provider)
+        # drop any unsupported parameters that were passed in
+        litellm.drop_params = True
+        super().__init__()
+
+    async def _create(
+        self,
+        input: ChatModelInput,
+        run: RunContext,
+    ) -> ChatModelOutput:
+        litellm_input = self._transform_input(input)
+        response = await acompletion(**litellm_input.model_dump())
+        response_message = response.get("choices", [{}])[0].get("message", {})
+        response_content = response_message.get("content", "")
+        tool_calls = response_message.tool_calls
+
+        if tool_calls:
+            litellm_input.messages.append({"role": Role.ASSISTANT, "content": response_content})
+            for tool_call in tool_calls:
+                function_name = tool_call.function.name
+                function_to_call: Tool = next(filter(lambda t: t.name == function_name, input.tools))
+
+                function_args = json.loads(tool_call.function.arguments)
+                function_response = function_to_call.run(input=function_args)
+                litellm_input.messages.append({"role": Role.TOOL, "content": function_response})
+
+                response = await acompletion(**litellm_input.model_dump())
+
+        response_output = self._transform_output(response)
+        logger.trace(f"Inference response output:\n{response_output}")
+        return response_output
+
+    async def _create_stream(self, input: ChatModelInput, _: RunContext) -> AsyncGenerator[ChatModelOutput]:
+        litellm_input = self._transform_input(input)
+        parameters = litellm_input.model_dump()
+        parameters["stream"] = True
+        response = await acompletion(**parameters)
+
+        # TODO: handle tool calling for streaming
+        async for chunk in response:
+            response_output = self._transform_output(chunk)
+            if not response_output:
+                continue
+            yield response_output
+
+    async def _create_structure(self, input: ChatModelStructureInput, run: RunContext) -> ChatModelStructureOutput:
+        if "response_format" not in self.supported_params:
+            logger.warning(f"{self.provider_id} model {self.model_id} does not support structured data.")
+            return await super()._create_structure(input, run)
+        else:
+            response = await self._create(
+                ChatModelInput(messages=input.messages, response_format=input.schema, abort_signal=input.abort_signal),
+                run,
+            )
+
+            logger.trace(f"Structured response received:\n{response}")
+
+            text_response = response.get_text_content()
+            result = parse_broken_json(text_response)
+            # TODO: validate result matches expected schema
+            return ChatModelStructureOutput(object=result)
+
+    def _get_model_name(self) -> str:
+        return f"{'ollama_chat' if self.provider_id == 'ollama' else self.provider_id}/{self.model_id}"
+
+    def _transform_input(self, input: ChatModelInput) -> LiteLLMParameters:
+        messages_list = [message.to_plain() for message in input.messages]
+
+        if input.tools:
+            prepared_tools_list = [{"type": "function", "function": tool.prompt_data()} for tool in input.tools]
+        else:
+            prepared_tools_list = None
+
+        model = self._get_model_name()
+
+        return LiteLLMParameters(
+            model=model,
+            messages=messages_list,
+            tools=prepared_tools_list,
+            response_format=input.response_format,
+            **self.settings,
+        )
+
+    def _transform_output(self, chunk: ModelResponse | ModelResponseStream) -> ChatModelOutput:
+        choice = chunk.get("choices", [{}])[0]
+        finish_reason = choice.get("finish_reason")
+        message: Message | None = None
+        usage = choice.get("usage")
+
+        if isinstance(chunk, ModelResponseStream):
+            if finish_reason:
+                return None
+            content = choice.get("delta", {}).get("content")
+            if choice.get("tool_calls"):
+                message = ToolMessage(content)
+            elif choice.get("delta"):
+                message = AssistantMessage(content)
+            else:
+                # TODO: handle other possible types
+                raise ChatModelError(f"Unhandled event: {choice}")
+        else:
+            response_message = choice.get("message")
+            content = response_message.get("content")
+            message = AssistantMessage(content)
+
+        return ChatModelOutput(messages=[message], finish_reason=finish_reason, usage=usage)

beeai_framework/adapters/ollama/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

beeai_framework/adapters/ollama/backend/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

beeai_framework/adapters/ollama/backend/chat.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Any
+
+from beeai_framework.adapters.litellm.chat import LiteLLMChatModel
+from beeai_framework.backend.constants import ProviderName
+from beeai_framework.utils.custom_logger import BeeLogger
+
+logger = BeeLogger(__name__)
+
+
+class OllamaChatModel(LiteLLMChatModel):
+    provider_id: ProviderName = "ollama"
+
+    def __init__(self, model_id: str | None = None, **settings: Any) -> None:
+        self._model_id = model_id if model_id else os.getenv("OLLAMA_CHAT_MODEL", "llama3.1:8b")
+        self.settings = {"base_url": "http://localhost:11434"} | settings
+        super().__init__()

beeai_framework/adapters/watsonx/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

beeai_framework/adapters/watsonx/backend/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

beeai_framework/adapters/watsonx/backend/chat.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Any
+
+from beeai_framework.adapters.litellm.chat import LiteLLMChatModel
+from beeai_framework.backend.constants import ProviderName
+from beeai_framework.utils.custom_logger import BeeLogger
+
+logger = BeeLogger(__name__)
+
+
+class WatsonxChatModel(LiteLLMChatModel):
+    provider_id: ProviderName = "watsonx"
+
+    def __init__(self, model_id: str | None = None, **settings: Any) -> None:
+        self._model_id = model_id if model_id else os.getenv("WATSONX_CHAT_MODEL", "ibm/granite-3-8b-instruct")
+        self.settings = settings
+        super().__init__()

beeai_framework/agents/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from beeai_framework.agents.base import BaseAgent
+from beeai_framework.agents.errors import AgentError
+
+__all__ = ["AgentError", "BaseAgent"]

beeai_framework/agents/base.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+
+from beeai_framework.agents.types import AgentMeta, BeeRunInput, BeeRunOptions, BeeRunOutput
+from beeai_framework.context import Run, RunContext, RunContextInput, RunInstance
+from beeai_framework.emitter import Emitter
+from beeai_framework.memory import BaseMemory
+from beeai_framework.utils.models import ModelLike, to_model, to_model_optional
+
+
+class BaseAgent(ABC):
+    is_running: bool = False
+    emitter: Emitter = None
+
+    def run(self, run_input: ModelLike[BeeRunInput], options: ModelLike[BeeRunOptions] | None = None) -> Run:
+        run_input = to_model(BeeRunInput, run_input)
+        options = to_model_optional(BeeRunOptions, options)
+
+        if self.is_running:
+            raise RuntimeError("Agent is already running!")
+
+        try:
+            return RunContext.enter(
+                RunInstance(emitter=self.emitter),
+                RunContextInput(signal=options.signal if options else None, params=(run_input, options)),
+                lambda context: self._run(run_input, options, context),
+            )
+        except Exception as e:
+            if isinstance(e, RuntimeError):
+                raise e
+            else:
+                raise RuntimeError("Error has occurred!") from e
+        finally:
+            self.is_running = False
+
+    @abstractmethod
+    async def _run(self, run_input: BeeRunInput, options: BeeRunOptions | None, context: RunContext) -> BeeRunOutput:
+        pass
+
+    def destroy(self) -> None:
+        self.emitter.destroy()
+
+    @property
+    @abstractmethod
+    def memory(self) -> BaseMemory:
+        pass
+
+    @memory.setter
+    @abstractmethod
+    def memory(self, memory: BaseMemory) -> None:
+        pass
+
+    @property
+    def meta(self) -> AgentMeta:
+        return AgentMeta(
+            name=self.__class__.__name__,
+            description="",
+            tools=[],
+        )

beeai_framework/agents/bee/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from beeai_framework.agents.bee.agent import BeeAgent
+
+__all__ = ["BeeAgent"]

beeai_framework/agents/bee/agent.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+from datetime import UTC, datetime
+
+from beeai_framework.agents.base import BaseAgent
+from beeai_framework.agents.runners.base import (
+    BaseRunner,
+    BeeRunnerToolInput,
+    BeeRunnerToolResult,
+    RunnerIteration,
+)
+from beeai_framework.agents.runners.default.runner import DefaultRunner
+from beeai_framework.agents.runners.granite.runner import GraniteRunner
+from beeai_framework.agents.types import (
+    AgentMeta,
+    BeeAgentExecutionConfig,
+    BeeInput,
+    BeeRunInput,
+    BeeRunOptions,
+    BeeRunOutput,
+)
+from beeai_framework.backend import Message
+from beeai_framework.backend.message import AssistantMessage, MessageMeta, UserMessage
+from beeai_framework.context import RunContext
+from beeai_framework.emitter import Emitter, EmitterInput
+from beeai_framework.memory import BaseMemory
+
+
+class BeeAgent(BaseAgent):
+    runner: Callable[..., BaseRunner]
+
+    def __init__(self, bee_input: BeeInput) -> None:
+        self.input = bee_input
+        if "granite" in self.input.llm.model_id:
+            self.runner = GraniteRunner
+        else:
+            self.runner = DefaultRunner
+        self.emitter = Emitter.root().child(
+            EmitterInput(
+                namespace=["agent", "bee"],
+                creator=self,
+            )
+        )
+
+    @property
+    def memory(self) -> BaseMemory:
+        return self.input.memory
+
+    @memory.setter
+    def memory(self, memory: BaseMemory) -> None:
+        self.input.memory = memory
+
+    @property
+    def meta(self) -> AgentMeta:
+        tools = self.input.tools[:]
+
+        if self.input.meta:
+            return AgentMeta(
+                name=self.input.meta.name,
+                description=self.input.meta.description,
+                extra_description=self.input.meta.extra_description,
+                tools=tools,
+            )
+
+        extra_description = ["Tools that I can use to accomplish given task."]
+        for tool in tools:
+            extra_description.append(f"Tool ${tool.name}': ${tool.description}.")
+
+        return AgentMeta(
+            name="BeeAI",
+            tools=tools,
+            description="The BeeAI framework demonstrates its ability to auto-correct and adapt in real-time, improving"
+            " the overall reliability and resilience of the system.",
+            extra_description="\n".join(extra_description) if len(tools) > 0 else None,
+        )
+
+    async def _run(self, run_input: BeeRunInput, options: BeeRunOptions | None, context: RunContext) -> BeeRunOutput:
+        runner = self.runner(
+            self.input,
+            (
+                options
+                if options
+                else BeeRunOptions(
+                    execution=self.input.execution
+                    or (options.execution if options is not None else None)
+                    or BeeAgentExecutionConfig(
+                        max_retries_per_step=3,
+                        total_max_retries=20,
+                        max_iterations=10,
+                    ),
+                    signal=None,
+                )
+            ),
+            context,
+        )
+        await runner.init(run_input)
+
+        final_message: Message | None = None
+        while not final_message:
+            iteration: RunnerIteration = await runner.create_iteration()
+
+            if iteration.state.tool_name and iteration.state.tool_input:
+                tool_result: BeeRunnerToolResult = await runner.tool(
+                    input=BeeRunnerToolInput(
+                        state=iteration.state,
+                        emitter=iteration.emitter,
+                        meta=iteration.meta,
+                        signal=iteration.signal,
+                    )
+                )
+                await runner.memory.add(
+                    AssistantMessage(
+                        content=runner.templates.assistant.render(
+                            {
+                                "thought": iteration.state.thought,
+                                "tool_name": iteration.state.tool_name,
+                                "tool_input": iteration.state.tool_input,
+                                "tool_output": tool_result.output.to_string(),
+                                "final_answer": iteration.state.final_answer,
+                            }
+                        ),
+                        meta=MessageMeta({"success": tool_result.success}),
+                    )
+                )
+                iteration.state.tool_output = tool_result.output.get_text_content()
+
+                for key in ["partialUpdate", "update"]:
+                    await iteration.emitter.emit(
+                        key,
+                        {
+                            "data": iteration.state,
+                            "update": {
+                                "key": "tool_output",
+                                "value": tool_result.output,
+                                "parsedValue": tool_result.output.to_string(),
+                            },
+                            "meta": {"success": tool_result.success},  # TODO deleted meta
+                            "memory": runner.memory,
+                        },
+                    )
+
+            if iteration.state.final_answer:
+                final_message = AssistantMessage(
+                    content=iteration.state.final_answer, meta=MessageMeta({"createdAt": datetime.now(tz=UTC)})
+                )
+                await runner.memory.add(final_message)
+                await iteration.emitter.emit(
+                    "success",
+                    {
+                        "data": final_message,
+                        "iterations": runner.iterations,
+                        "memory": runner.memory,
+                        "meta": iteration.meta,
+                    },
+                )
+
+        if run_input.prompt is not None:
+            await self.input.memory.add(
+                UserMessage(content=run_input.prompt, meta=MessageMeta({"createdAt": context.created_at}))
+            )
+
+        await self.input.memory.add(final_message)
+
+        return BeeRunOutput(result=final_message, iterations=runner.iterations, memory=runner.memory)

beeai_framework/agents/errors.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from beeai_framework.errors import FrameworkError
+
+
+class AgentError(FrameworkError):
+    """Raised for errors caused by agents."""
+
+    def __init__(self, message: str = "Agent error", *, cause: Exception | None = None) -> None:
+        super().__init__(message, is_fatal=True, is_retryable=False, cause=cause)

beeai_framework/agents/runners/base.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+from beeai_framework.agents.types import (
+    BeeAgentRunIteration,
+    BeeAgentTemplates,
+    BeeInput,
+    BeeIterationResult,
+    BeeMeta,
+    BeeRunInput,
+    BeeRunOptions,
+)
+from beeai_framework.cancellation import AbortSignal
+from beeai_framework.context import RunContext
+from beeai_framework.emitter.emitter import Emitter
+from beeai_framework.emitter.types import EmitterInput
+from beeai_framework.memory.base_memory import BaseMemory
+from beeai_framework.tools import ToolOutput
+from beeai_framework.utils.counter import RetryCounter
+
+
+@dataclass
+class BeeRunnerLLMInput:
+    meta: BeeMeta
+    signal: AbortSignal
+    emitter: Emitter
+
+
+@dataclass
+class RunnerIteration:
+    emitter: Emitter
+    state: BeeIterationResult
+    meta: BeeMeta
+    signal: AbortSignal
+
+
+@dataclass
+class BeeRunnerToolResult:
+    output: ToolOutput
+    success: bool
+
+
+@dataclass
+class BeeRunnerToolInput:
+    state: BeeIterationResult  # TODO BeeIterationToolResult
+    meta: BeeMeta
+    signal: AbortSignal
+    emitter: Emitter
+
+
+class BaseRunner(ABC):
+    def __init__(self, input: BeeInput, options: BeeRunOptions, run: RunContext) -> None:
+        self._input = input
+        self._options = options
+        self._failed_attempts_counter = RetryCounter(
+            max_retries=(
+                options.execution.max_iterations if options.execution and options.execution.max_iterations else 0
+            ),
+            error_type=Exception,  # TODO Specific error type
+        )
+
+        self._memory: BaseMemory | None = None
+
+        self._iterations: list[BeeAgentRunIteration] = []
+        self._failedAttemptsCounter: RetryCounter = RetryCounter(
+            error_type=Exception,  # TODO AgentError
+            max_retries=(
+                options.execution.total_max_retries if options.execution and options.execution.total_max_retries else 0
+            ),
+        )
+        self._run = run
+
+    @property
+    def iterations(self) -> list[BeeAgentRunIteration]:
+        return self._iterations
+
+    @property
+    def memory(self) -> BaseMemory:
+        if self._memory is not None:
+            return self._memory
+        raise Exception("Memory has not been initialized.")
+
+    async def create_iteration(self) -> RunnerIteration:
+        meta: BeeMeta = BeeMeta(iteration=len(self._iterations) + 1)
+        max_iterations = (
+            self._options.execution.max_iterations
+            if self._options.execution and self._options.execution.max_iterations
+            else 0
+        )
+
+        if meta.iteration > max_iterations:
+            # TODO: Raise Agent Error with metadata
+            # https://github.com/i-am-bee/beeai-framework/blob/aa4d5e6091ed3bab8096492707ceb03d3b03863b/src/agents/bee/runners/base.ts#L70
+            raise Exception(f"Agent was not able to resolve the task in {max_iterations} iterations.")
+
+        emitter = self._run.emitter.child(emitter_input=EmitterInput(group_id=f"`iteration-${meta.iteration}"))
+        iteration: BeeAgentRunIteration = await self.llm(
+            BeeRunnerLLMInput(emitter=emitter, signal=self._run.signal, meta=meta)
+        )
+        self._iterations.append(iteration)
+
+        return RunnerIteration(emitter=emitter, state=iteration.state, meta=meta, signal=self._run.signal)
+
+    async def init(self, input: BeeRunInput) -> None:
+        self._memory = await self.init_memory(input)
+
+    @abstractmethod
+    async def llm(self, input: BeeRunnerLLMInput) -> BeeAgentRunIteration:
+        pass
+
+    @abstractmethod
+    async def tool(self, input: BeeRunnerToolInput) -> BeeRunnerToolResult:
+        pass
+
+    @abstractmethod
+    def default_templates(self) -> BeeAgentTemplates:
+        pass
+
+    @abstractmethod
+    async def init_memory(self, input: BeeRunInput) -> BaseMemory:
+        pass
+
+    @property
+    def templates(self) -> BeeAgentTemplates:
+        # TODO: overrides
+        return self.default_templates()
+
+    # TODO: Serialization