wehub-resource-sync
373 lines
17 KiB
Python
373 lines
17 KiB
Python
# ===============================================================================
|
|
# IMPORTS AND SETUP
|
|
# ===============================================================================
|
|
|
|
# FastAPI framework for building the web API
|
|
from fastapi import FastAPI
|
|
from fastapi.responses import StreamingResponse # For streaming real-time responses to client
|
|
|
|
# Standard Python libraries
|
|
import uuid
|
|
from typing import Any
|
|
import os
|
|
import uvicorn
|
|
import asyncio
|
|
|
|
# AG UI core components for agent communication and event handling
|
|
from ag_ui.core import (
|
|
RunAgentInput,
|
|
StateSnapshotEvent,
|
|
EventType,
|
|
RunStartedEvent,
|
|
RunFinishedEvent,
|
|
TextMessageStartEvent,
|
|
TextMessageEndEvent,
|
|
TextMessageContentEvent,
|
|
ToolCallStartEvent,
|
|
ToolCallEndEvent,
|
|
ToolCallArgsEvent,
|
|
StateDeltaEvent
|
|
)
|
|
from ag_ui.encoder import EventEncoder # Encoder for converting events to streamable format
|
|
|
|
# Import our custom stock analysis workflow
|
|
from stock_analysis import StockAnalysisFlow
|
|
|
|
# CopilotKit state management
|
|
from copilotkit import CopilotKitState
|
|
|
|
# ===============================================================================
|
|
# APPLICATION SETUP
|
|
# ===============================================================================
|
|
|
|
# Initialize FastAPI application instance
|
|
app = FastAPI()
|
|
|
|
|
|
# ===============================================================================
|
|
# STATE MANAGEMENT
|
|
# ===============================================================================
|
|
|
|
class AgentState(CopilotKitState):
|
|
"""
|
|
Agent state class that manages the data throughout the stock analysis workflow.
|
|
|
|
This class extends CopilotKitState to provide structured state management
|
|
for the stock analysis agent. It tracks:
|
|
- Tool configurations and message history
|
|
- Stock data and analysis arguments
|
|
- Financial information (cash, investments, summaries)
|
|
- UI interaction logs
|
|
|
|
Inherits from CopilotKitState (which extends langgraph's MessagesState)
|
|
"""
|
|
|
|
tools: list
|
|
messages: list
|
|
be_stock_data: Any
|
|
be_arguments: dict
|
|
available_cash: int
|
|
investment_summary: dict
|
|
tool_logs: list
|
|
|
|
# ===============================================================================
|
|
# MAIN API ENDPOINT
|
|
# ===============================================================================
|
|
|
|
@app.post("/crewai-agent")
|
|
async def crewai_agent(input_data: RunAgentInput):
|
|
"""
|
|
Main API endpoint for processing stock analysis requests.
|
|
|
|
This endpoint:
|
|
1. Receives user input and current state from the frontend
|
|
2. Streams real-time events back to the client during processing
|
|
3. Runs the StockAnalysisFlow workflow asynchronously
|
|
4. Returns results via Server-Sent Events (SSE) streaming
|
|
|
|
Args:
|
|
input_data (RunAgentInput): Contains user messages, tools, state, thread/run IDs
|
|
|
|
Returns:
|
|
StreamingResponse: Real-time stream of events during agent execution
|
|
"""
|
|
try:
|
|
|
|
async def event_generator():
|
|
"""
|
|
Asynchronous generator that streams events to the client in real-time.
|
|
|
|
This function orchestrates the entire stock analysis workflow:
|
|
1. Sets up event streaming infrastructure
|
|
2. Emits initial state and run started events
|
|
3. Launches the StockAnalysisFlow workflow
|
|
4. Streams progress events as they occur
|
|
5. Handles final results (tool calls or text messages)
|
|
6. Emits run completion events
|
|
|
|
Yields:
|
|
Encoded events for Server-Sent Events (SSE) streaming
|
|
"""
|
|
# Step 1: Initialize event streaming infrastructure
|
|
encoder = EventEncoder() # Converts events to SSE format
|
|
event_queue = asyncio.Queue() # Queue for events from the workflow
|
|
|
|
def emit_event(event):
|
|
"""Callback function for the workflow to emit events"""
|
|
event_queue.put_nowait(event)
|
|
|
|
# Generate unique identifier for text messages
|
|
# Generate unique identifier for text messages
|
|
message_id = str(uuid.uuid4())
|
|
|
|
# Step 2: Emit run started event to notify client that processing has begun
|
|
yield encoder.encode(
|
|
RunStartedEvent(
|
|
type=EventType.RUN_STARTED,
|
|
thread_id=input_data.thread_id, # Conversation thread identifier
|
|
run_id=input_data.run_id, # Unique run identifier
|
|
)
|
|
)
|
|
|
|
# Step 3: Send initial state snapshot to client
|
|
# This provides the client with current financial state before analysis
|
|
yield encoder.encode(
|
|
StateSnapshotEvent(
|
|
type=EventType.STATE_SNAPSHOT,
|
|
snapshot={
|
|
"available_cash": input_data.state["available_cash"], # Current cash balance
|
|
"investment_summary": input_data.state["investment_summary"], # Previous analysis results
|
|
"investment_portfolio": input_data.state["investment_portfolio"], # Current holdings
|
|
"tool_logs": [] # Reset tool logs for new analysis
|
|
}
|
|
)
|
|
)
|
|
|
|
# Step 4: Initialize agent state with input data
|
|
state = AgentState(
|
|
tools=input_data.tools, # Available tools
|
|
messages=input_data.messages, # Conversation history
|
|
be_stock_data=None, # Will be populated during analysis
|
|
be_arguments=None, # Will be populated during analysis
|
|
available_cash=input_data.state["available_cash"], # Current cash
|
|
investment_portfolio=input_data.state["investment_portfolio"], # Current portfolio
|
|
tool_logs=[] # Progress tracking
|
|
)
|
|
|
|
# Step 5: Launch the stock analysis workflow asynchronously
|
|
# This creates a task that runs the StockAnalysisFlow in the background
|
|
agent_task = asyncio.create_task(
|
|
StockAnalysisFlow().kickoff_async(inputs={
|
|
"state": state, # Agent state
|
|
"emit_event": emit_event, # Event emission callback
|
|
"investment_portfolio": input_data.state["investment_portfolio"] # Portfolio data
|
|
})
|
|
)
|
|
|
|
# Step 6: Event streaming loop - relay events from workflow to client
|
|
chart_data_sent = False # Track if chart data has been sent
|
|
while True:
|
|
try:
|
|
# Try to get an event from the queue with a short timeout
|
|
event = await asyncio.wait_for(event_queue.get(), timeout=0.1)
|
|
|
|
# Only stream events that are NOT from the insights stage to prevent flickering
|
|
# Skip events that might interfere with the completed chart
|
|
should_stream = True
|
|
|
|
# Check if this is a tool log completion event (indicates allocation stage done)
|
|
if (hasattr(event, 'delta') and event.delta and
|
|
isinstance(event.delta, list) and len(event.delta) > 0 and
|
|
event.delta[0].get('path') == '/tool_logs' and
|
|
event.delta[0].get('op') == 'replace' and
|
|
event.delta[0].get('value') == 'completed'):
|
|
chart_data_sent = True
|
|
# Don't stop streaming yet - let the chart data go through first
|
|
|
|
# Check if this is the actual chart tool call being sent
|
|
if (hasattr(event, 'type') and event.type == EventType.TOOL_CALL_ARGS and
|
|
hasattr(event, 'delta') and 'render_standard_charts_and_table' in str(event.delta)):
|
|
# This is the chart data being sent - mark it and allow it through
|
|
chart_data_sent = True
|
|
should_stream = True
|
|
|
|
# After chart data is sent, block insights-related events but allow portfolio updates
|
|
if chart_data_sent and (hasattr(event, 'delta') and event.delta and
|
|
isinstance(event.delta, list) and len(event.delta) > 0):
|
|
# Allow portfolio updates to go through
|
|
if event.delta[0].get('path') == '/investment_portfolio':
|
|
should_stream = True
|
|
# Block insights-related events
|
|
elif ('insights' in str(event.delta[0]).lower() or
|
|
'processing' in str(event.delta[0]).lower() or
|
|
'extracting' in str(event.delta[0]).lower()):
|
|
should_stream = False
|
|
|
|
if should_stream:
|
|
yield encoder.encode(event) # Stream the event to client
|
|
|
|
except asyncio.TimeoutError:
|
|
# No events in queue - check if workflow is complete
|
|
if agent_task.done():
|
|
break # Exit loop when workflow finishes
|
|
|
|
# If chart data has been sent and workflow is still running (insights stage),
|
|
# we can break early to make graph interactive while insights generate
|
|
if chart_data_sent and not agent_task.done():
|
|
# Minimal delay to ensure chart data is fully processed
|
|
await asyncio.sleep(0.2)
|
|
break # Exit early to make graph interactive
|
|
|
|
# Step 7: Clear tool logs after workflow completion
|
|
# This prevents old progress logs from cluttering the UI
|
|
yield encoder.encode(
|
|
StateDeltaEvent(
|
|
type=EventType.STATE_DELTA,
|
|
delta=[
|
|
{
|
|
"op": "replace",
|
|
"path": "/tool_logs",
|
|
"value": []
|
|
}
|
|
]
|
|
)
|
|
)
|
|
|
|
# Step 8: Handle workflow results based on the final message type
|
|
# Check if the last message is from the assistant (AI agent)
|
|
if state["messages"][-1].role == "assistant":
|
|
|
|
# Step 8.1: Handle tool call results (charts, data visualizations)
|
|
if state["messages"][-1].tool_calls:
|
|
# The workflow generated a tool call (e.g., render charts)
|
|
# Stream tool call events to trigger UI rendering
|
|
|
|
yield encoder.encode(
|
|
ToolCallStartEvent(
|
|
type=EventType.TOOL_CALL_START,
|
|
tool_call_id=state["messages"][-1].tool_calls[0].id,
|
|
toolCallName=state["messages"][-1]
|
|
.tool_calls[0]
|
|
.function.name,
|
|
)
|
|
)
|
|
|
|
# Stream the tool call arguments (contains analysis results)
|
|
yield encoder.encode(
|
|
ToolCallArgsEvent(
|
|
type=EventType.TOOL_CALL_ARGS,
|
|
tool_call_id=state["messages"][-1].tool_calls[0].id,
|
|
delta=state["messages"][-1]
|
|
.tool_calls[0]
|
|
.function.arguments, # Contains investment summary and insights
|
|
)
|
|
)
|
|
|
|
# Signal end of tool call
|
|
yield encoder.encode(
|
|
ToolCallEndEvent(
|
|
type=EventType.TOOL_CALL_END,
|
|
tool_call_id=state["messages"][-1].tool_calls[0].id,
|
|
)
|
|
)
|
|
else:
|
|
# Step 8.2: Handle text message results (when no analysis was performed)
|
|
yield encoder.encode(
|
|
TextMessageStartEvent(
|
|
type=EventType.TEXT_MESSAGE_START,
|
|
message_id=message_id,
|
|
role="assistant",
|
|
)
|
|
)
|
|
|
|
# Step 8.2.1: Stream text content if available
|
|
if state["messages"][-1].content:
|
|
content = state["messages"][-1].content
|
|
|
|
# Split content into chunks for smooth streaming effect
|
|
n_parts = 100 # Number of chunks for streaming
|
|
part_length = max(1, len(content) // n_parts)
|
|
parts = [content[i:i+part_length] for i in range(0, len(content), part_length)]
|
|
|
|
# Ensure we don't exceed the target number of parts
|
|
if len(parts) > n_parts:
|
|
parts = parts[:n_parts-1] + [''.join(parts[n_parts-1:])]
|
|
|
|
# Stream each part with a small delay for smooth typing effect
|
|
for part in parts:
|
|
yield encoder.encode(
|
|
TextMessageContentEvent(
|
|
type=EventType.TEXT_MESSAGE_CONTENT,
|
|
message_id=message_id,
|
|
delta=part, # Chunk of text content
|
|
)
|
|
)
|
|
await asyncio.sleep(0.05) # Small delay for typing effect
|
|
else:
|
|
# Step 8.2.2: Handle case where no content is available (error scenario)
|
|
yield encoder.encode(
|
|
TextMessageContentEvent(
|
|
type=EventType.TEXT_MESSAGE_CONTENT,
|
|
message_id=message_id,
|
|
delta="Something went wrong! Please try again.",
|
|
)
|
|
)
|
|
|
|
# Step 8.2.3: Signal end of text message
|
|
yield encoder.encode(
|
|
TextMessageEndEvent(
|
|
type=EventType.TEXT_MESSAGE_END,
|
|
message_id=message_id,
|
|
)
|
|
)
|
|
|
|
# Step 9: Emit run finished event to signal completion
|
|
yield encoder.encode(
|
|
RunFinishedEvent(
|
|
type=EventType.RUN_FINISHED,
|
|
thread_id=input_data.thread_id, # Same thread ID from start
|
|
run_id=input_data.run_id, # Same run ID from start
|
|
)
|
|
)
|
|
|
|
except Exception as e:
|
|
# Step 10: Handle any unexpected errors during processing
|
|
print(e) # Log error for debugging
|
|
|
|
# Step 11: Return streaming response to client
|
|
# Step 11: Return streaming response to client
|
|
return StreamingResponse(event_generator(), media_type="text/event-stream")
|
|
|
|
|
|
# ===============================================================================
|
|
# SERVER STARTUP AND CONFIGURATION
|
|
# ===============================================================================
|
|
|
|
def main():
|
|
"""
|
|
Main function to start the uvicorn server.
|
|
|
|
This function:
|
|
- Reads the port from environment variables (defaults to 8000)
|
|
- Configures uvicorn server settings
|
|
- Starts the server with hot reload enabled for development
|
|
"""
|
|
# Get port from environment variable or use default
|
|
port = int(os.getenv("PORT", "8000"))
|
|
|
|
# Start uvicorn server with configuration
|
|
uvicorn.run(
|
|
"main:app", # Module and app instance
|
|
host="0.0.0.0", # Listen on all network interfaces
|
|
port=port, # Port number
|
|
reload=True, # Enable hot reload for development
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
"""
|
|
Entry point when script is run directly.
|
|
Starts the FastAPI server using uvicorn.
|
|
"""
|
|
main()
|