fix: add ty and fix types

flo_ai/examples/agent_builder_usage.py

@@ -1,7 +1,6 @@
 import asyncio
 from flo_ai import UserMessage
 from flo_ai.builder.agent_builder import AgentBuilder
-from flo_ai.models import TextMessageContent
 from flo_ai.tool.base_tool import Tool
 from flo_ai.models.base_agent import ReasoningPattern
 from flo_ai.llm.openai_llm import OpenAI
@@ -19,11 +18,7 @@ async def example_simple_agent():
     )
 
     response = await agent.run(
-        [
-            UserMessage(
-                TextMessageContent(text='What is the formula for the area of a circle?')
-            )
-        ]
+        [UserMessage('What is the formula for the area of a circle?')]
     )
     print(f'Simple Agent Response: {response}')
 
@@ -74,14 +69,10 @@ async def calculate(operation: str, x: float, y: float) -> float:
         .build()
     )
 
-    response = await agent_openai.run(
-        [UserMessage(TextMessageContent(text='Calculate 5 plus 3'))]
-    )
+    response = await agent_openai.run([UserMessage('Calculate 5 plus 3')])
     print(f'OpenAI Tool Agent Response: {response}')
 
-    response = await agent_claude.run(
-        [UserMessage(TextMessageContent(text='Calculate 5 plus 3'))]
-    )
+    response = await agent_claude.run([UserMessage('Calculate 5 plus 3')])
     print(f'Claude Tool Agent Response: {response}')
 
 
@@ -108,9 +99,7 @@ async def example_structured_output():
         .build()
     )
 
-    response = await agent.run(
-        [UserMessage(TextMessageContent(text='Solve: 2x + 5 = 15'))]
-    )
+    response = await agent.run([UserMessage('Solve: 2x + 5 = 15')])
     print(f'Structured Output Response: {response}')
 
 

flo_ai/examples/arium_examples.py

@@ -1,11 +1,14 @@
 """
 Examples demonstrating how to use the AriumBuilder pattern for creating and running Arium workflows.
+
+Note: Both string inputs and list of UserMessage inputs are supported:
+    - String: build_and_run('Hello') - simpler for single text inputs
+    - List: build_and_run([UserMessage('Hello')]) - required for multiple messages or complex message types
 """
 
 from typing import Literal
 from flo_ai.arium import AriumBuilder, create_arium
 from flo_ai.llm import OpenAI
-from flo_ai.models import TextMessageContent, UserMessage
 from flo_ai.models.agent import Agent
 from flo_ai.arium.nodes import FunctionNode
 from flo_ai.arium.memory import MessageMemory, MessageMemoryItem
@@ -47,7 +50,7 @@ async def example_linear_workflow():
         .connect(analyzer_agent, processing_function_node)
         .connect(processing_function_node, summarizer_agent)
         .end_with(summarizer_agent)
-        .build_and_run([UserMessage(TextMessageContent(text='Analyze this text'))])
+        .build_and_run('Analyze this text')
     )
 
     return result
@@ -58,14 +61,22 @@ async def example_branching_workflow():
     """Example of a branching workflow with conditional routing"""
 
     # Create agents and function nodes
-    classifier_agent = Agent(name='classifier', prompt='Classify the input type')
+    classifier_agent = Agent(
+        name='classifier',
+        system_prompt='Classify the input type',
+        llm=OpenAI(model='gpt-4o-mini'),
+    )
     text_processor_node = FunctionNode(
         name='text_processor', description='Process text', function=lambda x: x
     )
     image_processor_node = FunctionNode(
         name='image_processor', description='Process image', function=lambda x: x
     )
-    final_agent = Agent(name='final', prompt='Provide final response')
+    final_agent = Agent(
+        name='final',
+        system_prompt='Provide final response',
+        llm=OpenAI(model='gpt-4o-mini'),
+    )
 
     # Router function for conditional branching
     def content_router(memory) -> Literal['text_processor', 'image_processor']:
@@ -101,7 +112,7 @@ def content_router(memory) -> Literal['text_processor', 'image_processor']:
         .connect(text_processor_node, final_agent)
         .connect(image_processor_node, final_agent)
         .end_with(final_agent)
-        .build_and_run(['Process this content'])
+        .build_and_run('Process this content')
     )
     return result
 
@@ -111,10 +122,26 @@ async def example_complex_workflow():
     """Example of a more complex workflow with multiple agents and function nodes"""
 
     # Create multiple agents and function nodes
-    input_agent = Agent(name='input_handler', prompt='Handle initial input')
-    researcher_agent = Agent(name='researcher', prompt='Research the topic')
-    analyzer_agent = Agent(name='analyzer', prompt='Analyze findings')
-    writer_agent = Agent(name='writer', prompt='Write the final report')
+    input_agent = Agent(
+        name='input_handler',
+        system_prompt='Handle initial input',
+        llm=OpenAI(model='gpt-4o-mini'),
+    )
+    researcher_agent = Agent(
+        name='researcher',
+        system_prompt='Research the topic',
+        llm=OpenAI(model='gpt-4o-mini'),
+    )
+    analyzer_agent = Agent(
+        name='analyzer',
+        system_prompt='Analyze findings',
+        llm=OpenAI(model='gpt-4o-mini'),
+    )
+    writer_agent = Agent(
+        name='writer',
+        system_prompt='Write the final report',
+        llm=OpenAI(model='gpt-4o-mini'),
+    )
 
     search_function_node = FunctionNode(
         name='search_function', description='Search the web', function=lambda x: x
@@ -151,7 +178,7 @@ def analysis_router(memory) -> Literal['writer', 'researcher']:
     )
 
     # Run the workflow
-    result = await arium.run(['Research and write a report on AI trends'])
+    result = await arium.run('Research and write a report on AI trends')
     return result
 
 
@@ -166,15 +193,14 @@ async def example_convenience_function():
         name='agent2', system_prompt='Second agent', llm=OpenAI(model='gpt-4o-mini')
     )
 
-    # Fix: Use proper InputMessage format for consistency
     result = await (
         create_arium()
         .add_agent(agent1)
         .add_agent(agent2)
         .start_with(agent1)
         .connect(agent1, agent2)
         .end_with(agent2)
-        .build_and_run([UserMessage(TextMessageContent(text='Hello'))])
+        .build_and_run('Hello')
     )
 
     return result
@@ -184,14 +210,18 @@ async def example_convenience_function():
 async def example_build_and_reuse():
     """Example of building an Arium once and reusing it"""
 
-    agent = Agent(name='echo_agent', prompt='Echo the input')
+    agent = Agent(
+        name='echo_agent',
+        system_prompt='Echo the input',
+        llm=OpenAI(model='gpt-4o-mini'),
+    )
 
     # Build the Arium
     arium = AriumBuilder().add_agent(agent).start_with(agent).end_with(agent).build()
 
     # Run it multiple times with different inputs
-    result1 = await arium.run(['First input'])
-    result2 = await arium.run(['Second input'])
+    result1 = await arium.run('First input')
+    result2 = await arium.run('Second input')
 
     return result1, result2
 
@@ -201,18 +231,16 @@ async def example_function_nodes_with_filters():
     """Workflow of only FunctionNodes; each uses input_filter to read from specific nodes."""
 
     # Define simple functions as nodes
-    async def pass_through(inputs: List[BaseMessage] | str, variables=None, **kwargs):
+    async def pass_through(inputs: List[BaseMessage], variables=None, **kwargs):
         return inputs[-1].content
 
-    async def capitalize_last(
-        inputs: List[BaseMessage] | str, variables=None, **kwargs
-    ):
+    async def capitalize_last(inputs: List[BaseMessage], variables=None, **kwargs):
         return str(inputs[-1].content).capitalize()
 
-    async def uppercase_all(inputs: List[BaseMessage] | str, variables=None, **kwargs):
+    async def uppercase_all(inputs: List[BaseMessage], variables=None, **kwargs):
         return ' '.join([str(x.content).upper() for x in inputs])
 
-    async def summarize(inputs: List[BaseMessage] | str, variables=None, **kwargs):
+    async def summarize(inputs: List[BaseMessage], variables=None, **kwargs):
         return f"count={len(inputs or [])} last={(str(inputs[-1].content) if inputs else '')}"
 
     # Create four FunctionNodes with input filters
@@ -262,7 +290,7 @@ def router(memory: MessageMemory) -> Literal['function2', 'function4']:
         .connect(t2, t3)
         .connect(t3, t1)
         .end_with(t4)
-        .build_and_run(['hello world'])
+        .build_and_run('hello world')
     )
 
     return result

flo_ai/examples/arium_linear_usage.py

@@ -33,7 +33,7 @@ async def simple_example():
         .start_with(greeter)
         .connect(greeter, responder)  # Direct connection
         .end_with(responder)
-        .build_and_run(["Hello, I'd like to learn about Python programming!"])
+        .build_and_run("Hello, I'd like to learn about Python programming!")
     )
 
     print('Simple Example Result:')
@@ -87,14 +87,12 @@ async def main():
         .connect(content_analyst, summary_generator)  # Direct connection
         .end_with(summary_generator)
         .build_and_run(
-            [
-                'Machine learning is revolutionizing various industries. '
-                'From healthcare to finance, AI systems are being deployed '
-                'to automate processes, improve decision-making, and enhance '
-                'customer experiences. However, challenges remain around '
-                'data privacy, algorithmic bias, and the need for skilled '
-                'professionals to manage these systems effectively.'
-            ]
+            """Machine learning is revolutionizing various industries.
+                From healthcare to finance, AI systems are being deployed
+                to automate processes, improve decision-making, and enhance
+                customer experiences. However, challenges remain around
+                data privacy, algorithmic bias, and the need for skilled
+                professionals to manage these systems effectively."""
         )
     )
 

flo_ai/examples/arium_yaml_example.py

@@ -580,7 +580,7 @@ async def run_complex_example():
     # Test with mathematical content
     print('\nTesting with mathematical content:')
     result1 = await builder.build_and_run(
-        ['Please calculate the sum of 25 and 17, then multiply the result by 3.']
+        'Please calculate the sum of 25 and 17, then multiply the result by 3.'
     )
 
     print('Result:')
@@ -693,7 +693,7 @@ async def run_mixed_nodes_example():
     # Test with mathematical content
     print('\nTesting with mathematical content:')
     result1 = await builder.build_and_run(
-        ['Please calculate the sum of 15 and 27, then multiply by 2.']
+        'Please calculate the sum of 15 and 27, then multiply by 2.'
     )
 
     print('Result:')
@@ -820,14 +820,13 @@ async def run_prebuilt_agents_example():
 
     # Run the workflow
     result = await builder.build_and_run(
-        [
-            'The global renewable energy market reached $1.1 trillion in 2023, representing a 15% '
-            'increase from the previous year. Solar energy dominated with 45% market share, followed '
-            'by wind energy at 35%. Government incentives in Europe and Asia drove significant growth, '
-            'while corporate sustainability commitments increased private sector investment. However, '
-            'supply chain challenges and raw material costs remain key obstacles. Industry experts '
-            'predict continued expansion, with the market expected to reach $1.8 trillion by 2030.'
-        ]
+        """The global renewable energy market reached $1.1 trillion in 2023, representing a 15%
+            increase from the previous year. Solar energy dominated with 45% market share, followed by wind energy at 35%. 
+            Government incentives in Europe and Asia drove significant growth,
+            by wind energy at 35%. Government incentives in Europe and Asia drove significant growth
+            while corporate sustainability commitments increased private sector investment. However
+            supply chain challenges and raw material costs remain key obstacles. Industry experts
+            predict continued expansion, with the market expected to reach $1.8 trillion by 2030."""
     )
 
     print('Result:')

flo_ai/examples/cot_agent_example.py

@@ -5,7 +5,6 @@
 
 import asyncio
 from flo_ai import UserMessage
-from flo_ai import TextMessageContent
 from flo_ai.models.agent import Agent
 from flo_ai.models.base_agent import ReasoningPattern
 from flo_ai.llm.openai_llm import OpenAI
@@ -76,15 +75,11 @@ async def main():
 
     # Test questions
     questions = [
-        UserMessage(TextMessageContent(text='What is 15 + 27?')),
+        UserMessage('What is 15 + 27?'),
         UserMessage(
-            TextMessageContent(
-                text='If I have 100 apples and I give away 23, then buy 15 more, how many do I have?',
-            )
-        ),
-        UserMessage(
-            TextMessageContent(text='Calculate 8 * 7 and then add 12 to the result.')
+            'If I have 100 apples and I give away 23, then buy 15 more, how many do I have?'
         ),
+        UserMessage('Calculate 8 * 7 and then add 12 to the result.'),
     ]
 
     print('=== Chain of Thought (CoT) Reasoning Demo ===\n')

flo_ai/examples/cot_conversational_example.py

@@ -4,7 +4,7 @@
 """
 
 import asyncio
-from flo_ai.models.agent import Agent, TextMessageContent, UserMessage
+from flo_ai.models.agent import Agent, UserMessage
 from flo_ai.models.base_agent import ReasoningPattern
 from flo_ai.llm.openai_llm import OpenAI
 import os
@@ -34,19 +34,13 @@ async def main():
     # Test questions that require step-by-step reasoning
     questions = [
         UserMessage(
-            TextMessageContent(
-                text='If a train leaves station A at 2 PM traveling 60 mph and another train leaves station B at 3 PM traveling 80 mph, and the stations are 300 miles apart, when will they meet?',
-            )
+            'If a train leaves station A at 2 PM traveling 60 mph and another train leaves station B at 3 PM traveling 80 mph, and the stations are 300 miles apart, when will they meet?'
         ),
         UserMessage(
-            TextMessageContent(
-                text='A store has a 20% discount on all items. If a customer buys 3 items that originally cost $50, $30, and $20, what is the final total after the discount?',
-            )
+            'A store has a 20% discount on all items. If a customer buys 3 items that originally cost $50, $30, and $20, what is the final total after the discount?'
         ),
         UserMessage(
-            TextMessageContent(
-                text='Explain why the sky appears blue during the day but red during sunset.',
-            )
+            'Explain why the sky appears blue during the day but red during sunset.'
         ),
     ]
 

flo_ai/examples/custom_plan_execute_demo.py

@@ -11,7 +11,6 @@
 from flo_ai.arium.memory import PlanAwareMemory
 from flo_ai.arium.llm_router import create_plan_execute_router
 from flo_ai.arium import AriumBuilder
-from flo_ai.models import TextMessageContent, UserMessage
 from flo_ai.models.plan_agents import PlannerAgent, ExecutorAgent
 
 
@@ -87,7 +86,7 @@ async def main():
     # Build workflow
     arium = (
         AriumBuilder()
-        .with_memory(memory)
+        .with_memory(memory)  # type: ignore[arg-type]
         .add_agents(agents)
         .start_with(planner)
         .add_edge(planner, agents, router)
@@ -99,12 +98,7 @@ async def main():
     )
 
     # Execute task
-    task = UserMessage(
-        TextMessageContent(
-            type='text',
-            text='Research the impact of AI on software development productivity',
-        )
-    )
+    task = 'Research the impact of AI on software development productivity'
     print(f'📋 Task: {task}')
     print('🔄 Executing custom research workflow...\n')
 

flo_ai/examples/document_processing_example.py

@@ -13,7 +13,7 @@
 import os
 import base64
 from pathlib import Path
-from flo_ai.models import DocumentMessageContent, TextMessageContent, UserMessage
+from flo_ai.models import DocumentMessageContent, UserMessage
 from reportlab.lib.pagesizes import letter
 from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
@@ -288,7 +288,7 @@ async def example_2_document_workflow():
         result = await workflow.run(
             [
                 document,
-                UserMessage(TextMessageContent(text='process this document')),
+                UserMessage('process this document'),
             ]
         )