/ examples / js / embeddings / retrieval.ts
retrieval.ts
  1  /**
  2   * Retrieval Example - Simple Vector Search
  3   * 
  4   * Demonstrates building a simple retrieval system with embeddings
  5   * 
  6   * Run: npx ts-node examples/js/embeddings/retrieval.ts
  7   */
  8  
  9  import { embed, embedMany, cosineSimilarity } from '../../../src/praisonai-ts/dist/llm/embeddings';
 10  
 11  interface Document {
 12    id: string;
 13    content: string;
 14    embedding?: number[];
 15  }
 16  
 17  interface SearchResult {
 18    document: Document;
 19    score: number;
 20  }
 21  
 22  // Simple in-memory vector store
 23  class SimpleVectorStore {
 24    private documents: Document[] = [];
 25  
 26    async add(doc: Document): Promise<void> {
 27      const { embedding } = await embed(doc.content);
 28      this.documents.push({ ...doc, embedding });
 29    }
 30  
 31    async addBatch(docs: Document[]): Promise<void> {
 32      const contents = docs.map(d => d.content);
 33      const { embeddings } = await embedMany(contents);
 34      docs.forEach((doc, i) => {
 35        this.documents.push({ ...doc, embedding: embeddings[i] });
 36      });
 37    }
 38  
 39    async search(query: string, topK: number = 3): Promise<SearchResult[]> {
 40      const { embedding: queryEmb } = await embed(query);
 41      
 42      const results = this.documents
 43        .filter(doc => doc.embedding)
 44        .map(doc => ({
 45          document: doc,
 46          score: cosineSimilarity(queryEmb, doc.embedding!)
 47        }))
 48        .sort((a, b) => b.score - a.score)
 49        .slice(0, topK);
 50      
 51      return results;
 52    }
 53  
 54    get(id: string): Document | undefined {
 55      return this.documents.find(d => d.id === id);
 56    }
 57  
 58    delete(id: string): boolean {
 59      const idx = this.documents.findIndex(d => d.id === id);
 60      if (idx >= 0) {
 61        this.documents.splice(idx, 1);
 62        return true;
 63      }
 64      return false;
 65    }
 66  }
 67  
 68  async function main() {
 69    console.log('=== Simple Vector Retrieval ===\n');
 70  
 71    // Create vector store
 72    console.log('1. Creating vector store...');
 73    const store = new SimpleVectorStore();
 74  
 75    // Add documents
 76    console.log('2. Adding documents...');
 77    const docs: Document[] = [
 78      { id: '1', content: 'PraisonAI is an AI agent framework for building intelligent applications.' },
 79      { id: '2', content: 'Agents can use tools to interact with external systems and APIs.' },
 80      { id: '3', content: 'Workflows allow orchestrating multiple agents in sequence or parallel.' },
 81      { id: '4', content: 'Memory systems help agents maintain context across conversations.' },
 82      { id: '5', content: 'The AI SDK provides a unified interface for multiple LLM providers.' },
 83      { id: '6', content: 'Embeddings enable semantic search and similarity matching.' },
 84      { id: '7', content: 'Guardrails ensure agent outputs meet safety and quality standards.' },
 85      { id: '8', content: 'Sessions persist conversation state for multi-turn interactions.' }
 86    ];
 87  
 88    await store.addBatch(docs);
 89    console.log(`   Added ${docs.length} documents`);
 90  
 91    // Search queries
 92    console.log('\n3. Running search queries...\n');
 93  
 94    const queries = [
 95      'How do I build an AI application?',
 96      'Can agents call external APIs?',
 97      'How to maintain conversation context?',
 98      'What providers are supported?'
 99    ];
100  
101    for (const query of queries) {
102      console.log(`Query: "${query}"`);
103      const results = await store.search(query);
104      
105      if (results.length > 0) {
106        console.log('Results:');
107        results.forEach((r, i) => {
108          console.log(`  ${i + 1}. [${(r.score * 100).toFixed(1)}%] ${r.document.content.substring(0, 60)}...`);
109        });
110      } else {
111        console.log('  No results found');
112      }
113      console.log('');
114    }
115  
116    // Get specific document
117    console.log('4. Get document by ID...');
118    const doc = store.get('5');
119    if (doc) {
120      console.log(`   Found: ${doc.content}`);
121    }
122  
123    // Delete document
124    console.log('\n5. Delete document...');
125    const deleted = store.delete('8');
126    console.log(`   Deleted doc 8: ${deleted}`);
127  
128    console.log('\n=== Complete ===');
129  }
130  
131  main().catch(console.error);