Fix: Agent Configuration and UI Improvements (#36)

* update configs to increase speed

* fix keywords in utils

* fix chat fragment warnings

* increase max tokens for agents

Author: dmitrkozlovsk

assistant/core/agents.py

@@ -7,6 +7,7 @@
 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 assistant.core.models import ModelFactory
 from assistant.core.schemas import Deps, ExperimentDefinition, OrchestrationResult, RouterOutput
@@ -71,6 +72,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 +91,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 +111,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 +131,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 +144,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=10_000),
             tools=[
                 retrieve_relevant_docs,
                 retrieve_codebase_docs,
@@ -160,6 +166,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=10_000),
             tools=[
                 retrieve_metrics_docs,
                 retrieve_relevant_docs,

demo_and_setup/configs/config.toml

@@ -21,4 +21,9 @@ backgroundColor = "#F0F2F6"
 gatherUsageStats = false
 
 [server]
-showEmailPrompt = false
+showEmailPrompt = false
+fileWatcherType = "none"
+runOnSave = false
+folderWatchList = ["src"]
+folderWatchBlacklist = [".git", ".venv", "venv", "node_modules",
+    "__pycache__", ".pytest_cache", "data", "dist", "build"]

src/ui/chat/chat.py

@@ -59,22 +59,24 @@ def render() -> None:
             is_stream_output=assistant_cfg.enable_streaming,
         )
         agent_placeholder = history_container.agent_placeholder
+        if chat_state.future_result:
+            agent_placeholder.show_status()
+
         # TODO: feature - select mode for agent (auto, create, multipurpose)
         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 = 2 if chat_state.future_result else None
 
         @st.fragment(run_every=run_every)
-        def handle_future_response(placeholder):
+        def handle_future_response():
             state = ChatStateManager.get_or_create_state()
             if not state.future_result:
                 return
             if state.future_result.done():
                 try:
                     response = state.future_result.result()  # get result
                     state.future_result = None
-                    agent_placeholder.handle_response(response)
                     ChatStateManager.update_state(response)
                 except Exception as e:
                     ChatStateManager.add_message(
@@ -83,10 +85,10 @@ def handle_future_response(placeholder):
                         content=str(e),
                         thinking=None,
                     )
+                    chat_scroll()
                 finally:  # prevent eternal loop
                     st.rerun()
             else:
-                placeholder.show_status()
-                chat_scroll()
+                return
 
-        handle_future_response(agent_placeholder)
+        handle_future_response()

src/ui/chat/elements.py

@@ -174,7 +174,7 @@ def handle_response(self, response: ChatResponse) -> None:
 
     def show_status(self) -> None:
         """Show thinking status message."""
-        self.placeholder.status("Thinking... Please, don't change the page. It may take a while.")
+        self.placeholder.status("Thinking... ", expanded=False)
 
     def show_error(self, error_msg: str | None) -> None:
         """Show error message.

src/utils.py

@@ -122,7 +122,6 @@ def validate_sql_query(sql: str | None) -> None:
             "DO",
             "VACUUM",
             "ANALYZE",
-            "SET",
             "RESET",
             "LOCK",
             "UNLOCK",