from flo_ai import configure_telemetry, shutdown_telemetry

# Configure at startup
configure_telemetry(
    service_name="my_ai_app",
    service_version="1.0.0",
    console_export=True  # For debugging
)

# Your application code here...

# Shutdown to flush data
shutdown_telemetry()

# Export to OTLP collector (Jaeger, Prometheus, etc.)
configure_telemetry(
    service_name="my_ai_app",
    otlp_endpoint="http://localhost:4317"
)

import asyncio
from typing import Any
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.llm import OpenAI
from flo_ai.models.agent import Agent

async def main() -> None:
    # Create a simple conversational agent
    agent: Agent = (
        AgentBuilder()
        .with_name('Math Tutor')
        .with_prompt('You are a helpful math tutor.')
        .with_llm(OpenAI(model='gpt-4o-mini'))
        .build()
    )

    response: Any = await agent.run('What is the formula for the area of a circle?')
    print(f'Response: {response}')

asyncio.run(main())

import asyncio
from typing import Any, Dict, List, Union
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.tool.base_tool import Tool
from flo_ai.models.base_agent import ReasoningPattern
from flo_ai.models.agent import Agent
from flo_ai.llm import Anthropic

async def calculate(operation: str, x: float, y: float) -> float:
    if operation == 'add':
        return x + y
    elif operation == 'multiply':
        return x * y
    raise ValueError(f'Unknown operation: {operation}')

# Define a calculator tool
calculator_tool: Tool = Tool(
    name='calculate',
    description='Perform basic calculations',
    function=calculate,
    parameters={
        'operation': {
            'type': 'string',
            'description': 'The operation to perform (add or multiply)',
        },
        'x': {'type': 'number', 'description': 'First number'},
        'y': {'type': 'number', 'description': 'Second number'},
    },
)

# Create a tool-using agent with Claude
agent: Agent = (
    AgentBuilder()
    .with_name('Calculator Assistant')
    .with_prompt('You are a math assistant that can perform calculations.')
    .with_llm(Anthropic(model='claude-3-5-sonnet-20240620'))
    .with_tools([calculator_tool])
    .with_reasoning(ReasoningPattern.REACT)
    .with_retries(2)
    .build()
)

response: Any = await agent.run('Calculate 5 plus 3')
print(f'Response: {response}')

import asyncio
from typing import Any, Dict
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.llm import OpenAI
from flo_ai.models.agent import Agent

# Define output schema for structured responses
math_schema: Dict[str, Any] = {
    'type': 'object',
    'properties': {
        'solution': {'type': 'string', 'description': 'The step-by-step solution'},
        'answer': {'type': 'string', 'description': 'The final answer'},
    },
    'required': ['solution', 'answer'],
}

# Create an agent with structured output
agent: Agent = (
    AgentBuilder()
    .with_name('Structured Math Solver')
    .with_prompt('You are a math problem solver that provides structured solutions.')
    .with_llm(OpenAI(model='gpt-4o'))
    .with_output_schema(math_schema)
    .build()
)

response: Any = await agent.run('Solve: 2x + 5 = 15')
print(f'Structured Response: {response}')

from typing import Any, List
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.models.agent import Agent

# Create agent from YAML
yaml_config: str = """..."""  # Your YAML configuration string
email_thread: List[str] = ["Email thread content..."]

builder: AgentBuilder = AgentBuilder.from_yaml(yaml_str=yaml_config)
agent: Agent = builder.build()

# Use the agent
result: Any = await agent.run(email_thread)

import asyncio
from typing import Any, Dict
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.llm import OpenAI
from flo_ai.models.agent import Agent

async def main() -> None:
    # Create agent with variables in system prompt
    agent: Agent = (
        AgentBuilder()
        .with_name('Data Analyst')
        .with_prompt('Analyze <dataset_path> and focus on <key_metric>. Generate insights for <target_audience>.')
        .with_llm(OpenAI(model='gpt-4o-mini'))
        .build()
    )
    
    # Define variables at runtime
    variables: Dict[str, str] = {
        'dataset_path': '/data/sales_q4_2024.csv',
        'key_metric': 'revenue growth',
        'target_audience': 'executive team'
    }
    
    # Run agent with variable resolution
    result: Any = await agent.run(
        'Please provide a comprehensive analysis with actionable recommendations.',
        variables=variables
    )
    
    print(f'Analysis: {result}')

asyncio.run(main())

import asyncio
from typing import Any, Dict
from flo_ai.models.agent import Agent
from flo_ai.llm import OpenAI

async def input_variables_example() -> None:
    agent: Agent = Agent(
        name='content_creator',
        system_prompt='You are a content creator specializing in <content_type>.',
        llm=OpenAI(model='gpt-4o-mini')
    )
    
    variables: Dict[str, str] = {
        'content_type': 'technical blog posts',
        'topic': 'machine learning fundamentals',
        'word_count': '1500',
        'target_level': 'intermediate'
    }
    
    # Variables in both system prompt and user input
    result: Any = await agent.run(
        'Create a <word_count>-word article about <topic> for <target_level> readers.',
        variables=variables
    )
    
    print(f'Content: {result}')

asyncio.run(input_variables_example())

import asyncio
from typing import Any, Dict, List
from flo_ai.arium import AriumBuilder
from flo_ai.models.agent import Agent
from flo_ai.llm import OpenAI

async def multi_agent_variables() -> List[Any]:
    llm: OpenAI = OpenAI(model='gpt-4o-mini')
    
    # Agent 1: Research phase
    researcher: Agent = Agent(
        name='researcher',
        system_prompt='Research <research_topic> and focus on <research_depth> analysis.',
        llm=llm
    )
    
    # Agent 2: Writing phase  
    writer: Agent = Agent(
        name='writer',
        system_prompt='Write a <document_type> based on the research for <target_audience>.',
        llm=llm
    )
    
    # Agent 3: Review phase
    reviewer: Agent = Agent(
        name='reviewer',
        system_prompt='Review the <document_type> for <review_criteria> and provide feedback.',
        llm=llm
    )
    
    # Shared variables across all agents
    shared_variables: Dict[str, str] = {
        'research_topic': 'sustainable energy solutions',
        'research_depth': 'comprehensive',
        'document_type': 'white paper',
        'target_audience': 'policy makers',
        'review_criteria': 'accuracy and policy relevance'
    }
    
    # Run multi-agent workflow with shared variables
    result: List[Any] = await (
        AriumBuilder()
        .add_agents([researcher, writer, reviewer])
        .start_with(researcher)
        .connect(researcher, writer)
        .connect(writer, reviewer)
        .end_with(reviewer)
        .build_and_run(
            ['Begin comprehensive research and document creation process'],
            variables=shared_variables
        )
    )
    
    return result

asyncio.run(multi_agent_variables())

import asyncio
from typing import Any, Dict
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.models.agent import Agent

async def yaml_with_variables() -> None:
    yaml_config: str = """..."""  # Your YAML configuration
    
    # Variables for YAML agent
    variables: Dict[str, str] = {
        'user_role': 'data scientist',
        'domain_expertise': 'machine learning and statistical analysis', 
        'primary_objective': 'deriving actionable insights from data',
        'experience_level': 'senior',
        'priority_constraints': 'computational efficiency and model interpretability'
    }
    
    # Create agent from YAML with variables
    builder: AgentBuilder = AgentBuilder.from_yaml(yaml_str=yaml_config)
    agent: Agent = builder.build()
    
    result: Any = await agent.run(
        'Help me design an ML pipeline for <use_case> with <data_constraints>',
        variables={
            **variables,
            'use_case': 'customer churn prediction',
            'data_constraints': 'limited labeled data'
        }
    )
    
    print(f'ML Pipeline Advice: {result}')

asyncio.run(yaml_with_variables())

import asyncio
from typing import Any, Dict
from flo_ai.models.agent import Agent
from flo_ai.llm import OpenAI

async def variable_validation_example() -> None:
    agent: Agent = Agent(
        name='validator_example',
        system_prompt='Process <required_param> and <another_param> for analysis.',
        llm=OpenAI(model='gpt-4o-mini')
    )
    
    # Incomplete variables (missing 'another_param')
    incomplete_variables: Dict[str, str] = {
        'required_param': 'dataset.csv'
        # 'another_param' is missing
    }
    
    try:
        result: Any = await agent.run(
            'Analyze the data in <data_source>',
            variables=incomplete_variables  # Missing 'another_param' and 'data_source'
        )
    except ValueError as e:
        print(f'Variable validation error: {e}')
        # Error will list all missing variables with their locations

asyncio.run(variable_validation_example())

import asyncio
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.llm import OpenAI
from flo_ai.models.document import DocumentMessage, DocumentType

async def basic_document_analysis():
    # Create document message from file path
    document = DocumentMessage(
        document_type=DocumentType.PDF,
        document_file_path='path/to/your/document.pdf'
    )
    
    # Create document analysis agent
    agent = (
        AgentBuilder()
        .with_name('Document Analyzer')
        .with_prompt('Analyze the provided document and extract key insights, themes, and important information.')
        .with_llm(OpenAI(model='gpt-4o-mini'))
        .build()
    )
    
    # Process document with agent
    result = await agent.run([document])
    print(f'Analysis: {result}')

asyncio.run(basic_document_analysis())

document = DocumentMessage(
    document_type=DocumentType.PDF,
    document_file_path='/path/to/document.pdf'
)

# Read file as bytes
with open('document.pdf', 'rb') as f:
    pdf_bytes = f.read()

document = DocumentMessage(
    document_type=DocumentType.PDF,
    document_bytes=pdf_bytes,
    mime_type='application/pdf'
)

import base64

# Encode file to base64
with open('document.pdf', 'rb') as f:
    pdf_base64 = base64.b64encode(f.read()).decode('utf-8')

document = DocumentMessage(
    document_type=DocumentType.PDF,
    document_base64=pdf_base64,
    mime_type='application/pdf'
)

import asyncio
from flo_ai.arium import AriumBuilder
from flo_ai.models.document import DocumentMessage, DocumentType

async def document_workflow():
    # Create document message
    document = DocumentMessage(
        document_type=DocumentType.PDF,
        document_file_path='business_report.pdf'
    )
    
    # Define workflow YAML
    workflow_yaml = """
    metadata:
      name: document-analysis-workflow
      version: 1.0.0
      description: "Multi-agent document analysis pipeline"

    arium:
      agents:
        - name: intake_agent
          role: "Document Intake Specialist"
          job: "Process and assess document content for analysis."
          model:
            provider: openai
            name: gpt-4o-mini

        - name: content_analyzer
          role: "Content Analyst"
          job: "Analyze document content for themes, insights, and key information."
          model:
            provider: openai
            name: gpt-4o-mini

        - name: summary_generator
          role: "Summary Writer"
          job: "Create comprehensive summaries of analyzed content."
          model:
            provider: openai
            name: gpt-4o-mini

      workflow:
        start: intake_agent
        edges:
          - from: intake_agent
            to: [content_analyzer]
          - from: content_analyzer
            to: [summary_generator]
        end: [summary_generator]
    """
    
    # Run workflow with document
    result = await (
        AriumBuilder()
        .from_yaml(yaml_str=workflow_yaml)
        .build_and_run([document, 'Analyze this business report and provide insights'])
    )
    
    return result

asyncio.run(document_workflow())

document = DocumentMessage(
    document_type=DocumentType.PDF,
    document_file_path='report.pdf',
    metadata={
        'source': 'quarterly_reports',
        'department': 'finance',
        'priority': 'high',
        'tags': ['financial', 'q4-2024']
    }
)

# PDF Document
pdf_doc = DocumentMessage(
    document_type=DocumentType.PDF,
    document_file_path='presentation.pdf'
)

# Text Document
txt_doc = DocumentMessage(
    document_type=DocumentType.TXT,
    document_file_path='notes.txt'
)

# Process both with the same agent
agent = AgentBuilder().with_name('Multi-Format Analyzer').build()

pdf_result = await agent.run([pdf_doc])
txt_result = await agent.run([txt_doc])

from flo_ai.tool import flo_tool
from flo_ai.models.document import DocumentMessage, DocumentType

@flo_tool(description="Extract key information from documents")
async def extract_document_info(document_path: str, doc_type: str) -> str:
    """Extract key information from a document."""
    document_type = DocumentType.PDF if doc_type.lower() == 'pdf' else DocumentType.TXT
    
    document = DocumentMessage(
        document_type=document_type,
        document_file_path=document_path
    )
    
    # Use document processing agent
    agent = AgentBuilder().with_name('Info Extractor').build()
    result = await agent.run([document])
    
    return result

# Use in agent
agent = (
    AgentBuilder()
    .with_name('Document Processor')
    .with_tools([extract_document_info.tool])
    .build()
)

from flo_ai.utils.document_processor import DocumentProcessingError

try:
    document = DocumentMessage(
        document_type=DocumentType.PDF,
        document_file_path='nonexistent.pdf'
    )
    result = await agent.run([document])
except DocumentProcessingError as e:
    print(f'Document processing failed: {e}')
except FileNotFoundError:
    print('Document file not found')

from typing import List
from flo_ai.tool.base_tool import Tool
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.llm import OpenAI
from flo_ai.models.agent import Agent

async def weather_lookup(city: str) -> str:
    # Your weather API call here
    return f"Weather in {city}: Sunny, 25°C"

weather_tool: Tool = Tool(
    name='weather_lookup',
    description='Get current weather for a city',
    function=weather_lookup,
    parameters={
        'city': {
            'type': 'string',
            'description': 'City name to get weather for'
        }
    }
)

# Add to your agent
agent: Agent = (
    AgentBuilder()
    .with_name('Weather Assistant')
    .with_llm(OpenAI(model='gpt-4o-mini'))
    .with_tools([weather_tool])
    .build()
)

from typing import Any, Dict, Union
from flo_ai.tool import flo_tool
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.llm import OpenAI
from flo_ai.models.agent import Agent

@flo_tool(
    description="Perform mathematical calculations",
    parameter_descriptions={
        "operation": "The operation to perform (add, subtract, multiply, divide)",
        "x": "First number",
        "y": "Second number"
    }
)
async def calculate(operation: str, x: float, y: float) -> Union[float, str]:
    """Calculate mathematical operations between two numbers."""
    operations: Dict[str, callable] = {
        'add': lambda: x + y,
        'subtract': lambda: x - y,
        'multiply': lambda: x * y,
        'divide': lambda: x / y if y != 0 else 'Cannot divide by zero',
    }
    if operation not in operations:
        raise ValueError(f'Unknown operation: {operation}')
    return operations[operation]()

# Function can be called normally
result: Union[float, str] = await calculate("add", 5, 3)  # Returns 8

# Tool object is automatically available
agent: Agent = (
    AgentBuilder()
    .with_name('Calculator Agent')
    .with_llm(OpenAI(model='gpt-4o-mini'))
    .with_tools([calculate.tool])  # Access the tool via .tool attribute
    .build()
)

from flo_ai.tool import flo_tool

@flo_tool()
async def convert_units(value: float, from_unit: str, to_unit: str) -> str:
    """Convert between different units (km/miles, kg/lbs, celsius/fahrenheit)."""
    # Implementation here
    result: float = 0.0  # Your conversion logic here
    return f"{value} {from_unit} = {result} {to_unit}"

# Tool is automatically available as convert_units.tool

from typing import Optional
from flo_ai.tool import flo_tool

@flo_tool(
    name="weather_checker",
    description="Get current weather information for a city",
    parameter_descriptions={
        "city": "The city to get weather for",
        "country": "The country (optional)",
    }
)
async def get_weather(city: str, country: Optional[str] = None) -> str:
    """Get weather information for a specific city."""
    return f"Weather in {city}: sunny"

from flo_ai.models.base_agent import ReasoningPattern
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.llm import OpenAI
from flo_ai.models.agent import Agent

agent: Agent = (
    AgentBuilder()
    .with_name('Reasoning Agent')
    .with_llm(OpenAI(model='gpt-4o'))
    .with_reasoning(ReasoningPattern.COT)  # or REACT, DIRECT
    .build()
)

from flo_ai.llm import OpenAI

llm: OpenAI = OpenAI(
    model='gpt-4o',
    temperature=0.7,
    api_key='your-api-key'  # or set OPENAI_API_KEY env var
)

from flo_ai.llm import Anthropic

llm: Anthropic = Anthropic(
    model='claude-3-5-sonnet-20240620',
    temperature=0.7,
    api_key='your-api-key'  # or set ANTHROPIC_API_KEY env var
)

from flo_ai.llm import Gemini

llm: Gemini = Gemini(
    model='gemini-2.5-flash',  # or gemini-2.5-pro
    temperature=0.7,
    api_key='your-api-key'  # or set GOOGLE_API_KEY env var
)

from flo_ai.llm import VertexAI

llm: VertexAI = VertexAI(
    model='gemini-2.5-flash',  # or gemini-2.5-pro
    temperature=0.7,
    project='your-gcp-project-id',  # or set GOOGLE_CLOUD_PROJECT env var
    location='us-central1'  # or set GOOGLE_CLOUD_LOCATION env var
)

from flo_ai.llm import Ollama

llm: Ollama = Ollama(
    model='llama2',
    base_url='http://localhost:11434'
)

from flo_ai.llm import Gemini

llm: Gemini = Gemini(
    model='gemini-2.5-flash',  # or gemini-2.5-pro
    temperature=0.7,
    api_key='your-api-key'  # or set GOOGLE_API_KEY env var
)
messages=[{"role": "user", "content": "Stream a short sentence."}]
chunks: List[str] = []
    async for chunk in llm.stream(messages=messages):
        text = chunk.get('content', '')
        if text:
            chunks.append(text)
        if len(''.join(chunks)) >= max_chars:
            break
    return ''.join(chunks)

from pydantic import BaseModel, Field
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.llm import OpenAI
from flo_ai.models.agent import Agent

class MathSolution(BaseModel):
    solution: str = Field(description="Step-by-step solution")
    answer: str = Field(description="Final answer")
    confidence: float = Field(description="Confidence level (0-1)")

agent: Agent = (
    AgentBuilder()
    .with_name('Math Solver')
    .with_llm(OpenAI(model='gpt-4o'))
    .with_output_schema(MathSolution)
    .build()
)

from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.llm import OpenAI
from flo_ai.models.agent import Agent

agent: Agent = (
    AgentBuilder()
    .with_name('Robust Agent')
    .with_llm(OpenAI(model='gpt-4o'))
    .with_retries(3)  # Retry up to 3 times on failure
    .build()
)

from typing import Dict, Any
from flo_ai.tool.base_tool import Tool

async def custom_function(param1: str, param2: int) -> Dict[str, str]:
    # Your async logic here
    return {"result": f"Processed {param1} with {param2}"}

custom_tool: Tool = Tool(
    name='custom_function',
    description='A custom async tool',
    function=custom_function,
    parameters={
        'param1': {'type': 'string', 'description': 'First parameter'},
        'param2': {'type': 'integer', 'description': 'Second parameter'}
    }
)

from typing import Dict, Any
from flo_ai.formatter.yaml_format_parser import FloYamlParser
from flo_ai.builder.agent_builder import AgentBuilder
from flo_ai.llm import OpenAI
from flo_ai.models.agent import Agent

# Create parser from YAML definition
yaml_config: Dict[str, Any] = {}  # Your YAML configuration dict
parser: FloYamlParser = FloYamlParser.create(yaml_dict=yaml_config)
output_schema: Any = parser.get_format()

agent: Agent = (
    AgentBuilder()
    .with_name('YAML Configured Agent')
    .with_llm(OpenAI(model='gpt-4o'))
    .with_output_schema(output_schema)
    .build()
)

import asyncio
from typing import Any, List
from flo_ai.arium import AriumBuilder
from flo_ai.models.agent import Agent
from flo_ai.llm.openai_llm import OpenAI

async def simple_chain() -> List[Any]:
    llm: OpenAI = OpenAI(model='gpt-4o-mini')
    
    # Create agents
    analyst: Agent = Agent(
        name='content_analyst',
        system_prompt='Analyze the input and extract key insights.',
        llm=llm
    )
    
    summarizer: Agent = Agent(
        name='summarizer', 
        system_prompt='Create a concise summary based on the analysis.',
        llm=llm
    )
    
    # Build and run workflow
    result: List[Any] = await (
        AriumBuilder()
        .add_agents([analyst, summarizer])
        .start_with(analyst)
        .connect(analyst, summarizer)  # analyst → summarizer
        .end_with(summarizer)
        .build_and_run(["Analyze this complex business report..."])
    )
    
    return result

asyncio.run(simple_chain())

import asyncio
from typing import Any, List
from flo_ai.arium import AriumBuilder
from flo_ai.models.agent import Agent
from flo_ai.llm.openai_llm import OpenAI
from flo_ai.arium.memory import BaseMemory

async def conditional_workflow() -> List[Any]:
    llm: OpenAI = OpenAI(model='gpt-4o-mini')
    
    # Create specialized agents
    classifier: Agent = Agent(
        name='classifier',
        system_prompt='Classify the input as either "technical" or "business".',
        llm=llm
    )
    
    tech_specialist: Agent = Agent(
        name='tech_specialist',
        system_prompt='Provide technical analysis and solutions.',
        llm=llm
    )
    
    business_specialist: Agent = Agent(
        name='business_specialist', 
        system_prompt='Provide business analysis and recommendations.',
        llm=llm
    )
    
    final_agent: Agent = Agent(
        name='final_reviewer',
        system_prompt='Provide final review and conclusions.',
        llm=llm
    )
    
    # Define routing logic
    def route_by_type(memory: BaseMemory) -> str:
        """Route based on classification result"""
        messages: List[Any] = memory.get()
        last_message: str = str(messages[-1]) if messages else ""
        
        if "technical" in last_message.lower():
            return "tech_specialist"
        else:
            return "business_specialist"
    
    # Build workflow with conditional routing
    result: List[Any] = await (
        AriumBuilder()
        .add_agents([classifier, tech_specialist, business_specialist, final_agent])
        .start_with(classifier)
        .add_edge(classifier, [tech_specialist, business_specialist], route_by_type)
        .connect(tech_specialist, final_agent)
        .connect(business_specialist, final_agent)
        .end_with(final_agent)
        .build_and_run(["How can we optimize our database performance?"])
    )
    
    return result

import asyncio
from typing import Any, List
from flo_ai.tool import flo_tool
from flo_ai.arium import AriumBuilder
from flo_ai.models.agent import Agent
from flo_ai.llm.openai_llm import OpenAI

@flo_tool(description="Search for relevant information")
async def search_tool(query: str) -> str:
    # Your search implementation
    return f"Search results for: {query}"

@flo_tool(description="Perform calculations")
async def calculator(expression: str) -> float:
    # Your calculation implementation
    return eval(expression)  # Note: Use safely in production

async def agent_tool_workflow() -> List[Any]:
    llm: OpenAI = OpenAI(model='gpt-4o-mini')
    
    research_agent: Agent = Agent(
        name='researcher',
        system_prompt='Research topics and gather information.',
        llm=llm
    )
    
    analyst_agent: Agent = Agent(
        name='analyst',
        system_prompt='Analyze data and provide insights.',
        llm=llm
    )
    
    # Mix agents and tools in workflow
    result: List[Any] = await (
        AriumBuilder()
        .add_agent(research_agent)
        .add_tools([search_tool.tool, calculator.tool])
        .add_agent(analyst_agent)
        .start_with(research_agent)
        .connect(research_agent, search_tool.tool)
        .connect(search_tool.tool, calculator.tool)
        .connect(calculator.tool, analyst_agent)
        .end_with(analyst_agent)
        .build_and_run(["Research market trends for Q4 2024"])
    )
    
    return result

from typing import Any, List, Callable, Optional
from flo_ai.arium import AriumBuilder
from flo_ai.arium.arium import Arium
from flo_ai.models.agent import Agent
from flo_ai.tool.base_tool import Tool

# Assume these are defined elsewhere
agent1: Agent = ...  # Your agent definitions
agent2: Agent = ...
agent3: Agent = ...
tool1: Tool = ...    # Your tool definitions
tool2: Tool = ...
router_function: Callable = ...  # Your router function

# Build workflow and generate visual diagram
arium: Arium = (
    AriumBuilder()
    .add_agents([agent1, agent2, agent3])
    .add_tools([tool1, tool2])
    .start_with(agent1)
    .connect(agent1, tool1)
    .add_edge(tool1, [agent2, agent3], router_function)
    .end_with(agent2)
    .end_with(agent3)
    .visualize("my_workflow.png", "Customer Service Workflow")  # Generates PNG
    .build()
)

# Run the workflow
result: List[Any] = await arium.run(["Customer complaint about billing"])

from typing import Any, List
from flo_ai.arium import AriumBuilder
from flo_ai.arium.memory import MessageMemory
from flo_ai.arium.arium import Arium
from flo_ai.models.agent import Agent

# Assume these agents are defined elsewhere
agent1: Agent = ...
agent2: Agent = ...
agent3: Agent = ...

# Custom memory for persistent context
custom_memory: MessageMemory = MessageMemory()

result: List[Any] = await (
    AriumBuilder()
    .with_memory(custom_memory)  # Shared across all agents
    .add_agents([agent1, agent2, agent3])
    .start_with(agent1)
    .connect(agent1, agent2)
    .connect(agent2, agent3)
    .end_with(agent3)
    .build_and_run(["Initial context and instructions"])
)

# Build the arium for reuse
arium: Arium = (
    AriumBuilder()
    .with_memory(custom_memory)
    .add_agents([agent1, agent2, agent3])
    .start_with(agent1)
    .connect(agent1, agent2)
    .connect(agent2, agent3)
    .end_with(agent3)
    .build()
)

# Memory persists and can be reused
result2: List[Any] = await arium.run(["Follow-up question based on previous context"])

from typing import Any, List
from flo_ai.arium import AriumBuilder
from flo_ai.arium.arium import Arium
from flo_ai.models.agent import Agent
from flo_ai.tool.base_tool import Tool

# Assume these are defined elsewhere
agent1: Agent = ...
agent2: Agent = ...
tool1: Tool = ...
inputs: List[str] = ["Your input messages"]

# All builder methods return self for chaining
workflow: Arium = (
    AriumBuilder()
    .add_agent(agent1)           # Add components
    .add_tool(tool1)
    .start_with(agent1)          # Define flow
    .connect(agent1, tool1)
    .end_with(tool1)
    .build()                     # Create Arium instance
)

# Or build and run in one step
result: List[Any] = await (
    AriumBuilder()
    .add_agents([agent1, agent2])
    .start_with(agent1)
    .connect(agent1, agent2)
    .end_with(agent2)
    .build_and_run(inputs)       # Build + run together
)

import asyncio
from typing import Any, List
from flo_ai.arium import AriumBuilder

async def run_yaml_workflow() -> List[Any]:
    yaml_config = """..."""  # Your YAML configuration
    
    # Create workflow from YAML
    result: List[Any] = await (
        AriumBuilder()
        .from_yaml(yaml_config)
        .build_and_run(["Analyze this quarterly business report..."])
    )
    
    return result

asyncio.run(run_yaml_workflow())

import asyncio
from typing import Any, Dict, List, Literal
from flo_ai.arium import AriumBuilder
from flo_ai.tool.base_tool import Tool
from flo_ai.arium.memory import BaseMemory

# Define tools in Python (cannot be defined in YAML)
async def web_search(query: str) -> str:
    # Your search implementation
    return f"Search results for: {query}"

async def calculate(expression: str) -> str:
    # Your calculation implementation
    try:
        result = eval(expression)  # Note: Use safely in production
        return f"Calculation result: {result}"
    except:
        return "Invalid expression"

# Create tool objects
tools: Dict[str, Tool] = {
    "web_search": Tool(
        name="web_search",
        description="Search the web for current information",
        function=web_search,
        parameters={
            "query": {
                "type": "string",
                "description": "Search query"
            }
        }
    ),
    "calculator": Tool(
        name="calculator",
        description="Perform mathematical calculations",
        function=calculate,
        parameters={
            "expression": {
                "type": "string",
                "description": "Mathematical expression to calculate"
            }
        }
    )
}

# Define router functions in Python (cannot be defined in YAML)
def classification_router(memory: BaseMemory) -> Literal["researcher", "analyst"]:
    """Route based on task classification"""
    content = str(memory.get()[-1]).lower()
    if 'research' in content or 'investigate' in content:
        return 'researcher'
    elif 'calculate' in content or 'analyze data' in content:
        return 'analyst'
    return 'researcher'  # default

routers: Dict[str, callable] = {
    "classification_router": classification_router
}

async def run_workflow() -> List[Any]:
    yaml_config = """..."""  # Your YAML configuration from above
    
    # Create workflow with tools and routers provided as Python objects
    result: List[Any] = await (
        AriumBuilder()
        .from_yaml(
            yaml_str=yaml_config,
            tools=tools,      # Tools must be provided as Python objects
            routers=routers   # Routers must be provided as Python functions
        )
        .build_and_run(["Research the latest trends in renewable energy"])
    )
    
    return result

# Using LLM routers is incredibly simple!
async def run_intelligent_workflow():
    # No routers dictionary needed - they're defined in YAML!
    result = await (
        AriumBuilder()
        .from_yaml(yaml_str=intelligent_workflow_yaml)
        .build_and_run(["Write a technical tutorial on Docker containers"])
    )
    # The LLM will automatically route to technical_writer! ✨
    return result

import asyncio
from flo_ai.arium import AriumBuilder
from flo_ai.models.agent import Agent
from flo_ai.llm import OpenAI
from flo_ai.arium.llm_router import create_main_critic_reflection_router

async def reflection_workflow_example():
    llm = OpenAI(model='gpt-4o-mini', api_key='your-api-key')
    
    # Create agents
    main_agent = Agent(
        name='main_agent',
        system_prompt='Create solutions and improve them based on feedback.',
        llm=llm
    )
    
    critic = Agent(
        name='critic', 
        system_prompt='Provide constructive feedback for improvement.',
        llm=llm
    )
    
    final_agent = Agent(
        name='final_agent',
        system_prompt='Polish and finalize the work.',
        llm=llm
    )
    
    # Create reflection router - A → B → A → C pattern
    reflection_router = create_main_critic_reflection_router(
        main_agent='main_agent',
        critic_agent='critic',
        final_agent='final_agent',
        allow_early_exit=False,  # Strict pattern adherence
        llm=llm
    )
    
    # Build workflow
    result = await (
        AriumBuilder()
        .add_agents([main_agent, critic, final_agent])
        .start_with(main_agent)
        .add_edge(main_agent, [critic, final_agent], reflection_router)
        .add_edge(critic, [main_agent, final_agent], reflection_router)
        .end_with(final_agent)
        .build_and_run(["Create a comprehensive project proposal"])
    )
    
    return result

# Alternative: Direct factory usage
from flo_ai.arium.llm_router import create_llm_router

reflection_router = create_llm_router(
    'reflection',
    flow_pattern=['writer', 'editor', 'writer'],  # A → B → A
    allow_early_exit=False,
    llm=llm
)

# Minimal A → B → A pattern
yaml_config = """
arium:
  agents:
    - name: main_agent
      job: "Main work agent"
      model: {provider: openai, name: gpt-4o-mini}
    - name: critic
      job: "Feedback agent" 
      model: {provider: openai, name: gpt-4o-mini}

  routers:
    - name: reflection_router
      type: reflection
      flow_pattern: [main_agent, critic, main_agent]

  workflow:
    start: main_agent
    edges:
      - from: main_agent
        to: [critic, main_agent]
        router: reflection_router
      - from: critic
        to: [main_agent]
        router: reflection_router
    end: [main_agent]
"""

result = await AriumBuilder().from_yaml(yaml_str=yaml_config).build_and_run(["Your task"])

import asyncio
from flo_ai.arium import AriumBuilder
from flo_ai.arium.memory import PlanAwareMemory
from flo_ai.models.agent import Agent
from flo_ai.llm import OpenAI
from flo_ai.arium.llm_router import create_plan_execute_router

async def cursor_style_workflow():
    llm = OpenAI(model='gpt-4o-mini', api_key='your-api-key')
    
    # Create specialized agents
    planner = Agent(
        name='planner',
        system_prompt='Create detailed execution plans by breaking down tasks into sequential steps.',
        llm=llm
    )
    
    developer = Agent(
        name='developer', 
        system_prompt='Implement features step by step according to execution plans.',
        llm=llm
    )
    
    tester = Agent(
        name='tester',
        system_prompt='Test implementations and validate functionality thoroughly.',
        llm=llm
    )
    
    reviewer = Agent(
        name='reviewer',
        system_prompt='Review completed work and provide final approval.',
        llm=llm
    )
    
    # Create plan-execute router
    plan_router = create_plan_execute_router(
        planner_agent='planner',
        executor_agent='developer',
        reviewer_agent='reviewer',
        additional_agents={'tester': 'Tests implementations and validates quality'},
        llm=llm
    )
    
    # Use PlanAwareMemory for plan state persistence
    memory = PlanAwareMemory()
    
    # Build and run workflow
    result = await (
        AriumBuilder()
        .with_memory(memory)
        .add_agents([planner, developer, tester, reviewer])
        .start_with(planner)
        .add_edge(planner, [developer, tester, reviewer, planner], plan_router)
        .add_edge(developer, [developer, tester, reviewer, planner], plan_router)
        .add_edge(tester, [developer, tester, reviewer, planner], plan_router)
        .add_edge(reviewer, [developer, tester, reviewer, planner], plan_router)
        .end_with(reviewer)
        .build_and_run(["Create a REST API for user authentication with JWT tokens"])
    )
    
    return result

# Alternative: Factory function
from flo_ai.arium.llm_router import create_plan_execute_router

plan_router = create_plan_execute_router(
    planner_agent='planner',
    executor_agent='developer', 
    reviewer_agent='reviewer',
    llm=llm
)

# Minimal plan-execute workflow
yaml_config = """
arium:
  agents:
    - name: planner
      job: "Create execution plans"
      model: {provider: openai, name: gpt-4o-mini}
    - name: executor
      job: "Execute plan steps"
      model: {provider: openai, name: gpt-4o-mini}
    - name: reviewer
      job: "Review final results"
      model: {provider: openai, name: gpt-4o-mini}

  routers:
    - name: simple_plan_router
      type: plan_execute
      agents:
        planner: "Creates plans"
        executor: "Executes steps"
        reviewer: "Reviews results"
      settings:
        planner_agent: planner
        executor_agent: executor
        reviewer_agent: reviewer

  workflow:
    start: planner
    edges:
      - from: planner
        to: [executor, reviewer, planner]
        router: simple_plan_router
      - from: executor
        to: [executor, reviewer, planner]
        router: simple_plan_router
      - from: reviewer
        to: [end]
    end: [reviewer]
"""

result = await AriumBuilder().from_yaml(yaml_str=yaml_config).build_and_run(["Your complex task"])

import asyncio
from typing import Any, Dict, List
from flo_ai.arium import AriumBuilder

async def run_personalized_workflow() -> List[Any]:
    yaml_config = """..."""  # Your YAML configuration with variables
    
    # Define variables for the workflow
    variables: Dict[str, str] = {
        'expert_role': 'Data Scientist',
        'domain': 'machine learning and predictive analytics',
        'output_type': 'technical analysis report',
        'target_audience': 'engineering team',
        'preferred_llm_provider': 'anthropic',
        'model_name': 'claude-3-5-sonnet-20240620',
        'reasoning_style': 'COT',
        'quality_criteria': 'technical accuracy and clarity'
    }
    
    result: List[Any] = await (
        AriumBuilder()
        .from_yaml(yaml_config)
        .build_and_run(
            ["Analyze our customer churn prediction model performance"],
            variables=variables
        )
    )
    
    return result