Revert "small fixed for agent and chat (#34)" (#35)

This reverts commit aa39dcc.

Author: dmitrkozlovsk

assistant/core/agents.py

@@ -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
@@ -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,
             ],
@@ -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,
             ],
@@ -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,
             ],
@@ -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)],
         )
 
@@ -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,
@@ -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,
@@ -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")

src/ui/chat/chat.py

@@ -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
@@ -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):
@@ -87,5 +87,6 @@ def handle_future_response(placeholder):
                     st.rerun()
             else:
                 placeholder.show_status()
+                chat_scroll()
 
         handle_future_response(agent_placeholder)

tests/assistant/test_agents.py

@@ -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
@@ -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=[])