Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 1 addition & 19 deletions assistant/core/agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,6 @@
from pydantic_ai import Agent
from pydantic_ai.common_tools.tavily import tavily_search_tool
from pydantic_ai.messages import ModelMessage, ModelResponse, ThinkingPart
from pydantic_ai.settings import ModelSettings
from pydantic_ai.usage import UsageLimitExceeded, UsageLimits

from assistant.core.models import ModelFactory
from assistant.core.schemas import Deps, ExperimentDefinition, OrchestrationResult, RouterOutput
Expand Down Expand Up @@ -73,7 +71,6 @@ def _create_experiment_creator_agent(self) -> Agent[Deps, ExperimentDefinition]:
instructions=load_prompt("experiment_creator_instructions.md").render(),
output_type=ExperimentDefinition,
deps_type=Deps,
model_settings=ModelSettings(max_tokens=10_000),
tools=[
retrieve_metrics_docs,
],
Expand All @@ -92,7 +89,6 @@ def _create_internal_database_agent(self) -> Agent[Deps, str]:
instructions=load_prompt("sql_expert_instructions.md").render(),
output_type=str,
deps_type=Deps,
model_settings=ModelSettings(max_tokens=10_000),
tools=[
retrieve_internal_db,
],
Expand All @@ -112,7 +108,6 @@ def _create_experiment_analyst_agent(self) -> Agent[Deps, str]:
instructions=load_prompt("experiment_analyst_instructions.md").render(),
output_type=str,
deps_type=Deps,
model_settings=ModelSettings(max_tokens=20_000),
tools=[
get_expanto_app_context,
],
Expand All @@ -132,7 +127,6 @@ def _create_internet_search_agent(self) -> Agent[Deps, str]:
instructions=load_prompt("internet_search_instructions.md").render(),
output_type=str,
deps_type=Deps,
model_settings=ModelSettings(max_tokens=10_000),
tools=[tavily_search_tool(self.tavily_api_key)],
)

Expand All @@ -145,7 +139,6 @@ def _create_expanto_assistant(self) -> Agent[Deps, str]:
instructions="Use as many tool call as you needed",
output_type=str,
deps_type=Deps,
model_settings=ModelSettings(max_tokens=20_000),
tools=[
retrieve_relevant_docs,
retrieve_codebase_docs,
Expand All @@ -167,7 +160,6 @@ def _create_universal_agent(self) -> Agent[Deps, str]:
"Use any tools if you need to answer user question or execute user task",
output_type=str,
deps_type=Deps,
model_settings=ModelSettings(max_tokens=20_000),
tools=[
retrieve_metrics_docs,
retrieve_relevant_docs,
Expand Down Expand Up @@ -255,17 +247,7 @@ async def process(self, user_input: str, deps: Deps, message_history: list[Any])
selected_agent = self.agent_manager.get_agent(route_id=route_output.route_id)
logger.info(f"Router decision: {route_output.route_id} → Selected: {selected_agent.name}")
try:
usage_limits = UsageLimits(
request_limit=5,
)
response = await selected_agent.run(
user_input, deps=deps, message_history=message_history, usage_limits=usage_limits
)
except UsageLimitExceeded:
response = await selected_agent.run(
"You working too long. Return final answer", deps=deps, message_history=message_history
)

response = await selected_agent.run(user_input, deps=deps, message_history=message_history)
except Exception as e:
logger.error(f"Agent {selected_agent.name} failed: {e}")
logger.info("Falling back to Multipurpose agent")
Expand Down
5 changes: 3 additions & 2 deletions src/ui/chat/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import streamlit as st

from src.settings import AssistantServiceCfg
from src.ui.chat.elements import MessageHistoryContainer, TokenUsageBar, UserInputField
from src.ui.chat.elements import MessageHistoryContainer, TokenUsageBar, UserInputField, chat_scroll
from src.ui.chat.schemas import MessageType, Role
from src.ui.chat.services import ChatController, HttpAssistantService
from src.ui.resources import load_assistant_service_cfg
Expand Down Expand Up @@ -63,7 +63,7 @@ def render() -> None:
UserInputField.render(controller) # if user input: save to chat state and rerun

# Handle input and response logic
run_every = 1.5 if chat_state.future_result else None
run_every = 0.5 if chat_state.future_result else None

@st.fragment(run_every=run_every)
def handle_future_response(placeholder):
Expand All @@ -87,5 +87,6 @@ def handle_future_response(placeholder):
st.rerun()
else:
placeholder.show_status()
chat_scroll()

handle_future_response(agent_placeholder)
5 changes: 1 addition & 4 deletions tests/assistant/test_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
)
from pydantic_ai.models.test import TestModel
from pydantic_ai.result import StreamedRunResult
from pydantic_ai.usage import UsageLimits

from assistant.core.agents import AgentManager, AgentOrchestrator
from assistant.core.models import ModelFactory
Expand Down Expand Up @@ -339,6 +338,4 @@ def get_agent_side_effect(route_id):
assert result.output == mock_output_content
mock_agent_manager.get_agent.assert_any_call(route_id="route")
mock_agent_manager.get_agent.assert_any_call(route_id="universal")
mock_agent.run.assert_awaited_once_with(
mock_user_input, deps=deps, message_history=[], usage_limits=UsageLimits(request_limit=5)
)
mock_agent.run.assert_awaited_once_with(mock_user_input, deps=deps, message_history=[])