119 lines
3.9 KiB
Python
119 lines
3.9 KiB
Python
from typing import List, Dict, Any
|
|
from gateway.user_session import UserSession
|
|
|
|
|
|
def _get_chart_store_context() -> str:
|
|
"""Get current ChartStore state for context injection.
|
|
|
|
Returns:
|
|
Formatted string with ChartStore contents, or empty string if unavailable
|
|
"""
|
|
try:
|
|
from agent.tools import _registry
|
|
|
|
if not _registry:
|
|
return ""
|
|
|
|
chart_store = _registry.entries.get("ChartStore")
|
|
if not chart_store:
|
|
return ""
|
|
|
|
chart_state = chart_store.model.model_dump(mode="json")
|
|
chart_data = chart_state.get("chart_state", {})
|
|
|
|
# Only include if there's actual chart data
|
|
if not chart_data or not chart_data.get("symbol"):
|
|
return ""
|
|
|
|
# Format the chart information
|
|
symbol = chart_data.get("symbol", "N/A")
|
|
interval = chart_data.get("interval", "N/A")
|
|
start_time = chart_data.get("start_time")
|
|
end_time = chart_data.get("end_time")
|
|
selected_shapes = chart_data.get("selected_shapes", [])
|
|
|
|
selected_info = ""
|
|
if selected_shapes:
|
|
selected_info = f"\n- **Selected Shapes**: {len(selected_shapes)} shape(s) selected (IDs: {', '.join(selected_shapes)})"
|
|
|
|
chart_context = f"""
|
|
## Current Chart Context
|
|
|
|
The user is currently viewing a chart with the following settings:
|
|
- **Symbol**: {symbol}
|
|
- **Interval**: {interval}
|
|
- **Time Range**: {f"from {start_time} to {end_time}" if start_time and end_time else "not set"}{selected_info}
|
|
|
|
This information is automatically available because you're connected via websocket.
|
|
When the user refers to "the chart", "this chart", or "what I'm viewing", this is what they mean.
|
|
"""
|
|
return chart_context
|
|
|
|
except Exception:
|
|
# Silently fail - chart context is optional enhancement
|
|
return ""
|
|
|
|
|
|
def build_system_prompt(context: str, active_channels: List[str]) -> str:
|
|
"""Build the system prompt for the agent.
|
|
|
|
The main system prompt comes from system_prompt.md (loaded in context).
|
|
This function adds dynamic session information.
|
|
|
|
Args:
|
|
context: Context from loaded markdown documents (includes system_prompt.md)
|
|
active_channels: List of active channel IDs for this session
|
|
|
|
Returns:
|
|
Formatted system prompt
|
|
"""
|
|
channels_str = ", ".join(active_channels) if active_channels else "none"
|
|
|
|
# Check if user is connected via websocket - if so, inject chart context
|
|
# Note: We check for websocket by looking for "websocket" in channel IDs
|
|
# since WebSocketChannel uses channel_id like "websocket-{uuid}"
|
|
has_websocket = any("websocket" in channel_id.lower() for channel_id in active_channels)
|
|
|
|
chart_context = ""
|
|
if has_websocket:
|
|
chart_context = _get_chart_store_context()
|
|
|
|
# Context already includes system_prompt.md and other docs
|
|
# Just add current session information
|
|
prompt = f"""{context}
|
|
|
|
## Current Session Information
|
|
|
|
**Active Channels**: {channels_str}
|
|
|
|
Your responses will be sent to all active channels. Your responses are streamed back in real-time.
|
|
If the user sends a new message while you're responding, your current response will be interrupted
|
|
and you'll be re-invoked with the updated context.
|
|
{chart_context}"""
|
|
return prompt
|
|
|
|
|
|
def build_user_prompt_with_history(session: UserSession, current_message: str) -> str:
|
|
"""Build a user prompt including conversation history.
|
|
|
|
Args:
|
|
session: User session with conversation history
|
|
current_message: Current user message
|
|
|
|
Returns:
|
|
Formatted prompt with history
|
|
"""
|
|
messages = []
|
|
|
|
# Get recent history (last 10 messages)
|
|
history = session.get_history(limit=10)
|
|
|
|
for msg in history:
|
|
role_label = "User" if msg.role == "user" else "Assistant"
|
|
messages.append(f"{role_label}: {msg.content}")
|
|
|
|
# Add current message
|
|
messages.append(f"User: {current_message}")
|
|
|
|
return "\n\n".join(messages)
|