@neural-technologies-indonesia/nl2sql 0.0.1 → 0.0.3

package/README.md

@@ -23,7 +23,7 @@ Core reasoning flow for text-to-SQL that can be exposed by any transport (Stream
 ## What lives here (and how it works)
 - Pure orchestrator in TypeScript (Bun-friendly) with clean interfaces for prompt composing, retrieval, LLM reasoning, query execution, and feedback policy.
 - Transport-agnostic: callers plug in their own input/output (UI, REST, MCP, etc.).
-- Defaults: prompt composer (`DefaultPromptComposer`), feedback policy (`SimpleFeedbackPolicy`), MCP executor (`MCPQueryExecutor`), MCP retriever (`MCPMilvusRetriever`), and OpenAI LLM wrapper (`OpenAILLM`).
+- Defaults: prompt composer (`DefaultPromptComposer`), feedback policy (`SimpleFeedbackPolicy`), MCP executor (`MCPQueryExecutor`), Milvus retriever (`MilvusRetriever`), and OpenAI LLM wrapper (`OpenAILLM`).
 - Flow: retrieve context → compose prompt → ask LLM for SQL candidates → execute via MCP `run_query` → accept/retry via feedback policy.
 
 ## Example wiring
@@ -64,13 +64,69 @@ console.log(result);
 
 ## Env
 - Copy `.env.example` and set `OPENAI_API_KEY` (and optionally `OPENAI_BASE_URL`, `OPENAI_MODEL`).
+- Milvus retriever defaults: `MILVUS_ADDRESS=milvus:19530`, `MILVUS_COLLECTION=rag_tables`, `MILVUS_VECTOR_FIELD=embedding`, `OPENAI_EMBEDDING_MODEL=text-embedding-3-small`.
+
+## Quick start (run-node)
+- Make sure `../mcp-server` is available (defaults to `../mcp-server/index.ts`).
+- Add `OPENAI_API_KEY` to `example/.env`.
+- Run: `bun run example/run-node.ts`.
+
+```ts
+import {
+  DefaultPromptComposer,
+  OpenAILLM,
+  ReasoningPipeline,
+  SimpleFeedbackPolicy,
+  createMcpUseAdapters,
+} from "@neural-technologies-indonesia/nl2sql";
+import dotenv from "dotenv";
+import path from "node:path";
+import { fileURLToPath } from "node:url";
+
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+dotenv.config({ path: path.join(__dirname, ".env") });
+
+const serverCwd =
+  process.env.MCP_SERVER_CWD ??
+  path.resolve(__dirname, "..", "..", "mcp-server");
+const serverEntry = path.join(serverCwd, "index.ts");
+
+const { executor, retriever, listIndexes, close } = await createMcpUseAdapters({
+  serverCwd,
+  command: "bun",
+  args: ["run", serverEntry],
+  env: Object.fromEntries(
+    Object.entries(process.env).filter(
+      (entry): entry is [string, string] => typeof entry[1] === "string"
+    )
+  ),
+});
+
+const pipeline = new ReasoningPipeline({
+  composer: new DefaultPromptComposer(),
+  retriever,
+  llm: new OpenAILLM({ model: process.env.OPENAI_MODEL ?? "gpt-4o-mini" }),
+  executor,
+  indexLookup: { listIndexes },
+  policy: new SimpleFeedbackPolicy({ durationBudgetMs: 5000 }),
+});
+
+const result = await pipeline.run({
+  question:
+    "Show monthly entity counts per tier since 2020. For each tier and month, return total entities and how many match metadata tier. Ignore groups with total ≤ 1000, ordered by most recent month and highest totals.",
+});
+
+console.log(JSON.stringify(result, null, 2));
+
+await close();
+```
 
 ## Examples
-- `example/bun-server.ts`: starts a Bun HTTP server exposing POST `/query` that runs the pipeline (OpenAI + MCP-backed executor/retriever). Run with `bun run example/bun-server.ts`.
+- `example/bun-server.ts`: starts a Bun HTTP server exposing POST `/query` that runs the pipeline (OpenAI + Milvus retriever + MCP-backed executor). Run with `bun run example/bun-server.ts`.
 - `example/streamlit_app.py`: minimal Streamlit UI that POSTs to the Bun server. Configure `API_URL` in `.streamlit/secrets.toml` or edit the default (`http://localhost:8787/query`).
 
 ## Use the sibling MCP server (`../mcp-server`) via `mcp-use`
-- That folder ships a FastMCP server exposing `run_query` (Postgres latency) and `milvus_retrieve` (RAG search). Start it with your env: `DATABASE_URL=... OPENAI_API_KEY=... bun run start` from `../mcp-server`.
+- That folder ships a FastMCP server exposing `run_query`/`run_query_explain`/`list_indexes` for Postgres. Start it with your env: `DATABASE_URL=... bun run start` from `../mcp-server`.
 - Install the client dependency here: `bun add mcp-use` (or `npm install mcp-use`).
 - If your MCP server lives somewhere else, point the helper there with env vars: `MCP_SERVER_CWD=/path/to/server`, optionally `MCP_SERVER_CMD=bun` and `MCP_SERVER_ARGS="run start"`. The example Bun server picks these up automatically.
 - Wire the pipeline to that server with the provided helper:
@@ -102,7 +158,7 @@ console.log(result);
   console.log(result);
   await close();
   ```
-- `createMcpUseAdapters` calls the MCP tools (no duplicated logic here) and returns the MCP-backed executor/retriever pair.
+- `createMcpUseAdapters` calls the MCP tools for execution (no duplicated logic here) and wires the retriever directly to Milvus using the provided/env config.
 
 ## Next steps
 - Swap in real retriever and LLM clients.

package/dist/adapters/milvusRetriever.d.ts

@@ -0,0 +1,42 @@
+import type { ContextChunk, PromptInput, Retriever } from "../types.js";
+export type MilvusSearchResult = {
+    id?: string;
+    score?: number;
+    rag_text?: string;
+    metadata?: Record<string, unknown>;
+};
+export type MilvusRetrieverOptions = {
+    topK?: number;
+    collection?: string;
+    vectorField?: string;
+    milvusAddress?: string;
+    timeoutMs?: number;
+    loadRetries?: number;
+    openAIApiKey?: string;
+    openAIModel?: string;
+    openAIBaseUrl?: string;
+};
+export declare class MilvusRetriever implements Retriever {
+    private readonly topK;
+    private readonly collection;
+    private readonly vectorField;
+    private readonly milvusAddress;
+    private readonly timeoutMs;
+    private readonly loadRetries;
+    private readonly embeddingModel;
+    private readonly openAIBaseUrl?;
+    private readonly openAIApiKey?;
+    private milvusClient;
+    private openAIClient;
+    private collectionLoaded;
+    private collectionLoading;
+    constructor(options?: MilvusRetrieverOptions);
+    retrieve(input: PromptInput): Promise<ContextChunk[]>;
+    private getMilvusClient;
+    private getOpenAIClient;
+    private ensureCollectionLoaded;
+    private embedQuery;
+    private searchMilvus;
+}
+export { MilvusRetriever as MCPMilvusRetriever };
+//# sourceMappingURL=milvusRetriever.d.ts.map

package/dist/adapters/milvusRetriever.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"milvusRetriever.d.ts","sourceRoot":"","sources":["../../src/adapters/milvusRetriever.ts"],"names":[],"mappings":"AAEA,OAAO,KAAK,EAAE,YAAY,EAAE,WAAW,EAAE,SAAS,EAAE,MAAM,aAAa,CAAC;AAExE,MAAM,MAAM,kBAAkB,GAAG;IAC/B,EAAE,CAAC,EAAE,MAAM,CAAC;IACZ,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB,QAAQ,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;CACpC,CAAC;AAEF,MAAM,MAAM,sBAAsB,GAAG;IACnC,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,UAAU,CAAC,EAAE,MAAM,CAAC;IACpB,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,aAAa,CAAC,EAAE,MAAM,CAAC;IACvB,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,YAAY,CAAC,EAAE,MAAM,CAAC;IACtB,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,aAAa,CAAC,EAAE,MAAM,CAAC;CACxB,CAAC;AAEF,qBAAa,eAAgB,YAAW,SAAS;IAC/C,OAAO,CAAC,QAAQ,CAAC,IAAI,CAAS;IAC9B,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAS;IACpC,OAAO,CAAC,QAAQ,CAAC,WAAW,CAAS;IACrC,OAAO,CAAC,QAAQ,CAAC,aAAa,CAAS;IACvC,OAAO,CAAC,QAAQ,CAAC,SAAS,CAAS;IACnC,OAAO,CAAC,QAAQ,CAAC,WAAW,CAAS;IACrC,OAAO,CAAC,QAAQ,CAAC,cAAc,CAAS;IACxC,OAAO,CAAC,QAAQ,CAAC,aAAa,CAAC,CAAS;IACxC,OAAO,CAAC,QAAQ,CAAC,YAAY,CAAC,CAAS;IAEvC,OAAO,CAAC,YAAY,CAA6B;IACjD,OAAO,CAAC,YAAY,CAAuB;IAC3C,OAAO,CAAC,gBAAgB,CAAS;IACjC,OAAO,CAAC,iBAAiB,CAA8B;gBAE3C,OAAO,CAAC,EAAE,sBAAsB;IAiBtC,QAAQ,CAAC,KAAK,EAAE,WAAW,GAAG,OAAO,CAAC,YAAY,EAAE,CAAC;IAY3D,OAAO,CAAC,eAAe;IAWvB,OAAO,CAAC,eAAe;YAeT,sBAAsB;YA2BtB,UAAU;YAaV,YAAY;CA0B3B;AA0HD,OAAO,EAAE,eAAe,IAAI,kBAAkB,EAAE,CAAC"}

package/dist/adapters/milvusRetriever.js

@@ -0,0 +1,230 @@
+import { DataType, MilvusClient } from "@zilliz/milvus2-sdk-node";
+import OpenAI from "openai";
+export class MilvusRetriever {
+    topK;
+    collection;
+    vectorField;
+    milvusAddress;
+    timeoutMs;
+    loadRetries;
+    embeddingModel;
+    openAIBaseUrl;
+    openAIApiKey;
+    milvusClient = null;
+    openAIClient = null;
+    collectionLoaded = false;
+    collectionLoading = null;
+    constructor(options) {
+        this.topK = options?.topK ?? 5;
+        this.collection = options?.collection ?? process.env.MILVUS_COLLECTION ?? "rag_tables";
+        this.vectorField = options?.vectorField ?? process.env.MILVUS_VECTOR_FIELD ?? "embedding";
+        this.milvusAddress =
+            options?.milvusAddress ?? process.env.MILVUS_ADDRESS ?? "localhost:19530";
+        this.timeoutMs =
+            options?.timeoutMs ?? Number(process.env.MILVUS_TIMEOUT_MS ?? "120000");
+        this.loadRetries =
+            options?.loadRetries ?? Number(process.env.MILVUS_LOAD_RETRIES ?? "6");
+        this.embeddingModel =
+            options?.openAIModel ?? process.env.OPENAI_EMBEDDING_MODEL ?? "text-embedding-3-small";
+        this.openAIBaseUrl =
+            normalizeBaseUrl(options?.openAIBaseUrl ?? process.env.OPENAI_BASE_URL);
+        this.openAIApiKey = options?.openAIApiKey ?? process.env.OPENAI_API_KEY;
+    }
+    async retrieve(input) {
+        const query = buildSearchQuery(input);
+        if (!query)
+            return [];
+        try {
+            const results = await this.searchMilvus(query);
+            return normalizeResults(results);
+        }
+        catch {
+            return [];
+        }
+    }
+    getMilvusClient() {
+        if (this.milvusClient)
+            return this.milvusClient;
+        this.milvusClient = new MilvusClient({
+            address: this.milvusAddress,
+            timeout: this.timeoutMs,
+        });
+        return this.milvusClient;
+    }
+    getOpenAIClient() {
+        if (this.openAIClient)
+            return this.openAIClient;
+        if (!this.openAIApiKey) {
+            throw new Error("OPENAI_API_KEY is required for Milvus retrieval.");
+        }
+        this.openAIClient = new OpenAI({
+            apiKey: this.openAIApiKey,
+            baseURL: this.openAIBaseUrl,
+        });
+        return this.openAIClient;
+    }
+    async ensureCollectionLoaded(client) {
+        if (this.collectionLoaded)
+            return;
+        if (!this.collectionLoading) {
+            this.collectionLoading = (async () => {
+                for (let i = 0; i < this.loadRetries; i++) {
+                    try {
+                        await client.loadCollection({
+                            collection_name: this.collection,
+                            timeout: this.timeoutMs,
+                        });
+                        this.collectionLoaded = true;
+                        return;
+                    }
+                    catch {
+                        await sleep(Math.min(1000 * 2 ** i, 8000));
+                    }
+                }
+                throw new Error(`Milvus load failed after ${this.loadRetries} attempts for collection '${this.collection}'.`);
+            })();
+        }
+        await this.collectionLoading;
+    }
+    async embedQuery(query) {
+        const client = this.getOpenAIClient();
+        const response = await client.embeddings.create({
+            model: this.embeddingModel,
+            input: query,
+        });
+        const embedding = response.data?.[0]?.embedding;
+        if (!embedding) {
+            throw new Error("Failed to generate embedding for query.");
+        }
+        return embedding;
+    }
+    async searchMilvus(text) {
+        const client = this.getMilvusClient();
+        await this.ensureCollectionLoaded(client);
+        const embedding = await this.embedQuery(text);
+        const searchResponse = await client.search({
+            collection_name: this.collection,
+            vectors: [embedding],
+            vector_type: DataType.FloatVector,
+            search_params: {
+                anns_field: this.vectorField,
+                topk: String(this.topK),
+                metric_type: "L2",
+                params: JSON.stringify({ nprobe: 16 }),
+            },
+            output_fields: ["id", "rag_text", "metadata"],
+            timeout: this.timeoutMs,
+        });
+        const results = Array.isArray(searchResponse.results)
+            ? searchResponse.results
+            : [];
+        return results;
+    }
+}
+// Combine user input fields into a single search string.
+const buildSearchQuery = (input) => {
+    const queryParts = [input.question.trim()];
+    if (input.userContext?.trim()) {
+        queryParts.push(`Context: ${input.userContext.trim()}`);
+    }
+    if (input.constraints?.length) {
+        queryParts.push(`Constraints: ${input.constraints.filter(Boolean).join(" | ")}`);
+    }
+    return queryParts.filter(Boolean).join("\n").trim();
+};
+// Accepts any payload shape and converts it into ContextChunks.
+const normalizeResults = (results) => {
+    if (typeof results === "string") {
+        return [{ source: "milvus", snippet: results }];
+    }
+    if (!Array.isArray(results)) {
+        return [{ source: "milvus", snippet: contentToSnippet(results) }];
+    }
+    return results.map((item, index) => {
+        const snippetParts = [item.rag_text ?? "(no rag_text)"];
+        if (item.metadata) {
+            snippetParts.push(buildMetadataSnippet(item.metadata));
+        }
+        return {
+            source: item.id ? String(item.id) : `milvus-${index + 1}`,
+            snippet: snippetParts.join("\n"),
+            score: item.score,
+        };
+    });
+};
+const contentToSnippet = (content) => {
+    if (typeof content === "string")
+        return content;
+    if (Array.isArray(content))
+        return content.map((item) => JSON.stringify(item)).join("\n");
+    if (content && typeof content === "object")
+        return JSON.stringify(content);
+    return String(content ?? "");
+};
+const normalizeTableReference = (schema, table) => {
+    if (!table.includes(".")) {
+        return { schema, table };
+    }
+    const parts = table.split(".").filter(Boolean);
+    if (!parts.length) {
+        return { schema, table };
+    }
+    if (parts.length >= 2 && !schema) {
+        return {
+            schema: parts[parts.length - 2],
+            table: parts[parts.length - 1],
+        };
+    }
+    return { schema, table: parts[parts.length - 1] };
+};
+// Formats metadata for the prompt while avoiding fully qualified names that can confuse the model.
+const buildMetadataSnippet = (metadata) => {
+    const parts = [];
+    const schema = typeof metadata.schema === "string" ? metadata.schema : "";
+    const tableRaw = metadata.table ?? metadata.table_name ?? metadata.name;
+    const table = typeof tableRaw === "string" ? tableRaw : "";
+    const normalized = normalizeTableReference(schema, table);
+    const qualified = [normalized.schema, normalized.table].filter(Boolean).join(".");
+    if (qualified)
+        parts.push(`table=${qualified}`);
+    const tableType = typeof metadata.table_type === "string"
+        ? metadata.table_type
+        : typeof metadata.tableType === "string"
+            ? metadata.tableType
+            : "";
+    if (tableType)
+        parts.push(`type=${tableType}`);
+    const primaryKeys = Array.isArray(metadata.primary_keys) ? metadata.primary_keys : [];
+    if (primaryKeys.length)
+        parts.push(`primary_keys=${primaryKeys.map(String).join(",")}`);
+    const columns = Array.isArray(metadata.columns)
+        ? metadata.columns
+            .map((col) => col && typeof col === "object" && "name" in col ? String(col.name) : null)
+            .filter((name) => Boolean(name))
+        : [];
+    if (columns.length)
+        parts.push(`columns=${columns.join(", ")}`);
+    const description = typeof metadata.description === "string" ? metadata.description : "";
+    if (description)
+        parts.push(`desc=${description}`);
+    if (!parts.length) {
+        const sanitized = { ...metadata };
+        delete sanitized.fully_qualified_name;
+        delete sanitized.fullyQualifiedName;
+        return `metadata=${JSON.stringify(sanitized)}`;
+    }
+    return parts.join("; ");
+};
+const normalizeBaseUrl = (rawBaseUrl) => {
+    if (!rawBaseUrl)
+        return undefined;
+    const trimmed = rawBaseUrl.trim();
+    if (!trimmed)
+        return undefined;
+    if (!trimmed.startsWith("http://") && !trimmed.startsWith("https://")) {
+        return `https://${trimmed}`;
+    }
+    return trimmed;
+};
+const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
+export { MilvusRetriever as MCPMilvusRetriever };

package/dist/clients/mcpUseClient.d.ts

@@ -1,18 +1,20 @@
-import { MCPClient } from "mcp-use";
 import { MCPQueryExecutor } from "../adapters/mcpExecutor.js";
-import { MCPMilvusRetriever } from "../adapters/mcpRetriever.js";
-import type { MCPMilvusResult } from "../adapters/mcpRetriever.js";
+import { MilvusRetriever } from "../adapters/milvusRetriever.js";
+import type { MilvusRetrieverOptions } from "../adapters/milvusRetriever.js";
 export type MCPUseClientOptions = {
     serverName?: string;
     serverCwd?: string;
     command?: string;
     args?: string[];
     env?: Record<string, string>;
+    transportType?: "stdio" | "httpStream";
+    httpStreamUrl?: string;
+    milvus?: MilvusRetrieverOptions;
 };
 export declare const createMcpUseAdapters: (options?: MCPUseClientOptions) => Promise<{
     executor: MCPQueryExecutor;
-    retriever: MCPMilvusRetriever;
-    callTool: (name: "run_query" | "run_query_explain" | "list_indexes" | "milvus_retrieve", args: Record<string, unknown>) => Promise<string | MCPMilvusResult[] | Record<string, unknown>>;
+    retriever: MilvusRetriever;
+    callTool: (name: "run_query" | "run_query_explain" | "list_indexes", args: Record<string, unknown>) => Promise<string | Record<string, unknown>>;
     runQueryExplain: (sql: string) => Promise<{
         durationMs: number;
         explainLines: string[] | undefined;
@@ -33,7 +35,7 @@ export declare const createMcpUseAdapters: (options?: MCPUseClientOptions) => Pr
         tablename?: string;
     }) => Promise<string>;
     close: () => Promise<void>;
-    client: MCPClient;
-    session: import("mcp-use").MCPSession;
+    client: unknown;
+    session: unknown;
 }>;
 //# sourceMappingURL=mcpUseClient.d.ts.map

package/dist/clients/mcpUseClient.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"mcpUseClient.d.ts","sourceRoot":"","sources":["../../src/clients/mcpUseClient.ts"],"names":[],"mappings":"AAEA,OAAO,EAAE,SAAS,EAAE,MAAM,SAAS,CAAC;AACpC,OAAO,EAAE,gBAAgB,EAAE,MAAM,4BAA4B,CAAC;AAC9D,OAAO,EAAE,kBAAkB,EAAE,MAAM,6BAA6B,CAAC;AACjE,OAAO,KAAK,EAAE,eAAe,EAAE,MAAM,6BAA6B,CAAC;AAGnE,MAAM,MAAM,mBAAmB,GAAG;IAChC,UAAU,CAAC,EAAE,MAAM,CAAC;IACpB,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,OAAO,CAAC,EAAE,MAAM,CAAC;IACjB,IAAI,CAAC,EAAE,MAAM,EAAE,CAAC;IAChB,GAAG,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;CAC9B,CAAC;AAwEF,eAAO,MAAM,oBAAoB,GAAU,UAAU,mBAAmB;;;qBA2C9D,WAAW,GAAG,mBAAmB,GAAG,cAAc,GAAG,iBAAiB,QACtE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,KAC5B,OAAO,CAAC,MAAM,GAAG,eAAe,EAAE,GAAG,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;2BAuB5B,MAAM;;;;;;;;;;;;;;;4BAoDL;QACnC,UAAU,CAAC,EAAE,MAAM,CAAC;QACpB,SAAS,CAAC,EAAE,MAAM,CAAC;KACpB;;;;EA2BF,CAAC"}
+{"version":3,"file":"mcpUseClient.d.ts","sourceRoot":"","sources":["../../src/clients/mcpUseClient.ts"],"names":[],"mappings":"AAKA,OAAO,EAAE,gBAAgB,EAAE,MAAM,4BAA4B,CAAC;AAC9D,OAAO,EAAE,eAAe,EAAE,MAAM,gCAAgC,CAAC;AACjE,OAAO,KAAK,EAAE,sBAAsB,EAAE,MAAM,gCAAgC,CAAC;AAG7E,MAAM,MAAM,mBAAmB,GAAG;IAChC,UAAU,CAAC,EAAE,MAAM,CAAC;IACpB,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,OAAO,CAAC,EAAE,MAAM,CAAC;IACjB,IAAI,CAAC,EAAE,MAAM,EAAE,CAAC;IAChB,GAAG,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;IAC7B,aAAa,CAAC,EAAE,OAAO,GAAG,YAAY,CAAC;IACvC,aAAa,CAAC,EAAE,MAAM,CAAC;IACvB,MAAM,CAAC,EAAE,sBAAsB,CAAC;CACjC,CAAC;AA8FF,eAAO,MAAM,oBAAoB,GAAU,UAAU,mBAAmB;;;qBAa5D,WAAW,GAAG,mBAAmB,GAAG,cAAc,QAClD,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,KAC1B,OAAO,CAAC,MAAM,GAAG,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;2BAKV,MAAM;;;;;;;;;;;;;;;4BA4DL;QACnC,UAAU,CAAC,EAAE,MAAM,CAAC;QACpB,SAAS,CAAC,EAAE,MAAM,CAAC;KACpB;iBAnEY,OAAO,CAAC,IAAI,CAAC;;;EAuK7B,CAAC"}

package/dist/clients/mcpUseClient.js

@@ -1,8 +1,10 @@
 import fs from "node:fs";
 import path from "node:path";
 import { MCPClient } from "mcp-use";
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
 import { MCPQueryExecutor } from "../adapters/mcpExecutor.js";
-import { MCPMilvusRetriever } from "../adapters/mcpRetriever.js";
+import { MilvusRetriever } from "../adapters/milvusRetriever.js";
 // Accepts a whitespace-delimited env var and turns it into an args array.
 const parseArgs = (value) => {
     if (!value)
@@ -12,6 +14,24 @@ const parseArgs = (value) => {
         .map((part) => part.trim())
         .filter(Boolean);
 };
+const parseNumber = (value) => {
+    if (value === undefined)
+        return undefined;
+    const parsed = Number(value);
+    return Number.isNaN(parsed) ? undefined : parsed;
+};
+const resolveMilvusOptions = (options) => ({
+    ...options,
+    topK: options?.topK ?? parseNumber(process.env.MILVUS_TOP_K) ?? undefined,
+    collection: options?.collection ?? process.env.MILVUS_COLLECTION,
+    vectorField: options?.vectorField ?? process.env.MILVUS_VECTOR_FIELD,
+    milvusAddress: options?.milvusAddress ?? process.env.MILVUS_ADDRESS,
+    timeoutMs: options?.timeoutMs ?? parseNumber(process.env.MILVUS_TIMEOUT_MS),
+    loadRetries: options?.loadRetries ?? parseNumber(process.env.MILVUS_LOAD_RETRIES),
+    openAIApiKey: options?.openAIApiKey ?? process.env.OPENAI_API_KEY,
+    openAIBaseUrl: options?.openAIBaseUrl ?? process.env.OPENAI_BASE_URL,
+    openAIModel: options?.openAIModel ?? process.env.OPENAI_EMBEDDING_MODEL,
+});
 // Convert relative file-like args to absolute paths so Bun can locate the entrypoint
 // even if the MCP client spawns from a different working directory.
 const absolutizeArgs = (args, baseDir) => {
@@ -69,6 +89,96 @@ const parseDurationMs = (text) => {
 // Starts mcp-use against the sibling MCP server and returns MCP-backed adapters.
 export const createMcpUseAdapters = async (options) => {
     const serverName = options?.serverName ?? "query-runner";
+    const transportType = options?.transportType ??
+        (process.env.MCP_SERVER_TRANSPORT === "http-stream"
+            ? "httpStream"
+            : process.env.MCP_SERVER_TRANSPORT);
+    const httpStreamUrl = options?.httpStreamUrl ?? process.env.MCP_SERVER_URL ?? undefined;
+    const retriever = new MilvusRetriever(resolveMilvusOptions(options?.milvus));
+    const buildAdapters = (callTool, close, client, session) => {
+        const runQueryExplain = async (sql) => {
+            const content = await callTool("run_query_explain", { raw_query: sql });
+            const text = contentToText(content);
+            const durationMs = parseDurationMs(text) ?? Number.NaN;
+            const jsonStart = text.search(/JSON:\s*/i);
+            const explainMarker = text.search(/EXPLAIN ANALYZE:\s*/i);
+            let parsed = null;
+            if (jsonStart !== -1) {
+                const jsonBlock = text.slice(jsonStart).replace(/^JSON:\s*/i, "");
+                const jsonText = explainMarker !== -1
+                    ? jsonBlock.slice(0, explainMarker - jsonStart).trim()
+                    : jsonBlock.trim();
+                try {
+                    parsed = JSON.parse(jsonText);
+                }
+                catch {
+                    parsed = null;
+                }
+            }
+            if (parsed) {
+                const explainLines = Array.isArray(parsed.explainLines)
+                    ? parsed.explainLines.map((line) => String(line))
+                    : undefined;
+                return {
+                    durationMs: typeof parsed.durationMs === "number"
+                        ? parsed.durationMs
+                        : durationMs,
+                    explainLines,
+                    planningMs: typeof parsed.planningMs === "number"
+                        ? parsed.planningMs
+                        : undefined,
+                    executionMs: typeof parsed.executionMs === "number"
+                        ? parsed.executionMs
+                        : undefined,
+                    raw: text,
+                };
+            }
+            const explainSplit = text.split(/EXPLAIN ANALYZE:\s*/i)[1] ??
+                text.split(/EXPLAIN:\s*/i)[1] ??
+                "";
+            return { durationMs, explain: explainSplit.trim(), raw: text };
+        };
+        const executor = new MCPQueryExecutor({
+            runQuery: async (sql) => {
+                const content = await callTool("run_query", { raw_query: sql });
+                const text = contentToText(content);
+                const durationMs = parseDurationMs(text) ?? Number.NaN; // Keep numeric for typing even if parse fails.
+                return { durationMs };
+            },
+            runQueryExplain: async (sql) => runQueryExplain(sql),
+        });
+        const listIndexes = async (options) => {
+            const content = await callTool("list_indexes", {
+                schemaname: options?.schemaname,
+                tablename: options?.tablename,
+            });
+            return contentToText(content);
+        };
+        return {
+            executor,
+            retriever,
+            callTool,
+            runQueryExplain,
+            listIndexes,
+            close,
+            client,
+            session,
+        };
+    };
+    if (transportType === "httpStream" || httpStreamUrl) {
+        const url = httpStreamUrl ?? "http://localhost:8080/mcp";
+        const transport = new StreamableHTTPClientTransport(new URL(url));
+        const client = new Client({ name: serverName, version: "1.0.0" });
+        await client.connect(transport);
+        const callTool = async (name, args) => {
+            const result = await client.callTool({ name, arguments: args });
+            return toolResultPayload(result);
+        };
+        const close = async () => {
+            await transport.close();
+        };
+        return buildAdapters(callTool, close, client, null);
+    }
     console.log("mcp server cwd", process.env.MCP_SERVER_CWD);
     const serverCwd = options?.serverCwd ??
         process.env.MCP_SERVER_CWD ??
@@ -100,83 +210,11 @@ export const createMcpUseAdapters = async (options) => {
         const result = await session.callTool(name, args);
         return toolResultPayload(result);
     };
-    const executor = new MCPQueryExecutor({
-        runQuery: async (sql) => {
-            const content = await callTool("run_query", { raw_query: sql });
-            const text = contentToText(content);
-            const durationMs = parseDurationMs(text) ?? Number.NaN; // Keep numeric for typing even if parse fails.
-            return { durationMs };
-        },
-        runQueryExplain: async (sql) => runQueryExplain(sql),
-    });
-    const retriever = new MCPMilvusRetriever({
-        milvusRetrieve: async (query, topK) => callTool("milvus_retrieve", { query, top_k: topK ?? 5 }),
-    });
-    const runQueryExplain = async (sql) => {
-        const content = await callTool("run_query_explain", { raw_query: sql });
-        const text = contentToText(content);
-        const durationMs = parseDurationMs(text) ?? Number.NaN;
-        const jsonStart = text.search(/JSON:\s*/i);
-        const explainMarker = text.search(/EXPLAIN ANALYZE:\s*/i);
-        let parsed = null;
-        if (jsonStart !== -1) {
-            const jsonBlock = text
-                .slice(jsonStart)
-                .replace(/^JSON:\s*/i, "");
-            const jsonText = explainMarker !== -1
-                ? jsonBlock.slice(0, explainMarker - jsonStart).trim()
-                : jsonBlock.trim();
-            try {
-                parsed = JSON.parse(jsonText);
-            }
-            catch {
-                parsed = null;
-            }
-        }
-        if (parsed) {
-            const explainLines = Array.isArray(parsed.explainLines)
-                ? parsed.explainLines.map((line) => String(line))
-                : undefined;
-            return {
-                durationMs: typeof parsed.durationMs === "number"
-                    ? parsed.durationMs
-                    : durationMs,
-                explainLines,
-                planningMs: typeof parsed.planningMs === "number"
-                    ? parsed.planningMs
-                    : undefined,
-                executionMs: typeof parsed.executionMs === "number"
-                    ? parsed.executionMs
-                    : undefined,
-                raw: text,
-            };
-        }
-        const explainSplit = text.split(/EXPLAIN ANALYZE:\s*/i)[1] ??
-            text.split(/EXPLAIN:\s*/i)[1] ??
-            "";
-        return { durationMs, explain: explainSplit.trim(), raw: text };
-    };
-    const listIndexes = async (options) => {
-        const content = await callTool("list_indexes", {
-            schemaname: options?.schemaname,
-            tablename: options?.tablename,
-        });
-        return contentToText(content);
-    };
     const close = async () => {
         if (typeof client
             .closeAllSessions === "function") {
             await client.closeAllSessions();
         }
     };
-    return {
-        executor,
-        retriever,
-        callTool,
-        runQueryExplain,
-        listIndexes,
-        close,
-        client,
-        session,
-    };
+    return buildAdapters(callTool, close, client, session);
 };

package/dist/index.d.ts

@@ -2,7 +2,7 @@ export * from "./types.js";
 export * from "./pipeline.js";
 export * from "./defaults/promptComposer.js";
 export * from "./adapters/mcpExecutor.js";
-export * from "./adapters/mcpRetriever.js";
+export * from "./adapters/milvusRetriever.js";
 export * from "./adapters/openaiLLM.js";
 export * from "./clients/mcpUseClient.js";
 //# sourceMappingURL=index.d.ts.map

package/dist/index.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,cAAc,YAAY,CAAC;AAC3B,cAAc,eAAe,CAAC;AAC9B,cAAc,8BAA8B,CAAC;AAC7C,cAAc,2BAA2B,CAAC;AAC1C,cAAc,4BAA4B,CAAC;AAC3C,cAAc,yBAAyB,CAAC;AACxC,cAAc,2BAA2B,CAAC"}
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,cAAc,YAAY,CAAC;AAC3B,cAAc,eAAe,CAAC;AAC9B,cAAc,8BAA8B,CAAC;AAC7C,cAAc,2BAA2B,CAAC;AAC1C,cAAc,+BAA+B,CAAC;AAC9C,cAAc,yBAAyB,CAAC;AACxC,cAAc,2BAA2B,CAAC"}

package/dist/index.js

@@ -2,6 +2,6 @@ export * from "./types.js";
 export * from "./pipeline.js";
 export * from "./defaults/promptComposer.js";
 export * from "./adapters/mcpExecutor.js";
-export * from "./adapters/mcpRetriever.js";
+export * from "./adapters/milvusRetriever.js";
 export * from "./adapters/openaiLLM.js";
 export * from "./clients/mcpUseClient.js";

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@neural-technologies-indonesia/nl2sql",
-  "version": "0.0.1",
+  "version": "0.0.3",
   "description": "Core reasoning pipeline for text-to-SQL.",
   "private": false,
   "type": "module",
@@ -24,7 +24,9 @@
     "test": "echo \"no tests yet\""
   },
   "dependencies": {
+    "@zilliz/milvus2-sdk-node": "2.3.5",
     "langchain": "^1.2.1",
+    "@modelcontextprotocol/sdk": "^1.25.1",
     "mcp-use": "*",
     "openai": "^6.14.0"
   },