Skip to content

Commit

Permalink
fix(kb): Embedding mixed response handled
Browse files Browse the repository at this point in the history
  • Loading branch information
RezaRahemtola committed Jan 30, 2025
1 parent 8589beb commit 357526c
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 5 deletions.
2 changes: 1 addition & 1 deletion src/types/knowledge.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ export type KnowledgeSearchResult = {

const knowledgeDocumentChunk = z.object({
content: z.string(),
vector: z.array(z.number()),
vector: z.array(z.number()).optional(), // TODO: remove optional once issue fixed with Justin C.
});

export type KnowledgeDocumentChunk = z.infer<typeof knowledgeDocumentChunk>;
Expand Down
13 changes: 9 additions & 4 deletions src/utils/knowledge/embedding.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ export const generateChunks = async (
const embedding_vector = await embed(chunk.pageContent);
result.push({
content: chunk.pageContent,
vector: embedding_vector,
vector: embedding_vector ?? [],
});
}

Expand Down Expand Up @@ -68,7 +68,7 @@ export const searchDocuments = async (
// Iterate over all embeddings
documents.forEach((document) => {
document.chunks.forEach((chunk) => {
const euclidean_distance = distance.euclidean(query_vector, chunk.vector);
const euclidean_distance = distance.euclidean(query_vector, chunk.vector ?? []);

// If the distance is greater than the max_distance, skip it
if (euclidean_distance > max_distance) return;
Expand All @@ -88,11 +88,16 @@ async function embed(content: string): Promise<number[]> {
const errors = [];
for (let i = 0; i < tries; i++) {
try {
const response = await axios.post<{ embedding: number[] }>(DEFAULT_EMBEDDING_API_URL, {
const response = await axios.post<
| { embedding: number[][] }[]
| {
embedding: number[];
}
>(DEFAULT_EMBEDDING_API_URL, {
content,
});

return response.data.embedding;
return Array.isArray(response.data) ? response.data[0].embedding[0] : response.data.embedding;
} catch (error) {
errors.push(error);
console.error(`Error embedding text: ${error}`);
Expand Down

0 comments on commit 357526c

Please sign in to comment.