claude-brain 0.30.2 → 0.30.3

package/src/memory/chroma/search.ts

@@ -1,315 +1,319 @@
-import type { Where, WhereDocument } from 'chromadb'
-import type { Logger } from 'pino'
-import type { CollectionManager } from './collection-manager'
-import type { DecisionMetadata } from './schemas'
-
-export interface SearchOptions {
-  project?: string
-  limit?: number
-  minSimilarity?: number
-  type?: string
-  tags?: string[]
-  dateRange?: {
-    start?: string
-    end?: string
-  }
-}
-
-export interface SearchResult {
-  id: string
-  content: string
-  metadata: Record<string, any>
-  similarity: number
-}
-
-export class ChromaSearchEngine {
-  private logger: Logger
-  private collections: CollectionManager
-  private embeddings?: any
-
-  constructor(logger: Logger, collections: CollectionManager, embeddings?: any) {
-    this.logger = logger.child({ component: 'chroma-search' })
-    this.collections = collections
-    this.embeddings = embeddings
-  }
-
-  async searchDecisions(
-    query: string,
-    options: SearchOptions = {}
-  ): Promise<SearchResult[]> {
-    const {
-      project,
-      limit = 10,
-      minSimilarity = 0.5,
-      tags,
-      dateRange
-    } = options
-
-    try {
-      const collection = await this.collections.getDecisions()
-
-      const where = this.buildWhereClause({ project, tags, dateRange })
-
-      let results
-
-      if (this.embeddings) {
-        const embedding = await this.embeddings.generate(query)
-        results = await collection.query({
-          queryEmbeddings: [embedding],
-          nResults: limit,
-          where: Object.keys(where).length > 0 ? where : undefined,
-          include: ['documents', 'metadatas', 'distances']
-        })
-      } else {
-        results = await collection.query({
-          queryTexts: [query],
-          nResults: limit,
-          where: Object.keys(where).length > 0 ? where : undefined,
-          include: ['documents', 'metadatas', 'distances']
-        })
-      }
-
-      return this.processResults(results, minSimilarity)
-
-    } catch (error) {
-      this.logger.error({ error, query }, 'Decision search failed')
-      throw error
-    }
-  }
-
-  async searchMemories(
-    query: string,
-    options: SearchOptions = {}
-  ): Promise<SearchResult[]> {
-    const {
-      project,
-      limit = 10,
-      minSimilarity = 0.5,
-      type,
-      dateRange
-    } = options
-
-    try {
-      const collection = await this.collections.getMemories()
-
-      const where = this.buildWhereClause({ project, type, dateRange })
-
-      let results
-
-      if (this.embeddings) {
-        const embedding = await this.embeddings.generate(query)
-        results = await collection.query({
-          queryEmbeddings: [embedding],
-          nResults: limit,
-          where: Object.keys(where).length > 0 ? where : undefined,
-          include: ['documents', 'metadatas', 'distances']
-        })
-      } else {
-        results = await collection.query({
-          queryTexts: [query],
-          nResults: limit,
-          where: Object.keys(where).length > 0 ? where : undefined,
-          include: ['documents', 'metadatas', 'distances']
-        })
-      }
-
-      return this.processResults(results, minSimilarity)
-
-    } catch (error) {
-      this.logger.error({ error, query }, 'Memory search failed')
-      throw error
-    }
-  }
-
-  async hybridSearch(
-    query: string,
-    options: SearchOptions & {
-      fullTextQuery?: string
-    } = {}
-  ): Promise<SearchResult[]> {
-    const {
-      project,
-      limit = 10,
-      minSimilarity = 0.5,
-      fullTextQuery
-    } = options
-
-    try {
-      const collection = await this.collections.getDecisions()
-
-      const where = this.buildWhereClause({ project })
-
-      const whereDocument: WhereDocument | undefined = fullTextQuery
-        ? { '$contains': fullTextQuery }
-        : undefined
-
-      let results
-
-      if (this.embeddings) {
-        const embedding = await this.embeddings.generate(query)
-        results = await collection.query({
-          queryEmbeddings: [embedding],
-          nResults: limit,
-          where: Object.keys(where).length > 0 ? where : undefined,
-          whereDocument,
-          include: ['documents', 'metadatas', 'distances']
-        })
-      } else {
-        results = await collection.query({
-          queryTexts: [query],
-          nResults: limit,
-          where: Object.keys(where).length > 0 ? where : undefined,
-          whereDocument,
-          include: ['documents', 'metadatas', 'distances']
-        })
-      }
-
-      return this.processResults(results, minSimilarity)
-
-    } catch (error) {
-      this.logger.error({ error, query }, 'Hybrid search failed')
-      throw error
-    }
-  }
-
-  async findSimilar(id: string, limit: number = 5): Promise<SearchResult[]> {
-    try {
-      const collection = await this.collections.getDecisions()
-
-      const existing = await collection.get({
-        ids: [id],
-        include: ['documents', 'embeddings']
-      })
-
-      if (existing.ids.length === 0) {
-        throw new Error(`Decision ${id} not found`)
-      }
-
-      const results = await collection.query({
-        queryEmbeddings: existing.embeddings as number[][],
-        nResults: limit + 1,
-        include: ['documents', 'metadatas', 'distances']
-      })
-
-      const filtered = this.processResults(results, 0)
-        .filter(r => r.id !== id)
-        .slice(0, limit)
-
-      return filtered
-
-    } catch (error) {
-      this.logger.error({ error, id }, 'Find similar failed')
-      throw error
-    }
-  }
-
-  async searchByTags(
-    tags: string[],
-    options: SearchOptions = {}
-  ): Promise<SearchResult[]> {
-    const { project, limit = 10 } = options
-
-    try {
-      const collection = await this.collections.getDecisions()
-
-      const where: Where = {}
-
-      if (project) {
-        where['project'] = project
-      }
-
-      const results = await collection.get({
-        where: Object.keys(where).length > 0 ? where : undefined,
-        include: ['documents', 'metadatas'],
-        limit: limit * 3
-      })
-
-      const filtered: SearchResult[] = []
-
-      for (let i = 0; i < results.ids.length; i++) {
-        const metadata = results.metadatas![i] as unknown as DecisionMetadata
-        const hasMatchingTag = tags.some(tag =>
-          metadata.tags?.includes(tag)
-        )
-
-        if (hasMatchingTag) {
-          filtered.push({
-            id: results.ids[i]!,
-            content: (results.documents![i] ?? '') as string,
-            metadata: metadata as Record<string, any>,
-            similarity: 1.0
-          })
-        }
-
-        if (filtered.length >= limit) break
-      }
-
-      return filtered
-
-    } catch (error) {
-      this.logger.error({ error, tags }, 'Tag search failed')
-      throw error
-    }
-  }
-
-  private buildWhereClause(options: {
-    project?: string
-    type?: string
-    tags?: string[]
-    dateRange?: { start?: string; end?: string }
-  }): Where {
-    const where: Where = {}
-
-    if (options.project) {
-      where['project'] = options.project
-    }
-
-    if (options.type) {
-      where['type'] = options.type
-    }
-
-    if (options.dateRange?.start && options.dateRange?.end) {
-      (where as any)['$and'] = [
-        { 'created_at': { '$gte': options.dateRange.start } },
-        { 'created_at': { '$lte': options.dateRange.end } }
-      ]
-    } else if (options.dateRange?.start) {
-      where['created_at'] = { '$gte': options.dateRange.start }
-    } else if (options.dateRange?.end) {
-      where['created_at'] = { '$lte': options.dateRange.end }
-    }
-
-    return where
-  }
-
-  private processResults(
-    results: any,
-    minSimilarity: number
-  ): SearchResult[] {
-    if (!results.ids || results.ids.length === 0) {
-      return []
-    }
-
-    const processed: SearchResult[] = []
-
-    const ids = results.ids[0] || []
-    const documents = results.documents?.[0] || []
-    const metadatas = results.metadatas?.[0] || []
-    const distances = results.distances?.[0] || []
-
-    for (let i = 0; i < ids.length; i++) {
-      const similarity = 1 - (distances[i] || 0)
-
-      if (similarity >= minSimilarity) {
-        processed.push({
-          id: ids[i],
-          content: documents[i],
-          metadata: metadatas[i],
-          similarity
-        })
-      }
-    }
-
-    return processed.sort((a, b) => b.similarity - a.similarity)
-  }
-}
+import type { QueryResult, Where, WhereDocument } from 'chromadb'
+import type { Logger } from 'pino'
+import type { CollectionManager } from './collection-manager'
+import type { DecisionMetadata } from './schemas'
+
+export interface SearchOptions {
+  project?: string
+  limit?: number
+  minSimilarity?: number
+  type?: string
+  tags?: string[]
+  dateRange?: {
+    start?: string
+    end?: string
+  }
+}
+
+export interface SearchResult {
+  id: string
+  content: string
+  metadata: Record<string, unknown>
+  similarity: number
+}
+
+interface EmbeddingGenerator {
+  generate(text: string): Promise<number[]>
+}
+
+export class ChromaSearchEngine {
+  private logger: Logger
+  private collections: CollectionManager
+  private embeddings?: EmbeddingGenerator
+
+  constructor(logger: Logger, collections: CollectionManager, embeddings?: EmbeddingGenerator) {
+    this.logger = logger.child({ component: 'chroma-search' })
+    this.collections = collections
+    this.embeddings = embeddings
+  }
+
+  async searchDecisions(
+    query: string,
+    options: SearchOptions = {}
+  ): Promise<SearchResult[]> {
+    const {
+      project,
+      limit = 10,
+      minSimilarity = 0.5,
+      tags,
+      dateRange
+    } = options
+
+    try {
+      const collection = await this.collections.getDecisions()
+
+      const where = this.buildWhereClause({ project, tags, dateRange })
+
+      let results
+
+      if (this.embeddings) {
+        const embedding = await this.embeddings.generate(query)
+        results = await collection.query({
+          queryEmbeddings: [embedding],
+          nResults: limit,
+          where: Object.keys(where).length > 0 ? where : undefined,
+          include: ['documents', 'metadatas', 'distances']
+        })
+      } else {
+        results = await collection.query({
+          queryTexts: [query],
+          nResults: limit,
+          where: Object.keys(where).length > 0 ? where : undefined,
+          include: ['documents', 'metadatas', 'distances']
+        })
+      }
+
+      return this.processResults(results, minSimilarity)
+
+    } catch (error) {
+      this.logger.error({ error, query }, 'Decision search failed')
+      throw error
+    }
+  }
+
+  async searchMemories(
+    query: string,
+    options: SearchOptions = {}
+  ): Promise<SearchResult[]> {
+    const {
+      project,
+      limit = 10,
+      minSimilarity = 0.5,
+      type,
+      dateRange
+    } = options
+
+    try {
+      const collection = await this.collections.getMemories()
+
+      const where = this.buildWhereClause({ project, type, dateRange })
+
+      let results
+
+      if (this.embeddings) {
+        const embedding = await this.embeddings.generate(query)
+        results = await collection.query({
+          queryEmbeddings: [embedding],
+          nResults: limit,
+          where: Object.keys(where).length > 0 ? where : undefined,
+          include: ['documents', 'metadatas', 'distances']
+        })
+      } else {
+        results = await collection.query({
+          queryTexts: [query],
+          nResults: limit,
+          where: Object.keys(where).length > 0 ? where : undefined,
+          include: ['documents', 'metadatas', 'distances']
+        })
+      }
+
+      return this.processResults(results, minSimilarity)
+
+    } catch (error) {
+      this.logger.error({ error, query }, 'Memory search failed')
+      throw error
+    }
+  }
+
+  async hybridSearch(
+    query: string,
+    options: SearchOptions & {
+      fullTextQuery?: string
+    } = {}
+  ): Promise<SearchResult[]> {
+    const {
+      project,
+      limit = 10,
+      minSimilarity = 0.5,
+      fullTextQuery
+    } = options
+
+    try {
+      const collection = await this.collections.getDecisions()
+
+      const where = this.buildWhereClause({ project })
+
+      const whereDocument: WhereDocument | undefined = fullTextQuery
+        ? { '$contains': fullTextQuery }
+        : undefined
+
+      let results
+
+      if (this.embeddings) {
+        const embedding = await this.embeddings.generate(query)
+        results = await collection.query({
+          queryEmbeddings: [embedding],
+          nResults: limit,
+          where: Object.keys(where).length > 0 ? where : undefined,
+          whereDocument,
+          include: ['documents', 'metadatas', 'distances']
+        })
+      } else {
+        results = await collection.query({
+          queryTexts: [query],
+          nResults: limit,
+          where: Object.keys(where).length > 0 ? where : undefined,
+          whereDocument,
+          include: ['documents', 'metadatas', 'distances']
+        })
+      }
+
+      return this.processResults(results, minSimilarity)
+
+    } catch (error) {
+      this.logger.error({ error, query }, 'Hybrid search failed')
+      throw error
+    }
+  }
+
+  async findSimilar(id: string, limit: number = 5): Promise<SearchResult[]> {
+    try {
+      const collection = await this.collections.getDecisions()
+
+      const existing = await collection.get({
+        ids: [id],
+        include: ['documents', 'embeddings']
+      })
+
+      if (existing.ids.length === 0) {
+        throw new Error(`Decision ${id} not found`)
+      }
+
+      const results = await collection.query({
+        queryEmbeddings: existing.embeddings as number[][],
+        nResults: limit + 1,
+        include: ['documents', 'metadatas', 'distances']
+      })
+
+      const filtered = this.processResults(results, 0)
+        .filter(r => r.id !== id)
+        .slice(0, limit)
+
+      return filtered
+
+    } catch (error) {
+      this.logger.error({ error, id }, 'Find similar failed')
+      throw error
+    }
+  }
+
+  async searchByTags(
+    tags: string[],
+    options: SearchOptions = {}
+  ): Promise<SearchResult[]> {
+    const { project, limit = 10 } = options
+
+    try {
+      const collection = await this.collections.getDecisions()
+
+      const where: Where = {}
+
+      if (project) {
+        where['project'] = project
+      }
+
+      const results = await collection.get({
+        where: Object.keys(where).length > 0 ? where : undefined,
+        include: ['documents', 'metadatas'],
+        limit: limit * 3
+      })
+
+      const filtered: SearchResult[] = []
+
+      for (let i = 0; i < results.ids.length; i++) {
+        const metadata = results.metadatas![i] as unknown as DecisionMetadata
+        const hasMatchingTag = tags.some(tag =>
+          metadata.tags?.includes(tag)
+        )
+
+        if (hasMatchingTag) {
+          filtered.push({
+            id: results.ids[i]!,
+            content: (results.documents![i] ?? '') as string,
+            metadata: metadata as unknown as Record<string, unknown>,
+            similarity: 1.0
+          })
+        }
+
+        if (filtered.length >= limit) break
+      }
+
+      return filtered
+
+    } catch (error) {
+      this.logger.error({ error, tags }, 'Tag search failed')
+      throw error
+    }
+  }
+
+  private buildWhereClause(options: {
+    project?: string
+    type?: string
+    tags?: string[]
+    dateRange?: { start?: string; end?: string }
+  }): Where {
+    const where: Where = {}
+
+    if (options.project) {
+      where['project'] = options.project
+    }
+
+    if (options.type) {
+      where['type'] = options.type
+    }
+
+    if (options.dateRange?.start && options.dateRange?.end) {
+      (where as Record<string, unknown>)['$and'] = [
+        { 'created_at': { '$gte': options.dateRange.start } },
+        { 'created_at': { '$lte': options.dateRange.end } }
+      ]
+    } else if (options.dateRange?.start) {
+      where['created_at'] = { '$gte': options.dateRange.start }
+    } else if (options.dateRange?.end) {
+      where['created_at'] = { '$lte': options.dateRange.end }
+    }
+
+    return where
+  }
+
+  private processResults(
+    results: QueryResult,
+    minSimilarity: number
+  ): SearchResult[] {
+    if (!results.ids || results.ids.length === 0) {
+      return []
+    }
+
+    const processed: SearchResult[] = []
+
+    const ids = results.ids[0] || []
+    const documents = results.documents?.[0] || []
+    const metadatas = results.metadatas?.[0] || []
+    const distances = results.distances?.[0] || []
+
+    for (let i = 0; i < ids.length; i++) {
+      const similarity = 1 - (distances[i] || 0)
+
+      if (similarity >= minSimilarity) {
+        processed.push({
+          id: ids[i]!,
+          content: (documents[i] ?? '') as string,
+          metadata: (metadatas[i] ?? {}) as Record<string, unknown>,
+          similarity
+        })
+      }
+    }
+
+    return processed.sort((a, b) => b.similarity - a.similarity)
+  }
+}