@escher-dbai/rag-module
Version:
Enterprise RAG module with chat context storage, vector search, and session management. Complete chat history retrieval and streaming content extraction for Electron apps.
270 lines (235 loc) • 9.84 kB
JavaScript
/**
* ContextRetrievalService - Service for retrieving conversation context from Qdrant
*
* This service handles retrieval of query-response pairs from Qdrant vector store
* using context ID filters and other search criteria.
*/
class ContextRetrievalService {
/**
* @param {Object} qdrantClient - Qdrant client instance
* @param {string} collectionName - Default collection name
*/
constructor(qdrantClient, collectionName = 'chat-collection') {
this.qdrantClient = qdrantClient;
this.collectionName = collectionName;
}
/**
* Get organized query-response pairs for a specific context ID
* @param {string|Array} contextId - Context ID(s) to retrieve (string or array of strings)
* @param {number|Object} [limitOrOptions=100] - Maximum pairs to return OR options object
* @param {Object} [options={}] - Additional options (when second param is number)
* @param {string} [options.userId] - Optional user ID filter for security
* @param {string} [options.collectionName] - Override default collection name
* @returns {Promise<{contextId: string, chatTitle: string, queries: Array, responses: Array, pairs: Array, stats: Object}>} - Organized query-response data
*/
async getQueryResponsePairs(contextId, limitOrOptions = 100, options = {}) {
// Handle flexible parameter formats
let actualContextId, actualOptions, actualLimit;
if (Array.isArray(contextId)) {
actualContextId = contextId[0]; // Take first element if array
} else {
actualContextId = contextId;
}
if (typeof limitOrOptions === 'number') {
actualLimit = limitOrOptions;
actualOptions = options;
} else {
actualLimit = limitOrOptions.limit || 100;
actualOptions = limitOrOptions;
}
const { userId, collectionName = this.collectionName } = actualOptions;
try {
// Build filter conditions
const filterConditions = [
{
key: "context_id",
match: {
value: actualContextId
}
}
];
// Add user ID filter if provided for security
if (userId) {
filterConditions.push({
key: "user_id",
match: {
value: userId
}
});
}
// Search for documents with the context ID
const searchResults = await this.qdrantClient.scroll(collectionName, {
filter: {
must: filterConditions
},
limit: actualLimit,
with_payload: true,
with_vector: false
});
// Extract conversation messages and chat title
const messages = [];
let chatTitle = '';
console.log('🔍 ContextRetrieval: Processing searchResults.points:', searchResults.points ? searchResults.points.length : 0);
if (searchResults.points) {
for (const point of searchResults.points) {
const payload = point.payload;
console.log('🔍 ContextRetrieval: Processing payload for point:', point.id);
console.log('🔍 ContextRetrieval: Payload keys:', Object.keys(payload));
// Extract chat title
if (payload.chat_title) {
chatTitle = payload.chat_title;
console.log('🔍 ContextRetrieval: Found chat_title:', chatTitle);
} else if (payload.final_conversation_context && payload.final_conversation_context.ct) {
chatTitle = payload.final_conversation_context.ct;
console.log('🔍 ContextRetrieval: Found chat title from final_conversation_context.ct:', chatTitle);
}
// Check if this document contains conversation messages
if (payload.final_conversation_context && payload.final_conversation_context.m) {
console.log('🔍 ContextRetrieval: Found final_conversation_context.m with', payload.final_conversation_context.m.length, 'messages');
// Extract messages from the conversation context
for (const msg of payload.final_conversation_context.m) {
messages.push({
role: msg.r, // 0 = prompt/query, 1 = response
content: msg.c,
timestamp: payload.timestamp || payload.final_conversation_context.t,
messageId: `${contextId}_${msg.r}_${Date.now()}`
});
console.log('🔍 ContextRetrieval: Added message with role:', msg.r, 'content length:', msg.c.length);
}
} else if (payload.conversation_state_after_filtering && payload.conversation_state_after_filtering.total_messages > 0) {
console.log('🔍 ContextRetrieval: Using fallback - processing_results');
// Fallback: try to extract from processing results
if (payload.processing_results && payload.processing_results.clean_response) {
messages.push({
role: 1, // Response
content: payload.processing_results.clean_response,
timestamp: payload.timestamp,
messageId: `${contextId}_response_${Date.now()}`
});
console.log('🔍 ContextRetrieval: Added fallback response message');
}
} else {
console.log('🔍 ContextRetrieval: No conversation messages found in payload');
}
}
}
console.log('🔍 ContextRetrieval: Total messages extracted:', messages.length);
// Sort messages by timestamp if available
messages.sort((a, b) => {
if (a.timestamp && b.timestamp) {
return new Date(a.timestamp) - new Date(b.timestamp);
}
return 0;
});
const queries = []; // role: 0 (prompts/questions)
const responses = []; // role: 1 (assistant responses)
const pairs = []; // matched query-response pairs
// Separate queries and responses
for (const message of messages) {
if (message.role === 0) {
queries.push({
...message,
type: 'query'
});
} else if (message.role === 1) {
responses.push({
...message,
type: 'response'
});
}
}
// Create query-response pairs by matching sequential messages
for (let i = 0; i < Math.min(queries.length, responses.length); i++) {
pairs.push({
pairId: `${contextId}_pair_${i + 1}`,
query: queries[i],
response: responses[i],
pairIndex: i + 1
});
}
console.log('🔍 ContextRetrieval: Final results:');
console.log(' - Total messages:', messages.length);
console.log(' - Queries:', queries.length);
console.log(' - Responses:', responses.length);
console.log(' - Pairs:', pairs.length);
console.log(' - Chat title:', chatTitle);
if (pairs.length > 0) {
console.log('🔍 ContextRetrieval: Sample pair:');
console.log(' - Query:', pairs[0].query.content.substring(0, 100) + '...');
console.log(' - Response:', pairs[0].response.content.substring(0, 100) + '...');
}
return {
contextId: actualContextId,
chatTitle,
queries,
responses,
pairs,
stats: {
totalMessages: messages.length,
queryCount: queries.length,
responseCount: responses.length,
pairCount: pairs.length,
retrievedAt: new Date().toISOString()
}
};
} catch (error) {
throw new Error(`Failed to retrieve context messages from Qdrant: ${error.message}`);
}
}
/**
* Search for contexts by various criteria
* @param {Object} searchCriteria - Search criteria
* @param {string} [searchCriteria.userId] - User ID to filter by
* @param {string} [searchCriteria.agentType] - Agent type to filter by
* @param {Date} [searchCriteria.fromDate] - Start date filter
* @param {Date} [searchCriteria.toDate] - End date filter
* @param {number} [searchCriteria.limit=50] - Maximum contexts to return
* @param {string} [searchCriteria.collectionName] - Override default collection name
* @returns {Promise<Array<string>>} - List of context IDs matching criteria
*/
async findContextIds(searchCriteria = {}) {
const { userId, agentType, fromDate, toDate, limit = 50, collectionName = this.collectionName } = searchCriteria;
try {
const filterConditions = [];
// Add filters based on criteria
if (userId) {
filterConditions.push({
key: "user_id",
match: { value: userId }
});
}
if (agentType) {
filterConditions.push({
key: "agent_type",
match: { value: agentType }
});
}
// Date range filter (if timestamps are stored as Unix timestamps)
if (fromDate || toDate) {
const rangeFilter = { key: "timestamp", range: {} };
if (fromDate) rangeFilter.range.gte = Math.floor(fromDate.getTime() / 1000);
if (toDate) rangeFilter.range.lte = Math.floor(toDate.getTime() / 1000);
filterConditions.push(rangeFilter);
}
const scrollResults = await this.qdrantClient.scroll(collectionName, {
filter: filterConditions.length > 0 ? { must: filterConditions } : undefined,
limit: limit,
with_payload: true,
with_vector: false
});
// Extract unique context IDs
const contextIds = new Set();
if (scrollResults.points) {
for (const point of scrollResults.points) {
if (point.payload.context_id) {
contextIds.add(point.payload.context_id);
}
}
}
return Array.from(contextIds);
} catch (error) {
throw new Error(`Failed to find context IDs: ${error.message}`);
}
}
}
module.exports = ContextRetrievalService;