-
Notifications
You must be signed in to change notification settings - Fork 20
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: support bitdeer embedding and chat api (#107)
* feat: support bitdeer embedding and chat api * test
- Loading branch information
Showing
7 changed files
with
176 additions
and
22 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,32 @@ | ||
import type {ChatResponse} from "@/core/services/chating"; | ||
import {getLLM} from "@/lib/llamaindex/converters/llm"; | ||
import {baseRegistry} from "@/rag-spec/base"; | ||
import {getFlow} from "@/rag-spec/createFlow"; | ||
import {NextRequest, NextResponse} from 'next/server'; | ||
|
||
const flow = await getFlow(baseRegistry); | ||
|
||
export async function GET (req: NextRequest) { | ||
const url = new URL(req.url); | ||
const query = url.searchParams.get('query') || 'Why TiDB is the advanced MySQL alternative?'; | ||
const llm = getLLM(flow, 'bitdeer', { | ||
model: 'mistral', | ||
apiSecretAccessKey: process.env.BITDEER_API_SECRET_ACCESS_KEY! | ||
}) | ||
const res = await llm.chat({ | ||
messages: [ | ||
{ | ||
role: 'system', | ||
content: 'Keep your answers brief, less than 100 words, emoji is fine.' | ||
}, | ||
{ | ||
role: 'user', | ||
content: query | ||
} | ||
], | ||
stream: false | ||
}); | ||
return NextResponse.json(res.message); | ||
} | ||
|
||
export const dynamic = 'force-dynamic'; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,22 @@ | ||
import {getEmbedding} from "@/lib/llamaindex/converters/embedding"; | ||
import {BitdeerEmbedding, BitdeerEmbeddingModelType} from "@/lib/llamaindex/embeddings/BitdeerEmbedding"; | ||
import {baseRegistry} from "@/rag-spec/base"; | ||
import {getFlow} from "@/rag-spec/createFlow"; | ||
import {NextRequest, NextResponse} from 'next/server'; | ||
|
||
const flow = await getFlow(baseRegistry); | ||
|
||
export async function GET (req: NextRequest) { | ||
const url = new URL(req.url); | ||
const input = url.searchParams.get('input') || 'I want a database to replace MySQL.'; | ||
const bitdeerEmbedding = getEmbedding(flow, 'bitdeer', { | ||
model: BitdeerEmbeddingModelType.MISTRAL_EMBED_LARGE, | ||
apiSecretAccessKey: process.env.BITDEER_API_SECRET_ACCESS_KEY! | ||
}); | ||
|
||
const embedding = await bitdeerEmbedding.getQueryEmbedding(input); | ||
|
||
return NextResponse.json(embedding); | ||
} | ||
|
||
export const dynamic = 'force-dynamic'; |
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,60 @@ | ||
import {BaseEmbedding} from "llamaindex"; | ||
import util from "node:util"; | ||
|
||
|
||
export enum BitdeerEmbeddingModelType { | ||
MISTRAL_EMBED_LARGE = "mxbai-embed-large", | ||
} | ||
|
||
export class BitdeerEmbedding extends BaseEmbedding { | ||
baseURL: string = "https://www.bitdeer.ai/public/v1"; | ||
model: BitdeerEmbeddingModelType; | ||
|
||
apiSecretAccessKey: string; | ||
requestTimeout: number = 60 * 1000; // Default is 60 seconds | ||
|
||
constructor(init?: Partial<BitdeerEmbedding>) { | ||
super(); | ||
this.model = BitdeerEmbeddingModelType.MISTRAL_EMBED_LARGE; | ||
if (typeof init?.apiSecretAccessKey !== "string") { | ||
throw new Error("Bitdeer API secret access key is required."); | ||
} | ||
this.apiSecretAccessKey = init?.apiSecretAccessKey; | ||
} | ||
|
||
private async getBitdeerEmbedding(input: string) { | ||
const payload = { | ||
model: this.model, | ||
prompt: input | ||
}; | ||
const url = `${this.baseURL}/models/${this.model}/generate`; | ||
const response = await fetch(url, { | ||
body: JSON.stringify(payload), | ||
method: "POST", | ||
// signal: AbortSignal.timeout(this.requestTimeout), | ||
headers: { | ||
"Content-Type": "application/json", | ||
"X-Api-Key": this.apiSecretAccessKey, | ||
}, | ||
}); | ||
|
||
if (!response.ok) { | ||
throw new Error(util.format( | ||
'Failed to call Bitdeer embedding API (status: %d, statusText: %s).', | ||
response.status, | ||
response.statusText | ||
)); | ||
} | ||
|
||
const raw = await response.json(); | ||
return raw.data.embedding; | ||
} | ||
|
||
async getTextEmbedding(text: string): Promise<number[]> { | ||
return this.getBitdeerEmbedding(text); | ||
} | ||
|
||
async getQueryEmbedding(query: string): Promise<number[]> { | ||
return this.getBitdeerEmbedding(query); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters