@codai/memorai-core
Version:
Core memory engine with vector operations for AI agents
292 lines • 10.1 kB
JavaScript
import { QdrantClient } from '@qdrant/js-client-rest';
import { VectorStoreError } from '../types/index.js';
export class QdrantVectorStore {
client;
collection;
dimension;
constructor(url, collection, dimension, apiKey) {
const clientConfig = { url };
if (apiKey) {
clientConfig.apiKey = apiKey;
}
this.client = new QdrantClient(clientConfig);
this.collection = collection;
this.dimension = dimension;
}
async initialize() {
try {
// Check if collection exists
const collections = await this.client.getCollections();
const exists = collections.collections.some((c) => c.name === this.collection);
if (!exists) {
await this.client.createCollection(this.collection, {
vectors: {
size: this.dimension,
distance: 'Cosine',
},
optimizers_config: {
default_segment_number: 2,
max_segment_size: 20000,
memmap_threshold: 50000,
indexing_threshold: 20000,
flush_interval_sec: 5,
},
hnsw_config: {
m: 16,
ef_construct: 100,
full_scan_threshold: 10000,
},
});
// Create indexes for filtering
await this.client.createPayloadIndex(this.collection, {
field_name: 'tenant_id',
field_schema: 'keyword',
});
await this.client.createPayloadIndex(this.collection, {
field_name: 'type',
field_schema: 'keyword',
});
await this.client.createPayloadIndex(this.collection, {
field_name: 'created_at',
field_schema: 'datetime',
});
}
}
catch (error) {
if (error instanceof Error) {
throw new VectorStoreError(`Failed to initialize Qdrant collection: ${error.message}`);
}
throw new VectorStoreError('Unknown initialization error');
}
}
async upsert(points) {
if (points.length === 0) {
return;
}
try {
const qdrantPoints = points.map(point => ({
id: point.id,
vector: point.vector,
payload: point.payload,
}));
await this.client.upsert(this.collection, {
wait: true,
points: qdrantPoints,
});
}
catch (error) {
if (error instanceof Error) {
throw new VectorStoreError(`Failed to upsert points: ${error.message}`, {
point_count: points.length,
});
}
throw new VectorStoreError('Unknown upsert error');
}
}
async search(vector, query) {
try {
const filter = {
must: [
{
key: 'tenant_id',
match: { value: query.tenant_id },
},
],
};
// Add optional filters
if (query.type) {
filter.must.push({
key: 'type',
match: { value: query.type },
});
}
if (query.agent_id) {
filter.must.push({
key: 'agent_id',
match: { value: query.agent_id },
});
}
const searchResult = await this.client.search(this.collection, {
vector,
filter,
limit: query.limit,
score_threshold: query.threshold,
with_payload: true,
});
return searchResult.map((point) => ({
id: point.id,
score: point.score,
payload: point.payload ?? {},
}));
}
catch (error) {
if (error instanceof Error) {
throw new VectorStoreError(`Search failed: ${error.message}`, {
query: query.query.substring(0, 100),
tenant_id: query.tenant_id,
});
}
throw new VectorStoreError('Unknown search error');
}
}
async delete(ids) {
if (ids.length === 0) {
return;
}
try {
await this.client.delete(this.collection, {
wait: true,
points: ids,
});
}
catch (error) {
if (error instanceof Error) {
throw new VectorStoreError(`Failed to delete points: ${error.message}`, {
id_count: ids.length,
});
}
throw new VectorStoreError('Unknown delete error');
}
}
async count(tenantId) {
try {
const result = await this.client.count(this.collection, {
filter: {
must: [
{
key: 'tenant_id',
match: { value: tenantId },
},
],
},
});
return result.count;
}
catch (error) {
if (error instanceof Error) {
throw new VectorStoreError(`Failed to count points: ${error.message}`, {
tenant_id: tenantId,
});
}
throw new VectorStoreError('Unknown count error');
}
}
async healthCheck() {
try {
const collections = await this.client.getCollections();
return collections.collections.some((c) => c.name === this.collection);
}
catch {
return false;
}
}
}
export class MemoryVectorStore {
store;
isInitialized = false;
constructor(store) {
this.store = store;
}
async initialize() {
if (this.isInitialized) {
return;
}
await this.store.initialize();
this.isInitialized = true;
}
async storeMemory(memory, embedding) {
await this.ensureInitialized();
const point = {
id: memory.id,
vector: embedding,
payload: {
...memory,
created_at: memory.createdAt.toISOString(),
updated_at: memory.updatedAt.toISOString(),
last_accessed_at: memory.lastAccessedAt.toISOString(),
ttl: memory.ttl?.toISOString(),
},
};
await this.store.upsert([point]);
}
async storeMemories(memories, embeddings) {
await this.ensureInitialized();
if (memories.length !== embeddings.length) {
throw new VectorStoreError('Memories and embeddings count mismatch');
}
const points = memories.map((memory, index) => ({
id: memory.id,
vector: embeddings[index],
payload: {
...memory,
created_at: memory.createdAt.toISOString(),
updated_at: memory.updatedAt.toISOString(),
last_accessed_at: memory.lastAccessedAt.toISOString(),
ttl: memory.ttl?.toISOString(),
},
}));
await this.store.upsert(points);
}
async searchMemories(queryEmbedding, query) {
await this.ensureInitialized();
const results = await this.store.search(queryEmbedding, query);
return results.map(result => {
const payload = result.payload;
return {
memory: {
id: payload.id,
type: payload.type,
content: payload.content,
confidence: payload.confidence,
createdAt: new Date(payload.created_at),
updatedAt: new Date(payload.updated_at),
lastAccessedAt: new Date(payload.last_accessed_at),
accessCount: payload.accessCount,
importance: payload.importance,
emotional_weight: payload.emotional_weight,
tags: payload.tags,
context: payload.context,
tenant_id: payload.tenant_id,
agent_id: payload.agent_id,
ttl: payload.ttl ? new Date(payload.ttl) : undefined,
},
score: result.score,
relevance_reason: this.generateRelevanceReason(result.score, query.query),
};
});
}
async deleteMemories(ids) {
await this.ensureInitialized();
await this.store.delete(ids);
}
async getMemoryCount(tenantId) {
await this.ensureInitialized();
return this.store.count(tenantId);
}
async healthCheck() {
if (!this.isInitialized) {
return false;
}
return this.store.healthCheck();
}
async ensureInitialized() {
if (!this.isInitialized) {
await this.initialize();
}
}
generateRelevanceReason(score, query) {
if (score >= 0.9) {
return `Highly relevant to "${query}" with excellent semantic match`;
}
else if (score >= 0.8) {
return `Strong relevance to "${query}" with good semantic similarity`;
}
else if (score >= 0.7) {
return `Moderately relevant to "${query}" with decent semantic overlap`;
}
else {
return `Some relevance to "${query}" but weaker semantic connection`;
}
}
}
//# sourceMappingURL=VectorStore.js.map