feat(flo-ai): chat history feature for all agents

flo_ai/examples/agent_builder_usage.py

@@ -1,5 +1,7 @@
 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
@@ -16,7 +18,13 @@ async def example_simple_agent():
         .build()
     )
 
-    response = await agent.run('What is the formula for the area of a circle?')
+    response = await agent.run(
+        [
+            UserMessage(
+                TextMessageContent(text='What is the formula for the area of a circle?')
+            )
+        ]
+    )
     print(f'Simple Agent Response: {response}')
 
 
@@ -59,17 +67,21 @@ async def calculate(operation: str, x: float, y: float) -> float:
         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', temperature=0.7))
+        .with_llm(Anthropic(model='claude-sonnet-4-5', temperature=0.7))
         .with_tools([calculator_tool])
         .with_reasoning(ReasoningPattern.REACT)
         .with_retries(2)
         .build()
     )
 
-    response = await agent_openai.run('Calculate 5 plus 3')
+    response = await agent_openai.run(
+        [UserMessage(TextMessageContent(text='Calculate 5 plus 3'))]
+    )
     print(f'OpenAI Tool Agent Response: {response}')
 
-    response = await agent_claude.run('Calculate 5 plus 3')
+    response = await agent_claude.run(
+        [UserMessage(TextMessageContent(text='Calculate 5 plus 3'))]
+    )
     print(f'Claude Tool Agent Response: {response}')
 
 
@@ -96,7 +108,9 @@ async def example_structured_output():
         .build()
     )
 
-    response = await agent.run('Solve: 2x + 5 = 15')
+    response = await agent.run(
+        [UserMessage(TextMessageContent(text='Solve: 2x + 5 = 15'))]
+    )
     print(f'Structured Output Response: {response}')
 
 

flo_ai/examples/arium_examples.py

@@ -4,31 +4,52 @@
 
 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.tool import flo_tool
 from flo_ai.tool.base_tool import Tool
 from flo_ai.arium.memory import MessageMemory
 
 
+@flo_tool(
+    description='Calculate mathematical expressions',
+    parameter_descriptions={
+        'result': 'The result to print',
+    },
+)
+async def print_result(result: str) -> str:
+    print(f'Result: {result}')
+    return result
+
+
 # Example 1: Simple Linear Workflow
 async def example_linear_workflow():
     """Example of a simple linear workflow: Agent -> Tool -> Agent"""
 
     # Create some example agents and tools (these would be your actual implementations)
-    analyzer_agent = Agent(name='analyzer', prompt='Analyze the input')
-    processing_tool = Tool(name='processor')
-    summarizer_agent = Agent(name='summarizer', prompt='Summarize the results')
+    analyzer_agent = Agent(
+        name='analyzer',
+        system_prompt='Analyze the input',
+        llm=OpenAI(model='gpt-4o-mini'),
+    )
+    summarizer_agent = Agent(
+        name='summarizer',
+        system_prompt='Summarize the results',
+        llm=OpenAI(model='gpt-4o-mini'),
+    )
 
     # Build and run the workflow
     result = await (
         AriumBuilder()
         .add_agent(analyzer_agent)
-        .add_tool(processing_tool)
+        .add_tool(print_result.tool)
         .add_agent(summarizer_agent)
         .start_with(analyzer_agent)
-        .connect(analyzer_agent, processing_tool)
-        .connect(processing_tool, summarizer_agent)
+        .connect(analyzer_agent, print_result.tool)
+        .connect(print_result.tool, summarizer_agent)
         .end_with(summarizer_agent)
-        .build_and_run(['Analyze this text'])
+        .build_and_run([UserMessage(TextMessageContent(text='Analyze this text'))])
     )
 
     return result
@@ -119,17 +140,22 @@ def analysis_router(memory) -> Literal['writer', 'researcher']:
 async def example_convenience_function():
     """Example using the create_arium convenience function"""
 
-    agent1 = Agent(name='agent1', prompt='First agent')
-    agent2 = Agent(name='agent2', prompt='Second agent')
+    agent1 = Agent(
+        name='agent1', system_prompt='First agent', llm=OpenAI(model='gpt-4o-mini')
+    )
+    agent2 = Agent(
+        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(['Hello'])
+        .build_and_run([UserMessage(TextMessageContent(text='Hello'))])
     )
 
     return result
@@ -164,7 +190,6 @@ async def main():
         # result3 = await example_complex_workflow()
         # result4 = await example_convenience_function()
         # result5 = await example_build_and_reuse()
-
         print('Examples completed!')
 
     asyncio.run(main())

flo_ai/examples/arium_yaml_example.py

@@ -8,6 +8,7 @@
 import asyncio
 from typing import Dict, Literal
 from flo_ai.arium.builder import AriumBuilder
+from flo_ai.models import TextMessageContent, UserMessage
 from flo_ai.tool.base_tool import Tool
 from flo_ai.llm import OpenAI
 from flo_ai.arium.memory import BaseMemory
@@ -320,10 +321,14 @@ async def run_simple_example():
     # Run the workflow
     result = await builder.build_and_run(
         [
-            'Machine learning is transforming healthcare by enabling predictive analytics, '
-            'personalized treatment recommendations, and automated medical imaging analysis. '
-            'However, challenges include data privacy concerns, the need for regulatory approval, '
-            'and ensuring AI systems are transparent and unbiased in their decision-making.'
+            UserMessage(
+                TextMessageContent(
+                    text='Machine learning is transforming healthcare by enabling predictive analytics, '
+                    'personalized treatment recommendations, and automated medical imaging analysis. '
+                    'However, challenges include data privacy concerns, the need for regulatory approval, '
+                    'and ensuring AI systems are transparent and unbiased in their decision-making.',
+                )
+            ),
         ]
     )
 
@@ -370,8 +375,16 @@ async def run_complex_example():
 
     result2 = await builder.build_and_run(
         [
-            "Please analyze this text and process it: 'The quick brown fox jumps over the lazy dog. "
-            "This sentence contains every letter of the alphabet at least once.'"
+            UserMessage(
+                TextMessageContent(
+                    text="Please analyze this text and process it: 'The quick brown fox jumps over the lazy dog. ",
+                )
+            ),
+            UserMessage(
+                TextMessageContent(
+                    text="This sentence contains every letter of the alphabet at least once.'",
+                )
+            ),
         ]
     )
 
@@ -394,10 +407,14 @@ async def run_mixed_config_example():
     # Run the workflow
     result = await builder.build_and_run(
         [
-            'Please analyze this business report: Our Q3 revenue increased by 15% compared to Q2, '
-            'driven primarily by strong performance in the software division. However, hardware sales '
-            'declined by 8%. Customer satisfaction scores improved to 4.2/5.0. We recommend focusing '
-            'on digital transformation initiatives and reconsidering the hardware product line.'
+            UserMessage(
+                TextMessageContent(
+                    text='Please analyze this business report: Our Q3 revenue increased by 15% compared to Q2, '
+                    'driven primarily by strong performance in the software division. However, hardware sales '
+                    'declined by 8%. Customer satisfaction scores improved to 4.2/5.0. We recommend focusing '
+                    'on digital transformation initiatives and reconsidering the hardware product line.',
+                )
+            ),
         ]
     )
 

flo_ai/examples/chat_history.py

@@ -1,9 +1,26 @@
 import asyncio
 from typing import Any
 from flo_ai.builder.agent_builder import AgentBuilder
-from flo_ai.llm import OpenAI
+from flo_ai.llm import Gemini
 from flo_ai.models.agent import Agent
-from flo_ai.models.chat_message import ChatMessage
+from flo_ai.models import (
+    AssistantMessage,
+    UserMessage,
+    TextMessageContent,
+)
+from flo_ai.tool import flo_tool
+
+
+@flo_tool(
+    description='Calculate the area of a rectangle',
+    parameter_descriptions={
+        'length': 'Length of the rectangle',
+        'breadth': 'Breadth of the rectangle',
+    },
+)
+async def calculate(length: float, breadth: float) -> float:
+    """Calculate the area of a rectangle."""
+    return length * breadth
 
 
 async def main() -> None:
@@ -12,30 +29,40 @@ async def main() -> None:
         AgentBuilder()
         .with_name('Math Tutor')
         .with_prompt('You are a helpful math tutor.')
-        .with_llm(OpenAI(model='gpt-4o-mini'))
+        .with_llm(Gemini(model='gemini-2.5-flash'))
+        .add_tool(calculate.tool)
         .build()
     )
 
     response: Any = await agent.run(
         [
-            ChatMessage(
-                role='user', content='What is the formula for the area of a circle?'
+            UserMessage(
+                TextMessageContent(
+                    text='What is the formula for the area of a circle?'
+                ),
             ),
-            ChatMessage(
-                role='assistant',
-                content='The formula for the area of a circle is πr^2.',
+            AssistantMessage(
+                TextMessageContent(
+                    text='The formula for the area of a circle is πr^2.'
+                ),
             ),
-            ChatMessage(
-                role='user', content='What is the formula for the area of a rectangle?'
+            UserMessage(
+                TextMessageContent(
+                    text='What is the formula for the area of a rectangle?'
+                )
             ),
-            ChatMessage(
-                role='assistant',
-                content='The formula for the area of a rectangle is length * width.',
+            AssistantMessage(
+                TextMessageContent(
+                    text='The formula for the area of a rectangle is length * width.',
+                ),
             ),
-            ChatMessage(
-                role='user', content='What is the formula for the area of a triangle?'
+            UserMessage(
+                TextMessageContent(
+                    text='What is the area of a rectable of length <length> and breadth <breadth>',
+                ),
             ),
-        ]
+        ],
+        variables={'length': 10, 'breadth': 70},
     )
     print(f'Response: {response}')
 

flo_ai/examples/cot_agent_example.py

@@ -4,6 +4,8 @@
 """
 
 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
@@ -74,9 +76,15 @@ async def main():
 
     # Test questions
     questions = [
-        'What is 15 + 27?',
-        'If I have 100 apples and I give away 23, then buy 15 more, how many do I have?',
-        'Calculate 8 * 7 and then add 12 to the result.',
+        UserMessage(TextMessageContent(text='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.')
+        ),
     ]
 
     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
+from flo_ai.models.agent import Agent, TextMessageContent, UserMessage
 from flo_ai.models.base_agent import ReasoningPattern
 from flo_ai.llm.openai_llm import OpenAI
 import os
@@ -33,9 +33,21 @@ async def main():
 
     # Test questions that require step-by-step reasoning
     questions = [
-        '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?',
-        '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?',
-        'Explain why the sky appears blue during the day but red during sunset.',
+        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?',
+            )
+        ),
+        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?',
+            )
+        ),
+        UserMessage(
+            TextMessageContent(
+                text='Explain why the sky appears blue during the day but red during sunset.',
+            )
+        ),
     ]
 
     print('=== Conversational Chain of Thought (CoT) Reasoning Demo ===\n')

flo_ai/examples/custom_plan_execute_demo.py

@@ -11,6 +11,7 @@
 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
 
 
@@ -98,12 +99,17 @@ async def main():
     )
 
     # Execute task
-    task = 'Research the impact of AI on software development productivity'
+    task = UserMessage(
+        TextMessageContent(
+            type='text',
+            text='Research the impact of AI on software development productivity',
+        )
+    )
     print(f'📋 Task: {task}')
     print('🔄 Executing custom research workflow...\n')
 
     try:
-        result = await arium.run([task])
+        result = await arium.run(task)
 
         print('\n' + '=' * 40)
         print('🎉 CUSTOM WORKFLOW COMPLETED!')