Merge pull request #4 from RedPlanetHQ/manoj/mcp

Fix: Deduplication of Entites
2026-01-10 23:48:26 +00:00 · 2025-06-23 21:37:56 +05:30 · 2025-06-23 21:37:56 +05:30 · bbf63b6b11
commit bbf63b6b11
parent 4a91b281e5 3a3055e3cc
10 changed files with 189 additions and 296 deletions
--- a/apps/webapp/app/env.server.ts
+++ b/apps/webapp/app/env.server.ts
@ -1,6 +1,7 @@
 import { z } from "zod";
 import { isValidDatabaseUrl } from "./utils/db";
 import { isValidRegex } from "./utils/regex";
+import { LLMModelEnum } from "@core/types";

 const EnvironmentSchema = z.object({
  NODE_ENV: z.union([
@ -69,6 +70,10 @@ const EnvironmentSchema = z.object({
  SMTP_SECURE: z.coerce.boolean().optional(),
  SMTP_USER: z.string().optional(),
  SMTP_PASSWORD: z.string().optional(),
+
+  // Model envs
+  MODEL: z.string().default(LLMModelEnum.GPT41),
+  OLLAMA_URL: z.string().optional(),
 });

 export type Environment = z.infer<typeof EnvironmentSchema>;
--- a/apps/webapp/app/lib/model.server.ts
+++ b/apps/webapp/app/lib/model.server.ts
@ -7,48 +7,58 @@ import {
 } from "ai";
 import { openai } from "@ai-sdk/openai";
 import { logger } from "~/services/logger.service";
+import { env } from "~/env.server";
+import { createOllama } from "ollama-ai-provider";

 export async function makeModelCall(
  stream: boolean,
-  model: LLMModelEnum,
  messages: CoreMessage[],
  onFinish: (text: string, model: string) => void,
  options?: any,
 ) {
  let modelInstance;
+  const model = env.MODEL;
  let finalModel: string = "unknown";
+  const ollamaUrl = process.env.OLLAMA_URL;

-  switch (model) {
-    case LLMModelEnum.GPT35TURBO:
-    case LLMModelEnum.GPT4TURBO:
-    case LLMModelEnum.GPT4O:
-    case LLMModelEnum.GPT41:
-    case LLMModelEnum.GPT41MINI:
-    case LLMModelEnum.GPT41NANO:
-      finalModel = LLMMappings[model];
-      modelInstance = openai(finalModel, { ...options });
-      break;
+  if (ollamaUrl) {
+    const ollama = createOllama({
+      baseURL: ollamaUrl,
+    });
+    modelInstance = ollama(model); // Default to llama2 if no model specified
+  } else {
+    switch (model) {
+      case LLMModelEnum.GPT35TURBO:
+      case LLMModelEnum.GPT4TURBO:
+      case LLMModelEnum.GPT4O:
+      case LLMModelEnum.GPT41:
+      case LLMModelEnum.GPT41MINI:
+      case LLMModelEnum.GPT41NANO:
+        finalModel = LLMMappings[model];
+        modelInstance = openai(finalModel, { ...options });
+        break;

-    case LLMModelEnum.CLAUDEOPUS:
-    case LLMModelEnum.CLAUDESONNET:
-    case LLMModelEnum.CLAUDEHAIKU:
-      finalModel = LLMMappings[model];
-      break;
+      case LLMModelEnum.CLAUDEOPUS:
+      case LLMModelEnum.CLAUDESONNET:
+      case LLMModelEnum.CLAUDEHAIKU:
+        finalModel = LLMMappings[model];
+        break;

-    case LLMModelEnum.GEMINI25FLASH:
-    case LLMModelEnum.GEMINI25PRO:
-    case LLMModelEnum.GEMINI20FLASH:
-    case LLMModelEnum.GEMINI20FLASHLITE:
-      finalModel = LLMMappings[model];
-      break;
+      case LLMModelEnum.GEMINI25FLASH:
+      case LLMModelEnum.GEMINI25PRO:
+      case LLMModelEnum.GEMINI20FLASH:
+      case LLMModelEnum.GEMINI20FLASHLITE:
+        finalModel = LLMMappings[model];
+        break;

-    default:
-      logger.warn(`Unsupported model type: ${model}`);
-      break;
+      default:
+        logger.warn(`Unsupported model type: ${model}`);
+        break;
+    }
  }

  if (stream) {
-    return await streamText({
+    return streamText({
      model: modelInstance as LanguageModelV1,
      messages,
      onFinish: async ({ text }) => {
--- a/apps/webapp/app/services/graphModels/entity.ts
+++ b/apps/webapp/app/services/graphModels/entity.ts
@ -63,12 +63,14 @@ export async function findSimilarEntities(params: {
  queryEmbedding: number[];
  limit: number;
  threshold: number;
+  userId: string;
 }): Promise<EntityNode[]> {
  const query = `
          MATCH (entity:Entity)
          WHERE entity.nameEmbedding IS NOT NULL
          WITH entity, vector.similarity.cosine($queryEmbedding, entity.nameEmbedding) AS score
          WHERE score >= $threshold
+          AND entity.userId = $userId
          RETURN entity, score
          ORDER BY score DESC
        `;
--- a/apps/webapp/app/services/knowledgeGraph.server.ts
+++ b/apps/webapp/app/services/knowledgeGraph.server.ts
@ -196,37 +196,30 @@ export class KnowledgeGraphService {

    let responseText = "";

-    await makeModelCall(
-      false,
-      LLMModelEnum.GPT41,
-      messages as CoreMessage[],
-      (text) => {
-        responseText = text;
-      },
-    );
+    await makeModelCall(false, messages as CoreMessage[], (text) => {
+      responseText = text;
+    });

    // Convert to EntityNode objects
-    const entities: EntityNode[] = [];
+    let entities: EntityNode[] = [];

    const outputMatch = responseText.match(/<output>([\s\S]*?)<\/output>/);
    if (outputMatch && outputMatch[1]) {
      responseText = outputMatch[1].trim();
      const extractedEntities = JSON.parse(responseText || "{}").entities || [];

-      entities.push(
-        ...(await Promise.all(
-          extractedEntities.map(async (entity: any) => ({
-            uuid: crypto.randomUUID(),
-            name: entity.name,
-            type: entity.type,
-            attributes: entity.attributes || {},
-            nameEmbedding: await this.getEmbedding(
-              `${entity.type}: ${entity.name}`,
-            ),
-            createdAt: new Date(),
-            userId: episode.userId,
-          })),
-        )),
+      entities = await Promise.all(
+        extractedEntities.map(async (entity: any) => ({
+          uuid: crypto.randomUUID(),
+          name: entity.name,
+          type: entity.type,
+          attributes: entity.attributes || {},
+          nameEmbedding: await this.getEmbedding(
+            `${entity.type}: ${entity.name}`,
+          ),
+          createdAt: new Date(),
+          userId: episode.userId,
+        })),
      );
    }

@ -260,15 +253,11 @@ export class KnowledgeGraphService {
    const messages = extractStatements(context);

    let responseText = "";
-    await makeModelCall(
-      false,
-      LLMModelEnum.GPT41,
-      messages as CoreMessage[],
-      (text) => {
-        responseText = text;
-      },
-    );
+    await makeModelCall(false, messages as CoreMessage[], (text) => {
+      responseText = text;
+    });

+    console.log(responseText);
    const outputMatch = responseText.match(/<output>([\s\S]*?)<\/output>/);
    if (outputMatch && outputMatch[1]) {
      responseText = outputMatch[1].trim();
@ -422,6 +411,7 @@ export class KnowledgeGraphService {
          queryEmbedding: entity.nameEmbedding,
          limit: 5,
          threshold: 0.85,
+          userId: episode.userId,
        });
        return {
          entity,
@ -483,14 +473,9 @@ export class KnowledgeGraphService {
    const messages = dedupeNodes(dedupeContext);
    let responseText = "";

-    await makeModelCall(
-      false,
-      LLMModelEnum.GPT41,
-      messages as CoreMessage[],
-      (text) => {
-        responseText = text;
-      },
-    );
+    await makeModelCall(false, messages as CoreMessage[], (text) => {
+      responseText = text;
+    });

    // Step 5: Process LLM response
    const outputMatch = responseText.match(/<output>([\s\S]*?)<\/output>/);
@ -507,14 +492,14 @@ export class KnowledgeGraphService {
      const entityResolutionMap = new Map<string, EntityNode>();

      nodeResolutions.forEach((resolution: any, index: number) => {
-        const originalEntity = uniqueEntities[resolution.id ?? index];
+        const originalEntity = allEntityResults[resolution.id ?? index];
        if (!originalEntity) return;

        const duplicateIdx = resolution.duplicate_idx ?? -1;

        // Get the corresponding result from allEntityResults
        const resultEntry = allEntityResults.find(
-          (result) => result.entity.uuid === originalEntity.uuid,
+          (result) => result.entity.uuid === originalEntity.entity.uuid,
        );

        if (!resultEntry) return;
@ -523,7 +508,7 @@ export class KnowledgeGraphService {
        const resolvedEntity =
          duplicateIdx >= 0 && duplicateIdx < resultEntry.similarEntities.length
            ? resultEntry.similarEntities[duplicateIdx]
-            : originalEntity;
+            : originalEntity.entity;

        // Update name if provided
        if (resolution.name) {
@ -531,7 +516,7 @@ export class KnowledgeGraphService {
        }

        // Map original UUID to resolved entity
-        entityResolutionMap.set(originalEntity.uuid, resolvedEntity);
+        entityResolutionMap.set(originalEntity.entity.uuid, resolvedEntity);
      });

      // Step 7: Reconstruct triples with resolved entities
@ -673,7 +658,7 @@ export class KnowledgeGraphService {
      let responseText = "";

      // Call the LLM to analyze all statements at once
-      await makeModelCall(false, LLMModelEnum.GPT41, messages, (text) => {
+      await makeModelCall(false, messages, (text) => {
        responseText = text;
      });

@ -804,14 +789,9 @@ export class KnowledgeGraphService {
    let responseText = "";

    // Call the LLM to extract attributes
-    await makeModelCall(
-      false,
-      LLMModelEnum.GPT41,
-      messages as CoreMessage[],
-      (text) => {
-        responseText = text;
-      },
-    );
+    await makeModelCall(false, messages as CoreMessage[], (text) => {
+      responseText = text;
+    });

    try {
      const outputMatch = responseText.match(/<output>([\s\S]*?)<\/output>/);
@ -864,7 +844,7 @@ export class KnowledgeGraphService {
    };
    const messages = normalizePrompt(context);
    let responseText = "";
-    await makeModelCall(false, LLMModelEnum.GPT41, messages, (text) => {
+    await makeModelCall(false, messages, (text) => {
      responseText = text;
    });
    let normalizedEpisodeBody = "";
--- a/apps/webapp/app/services/search/rerank.ts
+++ b/apps/webapp/app/services/search/rerank.ts
@ -100,7 +100,6 @@ export async function applyCrossEncoderReranking(
      let responseText = "";
      await makeModelCall(
        false,
-        LLMModelEnum.GPT41NANO,
        messages,
        (text) => {
          responseText = text;
--- a/apps/webapp/app/utils/presets/nodes.ts
+++ b/apps/webapp/app/utils/presets/nodes.ts
@ -2,12 +2,14 @@ export enum Apps {
  LINEAR = "LINEAR",
  SLACK = "SLACK",
  SOL = "SOL",
+  GITHUB = "GITHUB",
 }

 export const AppNames = {
  [Apps.LINEAR]: "Linear",
  [Apps.SLACK]: "Slack",
  [Apps.SOL]: "Sol",
+  [Apps.GITHUB]: "GitHub",
 } as const;

 // Define attribute structure
@ -97,7 +99,7 @@ export const GENERAL_NODE_TYPES = {
  },
  ALIAS: {
    name: "Alias",
-    description: "An alternative name or identifier for an entity",
+    description: "An alternative name or identifier for nouns and pronouns",
    attributes: [
      {
        name: "originalName",
@ -111,6 +113,24 @@ export const GENERAL_NODE_TYPES = {
      },
    ],
  },
+  FILE: {
+    name: "File",
+    description: "A document, image or other file shared in an app",
+    attributes: [
+      {
+        name: "fileId",
+        description: "Unique identifier for the file",
+        type: "string",
+        required: true,
+      },
+      {
+        name: "source",
+        description: "The source of the file",
+        type: "string",
+        required: true,
+      },
+    ],
+  },
 } as const;

 // App-specific node types
@ -127,32 +147,6 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "title",
-          description: "The title of the task",
-          type: "string",
-          required: true,
-        },
-        {
-          name: "description",
-          description: "The description of the task",
-          type: "string",
-        },
-        {
-          name: "status",
-          description: "The current status of the task",
-          type: "string",
-        },
-        {
-          name: "dueDate",
-          description: "The due date of the task",
-          type: "date",
-        },
-        {
-          name: "priority",
-          description: "The priority level of the task",
-          type: "string",
-        },
      ],
    },
    LIST: {
@ -166,22 +160,6 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "title",
-          description: "The title of the list",
-          type: "string",
-          required: true,
-        },
-        {
-          name: "description",
-          description: "The description of the list",
-          type: "string",
-        },
-        {
-          name: "itemCount",
-          description: "The number of items in the list",
-          type: "number",
-        },
      ],
    },
    PREFERENCE: {
@ -203,25 +181,6 @@ export const APP_NODE_TYPES = {
        },
      ],
    },
-    COMMAND: {
-      name: "Sol Command",
-      description:
-        "A user-issued command or trigger phrase, often starting with '/', that directs the system or an app to perform a specific action. Commands should always be extracted as distinct, important user actions.",
-      attributes: [
-        {
-          name: "commandId",
-          description: "Unique identifier for the command",
-          type: "string",
-          required: true,
-        },
-        {
-          name: "commandName",
-          description: "The name of the command",
-          type: "string",
-          required: true,
-        },
-      ],
-    },
    AUTOMATION: {
      name: "Sol Automation",
      description:
@ -233,18 +192,6 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "trigger",
-          description: "The event that triggers this automation",
-          type: "string",
-          required: true,
-        },
-        {
-          name: "action",
-          description: "The action performed by this automation",
-          type: "string",
-          required: true,
-        },
      ],
    },
  },
@ -259,27 +206,6 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "title",
-          description: "The title of the issue",
-          type: "string",
-          required: true,
-        },
-        {
-          name: "status",
-          description: "The current status of the issue",
-          type: "string",
-        },
-        {
-          name: "priority",
-          description: "The priority level of the issue",
-          type: "number",
-        },
-        {
-          name: "assignee",
-          description: "The person assigned to the issue",
-          type: "string",
-        },
      ],
    },
    PROJECT: {
@ -292,27 +218,6 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "name",
-          description: "The name of the project",
-          type: "string",
-          required: true,
-        },
-        {
-          name: "status",
-          description: "The current status of the project",
-          type: "string",
-        },
-        {
-          name: "startDate",
-          description: "The start date of the project",
-          type: "date",
-        },
-        {
-          name: "targetDate",
-          description: "The target completion date of the project",
-          type: "date",
-        },
      ],
    },
    CYCLE: {
@ -325,24 +230,6 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "name",
-          description: "The name of the cycle",
-          type: "string",
-          required: true,
-        },
-        {
-          name: "startDate",
-          description: "The start date of the cycle",
-          type: "date",
-          required: true,
-        },
-        {
-          name: "endDate",
-          description: "The end date of the cycle",
-          type: "date",
-          required: true,
-        },
      ],
    },
    TEAM: {
@ -355,22 +242,6 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "name",
-          description: "The name of the team",
-          type: "string",
-          required: true,
-        },
-        {
-          name: "key",
-          description: "The team's key or shorthand",
-          type: "string",
-        },
-        {
-          name: "memberCount",
-          description: "Number of members in the team",
-          type: "number",
-        },
      ],
    },
    LABEL: {
@ -383,17 +254,6 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "name",
-          description: "The name of the label",
-          type: "string",
-          required: true,
-        },
-        {
-          name: "color",
-          description: "The color of the label",
-          type: "string",
-        },
      ],
    },
  },
@ -408,22 +268,12 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "name",
-          description: "The name of the channel",
-          type: "string",
-          required: true,
-        },
+
        {
          name: "isPrivate",
          description: "Whether the channel is private",
          type: "boolean",
        },
-        {
-          name: "memberCount",
-          description: "The number of members in the channel",
-          type: "number",
-        },
      ],
    },
    THREAD: {
@ -442,11 +292,6 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "replyCount",
-          description: "Number of replies in the thread",
-          type: "number",
-        },
      ],
    },
    MESSAGE: {
@ -459,23 +304,6 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "content",
-          description: "The content of the message",
-          type: "string",
-          required: true,
-        },
-        {
-          name: "timestamp",
-          description: "When the message was sent",
-          type: "date",
-          required: true,
-        },
-        {
-          name: "reactions",
-          description: "Reactions to the message",
-          type: "array",
-        },
      ],
    },
    REACTION: {
@ -488,39 +316,79 @@ export const APP_NODE_TYPES = {
          type: "string",
          required: true,
        },
-        {
-          name: "count",
-          description: "Number of users who reacted with this emoji",
-          type: "number",
-          required: true,
-        },
      ],
    },
-    FILE: {
-      name: "Slack File",
-      description: "A document, image or other file shared in Slack",
+  },
+  [Apps.GITHUB]: {
+    REPOSITORY: {
+      name: "GitHub Repository",
+      description: "A code repository hosted on GitHub",
      attributes: [
        {
-          name: "fileId",
-          description: "Unique identifier for the file",
+          name: "repoId",
+          description: "Unique identifier for the repository",
          type: "string",
          required: true,
        },
        {
          name: "name",
-          description: "The name of the file",
+          description: "The name of the repository",
          type: "string",
          required: true,
        },
        {
-          name: "type",
-          description: "The file type or format",
+          name: "owner",
+          description: "Owner (user or organization) of the repository",
          type: "string",
+          required: true,
        },
+      ],
+    },
+    ISSUE: {
+      name: "GitHub Issue",
+      description: "An issue created to track bugs, tasks, or feature requests",
+      attributes: [
        {
-          name: "size",
-          description: "The size of the file in bytes",
-          type: "number",
+          name: "issueId",
+          description: "Unique identifier for the issue",
+          type: "string",
+          required: true,
+        },
+      ],
+    },
+    PULL_REQUEST: {
+      name: "GitHub Pull Request",
+      description: "A pull request to propose changes to a repository",
+      attributes: [
+        {
+          name: "PR number",
+          description: "Unique number for the pull request",
+          type: "string",
+          required: true,
+        },
+      ],
+    },
+    COMMIT: {
+      name: "GitHub Commit",
+      description: "A commit representing a set of changes in a repository",
+      attributes: [
+        {
+          name: "commitSha",
+          description: "SHA hash of the commit",
+          type: "string",
+          required: true,
+        },
+      ],
+    },
+    BRANCH: {
+      name: "GitHub Branch",
+      description: "A branch in a GitHub repository",
+      attributes: [
+        {
+          name: "branchName",
+          description: "Name of the branch",
+          type: "string",
+          required: true,
        },
      ],
    },
--- a/apps/webapp/package.json
+++ b/apps/webapp/package.json
@ -68,6 +68,7 @@
    "nanoid": "3.3.8",
    "neo4j-driver": "^5.28.1",
    "non.geist": "^1.0.2",
+    "ollama-ai-provider": "1.2.0",
    "posthog-js": "^1.116.6",
    "react": "^18.2.0",
    "react-dom": "^18.2.0",
--- a/package.json
+++ b/package.json
@ -1,10 +1,11 @@
 {
  "name": "core",
  "private": true,
-  "version": "0.1.3",
-  "workspaces":
-  [ "apps/*", "packages/*" ]
-  ,
+  "version": "0.1.5",
+  "workspaces": [
+    "apps/*",
+    "packages/*"
+  ],
  "scripts": {
    "build": "dotenv -- turbo run build",
    "dev": "dotenv -- turbo run dev",
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@ -204,6 +204,9 @@ importers:
      non.geist:
        specifier: ^1.0.2
        version: 1.0.4
+      ollama-ai-provider:
+        specifier: 1.2.0
+        version: 1.2.0(zod@3.23.8)
      posthog-js:
        specifier: ^1.116.6
        version: 1.250.2
@ -6580,6 +6583,15 @@ packages:
    resolution: {integrity: sha512-gXah6aZrcUxjWg2zR2MwouP2eHlCBzdV4pygudehaKXSGW4v2AsRQUK+lwwXhii6KFZcunEnmSUoYp5CXibxtA==}
    engines: {node: '>= 0.4'}

+  ollama-ai-provider@1.2.0:
+    resolution: {integrity: sha512-jTNFruwe3O/ruJeppI/quoOUxG7NA6blG3ZyQj3lei4+NnJo7bi3eIRWqlVpRlu/mbzbFXeJSBuYQWF6pzGKww==}
+    engines: {node: '>=18'}
+    peerDependencies:
+      zod: ^3.0.0
+    peerDependenciesMeta:
+      zod:
+        optional: true
+
  on-finished@2.3.0:
    resolution: {integrity: sha512-ikqdkGAAyf/X/gPhXGvfgAytDZtDbr+bkNUJ0N9h5MI/dmdgCs3l6hoHrcUv41sRKew3jIwrp4qQDXiK99Utww==}
    engines: {node: '>= 0.8'}
@ -6685,6 +6697,9 @@ packages:
    resolution: {integrity: sha512-CiyeOxFT/JZyN5m0z9PfXw4SCBJ6Sygz1Dpl0wqjlhDEGGBP1GnsUVEL0p63hoG1fcj3fHynXi9NYO4nWOL+qQ==}
    engines: {node: '>= 0.8'}

+  partial-json@0.1.7:
+    resolution: {integrity: sha512-Njv/59hHaokb/hRUjce3Hdv12wd60MtM9Z5Olmn+nehe0QDAsRtRbJPvJ0Z91TusF0SuZRIvnM+S4l6EIP8leA==}
+
  path-exists@4.0.0:
    resolution: {integrity: sha512-ak9Qy5Q7jYb2Wwcey5Fpvg2KoAc/ZIhLSLOSBmRmygPsGwkVVt0fZa0qrtMz+m6tJTAHfZQ8FnmB4MG4LWy7/w==}
    engines: {node: '>=8'}
@ -15619,6 +15634,14 @@ snapshots:
      define-properties: 1.2.1
      es-object-atoms: 1.1.1

+  ollama-ai-provider@1.2.0(zod@3.23.8):
+    dependencies:
+      '@ai-sdk/provider': 1.1.3
+      '@ai-sdk/provider-utils': 2.2.8(zod@3.23.8)
+      partial-json: 0.1.7
+    optionalDependencies:
+      zod: 3.23.8
+
  on-finished@2.3.0:
    dependencies:
      ee-first: 1.1.1
@ -15737,6 +15760,8 @@ snapshots:

  parseurl@1.3.3: {}

+  partial-json@0.1.7: {}
+
  path-exists@4.0.0: {}

  path-is-absolute@1.0.1: {}
--- a/turbo.json
+++ b/turbo.json
@ -58,6 +58,8 @@
    "NEO4J_PASSWORD",
    "OPENAI_API_KEY",
    "MAGIC_LINK_SECRET",
-    "ENABLE_EMAIL_LOGIN"
+    "ENABLE_EMAIL_LOGIN",
+    "MODEL",
+    "OLLAMA_URL"
  ]
 }