fix: add ty and fix types

[vishnurk6247]

Summary by CodeRabbit

• Refactor

  • Updated API to accept string inputs directly instead of lists in workflow execution methods.
  • Changed agent and workflow return values from strings to message-based objects for richer context.
  • Standardized LLM configuration parameters across provider implementations for consistency.
  • Expanded node type system to support additional workflow patterns.

• Bug Fixes

  • Added validation to prevent missing LLM configuration errors early in initialization.
  • Enhanced error handling with stricter parameter requirements and guards.

• Chores

  • Updated development dependencies.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai[bot]]

Walkthrough

Consolidates input shapes to plain strings across examples and builders, introduces AriumNodeType and FunctionNode, tightens memory and execution-plan types, changes Agent.run to return List[BaseMessage], updates LLM/router signatures (functions/output_schema/optional execution_context), and adjusts builders, YAML parsing, and tests to match these API changes.

Changes

Cohort / File(s)	Summary
Example input & message updates flo_ai/examples/...	Replaced list-wrapped prompts and TextMessageContent wrappers with single str inputs; examples updated to construct UserMessage/inputs as plain strings and to use Agent(name=..., system_prompt=..., llm=...) where applicable.
Builder / Arium public API flo_ai/flo_ai/arium/builder.py, flo_ai/flo_ai/arium/arium.py	build_and_run()/Arium.run() now accept scalar str inputs (wrapped to UserMessage internally); memory narrowed to MessageMemory; variables defaults tightened to {}; improved YAML loading and robust node/string resolution.
Arium core types & nodes flo_ai/flo_ai/arium/base.py, flo_ai/flo_ai/arium/nodes.py	Added AriumNodeType (includes FunctionNode); self.nodes and many APIs updated to use AriumNodeType; stricter runtime checks for executable nodes and updated edge/start/end semantics.
Routing / routers flo_ai/flo_ai/arium/llm_router.py	execution_context made Optional[dict] across router APIs; route/get_routing_prompt signatures adjusted; create_llm_router/llm_router decorator propagate optional execution_context; internal prompts guarded against None.
Memory & execution plan flo_ai/flo_ai/arium/memory.py	MessageMemoryItem constructor changed (result required), MessageMemory moved off generic base, PlanAwareMemory messages now dicts, and ExecutionPlan step completion/failure signatures tightened.
Agent & base agent flo_ai/flo_ai/models/agent.py, flo_ai/flo_ai/models/base_agent.py	Agent.run, _run_conversational, _run_with_tools now return List[BaseMessage]; history and message-handling updated; system_prompt and initialization paths adjusted.
LLM implementations & factory helpers flo_ai/flo_ai/llm/*.py, flo_ai/flo_ai/helpers/llm_factory.py	Constructor params (api_key/base_url/project/location) made Optional[str]; generate/stream gained functions and output_schema params in many LLMs; image formatting return types relaxed to Any/provider-specific; factories enforce required endpoints/keys.
Utilities & tooling flo_ai/flo_ai/tool/flo_tool.py, flo_ai/flo_ai/utils/variable_extractor.py, flo_ai/flo_ai/telemetry/telemetry.py	Safer attribute access (getattr, type-ignore), resolve_variables(text: str, variables: Optional[dict]) -> str now strictly returns str and errors for missing vars, telemetry params typed Optional[str].
YAML/format parser & agent builder flo_ai/flo_ai/formatter/yaml_format_parser.py, flo_ai/flo_ai/builder/agent_builder.py	YAML path support added, file-read encoding enforced, array/literal helpers refactored, and AgentBuilder.from_yaml accepts yaml_str or yaml_file with input validation.
Mocks, tests & examples adjustments flo_ai/tests/..., flo_ai/examples/...	MockLLM updated to BaseLLM shape; tests and examples adjusted for new LLM/Agent signatures, added type-ignore comments, VLLM base_url guard helper added, and minor defaults (e.g., country: str = '').
Dev dependency flo_ai/pyproject.toml	Added dev dependency ty >= 0.0.1a28.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant AriumBuilder
    participant Arium
    participant Memory
    participant Agent
    participant LLM

    User->>AriumBuilder: build_and_run(inputs: str)
    AriumBuilder->>Arium: run(inputs: str, variables: {})
    Arium->>Arium: _resolve_inputs -> wrap into UserMessage
    Arium->>Memory: add(MessageMemoryItem(node, result=UserMessage))
    Arium->>Agent: run(messages: List[BaseMessage], variables: Dict)
    Agent->>LLM: generate/stream (functions?, output_schema?)
    LLM-->>Agent: List[BaseMessage] (conversation history / function results)
    Agent->>Memory: add(results as MessageMemoryItem)
    Agent-->>Arium: return List[BaseMessage]
    Arium-->>AriumBuilder: aggregated result
    AriumBuilder-->>User: final output (List[BaseMessage])

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Areas requiring extra attention:

• flo_ai/flo_ai/models/agent.py and base_agent.py — callers must adapt to List[BaseMessage] return values and any string-expected code.
• flo_ai/flo_ai/arium/memory.py and arium.py — ensure MessageMemoryItem constructor and dict-based PlanAwareMemory interop are correct everywhere.
• flo_ai/flo_ai/arium/base.py, builder.py, nodes.py — AriumNodeType usage and string-to-node resolution paths need careful validation (foreach/start/end/edges).
• LLM implementations (flo_ai/flo_ai/llm/*) — verify generate/stream functions/output_schema handling and initialization guards (api_key/base_url) for all providers.
• Examples & tests — confirm mocks and tests reflect new LLM signatures and Agent output shapes; watch for remaining TextMessageContent usages.

"🐰 I hopped through types and gentle refactors bright,
Strings now travel singly, messages take flight.
Nodes and memories neatly sorted in rows,
I munch small mismatches where the new code grows.
Hooray — a tiny rabbit wiggle under moonlight."

Pre-merge checks and finishing touches

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 58.45% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.
Title check	❓ Inconclusive	The title 'fix: add ty and fix types' is vague and non-descriptive. While it mentions 'fix' and 'types', it does not clearly communicate the main changes or purpose of this extensive pull request.	Provide a more specific title that captures the primary intent, such as 'refactor: update type annotations and API signatures across codebase' or 'fix: standardize input handling and improve type safety'.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

✨ Finishing touches

• [ ] 📝 Generate docstrings

🧪 Generate unit tests (beta)

• [ ] Create PR with unit tests
• [ ] Post copyable unit tests in a comment
• [ ] Commit unit tests in branch fix/types

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}