feat: Add LangGraph streaming with real-time UI updates

- Add streaming schemas and events
- Implement run_review_stream in ReviewerAgent
- Update task_worker to broadcast streaming events via WebSocket
- Create ReviewStream component for real-time progress visualization
- Integrate ReviewStream into ReviewDetail page
- Show agent steps, LLM messages, and progress in real-time
This commit is contained in:
Primakov Alexandr Alexandrovich 2025-10-13 01:00:49 +03:00
parent 2ad11142ad
commit 4ab6400a87
6 changed files with 383 additions and 3 deletions

View File

@ -486,3 +486,55 @@ class ReviewerAgent:
final_state = await self.graph.ainvoke(initial_state)
return final_state
async def run_review_stream(
self,
review_id: int,
pr_number: int,
repository_id: int,
on_event: callable = None
) -> Dict[str, Any]:
"""Run the review workflow with streaming events"""
initial_state: ReviewState = {
"review_id": review_id,
"pr_number": pr_number,
"repository_id": repository_id,
"status": "pending",
"files": [],
"analyzed_files": [],
"comments": [],
"error": None,
"git_service": None
}
final_state = None
# Stream through the graph
async for event in self.graph.astream(
initial_state,
stream_mode=["updates", "messages"]
):
# Handle different event types
if isinstance(event, dict):
# Node updates
for node_name, node_data in event.items():
if on_event:
await on_event({
"type": "agent_step",
"step": node_name,
"data": node_data
})
# Store final state
if isinstance(node_data, dict):
final_state = node_data
# Handle message events (LLM calls)
elif hasattr(event, '__class__') and 'message' in event.__class__.__name__.lower():
if on_event:
await on_event({
"type": "llm_message",
"message": str(event)
})
return final_state or initial_state

View File

@ -16,6 +16,13 @@ from app.schemas.webhook import (
GitHubWebhook,
BitbucketWebhook
)
from app.schemas.streaming import (
StreamEvent,
AgentStepEvent,
LLMStreamEvent,
ReviewProgressEvent,
StreamEventType
)
__all__ = [
"RepositoryCreate",
@ -28,5 +35,10 @@ __all__ = [
"GiteaWebhook",
"GitHubWebhook",
"BitbucketWebhook",
"StreamEvent",
"AgentStepEvent",
"LLMStreamEvent",
"ReviewProgressEvent",
"StreamEventType",
]

View File

@ -0,0 +1,55 @@
"""Streaming events schemas"""
from typing import Optional, Any, Dict, Literal
from pydantic import BaseModel
class StreamEventType:
"""Stream event types"""
AGENT_START = "agent_start"
AGENT_UPDATE = "agent_update"
AGENT_STEP = "agent_step"
LLM_START = "llm_start"
LLM_STREAM = "llm_stream"
LLM_END = "llm_end"
AGENT_ERROR = "agent_error"
AGENT_COMPLETE = "agent_complete"
class StreamEvent(BaseModel):
"""Base streaming event"""
type: str
review_id: int
timestamp: str
data: Dict[str, Any]
class AgentStepEvent(BaseModel):
"""Agent step event"""
type: Literal["agent_step"] = "agent_step"
review_id: int
step: str # fetch_pr_info, fetch_files, analyze_files, post_comments
status: str # started, completed, failed
message: str
data: Optional[Dict[str, Any]] = None
class LLMStreamEvent(BaseModel):
"""LLM streaming event"""
type: Literal["llm_stream"] = "llm_stream"
review_id: int
file_path: Optional[str] = None
chunk: str # Часть ответа от LLM
is_complete: bool = False
class ReviewProgressEvent(BaseModel):
"""Review progress event"""
type: Literal["review_progress"] = "review_progress"
review_id: int
total_files: int
analyzed_files: int
total_comments: int
current_step: str
message: str

View File

@ -159,14 +159,45 @@ class ReviewTaskWorker:
await db.commit()
await db.refresh(review)
# Run review agent
# Run review agent with streaming
logger.info(f" 🤖 Запуск AI review для PR #{pull_request.pr_number}")
# Import broadcast function
from app.main import manager
from datetime import datetime as dt
# Create event handler
async def on_review_event(event: dict):
"""Handle review events and broadcast to clients"""
try:
# Prepare event data
event_data = {
"type": event.get("type", "agent_update"),
"review_id": review.id,
"pr_number": pull_request.pr_number,
"timestamp": dt.utcnow().isoformat(),
"data": event
}
# Broadcast to all connected clients
await manager.broadcast(event_data)
# Log the event
if event.get("type") == "agent_step":
step = event.get("step", "unknown")
logger.info(f" 📍 Step: {step}")
elif event.get("type") == "llm_message":
message = event.get("message", "")[:100]
logger.debug(f" 💬 LLM: {message}...")
except Exception as e:
logger.error(f" ❌ Ошибка broadcast события: {e}")
agent = ReviewerAgent(db)
await agent.run_review(
await agent.run_review_stream(
review_id=review.id,
pr_number=pull_request.pr_number,
repository_id=repository.id
repository_id=repository.id,
on_event=on_review_event
)
logger.info(f" ✅ Review завершен для PR #{pull_request.pr_number}")

View File

@ -0,0 +1,224 @@
import React, { useEffect, useState } from 'react';
import { WS_URL } from '../api/websocket';
interface StreamEvent {
type: string;
review_id: number;
pr_number: number;
timestamp: string;
data: {
type?: string;
step?: string;
message?: string;
data?: any;
};
}
interface ReviewStreamProps {
reviewId: number;
}
const STEP_NAMES: Record<string, { name: string; icon: string }> = {
fetch_pr_info: { name: 'Получение информации о PR', icon: '📋' },
fetch_files: { name: 'Загрузка файлов', icon: '📂' },
analyze_files: { name: 'Анализ кода', icon: '🔍' },
post_comments: { name: 'Публикация комментариев', icon: '💬' },
complete_review: { name: 'Завершение review', icon: '✅' },
};
export const ReviewStream: React.FC<ReviewStreamProps> = ({ reviewId }) => {
const [events, setEvents] = useState<StreamEvent[]>([]);
const [currentStep, setCurrentStep] = useState<string>('');
const [isConnected, setIsConnected] = useState(false);
const [llmMessages, setLlmMessages] = useState<string[]>([]);
useEffect(() => {
const ws = new WebSocket(WS_URL);
ws.onopen = () => {
console.log('WebSocket connected for streaming');
setIsConnected(true);
};
ws.onmessage = (event) => {
try {
const data: StreamEvent = JSON.parse(event.data);
// Filter events for this review
if (data.review_id === reviewId) {
setEvents((prev) => [...prev, data]);
// Update current step
if (data.data.type === 'agent_step') {
setCurrentStep(data.data.step || '');
}
// Collect LLM messages
if (data.data.type === 'llm_message') {
setLlmMessages((prev) => [...prev, data.data.message || '']);
}
}
} catch (error) {
console.error('Error parsing WebSocket message:', error);
}
};
ws.onclose = () => {
console.log('WebSocket disconnected');
setIsConnected(false);
};
ws.onerror = (error) => {
console.error('WebSocket error:', error);
};
return () => {
ws.close();
};
}, [reviewId]);
const getStepInfo = (step: string) => {
return STEP_NAMES[step] || { name: step, icon: '⚙️' };
};
const renderStepProgress = () => {
const steps = Object.keys(STEP_NAMES);
const currentIndex = steps.indexOf(currentStep);
return (
<div className="mb-6">
<div className="flex items-center justify-between mb-2">
{steps.map((step, index) => {
const stepInfo = getStepInfo(step);
const isActive = index === currentIndex;
const isCompleted = index < currentIndex;
return (
<div key={step} className="flex items-center flex-1">
<div
className={`flex items-center justify-center w-10 h-10 rounded-full transition-all ${
isActive
? 'bg-blue-600 text-white scale-110'
: isCompleted
? 'bg-green-600 text-white'
: 'bg-dark-card border-2 border-dark-border text-dark-text-muted'
}`}
>
<span className="text-lg">{stepInfo.icon}</span>
</div>
{index < steps.length - 1 && (
<div
className={`flex-1 h-1 mx-2 transition-all ${
isCompleted ? 'bg-green-600' : 'bg-dark-border'
}`}
/>
)}
</div>
);
})}
</div>
{currentStep && (
<div className="text-center text-dark-text-primary font-medium mt-2">
{getStepInfo(currentStep).name}
</div>
)}
</div>
);
};
const renderEvents = () => {
if (events.length === 0) {
return (
<div className="text-center text-dark-text-muted py-4">
Ожидание событий...
</div>
);
}
return (
<div className="space-y-2 max-h-60 overflow-y-auto">
{events.map((event, index) => (
<div
key={index}
className="bg-dark-card border border-dark-border rounded-lg p-3 text-sm"
>
<div className="flex items-center justify-between mb-1">
<span className="text-dark-text-secondary text-xs">
{new Date(event.timestamp).toLocaleTimeString()}
</span>
<span className="text-xs bg-blue-900/30 text-blue-400 px-2 py-1 rounded">
{event.data.type}
</span>
</div>
{event.data.step && (
<div className="text-dark-text-primary">
{getStepInfo(event.data.step).icon} {getStepInfo(event.data.step).name}
</div>
)}
{event.data.message && (
<div className="text-dark-text-muted mt-1 text-xs">
{event.data.message}
</div>
)}
</div>
))}
</div>
);
};
const renderLLMMessages = () => {
if (llmMessages.length === 0) return null;
return (
<div className="mt-4">
<h3 className="text-lg font-semibold text-dark-text-primary mb-2 flex items-center gap-2">
🤖 Ответы LLM
</h3>
<div className="space-y-2 max-h-40 overflow-y-auto">
{llmMessages.slice(-5).map((message, index) => (
<div
key={index}
className="bg-dark-card border border-dark-border rounded-lg p-2 text-xs text-dark-text-secondary font-mono"
>
{message}
</div>
))}
</div>
</div>
);
};
return (
<div className="bg-dark-bg rounded-lg p-6 border border-dark-border">
<div className="flex items-center justify-between mb-4">
<h2 className="text-xl font-bold text-dark-text-primary">
🔄 Процесс Review
</h2>
<div
className={`flex items-center gap-2 px-3 py-1 rounded-full text-sm ${
isConnected
? 'bg-green-900/30 text-green-400'
: 'bg-red-900/30 text-red-400'
}`}
>
<span className={`w-2 h-2 rounded-full ${
isConnected ? 'bg-green-400 animate-pulse' : 'bg-red-400'
}`} />
{isConnected ? 'Подключено' : 'Отключено'}
</div>
</div>
{renderStepProgress()}
<div className="mt-6">
<h3 className="text-lg font-semibold text-dark-text-primary mb-3">
📝 События
</h3>
{renderEvents()}
</div>
{renderLLMMessages()}
</div>
);
};

View File

@ -5,6 +5,7 @@ import { getReview, retryReview } from '../api/client';
import { wsClient } from '../api/websocket';
import ReviewProgress from '../components/ReviewProgress';
import CommentsList from '../components/CommentsList';
import { ReviewStream } from '../components/ReviewStream';
import { formatDistance } from 'date-fns';
import { ru } from 'date-fns/locale';
@ -127,6 +128,11 @@ export default function ReviewDetail() {
</div>
</div>
{/* Real-time Review Stream (only during review) */}
{['pending', 'fetching', 'analyzing', 'commenting'].includes(review.status) && (
<ReviewStream reviewId={review.id} />
)}
{/* Progress */}
<div className="bg-gray-800 rounded-lg p-6 border border-gray-700">
<h2 className="text-xl font-semibold text-white mb-4">Прогресс</h2>