flo-ai 1.0.0.dev2__tar.gz → 1.0.0.dev4__tar.gz

{flo_ai-1.0.0.dev2 → flo_ai-1.0.0.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flo_ai
-Version: 1.0.0.dev2
+Version: 1.0.0.dev4
 Summary: A easy way to create structured AI agents
 License: MIT
 Author: rootflo

flo_ai-1.0.0.dev4/flo_ai/__init__.py

@@ -0,0 +1,66 @@
+"""
+flo_ai - A flexible agent framework for LLM-powered applications
+"""
+
+# Models package - Agent framework components
+from .models import Agent, AgentError, BaseAgent, AgentType, ReasoningPattern
+
+from .builder.agent_builder import AgentBuilder
+
+# LLM package - Language model integrations
+from .llm import BaseLLM, Anthropic, OpenAI, OllamaLLM, Gemini
+
+# Tool package - Tool framework components
+from .tool import Tool, ToolExecutionError, flo_tool, create_tool_from_function
+
+# Arium package - Workflow and memory components
+from .arium import (
+    Arium,
+    BaseArium,
+    create_arium,
+    MessageMemory,
+    BaseMemory,
+    StartNode,
+    EndNode,
+    Edge,
+)
+
+# Utils package - Utility functions
+from .utils import FloUtils
+
+__all__ = [
+    # Models
+    'Agent',
+    'AgentError',
+    'BaseAgent',
+    'AgentType',
+    'ReasoningPattern',
+    # Utils
+    'FloUtils',
+    # LLM
+    'BaseLLM',
+    'Anthropic',
+    'OpenAI',
+    'OllamaLLM',
+    'Gemini',
+    # Tools
+    'Tool',
+    'ToolExecutionError',
+    'flo_tool',
+    'create_tool_from_function',
+    # Arium
+    'Arium',
+    'BaseArium',
+    'AriumBuilder',
+    'create_arium',
+    'MessageMemory',
+    'BaseMemory',
+    'StartNode',
+    'EndNode',
+    'Edge',
+    # Builder
+    'AgentBuilder',
+    'AriumBuilder',
+]
+
+__version__ = '1.0.0'

flo_ai-1.0.0.dev4/flo_ai/arium/arium.py

@@ -0,0 +1,164 @@
+from flo_ai.arium.base import BaseArium
+from flo_ai.arium.memory import MessageMemory, BaseMemory
+from flo_ai.llm.base_llm import ImageMessage
+from typing import List, Dict, Any, Optional
+from flo_ai.models.agent import Agent
+from flo_ai.tool.base_tool import Tool
+from flo_ai.arium.models import StartNode, EndNode
+from flo_ai.utils.logger import logger
+from flo_ai.utils.variable_extractor import (
+    extract_variables_from_inputs,
+    extract_agent_variables,
+    validate_multi_agent_variables,
+    resolve_variables,
+)
+
+
+class Arium(BaseArium):
+    def __init__(self, memory: BaseMemory):
+        super().__init__()
+        self.is_compiled = False
+        self.memory = memory if memory else MessageMemory()
+
+    def compile(self):
+        self.validate_graph()
+        self.is_compiled = True
+
+    async def run(
+        self,
+        inputs: List[str | ImageMessage],
+        variables: Optional[Dict[str, Any]] = None,
+    ):
+        if not self.is_compiled:
+            raise ValueError('Arium is not compiled')
+
+        if not self.memory:
+            raise ValueError('Arium has no memory')
+
+        if not self.nodes:
+            raise ValueError('Arium has no nodes')
+
+        # Extract and validate variables from inputs and all agents
+        self._extract_and_validate_variables(inputs, variables)
+
+        # Resolve variables in inputs and agent prompts
+        resolved_inputs = self._resolve_inputs(inputs, variables)
+        self._resolve_agent_prompts(variables)
+
+        return await self._execute_graph(resolved_inputs)
+
+    async def _execute_graph(self, inputs: List[str | ImageMessage]):
+        [self.memory.add(msg) for msg in inputs]
+
+        current_node = self.nodes[self.start_node_name]
+        current_edge = self.edges[self.start_node_name]
+
+        logger.info(f'Executing graph from {current_node.name}')
+        while current_node.name != self.end_node_name:
+            # execute current node
+            result = await self._execute_node(current_node)
+
+            # update results to memory
+            self._add_to_memory(result)
+
+            # find next node post current node
+            next_node_name = current_edge.router_fn(memory=self.memory)
+
+            # find next edge
+            # TODO: next_node_name might not be in self.edges if it's the end node. Handle this case
+            next_edge = (
+                self.edges[next_node_name] if next_node_name in self.edges else None
+            )
+
+            # update current node
+            current_node = self.nodes[next_node_name]
+            current_edge = next_edge
+
+        return self.memory.get()
+
+    def _extract_and_validate_variables(
+        self, inputs: List[str | ImageMessage], variables: Dict[str, Any]
+    ) -> None:
+        """Extract variables from inputs and agents, then validate them.
+
+        Args:
+            inputs: List of input messages
+            variables: Dictionary of variable name to value mappings
+
+        Raises:
+            ValueError: If any required variables are missing
+        """
+        # Extract variables from inputs
+        input_variables = extract_variables_from_inputs(inputs)
+
+        # Extract variables from all agents in the workflow
+        agents_variables = {}
+        for node in self.nodes.values():
+            if isinstance(node, Agent):
+                agent_vars = extract_agent_variables(node)
+                if agent_vars:
+                    agents_variables[node.name] = agent_vars
+
+        # Validate input variables separately with cleaner error message
+        if input_variables:
+            missing_input_vars = input_variables - set(variables.keys())
+            if missing_input_vars:
+                provided_keys = sorted(variables.keys())
+                raise ValueError(
+                    f'Input contains missing variables: {sorted(missing_input_vars)}. '
+                    f'Provided variables: {provided_keys}'
+                )
+
+        # Validate agent variables with detailed agent breakdown
+        if agents_variables:
+            validate_multi_agent_variables(agents_variables, variables)
+
+    def _resolve_inputs(
+        self, inputs: List[str | ImageMessage], variables: Dict[str, Any]
+    ) -> List[str | ImageMessage]:
+        """Resolve variables in input messages.
+
+        Args:
+            inputs: List of input messages
+            variables: Dictionary of variable name to value mappings
+
+        Returns:
+            List of inputs with variables resolved
+        """
+        resolved_inputs = []
+        for input_item in inputs:
+            if isinstance(input_item, str):
+                # Resolve variables in text input
+                resolved_input = resolve_variables(input_item, variables)
+                resolved_inputs.append(resolved_input)
+            else:
+                # ImageMessage objects don't need variable resolution
+                resolved_inputs.append(input_item)
+        return resolved_inputs
+
+    def _resolve_agent_prompts(self, variables: Dict[str, Any]) -> None:
+        """Resolve variables in all agent system prompts and mark them as resolved.
+
+        Args:
+            variables: Dictionary of variable name to value mappings
+        """
+        for node in self.nodes.values():
+            if isinstance(node, Agent):
+                node.system_prompt = resolve_variables(node.system_prompt, variables)
+                node.resolved_variables = True
+
+    async def _execute_node(self, node: Agent | Tool | StartNode | EndNode):
+        if isinstance(node, Agent):
+            # Variables are already resolved, pass empty dict to avoid re-processing
+            return await node.run(self.memory.get(), variables={})
+        elif isinstance(node, Tool):
+            return await node.execute(self.memory.get())
+        elif isinstance(node, StartNode):
+            return None
+        elif isinstance(node, EndNode):
+            return None
+
+    def _add_to_memory(self, result: str):
+        # TODO result will be None for start and end nodes
+        if result:
+            self.memory.add(result)

{flo_ai-1.0.0.dev2 → flo_ai-1.0.0.dev4}/flo_ai/arium/builder.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Callable, Union
+from typing import List, Optional, Callable, Union, Dict, Any
 from flo_ai.arium.arium import Arium
 from flo_ai.arium.memory import MessageMemory, BaseMemory
 from flo_ai.models.agent import Agent
@@ -128,10 +128,14 @@ class AriumBuilder:
         self._arium = arium
         return arium
 
-    async def build_and_run(self, inputs: List[Union[str, ImageMessage]]) -> List[dict]:
-        """Build the Arium and run it with the given inputs."""
+    async def build_and_run(
+        self,
+        inputs: List[Union[str, ImageMessage]],
+        variables: Optional[Dict[str, Any]] = None,
+    ) -> List[dict]:
+        """Build the Arium and run it with the given inputs and optional runtime variables."""
         arium = self.build()
-        return await arium.run(inputs)
+        return await arium.run(inputs, variables=variables)
 
     def visualize(
         self, output_path: str = 'arium_graph.png', title: str = 'Arium Workflow'

{flo_ai-1.0.0.dev2 → flo_ai-1.0.0.dev4}/flo_ai/models/__init__.py

@@ -6,10 +6,4 @@ from .agent import Agent
 from .agent_error import AgentError
 from .base_agent import BaseAgent, AgentType, ReasoningPattern
 
-__all__ = [
-    "Agent",
-    "AgentError", 
-    "BaseAgent",
-    "AgentType",
-    "ReasoningPattern"
-]
+__all__ = ['Agent', 'AgentError', 'BaseAgent', 'AgentType', 'ReasoningPattern']

{flo_ai-1.0.0.dev2 → flo_ai-1.0.0.dev4}/flo_ai/models/agent.py

@@ -5,6 +5,12 @@ from flo_ai.llm.base_llm import BaseLLM, ImageMessage
 from flo_ai.tool.base_tool import Tool, ToolExecutionError
 from flo_ai.models.agent_error import AgentError
 from flo_ai.utils.logger import logger
+from flo_ai.utils.variable_extractor import (
+    extract_variables_from_inputs,
+    extract_agent_variables,
+    validate_multi_agent_variables,
+    resolve_variables,
+)
 
 
 class Agent(BaseAgent):
@@ -40,35 +46,79 @@ class Agent(BaseAgent):
         self.output_schema = output_schema
         self.role = role
 
-    async def run(self, inputs: List[str | ImageMessage] | str) -> str:
+    async def run(
+        self,
+        inputs: List[str | ImageMessage] | str,
+        variables: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        variables = variables or {}
+
         if isinstance(inputs, str):
             inputs = [inputs]
 
-        for input in inputs:
-            if isinstance(input, ImageMessage):
-                self.add_to_history('user', self.llm.format_image_in_message(input))
-            else:
-                self.add_to_history('user', input)
+        # Perform runtime variable validation if not already resolved (single agent usage)
+        if not self.resolved_variables:
+            # Extract variables from inputs and system prompt
+            input_variables = extract_variables_from_inputs(inputs)
+            agent_variables = extract_agent_variables(self)
+            all_required_variables = input_variables.union(agent_variables)
+
+            # Validate that all required variables are provided
+            if all_required_variables:
+                agents_variables = {self.name: all_required_variables}
+                validate_multi_agent_variables(agents_variables, variables)
+
+            # Resolve variables and mark as resolved
+            self.system_prompt = resolve_variables(self.system_prompt, variables)
+
+            # Process inputs and resolve variables in string inputs
+            for input in inputs:
+                if isinstance(input, ImageMessage):
+                    self.add_to_history('user', self.llm.format_image_in_message(input))
+                else:
+                    # Resolve variables in text input
+                    resolved_input = resolve_variables(input, variables)
+                    self.add_to_history('user', resolved_input)
+
+            # after resolving agent system prompts and inputs, mark variables as resolved
+            self.resolved_variables = True
+
+        else:
+            # Variables already resolved, process inputs without variable resolution
+            for input in inputs:
+                if isinstance(input, ImageMessage):
+                    self.add_to_history('user', self.llm.format_image_in_message(input))
+                else:
+                    self.add_to_history('user', input)
 
         retry_count = 0
 
         # If no tools, act as conversational agent
         if not self.tools:
-            return await self._run_conversational(retry_count)
+            return await self._run_conversational(retry_count, variables)
 
         # Otherwise, run as tool agent
-        return await self._run_with_tools(retry_count)
+        return await self._run_with_tools(retry_count, variables)
 
-    async def _run_conversational(self, retry_count: int) -> str:
+    async def _run_conversational(
+        self, retry_count: int, variables: Optional[Dict[str, Any]] = None
+    ) -> str:
         """Run as a conversational agent when no tools are provided"""
+        variables = variables or {}
+
         while retry_count < self.max_retries:
             try:
+                # Resolve variables in system prompt
+                system_content = (
+                    self._get_cot_prompt(variables)
+                    if self.reasoning_pattern == ReasoningPattern.COT
+                    else resolve_variables(self.system_prompt, variables)
+                )
+
                 messages = [
                     {
                         'role': 'system',
-                        'content': self._get_cot_prompt()
-                        if self.reasoning_pattern == ReasoningPattern.COT
-                        else self.system_prompt,
+                        'content': system_content,
                     }
                 ] + self.conversation_history
 
@@ -111,18 +161,26 @@ class Agent(BaseAgent):
                         original_error=e,
                     )
 
-    async def _run_with_tools(self, retry_count: int = 0) -> str:
+    async def _run_with_tools(
+        self, retry_count: int = 0, variables: Optional[Dict[str, Any]] = None
+    ) -> str:
         """Run as a tool-using agent when tools are provided"""
+        variables = variables or {}
+
         while retry_count < self.max_retries:
             try:
+                # Resolve variables in system prompt based on reasoning pattern
+                if self.reasoning_pattern == ReasoningPattern.REACT:
+                    system_content = self._get_react_prompt(variables)
+                elif self.reasoning_pattern == ReasoningPattern.COT:
+                    system_content = self._get_cot_prompt(variables)
+                else:
+                    system_content = resolve_variables(self.system_prompt, variables)
+
                 messages = [
                     {
                         'role': 'system',
-                        'content': self._get_react_prompt()
-                        if self.reasoning_pattern == ReasoningPattern.REACT
-                        else self._get_cot_prompt()
-                        if self.reasoning_pattern == ReasoningPattern.COT
-                        else self.system_prompt,
+                        'content': system_content,
                     }
                 ] + self.conversation_history
 
@@ -239,12 +297,18 @@ class Agent(BaseAgent):
 
         raise AgentError(f'Failed after maximum {self.max_retries} attempts.')
 
-    def _get_react_prompt(self) -> str:
+    def _get_react_prompt(self, variables: Optional[Dict[str, Any]] = None) -> str:
         """Get system prompt modified for ReACT pattern"""
+        variables = variables or {}
+
         tools_desc = '\n'.join(
             [f'- {tool.name}: {tool.description}' for tool in self.tools]
         )
-        react_prompt = f"""{self.system_prompt}
+
+        # Resolve variables in the base system prompt
+        resolved_system_prompt = resolve_variables(self.system_prompt, variables)
+
+        react_prompt = f"""{resolved_system_prompt}
             When solving tasks, follow this format:
 
             Thought: Analyze the situation and think about what to do
@@ -263,12 +327,18 @@ class Agent(BaseAgent):
 
         return react_prompt
 
-    def _get_cot_prompt(self) -> str:
+    def _get_cot_prompt(self, variables: Optional[Dict[str, Any]] = None) -> str:
         """Get system prompt modified for Chain of Thought pattern"""
+        variables = variables or {}
+
         tools_desc = '\n'.join(
             [f'- {tool.name}: {tool.description}' for tool in self.tools]
         )
-        cot_prompt = f"""{self.system_prompt}
+
+        # Resolve variables in the base system prompt
+        resolved_system_prompt = resolve_variables(self.system_prompt, variables)
+
+        cot_prompt = f"""{resolved_system_prompt}
             When solving tasks, follow this Chain of Thought reasoning format:
 
             Let me think through this step by step:

{flo_ai-1.0.0.dev2 → flo_ai-1.0.0.dev4}/flo_ai/models/base_agent.py

@@ -29,6 +29,7 @@ class BaseAgent(ABC):
         self.agent_type = agent_type
         self.llm = llm
         self.max_retries = max_retries
+        self.resolved_variables = False
         self.conversation_history: List[Dict[str, str]] = []
 
     @abstractmethod

{flo_ai-1.0.0.dev2 → flo_ai-1.0.0.dev4}/flo_ai/utils/__init__.py

@@ -1,5 +1,3 @@
 from .flo_utils import FloUtils
 
-__all__ = [
-    'FloUtils'
-]
+__all__ = ['FloUtils']

{flo_ai-1.0.0.dev2 → flo_ai-1.0.0.dev4}/flo_ai/utils/flo_utils.py

@@ -1,11 +1,9 @@
 from typing import Dict, Any
-import re
 import json
 from flo_ai.utils.logger import logger
 
 
 class FloUtils:
-
     @staticmethod
     def extract_jsons_from_string(data: str, strict: bool = False) -> Dict[str, Any]:
         """
@@ -14,6 +12,7 @@ class FloUtils:
         3) Merge into one dict (later keys override earlier)
         4) On strict mode, raise FloException if no JSON found
         """
+
         # Custom function to find balanced braces since (?R) is not supported in Python re
         def find_balanced_braces(text):
             matches = []
@@ -25,10 +24,10 @@ class FloUtils:
                     j = i + 1
                     in_string = False
                     escape_next = False
-                    
+
                     while j < len(text) and brace_count > 0:
                         char = text[j]
-                        
+
                         if escape_next:
                             escape_next = False
                         elif char == '\\' and in_string:
@@ -40,9 +39,9 @@ class FloUtils:
                                 brace_count += 1
                             elif char == '}':
                                 brace_count -= 1
-                        
+
                         j += 1
-                    
+
                     if brace_count == 0:
                         candidate = text[start:j]
                         # Try to validate it's actually JSON by attempting to parse
@@ -77,7 +76,7 @@ class FloUtils:
             raise ValueError(f'No JSON found in strict mode: {data}')
 
         return merged
-    
+
     @staticmethod
     def strip_comments_from_string(data: str) -> str:
         """Remove JS-style comments (// and /*…*/) so json.loads() will succeed."""

flo_ai-1.0.0.dev4/flo_ai/utils/variable_extractor.py

@@ -0,0 +1,157 @@
+"""
+Variable extraction utilities for the Flo AI framework.
+
+This module provides functions to extract variable placeholders from text,
+inputs, and agent configurations for runtime variable validation.
+"""
+
+import re
+from typing import List, Set, Dict, Any
+from flo_ai.llm.base_llm import ImageMessage
+
+
+def extract_variables_from_text(text: str) -> Set[str]:
+    """Extract variable placeholders from text using <variable_name> pattern.
+
+    Args:
+        text: Text containing variable placeholders like <variable_name>
+
+    Returns:
+        Set of variable names found in the text
+
+    Examples:
+        >>> extract_variables_from_text("Hello <name>, your <role> is important")
+        {'name', 'role'}
+        >>> extract_variables_from_text("No variables here")
+        set()
+    """
+    if not text:
+        return set()
+
+    # Use regex to find all <variable_name> patterns
+    # \w+ matches word characters (letters, digits, underscore)
+    variable_pattern = r'<(\w+)>'
+    matches = re.findall(variable_pattern, text)
+    return set(matches)
+
+
+def extract_variables_from_inputs(inputs: List[str | ImageMessage]) -> Set[str]:
+    """Extract variable placeholders from a list of input messages.
+
+    Args:
+        inputs: List of input strings or ImageMessages
+
+    Returns:
+        Set of variable names found across all string inputs
+
+    Note:
+        ImageMessage objects are skipped as they don't contain variable placeholders
+    """
+    all_variables = set()
+
+    for input_item in inputs:
+        if isinstance(input_item, str):
+            variables = extract_variables_from_text(input_item)
+            all_variables.update(variables)
+        # Skip ImageMessage objects as they don't contain text variables
+
+    return all_variables
+
+
+def extract_agent_variables(agent) -> Set[str]:
+    """Extract variable placeholders from an agent's system prompt.
+
+    Args:
+        agent: Agent instance with a system_prompt attribute
+
+    Returns:
+        Set of variable names found in the agent's system prompt
+
+    Note:
+        This function avoids importing Agent to prevent circular imports
+    """
+    if not hasattr(agent, 'system_prompt'):
+        return set()
+
+    return extract_variables_from_text(agent.system_prompt)
+
+
+def validate_variables(
+    required_variables: Set[str], provided_variables: Dict[str, Any], context: str = ''
+) -> None:
+    """Validate that all required variables are provided.
+
+    Args:
+        required_variables: Set of variable names that are required
+        provided_variables: Dictionary of variable name to value mappings
+        context: Optional context string for error messages (e.g., agent name)
+
+    Raises:
+        ValueError: If any required variables are missing
+    """
+    missing = required_variables - set(provided_variables.keys())
+    if missing:
+        available = list(provided_variables.keys()) if provided_variables else []
+        context_str = f' for {context}' if context else ''
+        raise ValueError(
+            f'Missing required variables{context_str}: {sorted(missing)}. '
+            f'Available variables: {available}'
+        )
+
+
+def resolve_variables(text: str, variables: Dict[str, Any]) -> str:
+    """Replace <variable_name> patterns with actual values
+
+    Args:
+        text: Text containing variable placeholders like <variable_name>
+        variables: Dictionary of variable name to value mappings
+
+    Returns:
+        Text with variables resolved
+
+    Raises:
+        ValueError: If a variable placeholder is found but not provided in variables
+    """
+    if not text or not variables:
+        return text
+
+    def replace_var(match):
+        var_name = match.group(1)
+        if var_name in variables:
+            return str(variables[var_name])
+        else:
+            available = list(variables.keys())
+            raise ValueError(
+                f"Variable '{var_name}' referenced in text but not provided. "
+                f'Available variables: {available}'
+            )
+
+    return re.sub(r'<(\w+)>', replace_var, text)
+
+
+def validate_multi_agent_variables(
+    agents_variables: Dict[str, Set[str]], provided_variables: Dict[str, Any]
+) -> None:
+    """Validate variables for multiple agents and provide detailed error messages.
+
+    Args:
+        agents_variables: Dictionary mapping agent names to their required variables
+        provided_variables: Dictionary of variable name to value mappings
+
+    Raises:
+        ValueError: If any agents are missing required variables, with detailed breakdown
+    """
+    missing_by_agent = {}
+    provided_keys = set(provided_variables.keys())
+
+    for agent_name, required_vars in agents_variables.items():
+        missing = required_vars - provided_keys
+        if missing:
+            missing_by_agent[agent_name] = sorted(missing)
+
+    if missing_by_agent:
+        error_msg = 'Missing required variables for agents:\n'
+        for agent_name, missing_vars in missing_by_agent.items():
+            error_msg += f"  - Agent '{agent_name}': {missing_vars}\n"
+        error_msg += f'Provided variables: {sorted(provided_keys)}'
+        raise ValueError(error_msg)

{flo_ai-1.0.0.dev2 → flo_ai-1.0.0.dev4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "flo_ai"
-version = "1.0.0-dev2"
+version = "1.0.0-dev4"
 description = "A easy way to create structured AI agents"
 authors = ["rootflo <*@rootflo.ai>"]
 license = "MIT"

flo_ai-1.0.0.dev2/flo_ai/__init__.py

@@ -1,92 +0,0 @@
-"""
-flo_ai - A flexible agent framework for LLM-powered applications
-"""
-
-# Models package - Agent framework components
-from .models import (
-    Agent,
-    AgentError,
-    BaseAgent,
-    AgentType,
-    ReasoningPattern
-)
-
-from .builder.agent_builder import (
-    AgentBuilder
-)
-
-# LLM package - Language model integrations
-from .llm import (
-    BaseLLM,
-    Anthropic,
-    OpenAI,
-    OllamaLLM,
-    Gemini
-)
-
-# Tool package - Tool framework components
-from .tool import (
-    Tool,
-    ToolExecutionError,
-    flo_tool,
-    create_tool_from_function
-)
-
-# Arium package - Workflow and memory components
-from .arium import (
-    Arium,
-    BaseArium,
-    create_arium,
-    MessageMemory,
-    BaseMemory,
-    StartNode,
-    EndNode,
-    Edge
-)
-
-# Utils package - Utility functions
-from .utils import (
-    FloUtils
-)
-
-__all__ = [
-    # Models
-    "Agent",
-    "AgentError",
-    "BaseAgent", 
-    "AgentType",
-    "ReasoningPattern",
-    
-    # Utils
-    "FloUtils",
-    
-    # LLM
-    "BaseLLM",
-    "Anthropic",
-    "OpenAI",
-    "OllamaLLM",
-    "Gemini",
-    
-    # Tools
-    "Tool",
-    "ToolExecutionError",
-    "flo_tool",
-    "create_tool_from_function",
-    
-    # Arium
-    "Arium",
-    "BaseArium",
-    "AriumBuilder",
-    "create_arium",
-    "MessageMemory",
-    "BaseMemory",
-    "StartNode",
-    "EndNode",
-    "Edge"
-
-    # Builder
-    "AgentBuilder",
-    "AriumBuilder"
-]
-
-__version__ = "1.0.0"

flo_ai-1.0.0.dev2/flo_ai/arium/arium.py

@@ -1,75 +0,0 @@
-from flo_ai.arium.base import BaseArium
-from flo_ai.arium.memory import MessageMemory, BaseMemory
-from flo_ai.llm.base_llm import ImageMessage
-from typing import List
-from flo_ai.models.agent import Agent
-from flo_ai.tool.base_tool import Tool
-from flo_ai.arium.models import StartNode, EndNode
-from flo_ai.utils.logger import logger
-
-
-class Arium(BaseArium):
-    def __init__(self, memory: BaseMemory):
-        super().__init__()
-        self.is_compiled = False
-        self.memory = memory if memory else MessageMemory()
-
-    def compile(self):
-        self.validate_graph()
-        self.is_compiled = True
-
-    async def run(self, inputs: List[str | ImageMessage]):
-        if not self.is_compiled:
-            raise ValueError('Arium is not compiled')
-
-        if not self.memory:
-            raise ValueError('Arium has no memory')
-
-        if not self.nodes:
-            raise ValueError('Arium has no nodes')
-
-        return await self._execute_graph(inputs)
-
-    async def _execute_graph(self, inputs: List[str | ImageMessage]):
-        [self.memory.add(msg) for msg in inputs]
-
-        current_node = self.nodes[self.start_node_name]
-        current_edge = self.edges[self.start_node_name]
-
-        logger.info(f'Executing graph from {current_node.name}')
-        while current_node.name != self.end_node_name:
-            # execute current node
-            result = await self._execute_node(current_node)
-
-            # update results to memory
-            self._add_to_memory(result)
-
-            # find next node post current node
-            next_node_name = current_edge.router_fn(memory=self.memory)
-
-            # find next edge
-            # TODO: next_node_name might not be in self.edges if it's the end node. Handle this case
-            next_edge = (
-                self.edges[next_node_name] if next_node_name in self.edges else None
-            )
-
-            # update current node
-            current_node = self.nodes[next_node_name]
-            current_edge = next_edge
-
-        return self.memory.get()
-
-    async def _execute_node(self, node: Agent | Tool | StartNode | EndNode):
-        if isinstance(node, Agent):
-            return await node.run(self.memory.get())
-        elif isinstance(node, Tool):
-            return await node.execute(self.memory.get())
-        elif isinstance(node, StartNode):
-            return None
-        elif isinstance(node, EndNode):
-            return None
-
-    def _add_to_memory(self, result: str):
-        # TODO result will be None for start and end nodes
-        if result:
-            self.memory.add(result)