context-efficient-ai/database/milvus.ts

import { MilvusClient, DataType } from "@zilliz/milvus2-sdk-node";
import { env } from "../server/env";

const address = env.MILVUS_ADDRESS as string; // milvus address
const username = env.MILVUS_USERNAME; // optional username
const password = env.MILVUS_PASSWORD; // optional password
const ssl = true; // secure or not

// connect to milvus
export const client = new MilvusClient({ address, ssl, username, password });

await client.createCollection({
  collection_name: "facts",
  fields: [
    {
      name: "id",
      data_type: DataType.Int64,
      is_primary_key: true,
      //@ts-ignore
      type_params: { auto_id: true },
    },
    {
      name: "user_id",
      data_type: DataType.Int64,
      description:
        "Foreign key linking to the Users Collection. Crucial if you have multiple users.",
    },
    {
      name: "source_message_id",
      data_type: DataType.Int64,
      description:
        "Foreign key linking to the ChatMessages Collection, indicating which message initially triggered the learning of this fact. Useful for debugging and traceability.",
    },
    {
      name: "fact_content",
      data_type: DataType.VarChar,
      description: `The actual natural language statement of the fact (e.g., "The user likes pizza."). This is what you'll inject into the LLM context.`,
    },
    {
      name: "fact_embedding",
      data_type: DataType.FloatVector,
      type_params: {
        dim: 3072,
      },
      description: `Vector for the fact_content itself. This allows you to retrieve facts based on semantic similarity to a user's new statement (e.g., user asks "What do I like?", you embed that and search fact_embedding).`,
    },
    /** Unix timestamp of when the fact was extracted/created. */
    { name: "created_at", data_type: DataType.Int64 },
  ],
});

/** Stores the semantic contexts (trigger_embeddings) for when a specific Fact
 * should be injected. This allows for many-to-one relationship (many triggers
 * for one fact) and efficient vector search on triggers only.
 *
 * It is important to note that in general, we will not be running queries
 * against Facts, only FactTriggers; because the FactTriggers describe different
 * situations in which the Facts should be invoked, each with a different
 * priority_multiplier.
 */
await client.createCollection({
  collection_name: "fact_triggers",
  fields: [
    { name: "id", data_type: DataType.Int64, is_primary_key: true },
    {
      name: "fact_id",
      data_type: DataType.Int64,
      description:
        "Foreign key linking to the Facts Collection. This is the fact that should be triggered.",
    },
    {
      name: "trigger_phrase",
      data_type: DataType.VarChar,
      description: `The natural language phrase describing the trigger (e.g., "When food preferences are discussed"). Useful for context and debugging during trigger generation.`,
    },
    {
      name: "trigger_phrase_embedding",
      data_type: DataType.FloatVector,
      type_params: {
        dim: 3072,
      },
      description: `The vector representation of the trigger phrase. This is what you'll query with the current conversation context, with a large/smart LLM.`,
    },
    {
      name: "priority_multiplier",
      data_type: DataType.Float,
      description:
        "A score multiplier indicating the perceived importance or confidence in this fact trigger, within this FactTrigger's scope. For example, if the FactTrigger is scoped to 'every Wednesday', then it would be very relevant on Wednesdays.",
    },
    {
      name: "priority_multiplier_reason",
      data_type: DataType.VarChar,
      description:
        "The natural-language reason for why this FactTrigger was given its priority multiplier. We store the reason here so that we can easily explain it to the user; but more importantly, the same FactTrigger might have different priorities in different situations that are only discovered from subsequent user input. In this case, we don't want to update the existing FactTrigger's priority multiplier, but instead create a new FactTrigger with a new priority multiplier, with different trigger_phrase and different scope_priority_multiplier_reason.",
    },
    {
      name: "scope_conversation_id",
      data_type: DataType.Int64,
      description:
        "Foreign key linking to the Conversations Collection. This is optional. It is set for scoping the FactTrigger to a specific conversation.",
    },
    /** Unix timestamp of when this specific trigger was generated. */
    { name: "created_at", data_type: DataType.Int64 },
  ],
});

/** Stores metadata about the conversations themselves. Right now, it's only the
 * title and creation time. */
await client.createCollection({
  collection_name: "conversations",
  fields: [
    { name: "id", data_type: DataType.Int64, is_primary_key: true },
    {
      name: "user_id",
      data_type: DataType.Int64,
      description:
        "Foreign key linking to the Users Collection. This is the user who sent this message.",
    },
    {
      name: "title",
      data_type: DataType.VarChar,
      description:
        "The title of the conversation, for friendly display in the UI.",
    },
    {
      name: "created_at",
      data_type: DataType.Int64,
      description: "Unix timestamp of when the conversation was created.",
    },
  ],
});

/** To store the full history of your conversations, enabling summarization, context reconstruction, and traceability.
 *
 * Purpose: To persist the raw chat history for each user.
 */
await client.createCollection({
  collection_name: "conversation_messages",
  fields: [
    { name: "id", data_type: DataType.Int64, is_primary_key: true },
    {
      name: "user_id",
      data_type: DataType.Int64,
      description:
        "Foreign key linking to the Users Collection. This is the user who sent this message.",
    },
    {
      name: "conversation_id",
      data_type: DataType.Int64,
      description:
        "(Optional) Unique identifier for a continuous chat session. Useful for analyzing specific conversational flows.",
    },
    {
      name: "index",
      data_type: DataType.Int32,
      description:
        "The position of this message in the conversation. Useful for loading portions of the conversation history in the UI.",
    },
    {
      name: "role",
      data_type: DataType.VarChar,
      description: `"user" or "assistant". This is what you'll inject into the LLM context.`,
    },
    {
      name: "message_content",
      data_type: DataType.VarChar,
      description:
        "The raw text of the message. This is what you'll inject into the LLM context.",
    },
    {
      name: "message_embedding",
      data_type: DataType.FloatVector,
      type_params: {
        dim: 3072,
      },
      description:
        "Embedding of the message content. Can be used for semantic search within chat history or for generating conversation summary embeddings.",
    },
    {
      name: "running_summary",
      data_type: DataType.VarChar,
      description:
        "The raw text of a running summary of the conversation until this message.",
    },
    {
      name: "running_summary_embedding",
      data_type: DataType.FloatVector,
      type_params: {
        dim: 3072,
      },
      description:
        "(Optional but useful) Embedding of the running summary. Can be used for semantic search within chat history.",
    },
    {
      name: "timestamp",
      data_type: DataType.Int64,
      description:
        "Unix timestamp when the message was sent/received. This is what you'll inject into the LLM context.",
    },
  ],
});

/**
 * To store basic user information. Even if you only have one user initially,
 * it's good practice to structure for multi-tenancy.
 *
 * Purpose: To store basic information about each registered user.
 */
await client.createCollection({
  collection_name: "users",
  fields: [
    { name: "id", data_type: DataType.Int64, is_primary_key: true },
    {
      name: "username",
      data_type: DataType.VarChar,
      description: "User's chosen username.",
    },
    {
      name: "password",
      data_type: DataType.VarChar,
      description: "(Optional) User's password.",
    },
    {
      name: "email",
      data_type: DataType.VarChar,
      description: "(Optional) User's email address.",
    },
    {
      name: "created_at",
      data_type: DataType.Int64,
      description: "Unix timestamp of user creation.",
    },
    {
      name: "last_login",
      data_type: DataType.Int64,
      description: "Unix timestamp of last user login.",
    },
  ],
});

await client.createCollection({
  collection_name: "tools",
  fields: [
    /** Primary key, unique identifier for each fact. */
    {
      name: "id",
      data_type: DataType.Int64,
      is_primary_key: true,
      //@ts-ignore
      type_params: { auto_id: true },
    },
    /** Foreign key linking to the Users Collection.
     * Crucial if you have multiple users. */
    {
      name: "user_id",
      data_type: DataType.Int64,
      description:
        "Foreign key linking to the Users Collection. Crucial if you have multiple users.",
    },
    /** Foreign key linking to the ChatMessages Collection, indicating which
     * message initially triggered the creation of this tool. Useful for
     * debugging and traceability. */
    {
      name: "source_message_id",
      data_type: DataType.Int64,
      description:
        "Foreign key linking to the ChatMessages Collection, indicating which message initially triggered the creation of this tool. Useful for debugging and traceability.",
    },
    {
      name: "name",
      data_type: DataType.VarChar,
      description: `The name of the tool (e.g., "weather").`,
    },
    {
      name: "description",
      data_type: DataType.VarChar,
      description: `The description of the tool (e.g., "Get the current weather in a given location.").`,
    },
    {
      name: "parameter_schema",
      data_type: DataType.JSON,
      description: `The JSON Schema of the parameters to be passed to the tool (e.g., {"location": "string"}).`,
    },
    {
      name: "implementation_language",
      data_type: DataType.VarChar,
      description: `The language of the tool (e.g., "Python"). This is so we know how to execute the tool's implementation code.`,
    },

    {
      name: "implementation_code",
      data_type: DataType.VarChar,
      description: `The actual code that implements the tool (e.g., "def get_weather(location): return 'Sunny'").`,
    },
    /** Unix timestamp of when the fact was extracted/created. */
    { name: "created_at", data_type: DataType.Int64 },
  ],
});

/** I'm still not sure if it's useful to have a `tool_trigger` collection. */

/** TODO: How do I store Agents? Should I store Agents, since the user is
 * supposed to interact it but one "entity"? Is it even useful, since each LLM
 * call will contain what it needs? */