adk-typescript
Version:
TypeScript port of Google's Agent Development Kit (ADK)
361 lines (360 loc) • 14.7 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.VertexAiRagMemoryService = void 0;
const fs = __importStar(require("fs"));
const os = __importStar(require("os"));
const path = __importStar(require("path"));
// Import the Vertex AI client library
const vertexai_1 = require("@google-cloud/vertexai");
// Import axios for making REST API calls
const axios_1 = __importDefault(require("axios"));
// Import Google Auth for authentication
const google_auth_library_1 = require("google-auth-library");
// Import FormData for multipart file uploads
const form_data_1 = __importDefault(require("form-data"));
/**
* A memory service that uses Vertex AI RAG capabilities for semantic search.
*
* This implementation uses the Vertex AI embedding models to compute
* embeddings for memory entries and perform semantic searches.
*/
class VertexAiRagMemoryService {
/**
* Creates a new VertexAiRagMemoryService
* @param config The configuration for the service
*/
constructor(config) {
// Store for key-value pairs
this.memoryStore = new Map();
this.config = config;
this.vertexAi = new vertexai_1.VertexAI({
project: config.project,
location: config.location,
});
this.vertexRagStore = {
ragResources: [{ ragCorpus: config.ragCorpus }],
similarityTopK: config.similarityTopK,
vectorDistanceThreshold: config.vectorDistanceThreshold || 10,
};
// Set base URL for REST API calls
this.baseUrl = `https://${config.location}-aiplatform.googleapis.com/v1`;
// Initialize Google Auth client
this.auth = new google_auth_library_1.GoogleAuth({
scopes: ['https://www.googleapis.com/auth/cloud-platform']
});
}
/**
* Adds a session to the memory service.
* @param session The session to add
*/
async addSessionToMemory(session) {
// Create a temporary file to store the session data
const tempDir = os.tmpdir();
const tempFilePath = path.join(tempDir, `session_${session.id}_${Date.now()}.txt`);
try {
const outputLines = [];
for (const event of session.events) {
if (!event.content || !event.content.parts || event.content.parts.length === 0) {
continue;
}
const textParts = event.content.parts
.filter(part => part.text)
.map(part => part.text.replace('\n', ' '));
if (textParts.length > 0) {
outputLines.push(JSON.stringify({
author: event.author,
timestamp: event.timestamp,
text: textParts.join('.'),
}));
}
}
const outputString = outputLines.join('\n');
fs.writeFileSync(tempFilePath, outputString);
// Upload the file to Vertex AI RAG corpus
for (const ragResource of this.vertexRagStore.ragResources) {
await this.uploadFile(ragResource.ragCorpus, tempFilePath, `${session.appName}.${session.userId}.${session.id}`);
}
}
catch (error) {
console.error('Error adding session to memory:', error);
throw error;
}
finally {
// Clean up the temporary file
if (fs.existsSync(tempFilePath)) {
fs.unlinkSync(tempFilePath);
}
}
}
/**
* Uploads a file to the Vertex AI RAG corpus.
*
* @param corpusName The name of the corpus
* @param filePath The path to the file
* @param displayName The display name to store session info
*/
async uploadFile(corpusName, filePath, displayName) {
try {
// Format the full corpus name based on project and location
const formattedCorpusName = this.formatRagCorpusName(corpusName);
// Get authentication token
const authClient = await this.auth.getClient();
const accessToken = await authClient.getAccessToken();
// Create a form data object for multipart upload using form-data
const formData = new form_data_1.default();
// Add metadata as a part of formData
formData.append('metadata', JSON.stringify({
displayName: displayName,
mimeType: 'text/plain'
}), {
contentType: 'application/json'
});
// Read the file and add to form data
formData.append('file', fs.createReadStream(filePath), {
filename: path.basename(filePath),
contentType: 'text/plain'
});
// Upload the file using the REST API
const uploadUrl = `${this.baseUrl}/${formattedCorpusName}/ragFiles:import`;
const response = await axios_1.default.post(uploadUrl, formData, {
headers: {
'Authorization': `Bearer ${accessToken}`,
// The form-data package automatically sets the correct content-type header with boundary
...formData.getHeaders()
}
});
console.log(`File uploaded successfully: ${response.data.name}`);
}
catch (error) {
console.error('Error uploading file:', error);
throw error;
}
}
/**
* Searches for sessions that match the query.
* @param appName The name of the application
* @param userId The id of the user
* @param query The query to search for
* @returns A SearchMemoryResponse containing the matching memories
*/
async searchMemory(appName, userId, query) {
try {
// Call Vertex AI RAG retrieval query
const response = await this.retrieveContexts(query);
const memoryResults = [];
const sessionEventsMap = new Map();
// Process response contexts
for (const context of response.contexts?.contexts || []) {
if (!context.text)
continue;
const sessionId = context.sourceDisplayName?.split('.').pop() || '';
const events = [];
const lines = context.text.split('\n');
for (const line of lines) {
const trimmedLine = line.trim();
if (!trimmedLine)
continue;
try {
// Parse JSON event data
const eventData = JSON.parse(trimmedLine);
const author = eventData.author || '';
const timestamp = Number(eventData.timestamp) || 0;
const text = eventData.text || '';
const event = {
invocationId: '', // We don't have this from the retrieved data
author,
timestamp,
content: {
role: author,
parts: [{ text }]
}
};
events.push(event);
}
catch (error) {
// Skip invalid JSON lines
continue;
}
}
if (sessionId) {
if (sessionEventsMap.has(sessionId)) {
sessionEventsMap.get(sessionId).push(events);
}
else {
sessionEventsMap.set(sessionId, [events]);
}
}
}
// Merge and sort events from the same session
for (const [sessionId, eventLists] of sessionEventsMap.entries()) {
for (const events of this.mergeEventLists(eventLists)) {
const sortedEvents = events.sort((a, b) => (a.timestamp || 0) - (b.timestamp || 0));
memoryResults.push({
sessionId,
events: sortedEvents
});
}
}
return { memories: memoryResults };
}
catch (error) {
console.error('Error searching memory:', error);
return { memories: [] };
}
}
/**
* Calls Vertex AI RAG retrieval query.
*
* @param text The text to search for
* @returns The retrieval query response
*/
async retrieveContexts(text) {
try {
// Format the location for the API call
const formattedParent = `projects/${this.config.project}/locations/${this.config.location}`;
// Get authentication token
const authClient = await this.auth.getClient();
const accessToken = await authClient.getAccessToken();
// Create the request body
const requestBody = {
query: text,
ragResources: this.vertexRagStore.ragResources.map(res => ({
ragCorpus: this.formatRagCorpusName(res.ragCorpus)
})),
similarityTopK: this.vertexRagStore.similarityTopK,
vectorDistanceThreshold: this.vertexRagStore.vectorDistanceThreshold
};
// Make the API call using REST
const url = `${this.baseUrl}/${formattedParent}:retrieveContexts`;
const response = await axios_1.default.post(url, requestBody, {
headers: {
'Authorization': `Bearer ${accessToken}`,
'Content-Type': 'application/json'
}
});
return response.data;
}
catch (error) {
console.error('Error retrieving contexts:', error);
return { contexts: { contexts: [] } };
}
}
/**
* Formats the RAG corpus name to include project and location if not already included.
*
* @param corpusName The corpus name to format
* @returns The formatted corpus name
*/
formatRagCorpusName(corpusName) {
if (corpusName.startsWith('projects/')) {
return corpusName;
}
return `projects/${this.config.project}/locations/${this.config.location}/ragCorpora/${corpusName}`;
}
/**
* Merges event lists that have overlapping timestamps.
*
* @param eventLists Lists of events to merge
* @returns Merged event lists
*/
mergeEventLists(eventLists) {
const merged = [];
while (eventLists.length > 0) {
const current = eventLists.shift();
const currentTimestamps = new Set(current.map(event => event.timestamp));
let mergeFound = true;
// Keep merging until no new overlap is found
while (mergeFound) {
mergeFound = false;
const remaining = [];
for (const other of eventLists) {
const otherTimestamps = new Set(other.map(event => event.timestamp));
// Check for overlap by finding common timestamps
const hasOverlap = [...otherTimestamps].some(ts => currentTimestamps.has(ts));
if (hasOverlap) {
// Add events from 'other' that aren't in 'current'
const newEvents = other.filter(e => !currentTimestamps.has(e.timestamp));
current.push(...newEvents);
newEvents.forEach(e => e.timestamp && currentTimestamps.add(e.timestamp));
mergeFound = true;
}
else {
remaining.push(other);
}
}
eventLists = remaining;
}
merged.push(current);
}
return merged;
}
/**
* Stores a memory entry.
* @param appName The application name
* @param userId The user ID
* @param key The memory key
* @param value The memory value
*/
async store(appName, userId, key, value) {
const storeKey = `${appName}/${userId}/${key}`;
this.memoryStore.set(storeKey, value);
}
/**
* Retrieves a memory entry.
* @param appName The application name
* @param userId The user ID
* @param key The memory key
* @returns The memory value, or undefined if not found
*/
async retrieve(appName, userId, key) {
const storeKey = `${appName}/${userId}/${key}`;
return this.memoryStore.get(storeKey);
}
/**
* Deletes a memory entry.
* @param appName The application name
* @param userId The user ID
* @param key The memory key
*/
async delete(appName, userId, key) {
const storeKey = `${appName}/${userId}/${key}`;
this.memoryStore.delete(storeKey);
}
}
exports.VertexAiRagMemoryService = VertexAiRagMemoryService;