@promptev/context-engine 0.0.0 → 0.0.2

package/dist/graph/index.d.cts

@@ -1,7 +1,7 @@
 import { Pool } from 'pg';
-import { G as GraphConfig, C as ContextEngineConfig } from '../config-Bl9U789m.cjs';
-import { H as Hooks } from '../redaction-BmDSWJ7h.cjs';
-import { E as Embedder } from '../embeddings-B-jZ42mk.cjs';
+import { G as GraphConfig, C as ContextEngineConfig } from '../config-C5RZ00W6.cjs';
+import { H as Hooks } from '../redaction-BqD_DEUQ.cjs';
+import { E as Embedder, A as AnnTunable } from '../storage-Dvpq2xAC.cjs';
 import 'zod';
 
 type Neo4jDriver = {
@@ -63,8 +63,10 @@ declare function buildGraphRanked$1(query: string, opts: {
     graphStore: GraphStore;
     hooks?: Hooks | null;
     sourceIds?: string[] | null;
+    documentIds?: string[] | null;
     principals?: string[] | null;
     maxDepth?: number;
+    backend?: Partial<AnnTunable> | null;
 }): Promise<string[]>;
 
 declare function runGraphStage$1(opts: {

package/dist/graph/index.d.ts

@@ -1,7 +1,7 @@
 import { Pool } from 'pg';
-import { G as GraphConfig, C as ContextEngineConfig } from '../config-Bt9bUQqU.js';
-import { H as Hooks } from '../redaction-BmDSWJ7h.js';
-import { E as Embedder } from '../embeddings-DaSdAZN3.js';
+import { G as GraphConfig, C as ContextEngineConfig } from '../config-BODDdXJ7.js';
+import { H as Hooks } from '../redaction-BqD_DEUQ.js';
+import { E as Embedder, A as AnnTunable } from '../storage-CJrKgJeJ.js';
 import 'zod';
 
 type Neo4jDriver = {
@@ -63,8 +63,10 @@ declare function buildGraphRanked$1(query: string, opts: {
     graphStore: GraphStore;
     hooks?: Hooks | null;
     sourceIds?: string[] | null;
+    documentIds?: string[] | null;
     principals?: string[] | null;
     maxDepth?: number;
+    backend?: Partial<AnnTunable> | null;
 }): Promise<string[]>;
 
 declare function runGraphStage$1(opts: {

package/dist/graph/index.js

@@ -59,9 +59,52 @@ var init_hooks = __esm({
   }
 });
 
+// src/providers/google.ts
+async function buildGenaiClient(cfg, timeoutMs, purpose) {
+  const specifier = "@google/genai";
+  let mod;
+  try {
+    mod = await import(specifier);
+  } catch {
+    throw new ExtraMissingError("gemini", specifier, purpose);
+  }
+  const Ctor = mod.GoogleGenAI ?? mod.Client;
+  if (!Ctor) {
+    throw new ExtraMissingError("gemini", specifier, purpose);
+  }
+  const opts = { httpOptions: { timeout: timeoutMs } };
+  if (cfg.provider === "vertex_ai") {
+    opts.vertexai = true;
+    if (cfg.project || cfg.location) {
+      if (cfg.project) opts.project = cfg.project;
+      if (cfg.location) opts.location = cfg.location;
+    } else if (cfg.apiKey) {
+      opts.apiKey = cfg.apiKey;
+    }
+  } else {
+    opts.apiKey = cfg.apiKey ?? null;
+  }
+  try {
+    return new Ctor(opts);
+  } catch (err) {
+    const message = err instanceof Error ? err.message : String(err);
+    if (!message.includes("Authentication is not set up")) throw err;
+    throw new Error(
+      `${message} Set \`project\` on the provider config, or export GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION. Credentials themselves come from Application Default Credentials.`
+    );
+  }
+}
+var init_google = __esm({
+  "src/providers/google.ts"() {
+    init_errors();
+  }
+});
+
 // src/providers/llm.ts
 var llm_exports = {};
 __export(llm_exports, {
+  CALL_TIMEOUT_MS: () => CALL_TIMEOUT_MS,
+  GEMINI_CALL_TIMEOUT_MS: () => GEMINI_CALL_TIMEOUT_MS,
   LLMClient: () => LLMClient,
   TIMEOUT_MS: () => TIMEOUT_MS,
   buildLlmClient: () => buildLlmClient,
@@ -92,20 +135,6 @@ function buildOpenAIChatClient(cfg) {
     maxRetries: 0
   });
 }
-async function loadGeminiChatClient(apiKey) {
-  const specifier = "@google/genai";
-  let mod;
-  try {
-    mod = await import(specifier);
-  } catch {
-    throw new ExtraMissingError("gemini", specifier, "gemini llm");
-  }
-  const Ctor = mod.GoogleGenAI ?? mod.Client;
-  if (!Ctor) {
-    throw new ExtraMissingError("gemini", specifier, "gemini llm");
-  }
-  return new Ctor({ apiKey: apiKey ?? null });
-}
 async function loadBedrockSdk() {
   const specifier = "@aws-sdk/client-bedrock-runtime";
   try {
@@ -124,7 +153,7 @@ function buildLlmClient(cfg, opts) {
   if (OPENAI_FAMILY.has(cfg.provider)) {
     return new LLMClient(cfg, { client: buildOpenAIChatClient(cfg) });
   }
-  if (cfg.provider === "gemini" || cfg.provider === "bedrock") {
+  if (GOOGLE_FAMILY.has(cfg.provider) || cfg.provider === "bedrock") {
     return new LLMClient(cfg);
   }
   throw new Error(`unknown llm provider: ${JSON.stringify(cfg.provider)}`);
@@ -141,14 +170,18 @@ async function callLlm(cfg, opts) {
     await owned.aclose();
   }
 }
-var TIMEOUT_MS, ANTHROPIC_VERSION, ANTHROPIC_MAX_TOKENS, OPENAI_FAMILY, LLMClient;
+var CALL_TIMEOUT_MS, TIMEOUT_MS, GEMINI_CALL_TIMEOUT_MS, ANTHROPIC_VERSION, ANTHROPIC_MAX_TOKENS, OPENAI_FAMILY, GOOGLE_FAMILY, LLMClient;
 var init_llm = __esm({
   "src/providers/llm.ts"() {
     init_errors();
-    TIMEOUT_MS = 3e4;
+    init_google();
+    CALL_TIMEOUT_MS = 24e4;
+    TIMEOUT_MS = CALL_TIMEOUT_MS;
+    GEMINI_CALL_TIMEOUT_MS = CALL_TIMEOUT_MS;
     ANTHROPIC_VERSION = "2023-06-01";
     ANTHROPIC_MAX_TOKENS = 4096;
     OPENAI_FAMILY = /* @__PURE__ */ new Set(["openai", "azure_openai", "custom"]);
+    GOOGLE_FAMILY = /* @__PURE__ */ new Set(["gemini", "vertex_ai"]);
     LLMClient = class {
       cfg;
       provider;
@@ -173,22 +206,30 @@ var init_llm = __esm({
         await this.aclose();
       }
       async call(opts) {
-        const { system, user, jsonMode = false, images = null } = opts;
+        const {
+          system,
+          user,
+          jsonMode = false,
+          images = null,
+          maxTokens = null,
+          thinkingBudget = null,
+          temperature = null
+        } = opts;
         if (this.provider === "anthropic") {
-          return this.callAnthropic(system, user, jsonMode, images);
+          return this.callAnthropic(system, user, jsonMode, images, maxTokens, thinkingBudget, temperature);
         }
         if (OPENAI_FAMILY.has(this.provider)) {
-          return this.callOpenAI(system, user, jsonMode, images);
+          return this.callOpenAI(system, user, jsonMode, images, maxTokens, temperature);
         }
-        if (this.provider === "gemini") {
-          return this.callGemini(system, user, jsonMode, images);
+        if (GOOGLE_FAMILY.has(this.provider)) {
+          return this.callGemini(system, user, jsonMode, images, maxTokens, thinkingBudget, temperature);
         }
         if (this.provider === "bedrock") {
-          return this.callBedrock(system, user, jsonMode, images);
+          return this.callBedrock(system, user, jsonMode, images, maxTokens, thinkingBudget, temperature);
         }
         throw new Error(`unknown llm provider: ${JSON.stringify(this.provider)}`);
       }
-      async callAnthropic(system, user, jsonMode, images) {
+      async callAnthropic(system, user, jsonMode, images, maxTokens, thinkingBudget = null, temperature = null) {
         if (!this.fetchImpl) {
           throw new Error("anthropic llm client has no fetch implementation");
         }
@@ -209,26 +250,27 @@ Respond with valid JSON only.`;
           });
         }
         content.push({ type: "text", text: user });
-        const data = await postJson(
-          this.fetchImpl,
-          "https://api.anthropic.com/v1/messages",
-          {
-            model: this.model,
-            max_tokens: ANTHROPIC_MAX_TOKENS,
-            system,
-            messages: [{ role: "user", content }]
-          },
-          {
-            "x-api-key": this.cfg.apiKey ?? "",
-            "anthropic-version": ANTHROPIC_VERSION,
-            "content-type": "application/json"
-          }
-        );
+        const body = {
+          model: this.model,
+          max_tokens: maxTokens ?? ANTHROPIC_MAX_TOKENS,
+          system,
+          messages: [{ role: "user", content }]
+        };
+        if (thinkingBudget) {
+          body.thinking = { type: "enabled", budget_tokens: thinkingBudget };
+        } else if (temperature !== null) {
+          body.temperature = temperature;
+        }
+        const data = await postJson(this.fetchImpl, "https://api.anthropic.com/v1/messages", body, {
+          "x-api-key": this.cfg.apiKey ?? "",
+          "anthropic-version": ANTHROPIC_VERSION,
+          "content-type": "application/json"
+        });
         const text = data.content[0].text;
         const usage = data.usage ?? {};
         return [text, { input: usage.input_tokens ?? 0, output: usage.output_tokens ?? 0 }];
       }
-      async callOpenAI(system, user, jsonMode, images) {
+      async callOpenAI(system, user, jsonMode, images, maxTokens, temperature = null) {
         if (!this.client) {
           throw new Error("openai-family llm client has no client");
         }
@@ -249,24 +291,54 @@ Respond with valid JSON only.`;
         if (jsonMode) {
           body.response_format = { type: "json_object" };
         }
+        if (maxTokens) {
+          body.max_completion_tokens = maxTokens;
+        }
+        if (temperature !== null) {
+          body.temperature = temperature;
+        }
         const resp = await this.client.chat.completions.create(body);
         const text = resp.choices[0]?.message?.content ?? "";
         const usage = resp.usage;
         return [text, { input: usage?.prompt_tokens ?? 0, output: usage?.completion_tokens ?? 0 }];
       }
-      async callGemini(system, user, jsonMode, images) {
+      async callGemini(system, user, jsonMode, images, maxTokens, thinkingBudget = null, temperature = null) {
         if (!this.genaiClient) {
-          this.genaiClient = await loadGeminiChatClient(this.cfg.apiKey);
+          this.genaiClient = await buildGenaiClient(
+            this.cfg,
+            GEMINI_CALL_TIMEOUT_MS,
+            "gemini llm"
+          );
         }
         const parts = [];
         for (const img of images ?? []) {
           parts.push({ inlineData: { mimeType: "image/png", data: toBase64(img) } });
         }
         parts.push({ text: user });
-        const config = { systemInstruction: system };
+        const config = {
+          systemInstruction: system,
+          // ALWAYS off, and not as a preference. Automatic function calling means
+          // the SDK itself EXECUTES a callable it was handed as a tool and loops
+          // on the result — up to ten round trips — before returning anything.
+          // This package passes declarations only, so today nothing is executable;
+          // but that depends on every future caller continuing to do the same, and
+          // an application that gates tool execution behind human approval would
+          // have that gate bypassed silently, by a library, with the loop already
+          // run before it could object.
+          automaticFunctionCalling: { disable: true }
+        };
         if (jsonMode) {
           config.responseMimeType = "application/json";
         }
+        if (maxTokens) {
+          config.maxOutputTokens = maxTokens;
+        }
+        if (temperature !== null) {
+          config.temperature = temperature;
+        }
+        if (thinkingBudget !== null) {
+          config.thinkingConfig = { thinkingBudget };
+        }
         const resp = await this.genaiClient.models.generateContent({
           model: this.model,
           contents: parts,
@@ -278,7 +350,7 @@ Respond with valid JSON only.`;
         const outputTokens = usageMeta?.candidatesTokenCount ?? usageMeta?.candidates_token_count ?? 0;
         return [text, { input: inputTokens || 0, output: outputTokens || 0 }];
       }
-      async callBedrock(system, user, jsonMode, images) {
+      async callBedrock(system, user, jsonMode, images, maxTokens, thinkingBudget = null, temperature = null) {
         const { BedrockRuntimeClient, ConverseCommand } = await loadBedrockSdk();
         if (jsonMode) {
           system = `${system}
@@ -302,7 +374,21 @@ Respond with valid JSON only.`;
             new ConverseCommand({
               modelId: this.model,
               system: [{ text: system }],
-              messages: [{ role: "user", content }]
+              messages: [{ role: "user", content }],
+              ...maxTokens || temperature !== null ? {
+                inferenceConfig: {
+                  ...maxTokens ? { maxTokens } : {},
+                  ...temperature !== null ? { temperature } : {}
+                }
+              } : {},
+              // Anthropic-style passthrough — Converse forwards it to the model.
+              // Zero (the vision transcription contract) sends nothing: thinking
+              // is opt-in for the anthropic models bedrock hosts.
+              ...thinkingBudget ? {
+                additionalModelRequestFields: {
+                  thinking: { type: "enabled", budget_tokens: thinkingBudget }
+                }
+              } : {}
             })
           );
           const text = result.output?.message?.content?.[0]?.text ?? "";
@@ -1144,7 +1230,8 @@ var retrieval_exports = {};
 __export(retrieval_exports, {
   buildGraphRanked: () => buildGraphRanked,
   corpusIsAclUniform: () => corpusIsAclUniform,
-  shouldUseCommunitySummaries: () => shouldUseCommunitySummaries
+  shouldUseCommunitySummaries: () => shouldUseCommunitySummaries,
+  vectorSeedIds: () => vectorSeedIds
 });
 function vecLiteral(vector) {
   return `[${vector.map((x) => Number(x)).join(",")}]`;
@@ -1170,8 +1257,8 @@ async function deriveQueryEntities(pool, chunkIds, opts) {
      JOIN context_engine_chunks c ON c.id = ce.chunk_id
      WHERE ce.chunk_id = ANY($1::uuid[]) ${SCOPE}
      GROUP BY e.normalized_name, e.name, e.type
-     ORDER BY freq DESC LIMIT $4`,
-    [chunkIds, opts.sourceIds, opts.principals, ENTITY_LIMIT]
+     ORDER BY freq DESC LIMIT $5`,
+    [chunkIds, opts.sourceIds, opts.documentIds ?? null, opts.principals, ENTITY_LIMIT]
   );
   return result.rows.map((r) => ({
     normalized_name: r.normalized_name,
@@ -1236,8 +1323,8 @@ async function computeCommunityScores(pool, chunkIds, vector, opts) {
     `SELECT ce.chunk_id::text, ce.entity_id::text
      FROM context_engine_chunk_entities ce
      JOIN context_engine_chunks c ON c.id = ce.chunk_id
-     WHERE ce.chunk_id = ANY($1::uuid[]) AND ce.entity_id = ANY($4::uuid[]) ${SCOPE}`,
-    [chunkIds, opts.sourceIds, opts.principals, Object.keys(entityScore)]
+     WHERE ce.chunk_id = ANY($1::uuid[]) AND ce.entity_id = ANY($5::uuid[]) ${SCOPE}`,
+    [chunkIds, opts.sourceIds, opts.documentIds ?? null, opts.principals, Object.keys(entityScore)]
   );
   const chunkScores = {};
   for (const row of result.rows) {
@@ -1246,13 +1333,29 @@ async function computeCommunityScores(pool, chunkIds, vector, opts) {
   return chunkScores;
 }
 async function vectorSeedIds(pool, vector, opts) {
-  const result = await pool.query(
-    `SELECT c.id::text FROM context_engine_chunks c
-     WHERE c.embedding IS NOT NULL ${SCOPE}
-     ORDER BY c.embedding <=> CAST($1 AS vector) LIMIT $4`,
-    [vecLiteral(vector), opts.sourceIds, opts.principals, SEED_LIMIT]
-  );
-  return result.rows.map((r) => String(r.id));
+  const binds = [opts.sourceIds, opts.documentIds ?? null, opts.principals];
+  const scoped = binds.some((v) => v !== null);
+  const client = await pool.connect();
+  try {
+    await client.query("BEGIN");
+    await opts.backend?.tuneAnnScan?.(client, SEED_LIMIT, scoped);
+    const result = await client.query(
+      `SELECT c.id::text FROM context_engine_chunks c
+       WHERE c.embedding IS NOT NULL ${SCOPE}
+       ORDER BY c.embedding <=> CAST($1 AS vector) LIMIT $5`,
+      [vecLiteral(vector), ...binds, SEED_LIMIT]
+    );
+    await client.query("COMMIT");
+    return result.rows.map((r) => String(r.id));
+  } catch (err) {
+    try {
+      await client.query("ROLLBACK");
+    } catch {
+    }
+    throw err;
+  } finally {
+    client.release();
+  }
 }
 async function buildGraphRanked(query, opts) {
   try {
@@ -1261,11 +1364,18 @@ async function buildGraphRanked(query, opts) {
     const vector = vectors[0] ? [...vectors[0]] : null;
     if (!vector) return [];
     const sourceIds = opts.sourceIds ?? null;
+    const documentIds = opts.documentIds ?? null;
     const principals = opts.principals ?? null;
-    const seeds = await vectorSeedIds(opts.pool, vector, { sourceIds, principals });
+    const seeds = await vectorSeedIds(opts.pool, vector, {
+      sourceIds,
+      documentIds,
+      principals,
+      backend: opts.backend
+    });
     if (!seeds.length) return [];
     const queryEntities = await deriveQueryEntities(opts.pool, seeds.slice(0, TOP_SEEDS_FOR_ENTITIES), {
       sourceIds,
+      documentIds,
       principals
     });
     const entityNorms = queryEntities.map((e) => String(e.normalized_name));
@@ -1287,6 +1397,7 @@ async function buildGraphRanked(query, opts) {
     const relScores = await computeRelationshipScores(opts.pool, allIds, entityNorms);
     const commScores = await computeCommunityScores(opts.pool, allIds, vector, {
       sourceIds,
+      documentIds,
       principals,
       useSummaries
     });
@@ -1338,7 +1449,8 @@ var init_retrieval = __esm({
     TOP_SEEDS_FOR_ENTITIES = 20;
     SCOPE = `
   AND ($2::text[] IS NULL OR c.source_id = ANY($2::text[]))
-  AND ($3::text[] IS NULL OR c.acl IS NULL OR c.acl && $3::text[])
+  AND ($3::uuid[] IS NULL OR c.document_id = ANY($3::uuid[]))
+  AND ($4::text[] IS NULL OR c.acl IS NULL OR c.acl && $4::text[])
 `;
   }
 });