Skip to content

Update chat-plugin to support model-proxy #101

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Nov 5, 2025
Merged

Update chat-plugin to support model-proxy #101

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4006e4b
make chat-plugin support model-proxy
zhogu Oct 15, 2025
ae7ce27
clean code
zhogu Oct 15, 2025
7f8911c
clean code
zhogu Oct 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 13 additions & 81 deletions contrib/chat-plugin/src/app/JobBar.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@

"use client";

import { fetchJobList, fetchModelsInCurrentJob } from "../libs/api";
import { fetchAllModels } from "../libs/api";
import { useEffect, useState } from "react";
import { Job, Status, useChatStore } from "../libs/state";
import { Status, useChatStore } from "../libs/state";
import { RefreshCw } from "lucide-react";


Expand All @@ -28,49 +28,28 @@ const statusText = {
export default function JobBar() {
const [status, setStatus] = useState<Status>("offline");
const {
allJobs,
currentJob,
setCurrentJob,
allModelsInCurrentJob,
allModels,
currentModel,
setCurrentModel
} = useChatStore();

// Fetch job list on component mount and set up status polling
// Loading states for better UI feedback
const [isJobsLoading, setIsJobsLoading] = useState(false);
const [isModelsLoading, setIsModelsLoading] = useState(false);

const [error, setError] = useState<string | null>(null);

// Fetch job list with loading indicator
const fetchJobs = async () => {
setError(null);
setIsJobsLoading(true);
setStatus("loading");
try {
await fetchJobList();
setStatus("unknown");
} catch (err) {
setError("Failed to fetch jobs. Please try again.");
console.error(err);
} finally {
setIsJobsLoading(false);
}
};

// Fetch models with loading indicator
const fetchModels = async () => {
if (!currentJob) return;
setError(null);
setIsModelsLoading(true);
setStatus("loading");
try {

await fetchModelsInCurrentJob();
await fetchAllModels();
await new Promise(res => setTimeout(res, 2500));
// Use the latest models from the store after fetching
const models = useChatStore.getState().allModelsInCurrentJob;
const models = useChatStore.getState().allModels;
if (models.length > 0) {
setStatus("online");
} else {
Expand All @@ -85,50 +64,21 @@ export default function JobBar() {
};

useEffect(() => {
fetchJobs(); // Initial job list fetch
fetchModels();

const interval = setInterval(async () => {
}, 500);

return () => clearInterval(interval);
}, []);

// Fetch models when current job changes
useEffect(() => {
if (currentJob) {
fetchModels();
} else {
// Clear models when no job is selected
useChatStore.getState().setAllModelsInCurrentJob([]);
if (currentModel) {
// Clear current model selection when job changes
setCurrentModel(null);
}
setStatus("unknown");
}
}, [currentJob]);

// Handle job selection change
const handleJobChange = (e: React.ChangeEvent<HTMLSelectElement>) => {
const selectedJobId = e.target.value;
const job = allJobs.find(job => job.id === selectedJobId) || null;
setCurrentJob(job);
};

// Handle model selection change
const handleModelChange = (e: React.ChangeEvent<HTMLSelectElement>) => {
const selectedModel = e.target.value;
setCurrentModel(selectedModel);
};

// Handle refresh button click
const handleRefresh = async () => {
await fetchJobs();
if (currentJob) {
await fetchModels();
}
};

return (
<div className="w-full h-full flex items-center justify-between p-2">
<div className="flex items-center gap-2 text-gray-500">
Expand All @@ -138,24 +88,6 @@ export default function JobBar() {
</div>

<div className="flex items-center gap-4">
{/* Job select */}
<div className="flex items-center gap-2">
<label htmlFor="job-select" className="block text-sm font-medium text-gray-700 mr-2">Serving Job:</label>
<select
id="job-select"
className="block w-64 rounded-md border-gray-300 shadow-sm focus:border-blue-500 focus:ring-blue-500 sm:text-sm px-3 py-2"
value={currentJob?.id || ""}
onChange={handleJobChange}
disabled={isJobsLoading}
size={1}
style={{ minWidth: "16rem", maxWidth: "30rem" }}
>
<option value="">Select a job ({allJobs.length})</option>
{allJobs.map((job) => (
<option key={job.id} value={job.id} title={`${job.name} (${job.username})`}>{job.name}</option>
))}
</select>
</div>

{/* Model select */}
<div className="flex items-center gap-2">
Expand All @@ -165,23 +97,23 @@ export default function JobBar() {
className="block w-48 rounded-md border-gray-300 shadow-sm focus:border-blue-500 focus:ring-blue-500 sm:text-sm"
value={currentModel || ""}
onChange={handleModelChange}
disabled={!currentJob || allModelsInCurrentJob.length === 0 || isModelsLoading}
disabled={allModels.length === 0 || isModelsLoading}
>
<option value="">Select a model ({allModelsInCurrentJob.length})</option>
{allModelsInCurrentJob.map((model) => (
<option value="">Select a model ({allModels.length})</option>
{allModels.map((model) => (
<option key={model} value={model} title={model}>{model}</option>
))}
</select>
</div>

{/* Refresh button */}
<button
onClick={handleRefresh}
onClick={fetchModels}
className="inline-flex items-center p-1.5 border border-transparent rounded-full shadow-sm text-white bg-blue-600 hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-blue-500"
title="Refresh job and model lists"
disabled={isJobsLoading || isModelsLoading}
title="Refresh model lists"
disabled={isModelsLoading}
>
<RefreshCw className={`h-5 w-5 ${(isJobsLoading || isModelsLoading) ? "animate-spin" : ""}`} aria-hidden="true" />
<RefreshCw className={`h-5 w-5 ${isModelsLoading ? "animate-spin" : ""}`} aria-hidden="true" />
</button>
</div>
</div>
Expand Down
8 changes: 3 additions & 5 deletions contrib/chat-plugin/src/app/index.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,15 +8,13 @@ import ChatBox from "./ChatBox";
import { useChatStore } from "../libs/state";
import { Toaster } from "../components/sonner";

export default function App({ restUrl, user, restToken, modelToken }:
{ restUrl: string; user: string; restToken: string; modelToken: string }) {
export default function App({ restUrl, user, restToken }:
{ restUrl: string; user: string; restToken: string }) {
// Initialize the chat store with the provided parameters
const { setRestServerPath, setUser, setRestServerToken, setJobServerToken } = useChatStore.getState();
const { setUser, setRestServerToken } = useChatStore.getState();

setRestServerPath(restUrl);
setUser(user);
setRestServerToken(restToken);
setJobServerToken(modelToken);

return (
<div className="w-full h-full flex flex-col overflow-hidden">
Expand Down
7 changes: 1 addition & 6 deletions contrib/chat-plugin/src/index.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ declare global {
id: string;
title: string;
uri: string;
token: string;
}];
}
}
Expand Down Expand Up @@ -58,13 +57,9 @@ class ProtocolPluginElement extends HTMLElement {
}
}
console.log("source", source);

const plugins = window.PAI_PLUGINS;
const pluginIndex = Number(params.get("index")) || 0;
const modelToken = plugins[pluginIndex].token || "";

const root = ReactDOM.createRoot(this);
root.render(React.createElement(App, {restUrl, user, restToken, modelToken}));
root.render(React.createElement(App, { restUrl, user, restToken }));
}

}
Expand Down
106 changes: 19 additions & 87 deletions contrib/chat-plugin/src/libs/api.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,10 @@
// Licensed under the MIT License.

import ky, { Options } from "ky";
import { Job, useChatStore } from "./state";
import { useChatStore } from "./state";
import { toast } from "sonner";

const API_BASE_URL = window.location.origin;
const TARGET_JOB_TAG = "model-serving";

export const api = {
post: async (path: string, key: string | null = null, options: Options = {}) => {
Expand Down Expand Up @@ -66,98 +65,36 @@ export const api = {

};

export async function fetchJobList(): Promise<void> {
const restServerPath = useChatStore.getState().restServerPath;
export async function fetchAllModels(): Promise<string[]> {
const modelProxyPath = useChatStore.getState().modelProxyPath;
const modelsUrl = `${modelProxyPath}/v1/models`;
const restServerToken = useChatStore.getState().restServerToken;

const getJobsUrl = `${restServerPath}/api/v2/jobs?state=RUNNING`;
try {
const data: any = await api.get(getJobsUrl, restServerToken);
const jobs = (data as { debugId: string, name: string, username: string, state: string }[])
.filter(job => job.name.includes(TARGET_JOB_TAG))
.map(job => ({
id: job.debugId,
name: job.name,
username: job.username,
status: job.state,
ip: "",
port: 0,
} as Job));

const jobsWithDetails = await Promise.all(jobs.map(async job => {
const jobDetailsUrl = `${restServerPath}/api/v2/jobs/${job.username}~${job.name}`;
try {
const jobInfo: any = await api.get(jobDetailsUrl, restServerToken);
const details = jobInfo as { taskRoles: { [key: string]: { taskStatuses: [{ containerIp: string, containerPorts: { http: string } }] } } };
if (!details || !details.taskRoles) {
console.warn(`No task roles found for job ${job.name}`);
return null;
}
const taskStatuses = Object.values(details.taskRoles)[0].taskStatuses;
if (!taskStatuses) {
console.warn(`No task statuses found for job ${job.name}`);
return null;
}
job.ip = taskStatuses[0].containerIp;
job.port = parseInt(taskStatuses[0].containerPorts.http);
return job;
} catch (error) {
console.warn(`Failed to fetch details for job ${job.name}:`, error);
return null;
}
}));

const filteredJobs = jobsWithDetails.filter((job): job is Job => job !== null);
useChatStore.getState().setAllJobs(filteredJobs);
console.log("Fetched job list:", filteredJobs);
} catch (error) {
console.error("Failed to fetch job list:", error);
useChatStore.getState().setAllJobs([]);
}
}

export async function fetchModelsInCurrentJob(): Promise<string[]> {
const currentJob = useChatStore.getState().currentJob;
if (!currentJob) {
console.warn("No current job selected");
return [];
}

const jobServerPath = useChatStore.getState().jobServerPath;
const jobServerToken = useChatStore.getState().jobServerToken;

const modelsUrl = `${jobServerPath}/${currentJob.ip}:${currentJob.port}/v1/models`;
try {
const data: any = await api.get(modelsUrl, jobServerToken);
const data: any = await api.get(modelsUrl, restServerToken);
const models = data as { data: [{ id: string }] };
if (!models || !models.data) {
console.warn(`No models found for job ${currentJob.name}`);
useChatStore.getState().setAllModelsInCurrentJob([]);
console.warn(`No models found from model proxy`);
useChatStore.getState().setAllModels([]);
return [];
}
const modelList = models.data.map((model: { id: string }) => model.id);
useChatStore.getState().setAllModelsInCurrentJob(modelList);
useChatStore.getState().setAllModels(modelList);
return modelList;
} catch (error) {
console.error("Failed to fetch models in current job:", error);
useChatStore.getState().setAllModelsInCurrentJob([]);
console.error("Failed to fetch all models:", error);
useChatStore.getState().setAllModels([]);
return [];
}
}

export async function chatRequest(abortSignal?: AbortSignal) {
const currentJob = useChatStore.getState().currentJob;
if (!currentJob) {
console.warn("No current job selected");
return;
}

const jobServerPath = useChatStore.getState().jobServerPath;
const jobServerToken = useChatStore.getState().jobServerToken;
const restServerToken = useChatStore.getState().restServerToken;
const modelProxyPath = useChatStore.getState().modelProxyPath;
const currentModel = useChatStore.getState().currentModel;
const chatMsgs = useChatStore.getState().chatMsgs;

const modelsUrl = `${jobServerPath}/${currentJob.ip}:${currentJob.port}/v1/chat/completions`;
const modelsUrl = `${modelProxyPath}/v1/chat/completions`;
const data = {
model: currentModel,
messages: chatMsgs.filter(
Expand All @@ -168,7 +105,7 @@ export async function chatRequest(abortSignal?: AbortSignal) {
})),
};
try {
const response: any = await api.post(modelsUrl, jobServerToken, {
const response: any = await api.post(modelsUrl, restServerToken, {
json: data,
signal: abortSignal,
});
Expand All @@ -195,18 +132,13 @@ export async function chatRequest(abortSignal?: AbortSignal) {
}

export async function chatStreamRequest(abortSignal?: AbortSignal) {
const currentJob = useChatStore.getState().currentJob;
if (!currentJob) {
console.warn("No current job selected");
return;
}

const jobServerPath = useChatStore.getState().jobServerPath;
const jobServerToken = useChatStore.getState().jobServerToken;
const restServerToken = useChatStore.getState().restServerToken;
const modelProxyPath = useChatStore.getState().modelProxyPath;
const currentModel = useChatStore.getState().currentModel;
const chatMsgs = useChatStore.getState().chatMsgs;

const modelsUrl = `${jobServerPath}/${currentJob.ip}:${currentJob.port}/v1/chat/completions`;
const modelsUrl = `${modelProxyPath}/v1/chat/completions`;

const data = {
model: currentModel,
stream: true,
Expand All @@ -228,7 +160,7 @@ export async function chatStreamRequest(abortSignal?: AbortSignal) {
useChatStore.getState().addChatMessage(newMsg);

try {
const response = await api.postStream(modelsUrl, jobServerToken, {
const response = await api.postStream(modelsUrl, restServerToken, {
json: data,
signal: abortSignal,
timeout: false, // Disable timeout for streaming
Expand Down
Loading