Skip to content

Commit

Permalink
fix not display
Browse files Browse the repository at this point in the history
  • Loading branch information
leehuwuj committed Jan 17, 2025
1 parent 90a55f3 commit acfc4ec
Show file tree
Hide file tree
Showing 4 changed files with 71 additions and 69 deletions.
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
from typing import Any, Dict, List, Literal

from app.api.routers.models import SourceNodes
from llama_index.core.schema import Node, NodeWithScore
from llama_index.core.schema import NodeWithScore
from llama_index.core.workflow import Event

from app.api.routers.models import SourceNodes


# Workflow events
class PlanResearchEvent(Event):
Expand All @@ -13,7 +14,7 @@ class PlanResearchEvent(Event):
class ResearchEvent(Event):
question_id: str
question: str
context_nodes: List[NodeWithScore | Node]
context_nodes: List[NodeWithScore]


class CollectAnswersEvent(Event):
Expand Down
36 changes: 18 additions & 18 deletions templates/components/agents/python/blog/app/workflows/writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,6 @@
import uuid
from typing import Any, Dict, List, Optional

from app.engine.index import IndexConfig, get_index
from app.workflows.agents import plan_research, research, write_report
from app.workflows.models import (
CollectAnswersEvent,
DataEvent,
PlanResearchEvent,
ResearchEvent,
SourceNodesEvent,
WriteReportEvent,
)
from llama_index.core.indices.base import BaseIndex
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.memory.simple_composable_memory import SimpleComposableMemory
Expand All @@ -26,7 +16,18 @@
step,
)

logger = logging.getLogger(__name__)
from app.engine.index import IndexConfig, get_index
from app.workflows.agents import plan_research, research, write_report
from app.workflows.models import (
CollectAnswersEvent,
DataEvent,
PlanResearchEvent,
ResearchEvent,
SourceNodesEvent,
WriteReportEvent,
)

logger = logging.getLogger("uvicorn")
logger.setLevel(logging.INFO)


Expand All @@ -45,7 +46,7 @@ def create_workflow(
return WriterWorkflow(
index=index,
chat_history=chat_history,
**kwargs,
timeout=120.0,
)


Expand All @@ -67,22 +68,20 @@ class WriterWorkflow(Workflow):
context_nodes: List[Node]
index: BaseIndex
user_request: str
stream: bool = False
stream: bool = True

def __init__(
self,
index: BaseIndex,
chat_history: Optional[List[ChatMessage]] = None,
stream: bool = False,
timeout: Optional[float] = 120.0,
stream: bool = True,
**kwargs,
):
super().__init__(**kwargs)
self.index = index
self.context_nodes = []
self.stream = stream
self.chat_history = chat_history
self.timeout = timeout
self.memory = SimpleComposableMemory.from_defaults(
primary_memory=ChatMemoryBuffer.from_defaults(
chat_history=chat_history,
Expand Down Expand Up @@ -142,7 +141,7 @@ async def analyze(
"""
Analyze the retrieved information
"""
print("Analyzing the retrieved information")
logger.info("Analyzing the retrieved information")
ctx.write_event_to_stream(
DataEvent(
type="analyze",
Expand Down Expand Up @@ -262,7 +261,7 @@ async def collect_answers(
content=f"<Question>{result.question}</Question>\n<Answer>{result.answer}</Answer>",
)
)
ctx.set("n_questions", 0)
await ctx.set("n_questions", 0)
self.memory.put(
message=ChatMessage(
role=MessageRole.ASSISTANT,
Expand All @@ -276,6 +275,7 @@ async def report(self, ctx: Context, ev: WriteReportEvent) -> StopEvent:
"""
Report the answers
"""
logger.info("Writing the report")
res = await write_report(
memory=self.memory,
user_request=self.user_request,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ export function ChatMessageContent() {
},
// add the writer card
{
position: ContentPosition.AFTER_EVENTS,
position: ContentPosition.CHAT_EVENTS,
component: <WriterCard message={message} />,
},
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import {
NotebookPen,
Search,
} from "lucide-react";
import { useEffect, useState } from "react";
import { useMemo } from "react";
import {
Collapsible,
CollapsibleContent,
Expand Down Expand Up @@ -45,42 +45,49 @@ type WriterState = {
};
};

// Update the state based on the event
const updateState = (state: WriterState, event: WriterEvent): WriterState => {
const stateIcon: Record<EventState, React.ReactNode> = {
pending: <Clock className="w-4 h-4 text-yellow-500" />,
inprogress: <CircleDashed className="w-4 h-4 text-blue-500 animate-spin" />,
done: <CheckCircle2 className="w-4 h-4 text-green-500" />,
error: <AlertCircle className="w-4 h-4 text-red-500" />,
};

// Transform the state based on the event without mutations
const transformState = (
state: WriterState,
event: WriterEvent,
): WriterState => {
switch (event.type) {
case "answer": {
const { id, question, answer } = event.data;
if (!id || !question) return state;

const questions = state.analyze.questions;
const existingQuestion = questions.find((q) => q.id === id);

const updatedQuestions = existingQuestion
? questions.map((q) =>
q.id === id
? {
...existingQuestion,
state: event.state,
answer: answer || existingQuestion.answer,
}
: q,
)
: [
...questions,
const updatedQuestions = state.analyze.questions.map((q) => {
if (q.id !== id) return q;
return {
...q,
state: event.state,
answer: answer ?? q.answer,
};
});

const newQuestion = !state.analyze.questions.some((q) => q.id === id)
? [
{
id,
question,
answer: answer || null,
answer: answer ?? null,
state: event.state,
isOpen: false,
},
];
]
: [];

return {
...state,
analyze: {
...state.analyze,
questions: updatedQuestions,
questions: [...updatedQuestions, ...newQuestion],
},
};
}
Expand All @@ -100,36 +107,30 @@ const updateState = (state: WriterState, event: WriterEvent): WriterState => {
}
};

export function WriterCard({ message }: { message: Message }) {
const [state, setState] = useState<WriterState>({
// Convert writer events to state
const writeEventsToState = (events: WriterEvent[] | undefined): WriterState => {
if (!events?.length) {
return {
retrieve: { state: null },
analyze: { state: null, questions: [] },
};
}

const initialState: WriterState = {
retrieve: { state: null },
analyze: { state: null, questions: [] },
});
};

return events.reduce(
(acc: WriterState, event: WriterEvent) => transformState(acc, event),
initialState,
);
};

export function WriterCard({ message }: { message: Message }) {
const writerEvents = message.annotations as WriterEvent[] | undefined;

useEffect(() => {
if (writerEvents?.length) {
writerEvents.forEach((event) => {
setState((currentState) => updateState(currentState, event));
});
}
}, [writerEvents]);

const getStateIcon = (state: EventState | null) => {
switch (state) {
case "pending":
return <Clock className="w-4 h-4 text-yellow-500" />;
case "inprogress":
return <CircleDashed className="w-4 h-4 text-blue-500 animate-spin" />;
case "done":
return <CheckCircle2 className="w-4 h-4 text-green-500" />;
case "error":
return <AlertCircle className="w-4 h-4 text-red-500" />;
default:
return null;
}
};
const state = useMemo(() => writeEventsToState(writerEvents), [writerEvents]);

if (!writerEvents?.length) {
return null;
Expand Down Expand Up @@ -167,7 +168,7 @@ export function WriterCard({ message }: { message: Message }) {
<CollapsibleTrigger className="w-full">
<div className="flex items-center gap-2 p-3 hover:bg-gray-50 transition-colors rounded-lg border border-gray-200">
<div className="flex-shrink-0">
{getStateIcon(question.state)}
{stateIcon[question.state]}
</div>
<span className="font-medium text-left flex-1">
{question.question}
Expand Down

0 comments on commit acfc4ec

Please sign in to comment.