diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/client/views/nodes/ChatBox/Agent.ts | 2 | ||||
| -rw-r--r-- | src/client/views/nodes/ChatBox/ChatBox.tsx | 2 | ||||
| -rw-r--r-- | src/client/views/nodes/ChatBox/tools/RAGTool.ts | 2 | ||||
| -rw-r--r-- | src/client/views/nodes/ChatBox/vectorstore/Vectorstore.ts (renamed from src/client/views/nodes/ChatBox/vectorstore/VectorstoreUpload.ts) | 38 | ||||
| -rw-r--r-- | src/server/ApiManagers/AssistantManager.ts | 54 |
5 files changed, 77 insertions, 21 deletions
diff --git a/src/client/views/nodes/ChatBox/Agent.ts b/src/client/views/nodes/ChatBox/Agent.ts index 04729414a..7b3703449 100644 --- a/src/client/views/nodes/ChatBox/Agent.ts +++ b/src/client/views/nodes/ChatBox/Agent.ts @@ -6,7 +6,7 @@ import { WikipediaTool } from './tools/WikipediaTool'; import { CalculateTool } from './tools/CalculateTool'; import { RAGTool } from './tools/RAGTool'; import { NoTool } from './tools/NoTool'; -import { Vectorstore } from './vectorstore/VectorstoreUpload'; +import { Vectorstore } from './vectorstore/Vectorstore'; import { ChatCompletionAssistantMessageParam, ChatCompletionMessageParam } from 'openai/resources'; import dotenv from 'dotenv'; import { ChatBox } from './ChatBox'; diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx index 3de5c1da3..13c418b32 100644 --- a/src/client/views/nodes/ChatBox/ChatBox.tsx +++ b/src/client/views/nodes/ChatBox/ChatBox.tsx @@ -12,7 +12,7 @@ import { FieldView, FieldViewProps } from '../FieldView'; import './ChatBox.scss'; import MessageComponentBox from './MessageComponent'; import { ASSISTANT_ROLE, AssistantMessage, AI_Document, Citation, CHUNK_TYPE, Chunk, getChunkType, TEXT_TYPE } from './types'; -import { Vectorstore } from './vectorstore/VectorstoreUpload'; +import { Vectorstore } from './vectorstore/Vectorstore'; import { Agent } from './Agent'; import dotenv from 'dotenv'; import { DocData, DocViews } from '../../../../fields/DocSymbols'; diff --git a/src/client/views/nodes/ChatBox/tools/RAGTool.ts b/src/client/views/nodes/ChatBox/tools/RAGTool.ts index 23b93b0f0..be591fa9a 100644 --- a/src/client/views/nodes/ChatBox/tools/RAGTool.ts +++ b/src/client/views/nodes/ChatBox/tools/RAGTool.ts @@ -1,5 +1,5 @@ import { BaseTool } from './BaseTool'; -import { Vectorstore } from '../vectorstore/VectorstoreUpload'; +import { Vectorstore } from '../vectorstore/Vectorstore'; import { Chunk } from '../types'; import * as fs from 'fs'; import { Networking } from '../../../../Network'; diff --git a/src/client/views/nodes/ChatBox/vectorstore/VectorstoreUpload.ts b/src/client/views/nodes/ChatBox/vectorstore/Vectorstore.ts index 787705bb6..25aec751f 100644 --- a/src/client/views/nodes/ChatBox/vectorstore/VectorstoreUpload.ts +++ b/src/client/views/nodes/ChatBox/vectorstore/Vectorstore.ts @@ -2,6 +2,7 @@ import { Pinecone, Index, IndexList, PineconeRecord, RecordMetadata, QueryRespon import { CohereClient } from 'cohere-ai'; import { EmbedResponse } from 'cohere-ai/api'; import dotenv from 'dotenv'; +import axios from 'axios'; import { Chunk, AI_Document, CHUNK_TYPE } from '../types'; import { Doc } from '../../../../../fields/Doc'; @@ -43,7 +44,7 @@ export class Vectorstore { if (!indexList.indexes?.some(index => index.name === this.indexName)) { await this.pinecone.createIndex({ name: this.indexName, - dimension: 1024, + dimension: 768, metric: 'cosine', spec: { serverless: { @@ -138,25 +139,26 @@ export class Vectorstore { async retrieve(query: string, topK: number = 10): Promise<Chunk[]> { console.log(`Retrieving chunks for query: ${query}`); try { - const queryEmbeddingResponse: EmbedResponse = await this.cohere.embed({ - texts: [query], - model: 'embed-english-v3.0', - inputType: 'search_query', - }); - - let queryEmbedding: number[]; - - if (Array.isArray(queryEmbeddingResponse.embeddings)) { - queryEmbedding = queryEmbeddingResponse.embeddings[0]; - } else if (queryEmbeddingResponse.embeddings && 'embeddings' in queryEmbeddingResponse.embeddings) { - queryEmbedding = (queryEmbeddingResponse.embeddings as { embeddings: number[][] }).embeddings[0]; - } else { - throw new Error('Invalid embedding response format'); + const url = 'https://api.jina.ai/v1/embeddings'; + const headers = { + 'Content-Type': 'application/json', + Authorization: `Bearer ${process.env.JINA_API_KEY}`, + }; + const data = { + model: 'jina-clip-v1', + normalized: true, + embedding_type: 'float', + input: [{ text: query }], + }; + + const response = await axios.post(url, data, { headers }); + const embeddings = response.data?.data?.[0]?.embedding; + + if (!embeddings || !Array.isArray(embeddings)) { + throw new Error('Invalid embedding response format from Jina API'); } - if (!Array.isArray(queryEmbedding)) { - throw new Error('Query embedding is not an array'); - } + const queryEmbedding = embeddings; const queryResponse: QueryResponse<RecordMetadata> = await this.index.query({ vector: queryEmbedding, diff --git a/src/server/ApiManagers/AssistantManager.ts b/src/server/ApiManagers/AssistantManager.ts index 36468157a..f69ca1383 100644 --- a/src/server/ApiManagers/AssistantManager.ts +++ b/src/server/ApiManagers/AssistantManager.ts @@ -9,6 +9,9 @@ import { Method } from '../RouteManager'; import ApiManager, { Registration } from './ApiManager'; import axios from 'axios'; import { Chunk } from '../../client/views/nodes/ChatBox/types'; +import { UnstructuredClient } from 'unstructured-client'; +import { PartitionResponse } from 'unstructured-client/sdk/models/operations'; +import { ChunkingStrategy, Strategy } from 'unstructured-client/sdk/models/shared'; export enum Directory { parsed_files = 'parsed_files', @@ -42,6 +45,11 @@ export default class AssistantManager extends ApiManager { apiKey: process.env._CLIENT_OPENAI_KEY, // Use client key so don't have to set key seperately for client and server. dangerouslyAllowBrowser: true, }); + const unstructuredClient = new UnstructuredClient({ + security: { + apiKeyAuth: process.env._CLIENT_UNSTRUCTURED_API_KEY!, + }, + }); register({ method: Method.POST, @@ -187,5 +195,51 @@ export default class AssistantManager extends ApiManager { res.send({ formattedChunks: content }); }, }); + + register({ + method: Method.POST, + subscription: '/chunkDocument', + secureHandler: async ({ req, res }) => { + const { file_path } = req.body; + const public_path = path.join(publicDirectory, file_path); + const file_name = path.basename(file_path); + + try { + // Read file data and convert to base64 + const file_data = await fs.promises.readFile(public_path); + + try { + const result = await unstructuredClient.general.partition({ + partitionParameters: { + files: { + content: file_data, + fileName: file_name, + }, + strategy: Strategy.Auto, + chunkingStrategy: ChunkingStrategy.ByTitle, + extractImageBlockTypes: ['Image', 'Table'], + }, + }); + + if (result.statusCode === 200) { + console.log(result.elements); + const jsonElements = JSON.stringify(result.elements, null, 2); + // Print the processed data. + console.log(jsonElements); + res.send({ document_json: jsonElements }); + } else { + console.error(`Unexpected status code: ${result.statusCode}`); + res.status(result.statusCode).send({ error: 'Failed to process the document', details: result }); + } + } catch (e: any) { + console.error('Error during partitioning:', e); + res.status(500).send({ error: 'Failed to partition the document', details: e.message }); + } + } catch (error: any) { + console.error('Error reading file:', error); + res.status(500).send({ error: 'Failed to read the file', details: error.message }); + } + }, + }); } } |