From 9f93824ba97b434de6c5a53df1cd1c9d5364c2a3 Mon Sep 17 00:00:00 2001
From: Harshith Mullapudi <harshithmullapudi@gmail.com>
Date: Mon, 23 Jun 2025 13:14:12 +0530
Subject: [PATCH] Fix: add source to env

---
 mcp/src/index.ts            | 9 +++++++++
 mcp/src/operations/graph.ts | 3 ++-
 mcp/src/types/graph.ts      | 1 -
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/mcp/src/index.ts b/mcp/src/index.ts
index 4eb0d15..e0b57ac 100644
--- a/mcp/src/index.ts
+++ b/mcp/src/index.ts
@@ -68,6 +68,15 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
 });
 
 async function runServer() {
+  // Check required environment variables
+  const requiredEnvVars = ["API_TOKEN", "API_BASE_URL", "SOURCE"];
+  const missingEnvVars = requiredEnvVars.filter((envVar) => !process.env[envVar]);
+
+  if (missingEnvVars.length > 0) {
+    console.error(`Missing required environment variables: ${missingEnvVars.join(", ")}`);
+    process.exit(1);
+  }
+
   const transport = new StdioServerTransport();
   await server.connect(transport);
   console.error("Sigma MCP Server running on stdio");
diff --git a/mcp/src/operations/graph.ts b/mcp/src/operations/graph.ts
index a7f79fd..c98568d 100644
--- a/mcp/src/operations/graph.ts
+++ b/mcp/src/operations/graph.ts
@@ -9,10 +9,11 @@ export async function searchKnowledgeGraph(args: SearchKG) {
   });
   return response.data;
 }
+
 export async function ingestKnowledgeGraph(args: IngestKG) {
   const response = await axiosInstance.post(`/ingest`, {
     data: args.data,
-    source: args.source,
+    source: process.env.SOURCE,
     referenceTime: args.referenceTime,
     sessionId: args.sessionId ? args.sessionId : undefined,
   });
diff --git a/mcp/src/types/graph.ts b/mcp/src/types/graph.ts
index 3587a38..4bc3f7c 100644
--- a/mcp/src/types/graph.ts
+++ b/mcp/src/types/graph.ts
@@ -9,7 +9,6 @@ export const SearchKGSchema = z.object({
 
 export const IngestKGSchema = z.object({
   data: z.string().describe("The data to ingest in text format"),
-  source: z.string().describe("The source of the data"),
   referenceTime: z.string().describe("The reference time in ISO format"),
   sessionId: z.string().optional().describe("The session id of the conversation"),
 });