ex-brain 0.1.0 → 0.1.1

package/src/ai/embed-factory.ts

@@ -0,0 +1,116 @@
+import { OpenAIEmbeddingFunction } from "@seekdb/openai";
+import type { EmbeddingFunction } from "seekdb";
+import type { ResolvedEmbed } from "../settings";
+import { LocalHashEmbeddingFunction } from "./hash-embed";
+
+/**
+ * 嵌入服务：与 seekdb 业务库（EBRAIN_SEEKDB_*）分离，仅由 EBRAIN_EMBED_* 控制。
+ * - `hash`（默认）：本地确定性向量，无网络。
+ * - `openai_compatible`：OpenAI 兼容 HTTP 端（如 DashScope compatible-mode）。
+ */
+export function createBrainEmbeddingFunction(cfg?: ResolvedEmbed): EmbeddingFunction {
+  // Fallback to env vars when no resolved settings passed
+  if (!cfg) {
+    return createFromEnv();
+  }
+
+  if (cfg.provider !== "openai_compatible") {
+    return new LocalHashEmbeddingFunction();
+  }
+
+  // Workaround: seekdb's Schema.fromJSON loads stored embedding function config
+  // (e.g. { name: "openai", properties: {} }) and instantiates it WITHOUT
+  // the API key. Setting OPENAI_API_KEY ensures seekdb can instantiate it.
+  if (cfg.apiKey) {
+    process.env.OPENAI_API_KEY = cfg.apiKey;
+  }
+
+  if (!cfg.apiKey) {
+    const fromEnv = process.env[cfg.apiKeyEnv]?.trim();
+    if (!fromEnv) {
+      console.warn(
+        `[ebrain] embed provider=openai_compatible but no API key; falling back to hash.`,
+      );
+      return new LocalHashEmbeddingFunction();
+    }
+    process.env.OPENAI_API_KEY = fromEnv;
+    return new OpenAIEmbeddingFunction({
+      baseURL: cfg.baseURL,
+      modelName: cfg.model,
+      dimensions: cfg.dimensions,
+      apiKeyEnvVar: cfg.apiKeyEnv,
+    });
+  }
+
+  return new OpenAIEmbeddingFunction({
+    baseURL: cfg.baseURL,
+    modelName: cfg.model,
+    dimensions: cfg.dimensions,
+    apiKey: cfg.apiKey,
+  });
+}
+
+// ---------------------------------------------------------------------------
+// Legacy fallback: read directly from env vars (backward compatible)
+// ---------------------------------------------------------------------------
+
+const DEFAULT_DASHSCOPE_COMPAT_URL =
+  "https://dashscope.aliyuncs.com/compatible-mode/v1";
+const DEFAULT_EMBED_MODEL = "text-embedding-v4";
+const DEFAULT_EMBED_DIMENSIONS = 1024;
+const DEFAULT_KEY_ENV = "DASHSCOPE_API_KEY";
+
+function createFromEnv(): EmbeddingFunction {
+  const provider = (process.env.EBRAIN_EMBED_PROVIDER ?? "hash")
+    .trim()
+    .toLowerCase();
+  if (provider !== "openai_compatible") {
+    return new LocalHashEmbeddingFunction();
+  }
+
+  const baseURL =
+    process.env.EBRAIN_EMBED_BASE_URL?.trim() || DEFAULT_DASHSCOPE_COMPAT_URL;
+  const modelName =
+    process.env.EBRAIN_EMBED_MODEL?.trim() || DEFAULT_EMBED_MODEL;
+  const dimensionsRaw = process.env.EBRAIN_EMBED_DIMENSIONS?.trim();
+  const dimensions = dimensionsRaw
+    ? Number(dimensionsRaw)
+    : DEFAULT_EMBED_DIMENSIONS;
+  if (!Number.isFinite(dimensions) || dimensions <= 0) {
+    throw new Error(
+      `[ebrain] EBRAIN_EMBED_DIMENSIONS must be a positive number, got: ${dimensionsRaw}`,
+    );
+  }
+
+  const directKey = process.env.EBRAIN_EMBED_API_KEY?.trim();
+  const keyEnv =
+    process.env.EBRAIN_EMBED_API_KEY_ENV?.trim() || DEFAULT_KEY_ENV;
+  const fromNamedEnv = process.env[keyEnv]?.trim();
+  const resolvedKey = directKey || fromNamedEnv;
+
+  if (!resolvedKey) {
+    console.warn(
+      `[ebrain] EBRAIN_EMBED_PROVIDER=openai_compatible but no API key (set EBRAIN_EMBED_API_KEY or ${keyEnv}); falling back to hash embedding.`,
+    );
+    return new LocalHashEmbeddingFunction();
+  }
+
+  // Set OPENAI_API_KEY for seekdb's Schema.fromJSON fallback
+  process.env.OPENAI_API_KEY = resolvedKey;
+
+  if (directKey) {
+    return new OpenAIEmbeddingFunction({
+      baseURL,
+      modelName,
+      dimensions,
+      apiKey: directKey,
+    });
+  }
+
+  return new OpenAIEmbeddingFunction({
+    baseURL,
+    modelName,
+    dimensions,
+    apiKeyEnvVar: keyEnv,
+  });
+}

package/src/ai/entity-link.ts

@@ -0,0 +1,226 @@
+import { ResolvedLLM } from "../settings";
+
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
+export type EntityType = "person" | "company" | "project" | "organization" | "event" | "other";
+
+export type RelationType = 
+  | "founder_of"
+  | "works_at"
+  | "leader_of"
+  | "collaborates_with"
+  | "competes_with"
+  | "acquired"
+  | "part_of"
+  | "invested_in"
+  | "mentioned_in"
+  | "related_to";
+
+export interface EntityRef {
+  name: string;
+  type: EntityType;
+}
+
+export interface EntityRelation {
+  type: "relation";
+  from: EntityRef;
+  to: EntityRef;
+  /** Semantic relation type. */
+  relation: RelationType;
+  /** The original sentence mentioning this relationship. */
+  context: string;
+  /** Confidence score 0.0 - 1.0. */
+  confidence: number;
+}
+
+export type ExtractionResult = EntityRelation[];
+
+// ---------------------------------------------------------------------------
+// Entity type mapping for slug prefix
+// ---------------------------------------------------------------------------
+
+const TYPE_PREFIX: Record<EntityType, string> = {
+  person: "people",
+  company: "companies",
+  project: "projects",
+  organization: "organizations",
+  event: "events",
+  other: "entities",
+};
+
+/**
+ * Convert an entity name to a slug: "Ali Partovi" → "ali-partovi"
+ */
+export function entityToSlug(name: string, type: EntityType): string {
+  const prefix = TYPE_PREFIX[type] ?? "entities";
+  const slugPart = name
+    .toLowerCase()
+    .replace(/[^a-z0-9\u4e00-\u9fff]+/g, "-")
+    .replace(/^-+|-+$/g, "");
+  return `${prefix}/${slugPart || "untitled"}`;
+}
+
+// ---------------------------------------------------------------------------
+// LLM extraction
+// ---------------------------------------------------------------------------
+
+const RELATION_TYPES = [
+  "founder_of", "works_at", "leader_of", 
+  "collaborates_with", "competes_with", "acquired", 
+  "part_of", "invested_in", "mentioned_in", "related_to"
+].join(", ");
+
+/**
+ * Use the configured LLM to extract entity relationships from text.
+ * Returns a list of relations with relation type, confidence, and context.
+ */
+export async function extractRelations(
+  content: string,
+  llm: ResolvedLLM,
+): Promise<ExtractionResult> {
+  const trimmed = content.trim();
+  if (!trimmed) return [];
+
+  // Truncate for API efficiency: first 4000 + last 1000 chars
+  let context: string;
+  if (trimmed.length <= 5000) {
+    context = trimmed;
+  } else {
+    context = trimmed.slice(0, 4000) + "\n\n...\n\n" + trimmed.slice(-1000);
+  }
+
+  const apiKey = resolveApiKey(llm);
+  if (!apiKey) return [];
+
+  const body = {
+    model: llm.model,
+    messages: [
+      {
+        role: "system",
+        content:
+          "You are a knowledge graph extraction assistant. " +
+          "Identify relationships between named entities. " +
+          "For each relationship, provide: from entity, to entity, relation type, confidence score, and exact context sentence. " +
+          `Allowed relation types: ${RELATION_TYPES}. ` +
+          "Output ONLY a JSON array. Schema: " +
+          '{ "type": "relation", "from": {"name": "...", "type": "..."}, ' +
+          '"to": {"name": "...", "type": "..."}, "relation": "...", "context": "...", "confidence": 0.9 }. ' +
+          "Output ONLY the JSON array. /no_think",
+      },
+      {
+        role: "user",
+        content: `Extract relationships from:\n\n${context}`,
+      },
+    ],
+    temperature: 0.1,
+    max_tokens: 1024,
+    enable_thinking: false,
+  };
+
+  try {
+    const resp = await fetch(
+      llm.baseURL.endsWith("/")
+        ? llm.baseURL + "chat/completions"
+        : llm.baseURL + "/chat/completions",
+      {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          Authorization: `Bearer ${apiKey}`,
+        },
+        body: JSON.stringify(body),
+      },
+    );
+
+    if (!resp.ok) {
+      const text = await resp.text();
+      console.warn(
+        `[ebrain] Entity extraction failed (${resp.status}): ${text.slice(0, 200)}`,
+      );
+      return [];
+    }
+
+    const data = await resp.json();
+    const raw = data.choices?.[0]?.message?.content?.trim();
+    if (!raw) return [];
+
+    const match = raw.match(/\[[\s\S]*\]/);
+    if (!match) return [];
+
+    const parsed = JSON.parse(match[0]) as unknown[];
+    const relations: ExtractionResult = [];
+
+    for (const item of parsed) {
+      if (typeof item !== "object" || item === null) continue;
+      const r = item as Record<string, unknown>;
+      if (r.type !== "relation") continue;
+
+      const fromRef = parseEntityRef(r.from);
+      const toRef = parseEntityRef(r.to);
+      const relation = String(r.relation || "related_to");
+      const contextStr = typeof r.context === "string" ? r.context.trim() : "";
+      const confidence = typeof r.confidence === "number" ? r.confidence : 0.8;
+
+      if (!fromRef || !toRef || !contextStr) continue;
+
+      relations.push({
+        type: "relation",
+        from: fromRef,
+        to: toRef,
+        relation: normalizeRelationType(relation),
+        context: contextStr,
+        confidence,
+      });
+    }
+
+    return relations;
+  } catch (error) {
+    const msg = error instanceof Error ? error.message : String(error);
+    console.warn(`[ebrain] Entity extraction error: ${msg}`);
+    return [];
+  }
+}
+
+function parseEntityRef(val: unknown): EntityRef | null {
+  if (typeof val !== "object" || val === null) return null;
+  const obj = val as Record<string, unknown>;
+  const name = typeof obj.name === "string" ? obj.name.trim() : "";
+  const rawType = typeof obj.type === "string" ? obj.type : "other";
+  if (!name) return null;
+  return { name, type: normalizeEntityType(rawType) };
+}
+
+function normalizeEntityType(raw: string): EntityType {
+  const lower = raw.toLowerCase().trim();
+  if (lower.includes("person") || lower.includes("people")) return "person";
+  if (lower.includes("company") || lower.includes("corp") || lower.includes("business")) return "company";
+  if (lower.includes("project")) return "project";
+  if (lower.includes("organization") || lower.includes("org") || lower.includes("ngo")) return "organization";
+  if (lower.includes("event")) return "event";
+  return "other";
+}
+
+export function normalizeRelationType(raw: string): RelationType {
+  const lower = raw.toLowerCase().trim().replace(/-/g, "_");
+  const validTypes = RELATION_TYPES.split(", ");
+  if (validTypes.includes(lower as RelationType)) return lower as RelationType;
+  // Fallbacks
+  if (lower.includes("founder") || lower.includes("create")) return "founder_of";
+  if (lower.includes("work") || lower.includes("join")) return "works_at";
+  if (lower.includes("lead") || lower.includes("head") || lower.includes("manage")) return "leader_of";
+  if (lower.includes("collabor") || lower.includes("partner")) return "collaborates_with";
+  if (lower.includes("compet")) return "competes_with";
+  if (lower.includes("acquir") || lower.includes("buy")) return "acquired";
+  if (lower.includes("invest")) return "invested_in";
+  if (lower.includes("part") || lower.includes("belong")) return "part_of";
+  if (lower.includes("mention") || lower.includes("refer")) return "mentioned_in";
+  return "related_to";
+}
+
+function resolveApiKey(llm: ResolvedLLM): string {
+  if (llm.apiKey) return llm.apiKey;
+  if (llm.apiKeyEnv) return process.env[llm.apiKeyEnv] ?? "";
+  return "";
+}

package/src/ai/hash-embed.ts

@@ -0,0 +1,30 @@
+import type { EmbeddingConfig, EmbeddingFunction } from "seekdb";
+
+const DIM = 384;
+
+function hashToVector(text: string): number[] {
+  const v = new Array<number>(DIM).fill(0);
+  for (let i = 0; i < text.length; i += 1) {
+    const j = i % DIM;
+    v[j] = (v[j]! + text.charCodeAt(i) * (i + 1)) / 1e6;
+  }
+  const norm = Math.sqrt(v.reduce((s, x) => s + x * x, 0)) || 1;
+  return v.map((x) => x / norm);
+}
+
+/**
+ * 零依赖、确定性的伪向量，满足 seekdb 集合对「文档 → 向量」的硬性要求；
+ * 不替代真实语义模型，仅用于嵌入模式本地可跑通全文 + 近似检索管线。
+ */
+export class LocalHashEmbeddingFunction implements EmbeddingFunction {
+  readonly name = "ebrain-local-hash";
+  readonly dimension = DIM;
+
+  async generate(texts: string[]): Promise<number[][]> {
+    return texts.map((t) => hashToVector(t));
+  }
+
+  getConfig(): EmbeddingConfig {
+    return { dimension: DIM };
+  }
+}