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
7 changes: 7 additions & 0 deletions assistant/core/agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
],
Expand All @@ -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,
],
Expand All @@ -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,
],
Expand All @@ -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)],
)

Expand All @@ -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,
Expand All @@ -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,
Expand Down
7 changes: 6 additions & 1 deletion demo_and_setup/configs/config.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
14 changes: 8 additions & 6 deletions src/ui/chat/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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()
2 changes: 1 addition & 1 deletion src/ui/chat/elements.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
1 change: 0 additions & 1 deletion src/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,6 @@ def validate_sql_query(sql: str | None) -> None:
"DO",
"VACUUM",
"ANALYZE",
"SET",
"RESET",
"LOCK",
"UNLOCK",
Expand Down