@agentmemory/agentmemory 0.7.0

package/src/providers/embedding/openai.ts

@@ -0,0 +1,45 @@
+import type { EmbeddingProvider } from "../../types.js";
+import { getEnvVar } from "../../config.js";
+
+const API_URL = "https://api.openai.com/v1/embeddings";
+
+export class OpenAIEmbeddingProvider implements EmbeddingProvider {
+  readonly name = "openai";
+  readonly dimensions = 1536;
+  private apiKey: string;
+
+  constructor(apiKey?: string) {
+    this.apiKey = apiKey || getEnvVar("OPENAI_API_KEY") || "";
+    if (!this.apiKey) throw new Error("OPENAI_API_KEY is required");
+  }
+
+  async embed(text: string): Promise<Float32Array> {
+    const [result] = await this.embedBatch([text]);
+    return result;
+  }
+
+  async embedBatch(texts: string[]): Promise<Float32Array[]> {
+    const response = await fetch(API_URL, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${this.apiKey}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        model: "text-embedding-3-small",
+        input: texts,
+      }),
+    });
+
+    if (!response.ok) {
+      const err = await response.text();
+      throw new Error(`OpenAI embedding failed (${response.status}): ${err}`);
+    }
+
+    const data = (await response.json()) as {
+      data: Array<{ embedding: number[] }>;
+    };
+
+    return data.data.map((d) => new Float32Array(d.embedding));
+  }
+}

package/src/providers/embedding/openrouter.ts

@@ -0,0 +1,51 @@
+import type { EmbeddingProvider } from "../../types.js";
+import { getEnvVar } from "../../config.js";
+
+const API_URL = "https://openrouter.ai/api/v1/embeddings";
+
+export class OpenRouterEmbeddingProvider implements EmbeddingProvider {
+  readonly name = "openrouter";
+  readonly dimensions = 1536;
+  private apiKey: string;
+  private model: string;
+
+  constructor(apiKey?: string) {
+    this.apiKey = apiKey || getEnvVar("OPENROUTER_API_KEY") || "";
+    if (!this.apiKey) throw new Error("OPENROUTER_API_KEY is required");
+    this.model =
+      getEnvVar("OPENROUTER_EMBEDDING_MODEL") ||
+      "openai/text-embedding-3-small";
+  }
+
+  async embed(text: string): Promise<Float32Array> {
+    const [result] = await this.embedBatch([text]);
+    return result;
+  }
+
+  async embedBatch(texts: string[]): Promise<Float32Array[]> {
+    const response = await fetch(API_URL, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${this.apiKey}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        model: this.model,
+        input: texts,
+      }),
+    });
+
+    if (!response.ok) {
+      const err = await response.text();
+      throw new Error(
+        `OpenRouter embedding failed (${response.status}): ${err}`,
+      );
+    }
+
+    const data = (await response.json()) as {
+      data: Array<{ embedding: number[] }>;
+    };
+
+    return data.data.map((d) => new Float32Array(d.embedding));
+  }
+}

package/src/providers/embedding/voyage.ts

@@ -0,0 +1,46 @@
+import type { EmbeddingProvider } from "../../types.js";
+import { getEnvVar } from "../../config.js";
+
+const API_URL = "https://api.voyageai.com/v1/embeddings";
+
+export class VoyageEmbeddingProvider implements EmbeddingProvider {
+  readonly name = "voyage";
+  readonly dimensions = 1024;
+  private apiKey: string;
+
+  constructor(apiKey?: string) {
+    this.apiKey = apiKey || getEnvVar("VOYAGE_API_KEY") || "";
+    if (!this.apiKey) throw new Error("VOYAGE_API_KEY is required");
+  }
+
+  async embed(text: string): Promise<Float32Array> {
+    const [result] = await this.embedBatch([text]);
+    return result;
+  }
+
+  async embedBatch(texts: string[]): Promise<Float32Array[]> {
+    const response = await fetch(API_URL, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${this.apiKey}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        model: "voyage-code-3",
+        input: texts,
+        input_type: "document",
+      }),
+    });
+
+    if (!response.ok) {
+      const err = await response.text();
+      throw new Error(`Voyage embedding failed (${response.status}): ${err}`);
+    }
+
+    const data = (await response.json()) as {
+      data: Array<{ embedding: number[] }>;
+    };
+
+    return data.data.map((d) => new Float32Array(d.embedding));
+  }
+}

package/src/providers/fallback-chain.ts

@@ -0,0 +1,31 @@
+import type { MemoryProvider } from "../types.js";
+
+export class FallbackChainProvider implements MemoryProvider {
+  name: string;
+
+  constructor(private providers: MemoryProvider[]) {
+    this.name = `fallback(${providers.map((p) => p.name).join(" -> ")})`;
+  }
+
+  async compress(systemPrompt: string, userPrompt: string): Promise<string> {
+    return this.tryAll((p) => p.compress(systemPrompt, userPrompt));
+  }
+
+  async summarize(systemPrompt: string, userPrompt: string): Promise<string> {
+    return this.tryAll((p) => p.summarize(systemPrompt, userPrompt));
+  }
+
+  private async tryAll(
+    fn: (p: MemoryProvider) => Promise<string>,
+  ): Promise<string> {
+    let lastError: Error | null = null;
+    for (const provider of this.providers) {
+      try {
+        return await fn(provider);
+      } catch (err) {
+        lastError = err instanceof Error ? err : new Error(String(err));
+      }
+    }
+    throw lastError || new Error("No providers available");
+  }
+}

package/src/providers/index.ts

@@ -0,0 +1,84 @@
+import type {
+  MemoryProvider,
+  ProviderConfig,
+  FallbackConfig,
+} from "../types.js";
+import { AgentSDKProvider } from "./agent-sdk.js";
+import { AnthropicProvider } from "./anthropic.js";
+import { OpenRouterProvider } from "./openrouter.js";
+import { ResilientProvider } from "./resilient.js";
+import { FallbackChainProvider } from "./fallback-chain.js";
+import { getEnvVar } from "../config.js";
+
+export { createEmbeddingProvider } from "./embedding/index.js";
+
+function requireEnvVar(key: string): string {
+  const value = getEnvVar(key);
+  if (!value) {
+    throw new Error(
+      `Missing required environment variable: ${key}. Set it in ~/.agentmemory/.env or as an environment variable.`,
+    );
+  }
+  return value;
+}
+
+export function createProvider(config: ProviderConfig): ResilientProvider {
+  return new ResilientProvider(createBaseProvider(config));
+}
+
+export function createFallbackProvider(
+  config: ProviderConfig,
+  fallbackConfig: FallbackConfig,
+): ResilientProvider {
+  if (fallbackConfig.providers.length === 0) {
+    return createProvider(config);
+  }
+
+  const providers: MemoryProvider[] = [createBaseProvider(config)];
+  for (const providerType of fallbackConfig.providers) {
+    if (providerType === config.provider) continue;
+    try {
+      const fbConfig: ProviderConfig = {
+        provider: providerType,
+        model: config.model,
+        maxTokens: config.maxTokens,
+      };
+      providers.push(createBaseProvider(fbConfig));
+    } catch {
+      // skip unavailable fallback providers
+    }
+  }
+
+  if (providers.length > 1) {
+    return new ResilientProvider(new FallbackChainProvider(providers));
+  }
+  return new ResilientProvider(providers[0]);
+}
+
+function createBaseProvider(config: ProviderConfig): MemoryProvider {
+  switch (config.provider) {
+    case "anthropic":
+      return new AnthropicProvider(
+        requireEnvVar("ANTHROPIC_API_KEY"),
+        config.model,
+        config.maxTokens,
+      );
+    case "gemini":
+      return new OpenRouterProvider(
+        requireEnvVar("GEMINI_API_KEY"),
+        config.model,
+        config.maxTokens,
+        "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
+      );
+    case "openrouter":
+      return new OpenRouterProvider(
+        requireEnvVar("OPENROUTER_API_KEY"),
+        config.model,
+        config.maxTokens,
+        "https://openrouter.ai/api/v1/chat/completions",
+      );
+    case "agent-sdk":
+    default:
+      return new AgentSDKProvider();
+  }
+}

package/src/providers/openrouter.ts

@@ -0,0 +1,71 @@
+import type { MemoryProvider } from "../types.js";
+
+export class OpenRouterProvider implements MemoryProvider {
+  name: string;
+  private apiKey: string;
+  private model: string;
+  private maxTokens: number;
+  private baseUrl: string;
+
+  constructor(
+    apiKey: string,
+    model: string,
+    maxTokens: number,
+    baseUrl: string,
+  ) {
+    this.apiKey = apiKey;
+    this.model = model;
+    this.maxTokens = maxTokens;
+    this.baseUrl = baseUrl;
+    this.name = baseUrl.includes("openrouter") ? "openrouter" : "gemini";
+  }
+
+  async compress(systemPrompt: string, userPrompt: string): Promise<string> {
+    return this.call(systemPrompt, userPrompt);
+  }
+
+  async summarize(systemPrompt: string, userPrompt: string): Promise<string> {
+    return this.call(systemPrompt, userPrompt);
+  }
+
+  private async call(
+    systemPrompt: string,
+    userPrompt: string,
+  ): Promise<string> {
+    const response = await fetch(this.baseUrl, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: `Bearer ${this.apiKey}`,
+        ...(this.baseUrl.includes("openrouter")
+          ? { "HTTP-Referer": "https://github.com/rohitg00/agentmemory" }
+          : {}),
+      },
+      body: JSON.stringify({
+        model: this.model,
+        max_tokens: this.maxTokens,
+        messages: [
+          { role: "system", content: systemPrompt },
+          { role: "user", content: userPrompt },
+        ],
+      }),
+    });
+
+    if (!response.ok) {
+      const text = await response.text();
+      throw new Error(`${this.name} API error (${response.status}): ${text}`);
+    }
+
+    const data = (await response.json()) as Record<string, unknown>;
+    const choices = data.choices as
+      | Array<{ message: { content: string } }>
+      | undefined;
+    const content = choices?.[0]?.message?.content;
+    if (!content) {
+      throw new Error(
+        `${this.name} returned unexpected response: ${JSON.stringify(data).slice(0, 200)}`,
+      );
+    }
+    return content;
+  }
+}

package/src/providers/resilient.ts

@@ -0,0 +1,37 @@
+import type { MemoryProvider, CircuitBreakerState } from "../types.js";
+import { CircuitBreaker } from "./circuit-breaker.js";
+
+export class ResilientProvider implements MemoryProvider {
+  private breaker = new CircuitBreaker();
+  name: string;
+
+  constructor(private inner: MemoryProvider) {
+    this.name = `resilient(${inner.name})`;
+  }
+
+  private async call(fn: () => Promise<string>): Promise<string> {
+    if (!this.breaker.isAllowed) {
+      throw new Error("circuit_breaker_open");
+    }
+    try {
+      const result = await fn();
+      this.breaker.recordSuccess();
+      return result;
+    } catch (err) {
+      this.breaker.recordFailure();
+      throw err;
+    }
+  }
+
+  async compress(systemPrompt: string, userPrompt: string): Promise<string> {
+    return this.call(() => this.inner.compress(systemPrompt, userPrompt));
+  }
+
+  async summarize(systemPrompt: string, userPrompt: string): Promise<string> {
+    return this.call(() => this.inner.summarize(systemPrompt, userPrompt));
+  }
+
+  get circuitState(): CircuitBreakerState {
+    return this.breaker.getState();
+  }
+}

package/src/state/hybrid-search.ts

@@ -0,0 +1,295 @@
+import { SearchIndex } from "./search-index.js";
+import { VectorIndex } from "./vector-index.js";
+import type {
+  EmbeddingProvider,
+  HybridSearchResult,
+  CompressedObservation,
+  QueryExpansion,
+} from "../types.js";
+import type { StateKV } from "./kv.js";
+import { KV } from "./schema.js";
+import {
+  GraphRetrieval,
+  type GraphRetrievalResult,
+} from "../functions/graph-retrieval.js";
+import { extractEntitiesFromQuery } from "../functions/query-expansion.js";
+
+const RRF_K = 60;
+
+export class HybridSearch {
+  private graphRetrieval: GraphRetrieval;
+
+  constructor(
+    private bm25: SearchIndex,
+    private vector: VectorIndex | null,
+    private embeddingProvider: EmbeddingProvider | null,
+    private kv: StateKV,
+    private bm25Weight = 0.4,
+    private vectorWeight = 0.6,
+    private graphWeight = 0.3,
+  ) {
+    this.graphRetrieval = new GraphRetrieval(kv);
+  }
+
+  async search(query: string, limit = 20): Promise<HybridSearchResult[]> {
+    return this.tripleStreamSearch(query, limit);
+  }
+
+  async searchWithExpansion(
+    query: string,
+    limit: number,
+    expansion: QueryExpansion,
+  ): Promise<HybridSearchResult[]> {
+    const allQueries = [
+      query,
+      ...expansion.reformulations,
+      ...expansion.temporalConcretizations,
+    ];
+
+    const allEntities = [
+      ...expansion.entityExtractions,
+      ...extractEntitiesFromQuery(query),
+    ];
+
+    const resultSets = await Promise.all(
+      allQueries.map((q) => this.tripleStreamSearch(q, limit, allEntities)),
+    );
+
+    const merged = new Map<string, HybridSearchResult>();
+    for (const results of resultSets) {
+      for (const r of results) {
+        const existing = merged.get(r.observation.id);
+        if (!existing || r.combinedScore > existing.combinedScore) {
+          merged.set(r.observation.id, r);
+        }
+      }
+    }
+
+    return Array.from(merged.values())
+      .sort((a, b) => b.combinedScore - a.combinedScore)
+      .slice(0, limit);
+  }
+
+  private async tripleStreamSearch(
+    query: string,
+    limit: number,
+    entityHints?: string[],
+  ): Promise<HybridSearchResult[]> {
+    const bm25Results = this.bm25.search(query, limit * 2);
+
+    let vectorResults: Array<{
+      obsId: string;
+      sessionId: string;
+      score: number;
+    }> = [];
+    let queryEmbedding: Float32Array | null = null;
+
+    if (this.vector && this.embeddingProvider && this.vector.size > 0) {
+      try {
+        queryEmbedding = await this.embeddingProvider.embed(query);
+        vectorResults = this.vector.search(queryEmbedding, limit * 2);
+      } catch {
+        // fall through to BM25-only
+      }
+    }
+
+    const entities =
+      entityHints && entityHints.length > 0
+        ? entityHints
+        : extractEntitiesFromQuery(query);
+    let graphResults: GraphRetrievalResult[] = [];
+    if (entities.length > 0) {
+      try {
+        graphResults = await this.graphRetrieval.searchByEntities(
+          entities,
+          2,
+          limit,
+        );
+      } catch {
+        // graph search is best-effort
+      }
+    }
+
+    const topVectorObs = vectorResults.slice(0, 5).map((r) => r.obsId);
+    if (topVectorObs.length > 0) {
+      try {
+        const expansionResults =
+          await this.graphRetrieval.expandFromChunks(topVectorObs, 1, 5);
+        graphResults = [...graphResults, ...expansionResults];
+      } catch {
+        // expansion is best-effort
+      }
+    }
+
+    const scores = new Map<
+      string,
+      {
+        bm25Rank: number;
+        vectorRank: number;
+        graphRank: number;
+        sessionId: string;
+        bm25Score: number;
+        vectorScore: number;
+        graphScore: number;
+        graphContext?: string;
+      }
+    >();
+
+    bm25Results.forEach((r, i) => {
+      scores.set(r.obsId, {
+        bm25Rank: i + 1,
+        vectorRank: Infinity,
+        graphRank: Infinity,
+        sessionId: r.sessionId,
+        bm25Score: r.score,
+        vectorScore: 0,
+        graphScore: 0,
+      });
+    });
+
+    vectorResults.forEach((r, i) => {
+      const existing = scores.get(r.obsId);
+      if (existing) {
+        existing.vectorRank = i + 1;
+        existing.vectorScore = r.score;
+      } else {
+        scores.set(r.obsId, {
+          bm25Rank: Infinity,
+          vectorRank: i + 1,
+          graphRank: Infinity,
+          sessionId: r.sessionId,
+          bm25Score: 0,
+          vectorScore: r.score,
+          graphScore: 0,
+        });
+      }
+    });
+
+    graphResults.forEach((r, i) => {
+      const existing = scores.get(r.obsId);
+      if (existing) {
+        existing.graphRank = Math.min(existing.graphRank, i + 1);
+        existing.graphScore = Math.max(existing.graphScore, r.score);
+        if (r.graphContext && !existing.graphContext) {
+          existing.graphContext = r.graphContext;
+        }
+      } else {
+        scores.set(r.obsId, {
+          bm25Rank: Infinity,
+          vectorRank: Infinity,
+          graphRank: i + 1,
+          sessionId: r.sessionId,
+          bm25Score: 0,
+          vectorScore: 0,
+          graphScore: r.score,
+          graphContext: r.graphContext,
+        });
+      }
+    });
+
+    const hasVector = vectorResults.length > 0;
+    const hasGraph = graphResults.length > 0;
+
+    let effectiveBm25W = this.bm25Weight;
+    let effectiveVectorW = hasVector ? this.vectorWeight : 0;
+    let effectiveGraphW = hasGraph ? this.graphWeight : 0;
+
+    const totalW = effectiveBm25W + effectiveVectorW + effectiveGraphW;
+    if (totalW > 0) {
+      effectiveBm25W /= totalW;
+      effectiveVectorW /= totalW;
+      effectiveGraphW /= totalW;
+    }
+
+    const combined = Array.from(scores.entries()).map(([obsId, s]) => ({
+      obsId,
+      sessionId: s.sessionId,
+      bm25Score: s.bm25Score,
+      vectorScore: s.vectorScore,
+      graphScore: s.graphScore,
+      graphContext: s.graphContext,
+      combinedScore:
+        effectiveBm25W * (1 / (RRF_K + s.bm25Rank)) +
+        effectiveVectorW * (1 / (RRF_K + s.vectorRank)) +
+        effectiveGraphW * (1 / (RRF_K + s.graphRank)),
+    }));
+
+    combined.sort((a, b) => b.combinedScore - a.combinedScore);
+    const diversified = this.diversifyBySession(combined, limit);
+    return this.enrichResults(diversified, limit);
+  }
+
+  private diversifyBySession(
+    results: Array<{
+      obsId: string;
+      sessionId: string;
+      bm25Score: number;
+      vectorScore: number;
+      graphScore: number;
+      combinedScore: number;
+      graphContext?: string;
+    }>,
+    limit: number,
+    maxPerSession = 3,
+  ): typeof results {
+    const selected: typeof results = [];
+    const sessionCounts = new Map<string, number>();
+
+    for (const r of results) {
+      const count = sessionCounts.get(r.sessionId) || 0;
+      if (count >= maxPerSession) continue;
+      selected.push(r);
+      sessionCounts.set(r.sessionId, count + 1);
+      if (selected.length >= limit) break;
+    }
+
+    if (selected.length < limit) {
+      for (const r of results) {
+        if (selected.length >= limit) break;
+        if (!selected.some(s => s.obsId === r.obsId)) {
+          selected.push(r);
+        }
+      }
+    }
+
+    return selected;
+  }
+
+  private async enrichResults(
+    results: Array<{
+      obsId: string;
+      sessionId: string;
+      bm25Score: number;
+      vectorScore: number;
+      graphScore: number;
+      combinedScore: number;
+      graphContext?: string;
+    }>,
+    limit: number,
+  ): Promise<HybridSearchResult[]> {
+    const sliced = results.slice(0, limit);
+    const observations = await Promise.all(
+      sliced.map((r) =>
+        this.kv
+          .get<CompressedObservation>(KV.observations(r.sessionId), r.obsId)
+          .catch(() => null),
+      ),
+    );
+    const enriched: HybridSearchResult[] = [];
+    for (let i = 0; i < sliced.length; i++) {
+      const obs = observations[i];
+      if (obs) {
+        enriched.push({
+          observation: obs,
+          bm25Score: sliced[i].bm25Score,
+          vectorScore: sliced[i].vectorScore,
+          graphScore: sliced[i].graphScore,
+          combinedScore: sliced[i].combinedScore,
+          sessionId: sliced[i].sessionId,
+          graphContext: sliced[i].graphContext,
+        });
+      }
+    }
+    return enriched;
+  }
+}