Refactor vectordb init to support zilliz auto create cluster and env variable management

Author: zc277584121

.env.example

@@ -0,0 +1,72 @@
+# CodeContext Environment Variables Example
+# 
+# Copy this file to ~/.codecontext/.env and modify the values as needed
+#
+# Usage: cp env.example ~/.codecontext/.env
+
+# =============================================================================
+# Embedding Provider Configuration
+# =============================================================================
+
+# Embedding provider: OpenAI, VoyageAI, Gemini, Ollama
+EMBEDDING_PROVIDER=OpenAI
+
+# Embedding model (provider-specific)
+EMBEDDING_MODEL=text-embedding-3-small
+
+# Embedding batch size for processing (default: 100)
+EMBEDDING_BATCH_SIZE=100
+
+# =============================================================================
+# OpenAI Configuration
+# =============================================================================
+
+# OpenAI API key
+OPENAI_API_KEY=your-openai-api-key-here
+
+# OpenAI base URL (optional, for custom endpoints)
+# OPENAI_BASE_URL=https://api.openai.com/v1
+
+# =============================================================================
+# VoyageAI Configuration
+# =============================================================================
+
+# VoyageAI API key
+# VOYAGEAI_API_KEY=your-voyageai-api-key-here
+
+# =============================================================================
+# Gemini Configuration
+# =============================================================================
+
+# Google Gemini API key
+# GEMINI_API_KEY=your-gemini-api-key-here
+
+# =============================================================================
+# Ollama Configuration
+# =============================================================================
+
+# Ollama model name
+# OLLAMA_MODEL=
+
+# Ollama host (default: http://localhost:11434)
+# OLLAMA_HOST=http://localhost:11434
+
+# =============================================================================
+# Vector Database Configuration (Milvus/Zilliz)
+# =============================================================================
+
+# Milvus server address
+MILVUS_ADDRESS=your-zilliz-cloud-public-endpoint
+
+# Milvus authentication token
+# MILVUS_TOKEN=your-zilliz-cloud-api-key
+
+# Zilliz Cloud base URL (optional, default: https://api.cloud.zilliz.com)
+# ZILLIZ_BASE_URL=https://api.cloud.zilliz.com
+
+# =============================================================================
+# Code Splitter Configuration
+# =============================================================================
+
+# Code splitter type: ast, langchain
+SPLITTER_TYPE=ast

examples/basic-usage/index.ts

@@ -1,4 +1,5 @@
 import { CodeContext, MilvusVectorDatabase, MilvusRestfulVectorDatabase, AstCodeSplitter, LangChainCodeSplitter } from '@zilliz/code-context-core';
+import { envManager } from '@zilliz/code-context-core';
 import * as path from 'path';
 
 // Try to load .env file
@@ -17,9 +18,9 @@ async function main() {
         // Set to true to use RESTful API (for environments without gRPC support)
         // Set to false to use gRPC (default, more efficient)
         const useRestfulApi = false;
-        const milvusAddress = process.env.MILVUS_ADDRESS || 'localhost:19530';
-        const milvusToken = process.env.MILVUS_TOKEN;
-        const splitterType = process.env.SPLITTER_TYPE?.toLowerCase() || 'ast';
+        const milvusAddress = envManager.get('MILVUS_ADDRESS') || 'localhost:19530';
+        const milvusToken = envManager.get('MILVUS_TOKEN');
+        const splitterType = envManager.get('SPLITTER_TYPE')?.toLowerCase() || 'ast';
 
         console.log(`🔧 Using ${useRestfulApi ? 'RESTful API' : 'gRPC'} implementation`);
         console.log(`🔌 Connecting to Milvus at: ${milvusAddress}`);

packages/core/src/context.ts

@@ -14,6 +14,7 @@ import {
     VectorSearchResult
 } from './vectordb';
 import { SemanticSearchResult } from './types';
+import { envManager } from './utils/env-manager';
 import * as fs from 'fs';
 import * as path from 'path';
 import * as crypto from 'crypto';
@@ -101,9 +102,9 @@ export class CodeContext {
     constructor(config: CodeContextConfig = {}) {
         // Initialize services
         this.embedding = config.embedding || new OpenAIEmbedding({
-            apiKey: process.env.OPENAI_API_KEY || 'your-openai-api-key',
+            apiKey: envManager.get('OPENAI_API_KEY') || 'your-openai-api-key',
             model: 'text-embedding-3-small',
-            ...(process.env.OPENAI_BASE_URL && { baseURL: process.env.OPENAI_BASE_URL })
+            ...(envManager.get('OPENAI_BASE_URL') && { baseURL: envManager.get('OPENAI_BASE_URL') })
         });
 
         if (!config.vectorDatabase) {
@@ -461,7 +462,7 @@ export class CodeContext {
         codebasePath: string,
         onFileProcessed?: (filePath: string, fileIndex: number, totalFiles: number) => void
     ): Promise<{ processedFiles: number; totalChunks: number; status: 'completed' | 'limit_reached' }> {
-        const EMBEDDING_BATCH_SIZE = Math.max(1, parseInt(process.env.EMBEDDING_BATCH_SIZE || '100', 10));
+        const EMBEDDING_BATCH_SIZE = Math.max(1, parseInt(envManager.get('EMBEDDING_BATCH_SIZE') || '100', 10));
         const CHUNK_LIMIT = 450000;
         console.log(`🔧 Using EMBEDDING_BATCH_SIZE: ${EMBEDDING_BATCH_SIZE}`);
 

packages/core/src/index.ts

@@ -3,4 +3,5 @@ export * from './embedding';
 export * from './vectordb';
 export * from './types';
 export * from './context';
-export * from './sync/synchronizer';
+export * from './sync/synchronizer';
+export * from './utils';

packages/core/src/utils/env-manager.ts

@@ -0,0 +1,101 @@
+import * as fs from 'fs';
+import * as path from 'path';
+import * as os from 'os';
+
+export class EnvManager {
+    private envFilePath: string;
+
+    constructor() {
+        const homeDir = os.homedir();
+        this.envFilePath = path.join(homeDir, '.codecontext', '.env');
+    }
+
+    /**
+     * Get environment variable by name
+     * Priority: process.env > .env file > undefined
+     */
+    get(name: string): string | undefined {
+        // First try to get from process environment variables
+        if (process.env[name]) {
+            return process.env[name];
+        }
+
+        // If not found in process env, try to read from .env file
+        try {
+            if (fs.existsSync(this.envFilePath)) {
+                const content = fs.readFileSync(this.envFilePath, 'utf-8');
+                const lines = content.split('\n');
+
+                for (const line of lines) {
+                    const trimmedLine = line.trim();
+                    if (trimmedLine.startsWith(`${name}=`)) {
+                        return trimmedLine.substring(name.length + 1);
+                    }
+                }
+            }
+        } catch (error) {
+            // Ignore file read errors
+        }
+
+        return undefined;
+    }
+
+    /**
+     * Set environment variable to the .env file
+     */
+    set(name: string, value: string): void {
+        try {
+            // Ensure directory exists
+            const envDir = path.dirname(this.envFilePath);
+            if (!fs.existsSync(envDir)) {
+                fs.mkdirSync(envDir, { recursive: true });
+            }
+
+            let content = '';
+            let found = false;
+
+            // Read existing content if file exists
+            if (fs.existsSync(this.envFilePath)) {
+                content = fs.readFileSync(this.envFilePath, 'utf-8');
+
+                // Update existing variable
+                const lines = content.split('\n');
+                for (let i = 0; i < lines.length; i++) {
+                    if (lines[i].trim().startsWith(`${name}=`)) {
+                        // Replace the existing value
+                        lines[i] = `${name}=${value}`;
+                        found = true;
+                        console.log(`✅ Updated ${name} in ${this.envFilePath}`);
+                        break;
+                    }
+                }
+                content = lines.join('\n');
+            }
+
+            // If variable not found, append it
+            if (!found) {
+                if (content && !content.endsWith('\n')) {
+                    content += '\n';
+                }
+                content += `${name}=${value}\n`;
+                console.log(`✅ Added ${name} to ${this.envFilePath}`);
+            }
+
+            fs.writeFileSync(this.envFilePath, content, 'utf-8');
+
+        } catch (error) {
+            console.error(`Failed to write env file: ${error}`);
+            throw error;
+        }
+    }
+
+    /**
+     * Get the path to the .env file
+     */
+    getEnvFilePath(): string {
+        return this.envFilePath;
+    }
+}
+
+// Export a default instance for convenience
+export const envManager = new EnvManager(); 

packages/core/src/utils/index.ts

@@ -0,0 +1 @@
+export { EnvManager, envManager } from './env-manager'; 

packages/core/src/vectordb/abstract-milvus-vectordb.ts

@@ -0,0 +1,72 @@
+import { VectorDatabase, VectorDocument, SearchOptions, VectorSearchResult } from './types';
+
+// Common configuration interface for Milvus implementations
+export interface BaseMilvusConfig {
+    address?: string;
+    token?: string;
+    username?: string;
+    password?: string;
+}
+
+/**
+ * Abstract base class for Milvus vector database implementations
+ * Provides common initialization logic for address resolution and environment management
+ */
+export abstract class AbstractMilvusVectorDatabase implements VectorDatabase {
+    protected config: BaseMilvusConfig;
+    protected initializationPromise: Promise<void>;
+
+    constructor(config: BaseMilvusConfig) {
+        this.config = config;
+
+        // Start initialization asynchronously without waiting
+        this.initializationPromise = this.initialize();
+    }
+
+    private async initialize(): Promise<void> {
+        const resolvedAddress = await this.resolveAddress();
+        await this.initializeClient(resolvedAddress);
+    }
+
+    /**
+     * Resolve address from config or token
+     * Common logic for both gRPC and REST implementations
+     */
+    protected async resolveAddress(): Promise<string> {
+        let finalConfig = { ...this.config };
+
+        // If address is not provided, get it using token
+        if (!finalConfig.address && finalConfig.token) {
+            const { ClusterManager } = await import('./zilliz-utils');
+            finalConfig.address = await ClusterManager.getAddressFromToken(finalConfig.token);
+        }
+
+        if (!finalConfig.address) {
+            throw new Error('Address is required and could not be resolved from token');
+        }
+
+        return finalConfig.address;
+    }
+
+    /**
+     * Initialize the specific client implementation
+     * Must be implemented by subclasses
+     */
+    protected abstract initializeClient(address: string): Promise<void>;
+
+    /**
+     * Ensure initialization is complete before method execution
+     */
+    protected async ensureInitialized(): Promise<void> {
+        await this.initializationPromise;
+    }
+
+    // Abstract methods that must be implemented by subclasses
+    abstract createCollection(collectionName: string, dimension: number, description?: string): Promise<void>;
+    abstract dropCollection(collectionName: string): Promise<void>;
+    abstract hasCollection(collectionName: string): Promise<boolean>;
+    abstract insert(collectionName: string, documents: VectorDocument[]): Promise<void>;
+    abstract search(collectionName: string, queryVector: number[], options?: SearchOptions): Promise<VectorSearchResult[]>;
+    abstract delete(collectionName: string, ids: string[]): Promise<void>;
+    abstract query(collectionName: string, filter: string, outputFields: string[]): Promise<Record<string, any>[]>;
+}