small fixed for agent and chat (#34)

Author: dmitrkozlovsk

assistant/core/agents.py

@@ -7,6 +7,8 @@
 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
@@ -71,6 +73,7 @@ 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,
             ],
@@ -89,6 +92,7 @@ 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,
             ],
@@ -108,6 +112,7 @@ 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,
             ],
@@ -127,6 +132,7 @@ 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)],
         )
 
@@ -139,6 +145,7 @@ 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,
@@ -160,6 +167,7 @@ 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,
@@ -247,7 +255,17 @@ 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:
-            response = await selected_agent.run(user_input, deps=deps, message_history=message_history)
+            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
+            )
+
         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, chat_scroll
+from src.ui.chat.elements import MessageHistoryContainer, TokenUsageBar, UserInputField
 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 = 0.5 if chat_state.future_result else None
+        run_every = 1.5 if chat_state.future_result else None
 
         @st.fragment(run_every=run_every)
         def handle_future_response(placeholder):
@@ -87,6 +87,5 @@ def handle_future_response(placeholder):
                     st.rerun()
             else:
                 placeholder.show_status()
-                chat_scroll()
 
         handle_future_response(agent_placeholder)

tests/assistant/test_agents.py

@@ -16,6 +16,7 @@
 )
 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
@@ -338,4 +339,6 @@ 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=[])
+    mock_agent.run.assert_awaited_once_with(
+        mock_user_input, deps=deps, message_history=[], usage_limits=UsageLimits(request_limit=5)
+    )