flo-ai 1.0.0.dev3__tar.gz → 1.0.0.dev5__tar.gz

{flo_ai-1.0.0.dev3 → flo_ai-1.0.0.dev5}/PKG-INFO

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

{flo_ai-1.0.0.dev3 → flo_ai-1.0.0.dev5}/flo_ai/__init__.py

@@ -8,7 +8,7 @@ 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
+from .llm import BaseLLM, Anthropic, OpenAI, OllamaLLM, Gemini, ImageMessage
 
 # Tool package - Tool framework components
 from .tool import Tool, ToolExecutionError, flo_tool, create_tool_from_function
@@ -43,6 +43,8 @@ __all__ = [
     'OpenAI',
     'OllamaLLM',
     'Gemini',
+    # LLM DataClass
+    'ImageMessage',
     # Tools
     'Tool',
     'ToolExecutionError',

flo_ai-1.0.0.dev5/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.dev3 → flo_ai-1.0.0.dev5}/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.dev5/flo_ai/llm/__init__.py

@@ -0,0 +1,16 @@
+from .base_llm import BaseLLM, ImageMessage
+from .anthropic_llm import Anthropic
+from .openai_llm import OpenAI
+from .ollama_llm import OllamaLLM
+from .gemini_llm import Gemini
+from .openai_vllm import OpenAIVLLM
+
+__all__ = [
+    'BaseLLM',
+    'Anthropic',
+    'OpenAI',
+    'OllamaLLM',
+    'Gemini',
+    'OpenAIVLLM',
+    'ImageMessage',
+]

flo_ai-1.0.0.dev5/flo_ai/llm/openai_vllm.py

@@ -0,0 +1,58 @@
+from typing import Any
+from .openai_llm import OpenAI
+
+
+class OpenAIVLLM(OpenAI):
+    def __init__(
+        self,
+        base_url: str,
+        model: str,
+        api_key: str = None,
+        temperature: float = 0.7,
+        **kwargs,
+    ):
+        super().__init__(
+            model=model,
+            api_key=api_key,
+            temperature=temperature,
+            base_url=base_url,
+            **kwargs,
+        )
+
+    # overriden
+    async def generate(
+        self, messages: list[dict], output_schema: dict = None, **kwargs
+    ) -> Any:
+        # Convert output_schema to OpenAI format if provided
+        if output_schema:
+            kwargs['extra_body'] = {'guided_json': output_schema.get('schema')}
+
+            # Add JSON format instruction to the system prompt
+            if messages and messages[0]['role'] == 'system':
+                messages[0]['content'] = (
+                    messages[0]['content']
+                    + '\n\nPlease provide your response in JSON format according to the specified schema.'
+                )
+            else:
+                messages.insert(
+                    0,
+                    {
+                        'role': 'system',
+                        'content': 'Please provide your response in JSON format according to the specified schema.',
+                    },
+                )
+
+        # Prepare OpenAI API parameters
+        vllm_openai_kwargs = {
+            'model': self.model,
+            'messages': messages,
+            **kwargs,
+            **self.kwargs,
+        }
+
+        # Make the API call
+        response = await self.client.chat.completions.create(**vllm_openai_kwargs)
+        message = response.choices[0].message
+
+        # Return the full message object instead of just the content
+        return message

{flo_ai-1.0.0.dev3 → flo_ai-1.0.0.dev5}/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"""
-        while retry_count < self.max_retries:
+        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
 
@@ -100,7 +150,7 @@ class Agent(BaseAgent):
 
                 should_retry, analysis = await self.handle_error(e, context)
 
-                if should_retry and retry_count < self.max_retries:
+                if should_retry and retry_count <= self.max_retries:
                     self.add_to_history(
                         'system', f'Error occurred. Analysis: {analysis}'
                     )
@@ -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"""
-        while retry_count < self.max_retries:
+        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
 
@@ -190,7 +248,7 @@ class Agent(BaseAgent):
                             'attempt': retry_count,
                         }
                         should_retry, analysis = await self.handle_error(e, context)
-                        if should_retry and retry_count < self.max_retries:
+                        if should_retry and retry_count <= self.max_retries:
                             self.add_to_history(
                                 'system', f'Tool execution error: {analysis}'
                             )
@@ -226,7 +284,7 @@ class Agent(BaseAgent):
                 }
 
                 should_retry, analysis = await self.handle_error(e, context)
-                if should_retry and retry_count < self.max_retries:
+                if should_retry and retry_count <= self.max_retries:
                     self.add_to_history(
                         'system', f'Error occurred. Analysis: {analysis}'
                     )
@@ -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.dev3 → flo_ai-1.0.0.dev5}/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.dev5/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.dev3 → flo_ai-1.0.0.dev5}/pyproject.toml

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

flo_ai-1.0.0.dev3/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)

flo_ai-1.0.0.dev3/flo_ai/llm/__init__.py

@@ -1,7 +0,0 @@
-from .base_llm import BaseLLM
-from .anthropic_llm import Anthropic
-from .openai_llm import OpenAI
-from .ollama_llm import OllamaLLM
-from .gemini_llm import Gemini
-
-__all__ = ['BaseLLM', 'Anthropic', 'OpenAI', 'OllamaLLM', 'Gemini']