Skip to content

Commit

Permalink
add ui
Browse files Browse the repository at this point in the history
  • Loading branch information
leehuwuj committed Jan 17, 2025
1 parent 6a36f44 commit f38e07f
Show file tree
Hide file tree
Showing 2 changed files with 199 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import {
useChatUI,
} from "@llamaindex/chat-ui";
import { Markdown } from "./custom/markdown";
import { WriterCard } from "./custom/writer-card";
import { ToolAnnotations } from "./tools/chat-tools";

export function ChatMessageContent() {
Expand All @@ -22,6 +23,11 @@ export function ChatMessageContent() {
/>
),
},
// add the writer card
{
position: ContentPosition.AFTER_EVENTS,
component: <WriterCard message={message} />,
},
{
// add the tool annotations after events
position: ContentPosition.AFTER_EVENTS,
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,193 @@
import { Message } from "@llamaindex/chat-ui";
import {
AlertCircle,
CheckCircle2,
ChevronDown,
CircleDashed,
Clock,
NotebookPen,
Search,
} from "lucide-react";
import { useEffect, useState } from "react";
import {
Collapsible,
CollapsibleContent,
CollapsibleTrigger,
} from "../../collapsible";

type EventState = "pending" | "inprogress" | "done" | "error";

type WriterEvent = {
type: "retrieve" | "analyze" | "answer";
state: EventState;
data: {
id?: string;
question?: string;
answer?: string | null;
};
};

type QuestionState = {
id: string;
question: string;
answer: string | null;
state: EventState;
isOpen: boolean;
};

type WriterState = {
retrieve: {
state: EventState | null;
};
analyze: {
state: EventState | null;
questions: QuestionState[];
};
};

// Update the state based on the event
const updateState = (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,
{
id,
question,
answer: answer || null,
state: event.state,
isOpen: false,
},
];

return {
...state,
analyze: {
...state.analyze,
questions: updatedQuestions,
},
};
}

case "retrieve":
case "analyze":
return {
...state,
[event.type]: {
...state[event.type],
state: event.state,
},
};

default:
return state;
}
};

export function WriterCard({ message }: { message: Message }) {
const [state, setState] = useState<WriterState>({
retrieve: { state: null },
analyze: { state: null, questions: [] },
});

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;
}
};

if (!writerEvents?.length) {
return null;
}

return (
<div className="bg-white rounded-2xl shadow-xl p-5 space-y-6 text-gray-800 w-full">
{state.retrieve.state !== null && (
<div className="border-t border-gray-200 pt-4">
<h3 className="font-bold text-lg mb-2 flex items-center gap-2">
<Search className="w-5 h-5" />
<span>
{state.retrieve.state === "inprogress"
? "Searching..."
: "Search completed"}
</span>
</h3>
</div>
)}

{state.analyze.state !== null && (
<div className="border-t border-gray-200 pt-4">
<h3 className="font-bold text-lg mb-2 flex items-center gap-2">
<NotebookPen className="w-5 h-5" />
<span>
{state.analyze.state === "inprogress"
? "Analyzing..."
: "Analysis"}
</span>
</h3>
{state.analyze.questions.length > 0 && (
<div className="space-y-2">
{state.analyze.questions.map((question) => (
<Collapsible key={question.id}>
<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)}
</div>
<span className="font-medium text-left flex-1">
{question.question}
</span>
<ChevronDown className="w-5 h-5 transition-transform ui-expanded:rotate-180" />
</div>
</CollapsibleTrigger>
{question.answer && (
<CollapsibleContent>
<div className="p-3 text-gray-600 text-left border border-t-0 border-gray-200 rounded-b-lg">
{question.answer}
</div>
</CollapsibleContent>
)}
</Collapsible>
))}
</div>
)}
</div>
)}
</div>
);
}

0 comments on commit f38e07f

Please sign in to comment.