Added streaming support for tool call agent #189
Conversation
WalkthroughThis pull request updates two primary components. In the ResearchAgent class, the Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Agent as ToolCallAgent
participant Logger
Client->>Agent: execute_stream(input_data)
Agent->>Logger: Log initialization & dependency checks
Agent->>Agent: Initialize LLM and Composio tools
Agent->>Logger: Log action processing stages
Agent-->>Client: Yield JSON response (status & logs)
Possibly related PRs
Suggested reviewers
Poem
Tip ⚡🧪 Multi-step agentic review comment chat (experimental)
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
athina/steps/tool_call_agent.py (1)
1-2: Remove unused imports for cleanliness.
Static analysis indicatesUnionandTypedDictare never used. Consider removing them:- from typing import Any, Dict, Union, Optional, List, AsyncGenerator, Literal, TypedDict + from typing import Any, Dict, Optional, List, AsyncGenerator, Literal🧰 Tools
🪛 Ruff (0.8.2)
2-2:
typing.Unionimported but unusedRemove unused import
(F401)
2-2:
typing.TypedDictimported but unusedRemove unused import
(F401)
🪛 GitHub Actions: Flake8, Pyflakes and Compileall Linter
[warning] 2-1: Imported modules are unused: 'typing.Union', 'typing.TypedDict', 'os'.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
athina/steps/research_agent_step.py(3 hunks)athina/steps/tool_call_agent.py(12 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
athina/steps/tool_call_agent.py
2-2: typing.Union imported but unused
Remove unused import
(F401)
2-2: typing.TypedDict imported but unused
Remove unused import
(F401)
73-73: Local variable execution_time is assigned to but never used
Remove assignment to unused variable execution_time
(F841)
🪛 GitHub Actions: MyPy static type checker
athina/steps/tool_call_agent.py
[error] 64-64: Signature of "_create_step_result" incompatible with supertype "Step" [override]
[error] 64-64: This violates the Liskov substitution principle
[error] 64-64: See https://mypy.readthedocs.io/en/stable/common_issues.html#incompatible-overrides
[error] 161-161: Signature of "answer" incompatible with supertype "QuestionAnswerer" [override]
[error] 161-161: This violates the Liskov substitution principle
[error] 161-161: See https://mypy.readthedocs.io/en/stable/common_issues.html#incompatible-overrides
[error] 188-188: Signature of "_create_step_result" incompatible with supertype "Step" [override]
[error] 188-188: This violates the Liskov substitution principle
[error] 188-188: See https://mypy.readthedocs.io/en/stable/common_issues.html#incompatible-overrides
[error] 310-310: Library stubs not installed for "requests" [import-untyped]
🪛 GitHub Actions: Flake8, Pyflakes and Compileall Linter
athina/steps/tool_call_agent.py
[warning] 2-1: Imported modules are unused: 'typing.Union', 'typing.TypedDict', 'os'.
[warning] 73-9: Local variable 'execution_time' is assigned to but never used.
athina/steps/research_agent_step.py
[error] 331-331: F541 f-string is missing placeholders
[error] 407-407: F541 f-string is missing placeholders
[error] 448-448: F541 f-string is missing placeholders
[error] 659-659: F541 f-string is missing placeholders
[error] 969-969: F541 f-string is missing placeholders
🔇 Additional comments (19)
athina/steps/research_agent_step.py (3)
5-5: No concerns regarding added import types.
This import addition appears to be used later in the code for asynchronous operations, and it aligns with the function signatures referencingAsyncGenerator.
449-449: Safer iteration using.get().
Usingupdated_statements.get("evaluation", [])helps avoidKeyErrorif"evaluation"is absent.
512-512: Factual integrity guideline is clear.
This explicit requirement to avoid inventing information aligns with the broader project goals.athina/steps/tool_call_agent.py (16)
4-4: No issues with JSON import.
This import is legitimately used later for serialization.
15-17: Logger setup looks good.
Defining a module-level logger early is a standard best practice.🧰 Tools
🪛 GitHub Actions: Flake8, Pyflakes and Compileall Linter
[warning] 73-9: Local variable 'execution_time' is assigned to but never used.
19-32: StreamLogHandler implementation is straightforward.
This class correctly accumulates log messages for retrieval, which is helpful for streaming logs.🧰 Tools
🪛 GitHub Actions: Flake8, Pyflakes and Compileall Linter
[warning] 73-9: Local variable 'execution_time' is assigned to but never used.
52-52: New attribute for log handler is appropriate.
Storing thestream_log_handlerat the class level ensures easy access throughout.🧰 Tools
🪛 GitHub Actions: Flake8, Pyflakes and Compileall Linter
[warning] 73-9: Local variable 'execution_time' is assigned to but never used.
54-56: Pydantic Config update is fine.
Allowing arbitrary types is often necessary when dealing with complex objects like log handlers.🧰 Tools
🪛 GitHub Actions: Flake8, Pyflakes and Compileall Linter
[warning] 73-9: Local variable 'execution_time' is assigned to but never used.
57-62: Logging initialization in init method makes sense.
Ensures the log handler is attached once per agent instance.🧰 Tools
🪛 GitHub Actions: Flake8, Pyflakes and Compileall Linter
[warning] 73-9: Local variable 'execution_time' is assigned to but never used.
82-94: Robust import error handling.
Importingcomposio_llamaindexinside this block ensures graceful failures when the dependency is missing.
122-132: Clear “no actions” error path.
Returning an error early prevents malformed agent executions.
135-145: Prompt validation is solid.
This block properly halts execution if no prompt is provided.
149-150: Environment loading is acceptable.
Loading environment variables here is consistent with typical usage.
154-157: Entity-specific toolset initialization.
Allowing a customentity_idfurther tailors the toolset usage; the logging is clear.
180-184: Verbose logging for actions.
Notifying how many actions/tools are retrieved helps with debugging the agent’s tool usage.
188-189: Graceful failure when no valid tools are found.
Immediate error response ensures no futile attempts to run the agent without tools.🧰 Tools
🪛 GitHub Actions: MyPy static type checker
[error] 188-188: Signature of "_create_step_result" incompatible with supertype "Step" [override]
[error] 188-188: This violates the Liskov substitution principle
[error] 188-188: See https://mypy.readthedocs.io/en/stable/common_issues.html#incompatible-overrides
244-253: Successful result encapsulation.
Returning a standardized payload fosters consistency across step results.
261-264: Comprehensive error response.
Capturing stack traces inmetadatais helpful for debugging.
276-585: Asynchronous streaming method is well-structured.
This large addition implements detailed log streaming and status updates. The incremental yields effectively inform callers about the agent’s progress.🧰 Tools
🪛 GitHub Actions: MyPy static type checker
[error] 310-310: Library stubs not installed for "requests" [import-untyped]
| def _create_step_result( | ||
| self, | ||
| status: Literal["success", "error"], | ||
| data: Any, | ||
| start_time: float, | ||
| metadata: Optional[Dict[str, Any]] = None, | ||
| ) -> StepResult: | ||
| """Create a standardized step result object.""" |
There was a problem hiding this comment.
Address signature incompatibility & unused variable.
MyPy flags this override as incompatible with the parent Step class. Verify the parent’s required signature or rename this method if it's not meant to override. Also, consider using or removing execution_time at line 73:
def _create_step_result(
self,
status: Literal["success", "error"],
data: Any,
start_time: float,
metadata: Optional[Dict[str, Any]] = None,
) -> StepResult:
end_time = time.perf_counter()
- execution_time = end_time - start_time
+ # Remove or utilize execution_time if required by the design
return {
"status": status,
"data": str(data),
"metadata": metadata or {},
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _create_step_result( | |
| self, | |
| status: Literal["success", "error"], | |
| data: Any, | |
| start_time: float, | |
| metadata: Optional[Dict[str, Any]] = None, | |
| ) -> StepResult: | |
| """Create a standardized step result object.""" | |
| def _create_step_result( | |
| self, | |
| status: Literal["success", "error"], | |
| data: Any, | |
| start_time: float, | |
| metadata: Optional[Dict[str, Any]] = None, | |
| ) -> StepResult: | |
| """Create a standardized step result object.""" | |
| end_time = time.perf_counter() | |
| # Remove or utilize execution_time if required by the design | |
| return { | |
| "status": status, | |
| "data": str(data), | |
| "metadata": metadata or {}, | |
| } |
🧰 Tools
🪛 GitHub Actions: MyPy static type checker
[error] 64-64: Signature of "_create_step_result" incompatible with supertype "Step" [override]
[error] 64-64: This violates the Liskov substitution principle
[error] 64-64: See https://mypy.readthedocs.io/en/stable/common_issues.html#incompatible-overrides
🪛 GitHub Actions: Flake8, Pyflakes and Compileall Linter
[warning] 73-9: Local variable 'execution_time' is assigned to but never used.
Summary by CodeRabbit
New Features
Bug Fixes