changes

Author: Veikkosuhonen

src/client/components/Rag/Rag.tsx

@@ -1,9 +1,7 @@
 import React, { useState } from 'react'
-import { TextField, Button, Box, Typography, Table, TableHead, TableBody, TableRow, TableCell, Paper, IconButton, Link } from '@mui/material'
+import { TextField, Button, Box, Typography, Table, TableHead, TableBody, TableRow, TableCell, Paper, Link } from '@mui/material'
 import apiClient from '../../util/apiClient'
 import { useMutation, useQuery } from '@tanstack/react-query'
-import { Settings } from '@mui/icons-material'
-import { useSnackbar } from 'notistack'
 import { useNavigate, Link as RouterLink } from 'react-router-dom'
 import { Chunk } from './Chunk'
 
@@ -50,9 +48,8 @@ const useCreateRagIndexMutation = () => {
 }
 
 const Rag: React.FC = () => {
-  const { enqueueSnackbar } = useSnackbar()
   const navigate = useNavigate()
-  const { data: indices, refetch } = useRagIndices()
+  const { data: indices } = useRagIndices()
   const createIndexMutation = useCreateRagIndexMutation()
   const [indexName, setIndexName] = useState('')
   const [selectedIndex, setSelectedIndex] = useState<RagIndexAttributes>(null)

src/config.ts

@@ -12,9 +12,12 @@ export const PUBLIC_URL = process.env.PUBLIC_URL || ''
 
 export const UPDATER_CRON_ENABLED = process.env.UPDATER_CRON_ENABLED === 'true'
 
-export const OLLAMA_URL = process.env.OLLAMA_HOST || 'http://ollama:11434/v1/'
+export const OLLAMA_URL = process.env.OLLAMA_URL || 'http://ollama:11434/v1/'
 export const RAG_ENABLED = process.env.RAG_ENABLED === 'true'
 
+export const LAAMA_API_URL = process.env.LAAMA_API_URL || ''
+export const LAAMA_API_TOKEN = process.env.LAAMA_API_TOKEN || ''
+
 export const DEFAULT_TOKEN_LIMIT = Number(process.env.DEFAULT_TOKEN_LIMIT) || 150_000
 
 export const FREE_MODEL = process.env.FREE_MODEL || 'gpt-4o-mini' // as it was decided in 23th Sept 2024 meeting

src/server/services/rag/embed.ts

@@ -1,5 +1,6 @@
 import { EMBED_DIM, EMBED_MODEL } from '../../../config'
 import type OpenAI from 'openai'
+import { getOllamaClient } from '../../util/ollama'
 
 export const getEmbeddingVector = async (client: OpenAI, query: string) => {
   const response = await client.embeddings.create({
@@ -22,3 +23,13 @@ export const getEmbeddingVectorBatch = async (client: OpenAI, queries: string[])
 
   return response.data
 }
+
+export const getOllamaEmbeddingVectorBatch = async (query: string[]) => {
+  const client = getOllamaClient()
+  const response = await client.embed({
+    model: EMBED_MODEL,
+    input: query,
+  })
+
+  return response.embeddings
+}

src/server/services/rag/ingestion/embedder.ts

@@ -1,7 +1,7 @@
 import type { Chunk } from './chunkingAlgorithms.ts'
 import { mkdirSync } from 'node:fs'
 import { writeFile } from 'node:fs/promises'
-import { getEmbeddingVectorBatch } from '../embed'
+import { getEmbeddingVectorBatch, getOllamaEmbeddingVectorBatch } from '../embed'
 import OpenAI from 'openai'
 
 export type EmbeddedChunk = Chunk & {
@@ -42,13 +42,13 @@ export class Embedder {
   private async embedBatch() {
     const chunkContents = this.currentBatch.map((chunk) => chunk.content.join('\n'))
     const startedAt = Date.now()
-    const result = await getEmbeddingVectorBatch(this.client, chunkContents)
+    const result = await getOllamaEmbeddingVectorBatch(chunkContents)
     const elapsed = Date.now() - startedAt
     console.log(`Embedded ${chunkContents.length} chunks in ${elapsed}ms`)
 
     const embeddedChunks: EmbeddedChunk[] = this.currentBatch.map((chunk, index) => ({
       ...chunk,
-      embedding: result[index].embedding,
+      embedding: result[index],
     }))
 
     await Promise.all(

src/server/util/ollama.ts

@@ -1,6 +1,15 @@
-import { OLLAMA_URL } from '../../config'
+import { Ollama } from 'ollama'
+import { LAAMA_API_TOKEN, LAAMA_API_URL, OLLAMA_URL } from '../../config'
 import OpenAI from 'openai'
 
+export const getOllamaClient = () =>
+  new Ollama({
+    host: LAAMA_API_URL,
+    headers: {
+      token: LAAMA_API_TOKEN,
+    },
+  })
+
 export const getOllamaOpenAIClient = () =>
   new OpenAI({
     apiKey: 'NOT_NEEDED',