This commit is contained in:
Hk-Gosuto
2024-04-07 18:00:21 +08:00
parent 7382ce48bb
commit b00e9f0c79
17 changed files with 307 additions and 122 deletions

View File

@@ -4,6 +4,7 @@ import { ACCESS_CODE_PREFIX, ModelProvider } from "@/app/constant";
import { OpenAIEmbeddings } from "@langchain/openai";
import { Pinecone } from "@pinecone-database/pinecone";
import { PineconeStore } from "@langchain/pinecone";
import { QdrantVectorStore } from "@langchain/community/vectorstores/qdrant";
import { getServerSideConfig } from "@/app/config/server";
interface RequestBody {
@@ -27,26 +28,40 @@ async function handle(req: NextRequest) {
const reqBody: RequestBody = await req.json();
const authToken = req.headers.get("Authorization") ?? "";
const token = authToken.trim().replaceAll("Bearer ", "").trim();
const pinecone = new Pinecone();
const pineconeIndex = pinecone.Index(process.env.PINECONE_INDEX!);
const serverConfig = getServerSideConfig();
// const pinecone = new Pinecone();
// const pineconeIndex = pinecone.Index(serverConfig.pineconeIndex!);
const apiKey = getOpenAIApiKey(token);
const baseUrl = getOpenAIBaseUrl(reqBody.baseUrl);
const embeddings = new OpenAIEmbeddings(
{
modelName: process.env.RAG_EMBEDDING_MODEL ?? "text-embedding-3-large",
modelName: serverConfig.ragEmbeddingModel ?? "text-embedding-3-large",
openAIApiKey: apiKey,
},
{ basePath: baseUrl },
);
const vectorStore = await PineconeStore.fromExistingIndex(embeddings, {
pineconeIndex,
});
const results = await vectorStore.similaritySearch(reqBody.query, 1, {
sessionId: reqBody.sessionId,
});
console.log(results);
return NextResponse.json(results, {
// const vectorStore = await PineconeStore.fromExistingIndex(embeddings, {
// pineconeIndex,
// });
// const results = await vectorStore.similaritySearch(reqBody.query, 4, {
// sessionId: reqBody.sessionId,
// });
const vectorStore = await QdrantVectorStore.fromExistingCollection(
embeddings,
{
url: process.env.QDRANT_URL,
apiKey: process.env.QDRANT_API_KEY,
collectionName: reqBody.sessionId,
},
);
const returnCunt = serverConfig.ragReturnCount
? parseInt(serverConfig.ragReturnCount, 10)
: 4;
const response = await vectorStore.similaritySearch(
reqBody.query,
returnCunt,
);
return NextResponse.json(response, {
status: 200,
});
} catch (e) {