Skip to content

Commit 2efb856

Browse files
maamokunmaamokun
committed
feat: Ollama support
1 parent 4c86659 commit 2efb856

6 files changed

Lines changed: 508 additions & 5 deletions

File tree

‎src/actions/ollama.ts‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
"use server";
2+
3+
import { getLogtoContext } from "@logto/next/server-actions";
4+
import { logtoConfig } from "@/lib/auth";
5+
import { prisma } from "@/lib/prisma";
6+
7+
export const validateOllamaHost = async (host: string) => {
8+
const { claims } = await getLogtoContext(logtoConfig);
9+
10+
if (!claims) {
11+
throw new Error("User not authenticated");
12+
}
13+
14+
try {
15+
const response = await fetch(`${host}/v1/models`);
16+
if (!response.ok) {
17+
throw new Error("Invalid Ollama host");
18+
}
19+
const data = await response.json();
20+
// @ts-ignore
21+
if (!data || !data.data) {
22+
throw new Error("No data returned from Ollama host");
23+
}
24+
return {
25+
isValid: true,
26+
// @ts-ignore
27+
modelCount: data.data.length,
28+
// @ts-ignore
29+
models: data.data,
30+
};
31+
} catch (error) {
32+
console.error("Error validating Ollama host:", error);
33+
return {
34+
isValid: false,
35+
modelCount: 0,
36+
models: [],
37+
};
38+
}
39+
};
40+
41+
export const addModel = async (host: string) => {
42+
const { claims } = await getLogtoContext(logtoConfig);
43+
44+
if (!claims) {
45+
throw new Error("User not authenticated");
46+
}
47+
48+
const validation = await validateOllamaHost(host);
49+
if (!validation.isValid) {
50+
throw new Error("Invalid Ollama host");
51+
}
52+
53+
const existingModel = await prisma.customProvider.findFirst({
54+
where: {
55+
endpoint: host,
56+
userId: claims.sub,
57+
type: "ollama",
58+
},
59+
});
60+
61+
if (existingModel) {
62+
throw new Error("Model already exists");
63+
}
64+
65+
return prisma.customProvider.create({
66+
data: {
67+
endpoint: host,
68+
userId: claims.sub,
69+
type: "ollama",
70+
},
71+
});
72+
};
73+
74+
export const deleteModel = async (id: string) => {
75+
const { claims } = await getLogtoContext(logtoConfig);
76+
77+
if (!claims) {
78+
throw new Error("User not authenticated");
79+
}
80+
81+
const model = await prisma.customProvider.findUnique({
82+
where: { id },
83+
});
84+
85+
if (!model || model.userId !== claims.sub) {
86+
throw new Error("Model not found or access denied");
87+
}
88+
89+
return prisma.customProvider.delete({
90+
where: { id },
91+
});
92+
};

‎src/app/(chat)/chat/page.tsx‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,15 @@ import { ChatInput } from "@/components/chatInput";
55
import * as React from "react";
66
import { prisma } from "@/lib/prisma";
77
import models from "@/consts/models.json";
8+
import { validateOllamaHost } from "@/actions/ollama";
9+
10+
//@ts-ignore
11+
interface OllamaModelWithStatus extends prisma.customProvider {
12+
isAvailable: boolean;
13+
endpoint: string;
14+
models: string[];
15+
id: string;
16+
}
817

918
export default async function Home() {
1019
const { claims } = await getLogtoContext(logtoConfig);
@@ -20,6 +29,43 @@ export default async function Home() {
2029
},
2130
});
2231

32+
const ollamaHosts = await prisma.customProvider.findMany({
33+
where: {
34+
userId: claims?.sub || "",
35+
type: "ollama",
36+
},
37+
});
38+
39+
const ollamaModels: OllamaModelWithStatus[] = [];
40+
41+
for (const host of ollamaHosts) {
42+
if (!host.endpoint) {
43+
ollamaModels.push({
44+
...host,
45+
endpoint: "",
46+
isAvailable: false,
47+
models: [],
48+
});
49+
continue;
50+
}
51+
try {
52+
const hostInfo = await validateOllamaHost(host.endpoint);
53+
ollamaModels.push({
54+
...host,
55+
endpoint: host.endpoint,
56+
isAvailable: hostInfo.isValid,
57+
models: hostInfo.models,
58+
});
59+
} catch (error) {
60+
ollamaModels.push({
61+
...host,
62+
endpoint: host.endpoint,
63+
isAvailable: false,
64+
models: [],
65+
});
66+
}
67+
}
68+
2369
const availableModels = Object.entries(models)
2470
.filter(([providerKey]) =>
2571
userKeys.some((key) => key.providerId === providerKey),
@@ -57,6 +103,7 @@ export default async function Home() {
57103
status={"ready"}
58104
openRouterModels={openRouterModels}
59105
openRouterEnabled={!!openRouterKey}
106+
ollamaModels={ollamaModels}
60107
/>
61108
</div>
62109
</main>

‎src/app/(settings)/settings/models/page.tsx‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,19 @@ import { redirect } from "next/navigation";
88

99
import { ModelCard } from "@/components/modelCard";
1010
import { OpenRouterModelList } from "@/components/OpenRouterModelList";
11+
import { OllamaModelList } from "@/components/OllamaModelList";
1112
import { OpenRouterBadge } from "@/components/OpenRouterConfigButton";
1213

14+
import { validateOllamaHost } from "@/actions/ollama";
15+
16+
//@ts-ignore
17+
interface OllamaModelWithStatus extends prisma.customProvider {
18+
isAvailable: boolean;
19+
endpoint: string;
20+
modelCount: number;
21+
id: string;
22+
}
23+
1324
export default async function Home() {
1425
const { claims } = await getLogtoContext(logtoConfig);
1526

@@ -23,12 +34,12 @@ export default async function Home() {
2334
},
2435
});
2536

26-
const userOllamaModels = await prisma.customProvider.findMany({
37+
let userOllamaModels = (await prisma.customProvider.findMany({
2738
where: {
2839
userId: claims?.sub,
2940
type: "ollama",
3041
},
31-
});
42+
})) as unknown as OllamaModelWithStatus[];
3243

3344
const userOpenRouterModels = await prisma.customProvider.findMany({
3445
where: {
@@ -37,6 +48,16 @@ export default async function Home() {
3748
},
3849
});
3950

51+
for (const model of userOllamaModels) {
52+
try {
53+
const hostInfo = await validateOllamaHost(model.endpoint);
54+
model.isAvailable = hostInfo.isValid;
55+
model.modelCount = hostInfo.modelCount;
56+
} catch (error) {
57+
model.isAvailable = false;
58+
}
59+
}
60+
4061
return (
4162
<div className="container mx-auto p-4">
4263
<h1 className={"text-base-content text-4xl"}>Model Settings</h1>
@@ -71,6 +92,9 @@ export default async function Home() {
7192
<h2 className={"text-base-content text-xl mt-5 mb-2"}>
7293
Custom Provider
7394
</h2>
95+
<div className="grid space-y-2">
96+
<OllamaModelList models={userOllamaModels} />
97+
</div>
7498
</div>
7599
);
76100
}

‎src/app/api/chat/route.ts‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { createAnthropic } from "@ai-sdk/anthropic";
1010
import { createXai } from "@ai-sdk/xai";
1111
import { createGroq } from "@ai-sdk/groq";
1212
import { createOpenRouter } from "@openrouter/ai-sdk-provider";
13+
import { createOllama } from "ollama-ai-provider";
1314
import { streamText, generateId, createDataStream } from "ai";
1415
import Models from "@/consts/models.json";
1516

@@ -80,6 +81,45 @@ export async function POST(req: NextRequest) {
8081
const modelId = chatData.model;
8182
const modelType = chatData.modelType;
8283

84+
if (modelType === "ollama") {
85+
try {
86+
const endpoint = modelId.split(" - ")[0];
87+
const modelName = modelId.split(" - ")[1];
88+
const provider = createOllama({ baseURL: `${endpoint}/api` });
89+
const streamId = generateId();
90+
await prisma.stream.create({
91+
data: {
92+
chatId: id,
93+
streamId,
94+
},
95+
});
96+
97+
const result = streamText({
98+
model: provider(modelName),
99+
messages,
100+
onFinish: async (message) => {
101+
await addMessage({
102+
message: {
103+
content: message.text,
104+
role: "assistant",
105+
},
106+
id,
107+
});
108+
await prisma.stream.deleteMany({
109+
where: { streamId: streamId },
110+
});
111+
},
112+
});
113+
return result.toDataStreamResponse();
114+
} catch (error) {
115+
console.error("Error in Ollama chat:", error);
116+
return new Response(
117+
"Failed to connect to Ollama model. Please check the endpoint and model name.",
118+
{ status: 500 },
119+
);
120+
}
121+
}
122+
83123
if (modelType === "openrouter") {
84124
const providerKey = await getUserORKey(claims?.sub);
85125
if (!providerKey) {

0 commit comments

Comments
 (0)