8000 feat: add AgentBrain abstract base class by AkashGowdaNC · Pull Request #73 · ARPAHLS/rooms · GitHub
[go: up one dir, main page]

Skip to content

feat: add AgentBrain abstract base class - #73

Open
AkashGowdaNC wants to merge 4 commits into
ARPAHLS:mainfrom
AkashGowdaNC:feat/issue-15-agentbrain-abc
Open

AkashGowdaNC wants to merge 4 commits into
ARPAHLS:mainfrom
AkashGowdaNC:feat/issue-15-agentbrain-abc

Conversation

@AkashGowdaNC
Copy link
Copy Markdown

Refs #15

What this PR does

Adds the AgentBrain abstract base class described in #15, plus tests.

This is step 1 of the plan discussed in the issue — the ABC and its tests only, with no production refactor yet. Follow-ups will move the LiteLLM path, then the custom-function path, behind this interface, and wire AgentConfig to pick a brain.

Why this is useful

The Agent class currently handles LiteLLM routing, custom-function dispatch, and Skillware tool-use in one place. Introducing the ABC first lets later PRs move each backend behind it without touching callers, keeping each change reviewable.

Changes

  • rooms/brains/__init__.py: new AgentBrain ABC with a generate() method matching the return shape of Agent.generate_response_with_events today
  • tests/test_agent_brain.py: 5 tests covering abstractness, minimal subclasses, result shape, and override-param pass-through

How to verify

@AkashGowdaNC
Copy link
Copy Markdown
Author

@rosspeili PR 1 is up: the AgentBrain ABC + tests. Ready when you are.

@rosspeili
Copy link
Copy Markdown
Contributor

Nice work on step1 @AkashGowdaNC!

A few small polish items before we land this and start Step 2:

  1. EOF newlines: Add a trailing newline to rooms/brains/__init__.py and tests/test_agent_brain.py (resolves Flake8 W292).
  2. Module exports: Add __all__ = ["AgentBrain"] to rooms/brains/__init__.py.
  3. Constructor & typing contract for Steps 2–4:
    • Document or declare the expected __init__ signature (e.g. accepting config: AgentConfig and optional skill_runtime) so concrete brains (LiteLLMBrain, CustomFunctionBrain) and external plugins have a uniform initialization contract.
    • (Optional/Recommended) Consider defining a TypedDict for the generate() return dict (content: str, skill_events: List[Dict[str, Any]]) to enforce typing across all future brains.

No changelog entry needed for this PR since it's internal scaffolding, we'll add one once the brains are wired to AgentConfig in Step 4. Push those tweaks and we're ready to merge.

@AkashGowdaNC
Copy link
Copy Markdown
Author

Thanks @rosspeili — pushed all four:

  • Added trailing newlines to both files (Flake8 W292 is clean)
  • Added __all__ = ["AgentBrain", "BrainResult"] to rooms/brains/__init__.py
  • Documented the constructor contract (config: AgentConfig, optional skill_runtime) in the class docstring so Step 2–4 brains and external plugins have a uniform init signature
  • Added a BrainResult TypedDict for the generate() return shape

Tests still pass (5/5). Ready when you are.

@rosspeili
Copy link
Copy Markdown
Contributor

Looks great @AkashGowdaNC! All 4 review items were addressed cleanly, the constructor contract and BrainResult TypedDict set up Steps 2–4 nicely, and all 82 tests pass.

Just one tiny formatting nit on line 5 of tests/test_agent_brain.py:

from rooms.brains import AgentBrain , BrainResult

There's a stray space before the comma (E203), and BrainResult is imported but unused in the test file (F401).

Please change it to:

from rooms.brains import AgentBrain

Once that's pushed, we'll merge this and open the floor for Step 2 (LiteLLMBrain).

@AkashGowdaNC
Copy link
Copy Markdown
Author

Thanks @rosspeili — fixed:

  • Removed the stray space before the comma (E203)
  • Dropped the unused BrainResult import (F401)

The file now imports from rooms.brains import AgentBrain. Tests still pass (5/5) and Flake8 is clean.

@AkashGowdaNC
Copy link
Copy Markdown
Author

@rosspeili — friendly nudge on this one. The last review round was addressed (import fix pushed in 7f53bb5) and the workflow is waiting on approval. Happy to wait if you're busy — just flagging in case the notification got buried.

@rosspeili
Copy link
Copy Markdown
Contributor

Thanks @AkashGowdaNC, small change of plan: all of #15 lands in this one PR. Use one commit per step so we can review commit by commit if it's easier for you. We'll review and merge it all together once it's complete.

1. Base contract (rooms/brains/base.py)

  • Move the AgentBrain and BrainResult code to base.py. Keep __init__.py for re-exports only. Why: concrete brains will import AgentConfig and SkillRuntime, and keeping logic out of __init__ avoids circular imports.
  • Put the real __init__(self, config: AgentConfig, skill_runtime: Optional[SkillRuntime] = None) on the base class, not just in the docstring.
  • generate(context_messages, override_params) receives the conversation without the system message. Each brain builds its own system prompt from the persona, custom_instructions and skill instructions. You can put a shared _build_system_prompt() helper on the base class. Why: that's how Agent works today, and skill instructions depend on the runtime the brain owns.

2. LiteLLMBrain (rooms/brains/litellm_brain.py)

  • Move the LiteLLM path out of Agent.generate_response_with_events as-is: params, timeout, tool calls, the max_skill_calls_per_turn limit, the second completion call, and the error/timeout messages.
  • Skillware tool use must keep working exactly as it does now.

3. CustomFunctionBrain (rooms/brains/custom_function_brain.py)

  • Move _execute_custom_function here, with the same timeout and error strings. It returns skill_events: [].

4. Choosing a brain + refactoring Agent

  • Add a small resolver (resolve_brain(config)) that maps ModelType to a brain class. Optionally let config name a plugin brain by dotted path (my_pkg.brains:MyBrain). Why: feat: plugin-based architecture for agent inference #15's goal is that third-party brains work without touching core code.
  • Agent becomes a thin wrapper that hands the call to self.brain. Keep generate_response, generate_response_with_events, consume_last_skill_events and __repr__ working exactly as now. Session, CLI and storage code must not change.

Tests

  • Keep the existing suite green without editing existing tests. That's our proof nothing broke for callers.
  • Add tests per brain, with litellm.completion mocked. Cover the tool-call flow, the too-many-calls guard, timeouts, the custom-function timeout and the missing file/function cases, the resolver mapping and the plugin path, and the error when the plugin path is wrong.

Docs & complementary files

  • README.md: add rooms/brains/ to the project tree.
  • CONTRIBUTING.md: add rooms/brains/ to the "Core framework" row.
  • docs/: a short BRAINS.md covering the contract, the built-in brains and how to write a plugin. Link it from docs/README.md.
  • docs/SETTINGS.md: document any new config/YAML key.
  • CHANGELOG.md → [Unreleased]: under Added, the pluggable AgentBrain architecture (feat: plugin-based architecture for agent inference #15); under Changed, Agent now hands inference to a brain.

Tips

  • Do a pure move first, then clean up. Don't change behaviour in the same commit as a move.
  • Run flake8 and pytest before every push. Make sure every file ends with a newline.
  • Update the PR description with a checklist of the 4 steps above.

Once it's all in, we review the whole thing and merge. 🚀

@AkashGowdaNC
Copy link
Copy Markdown
Author

@rosspeili Thanks for the plan — that's a much better structure. I want to be upfront: this is a multi-day refactor for me, so I'll push one commit per step.

Commit 1 is up: refactor: move AgentBrain and BrainResult to brains/base.py. Pure move — base.py has the ABC, __init__.py re-exports only, no behaviour change, 5 tests pass.

Starting on Commit 2 (__init__ + _build_system_prompt on the base class) next.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

2C16
0