diff --git a/ts/entrypoints/chunker.mjs b/ts/entrypoints/chunker.mjs new file mode 100644 index 00000000..5e377208 --- /dev/null +++ b/ts/entrypoints/chunker.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/chunking/service.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/doc-embeddings-query.mjs b/ts/entrypoints/doc-embeddings-query.mjs new file mode 100644 index 00000000..02424210 --- /dev/null +++ b/ts/entrypoints/doc-embeddings-query.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/query/embeddings/qdrant-doc-service.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/document-rag.mjs b/ts/entrypoints/document-rag.mjs new file mode 100644 index 00000000..fcb6b4b1 --- /dev/null +++ b/ts/entrypoints/document-rag.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/retrieval/document-rag-service.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/extractor.mjs b/ts/entrypoints/extractor.mjs new file mode 100644 index 00000000..a58f6aa1 --- /dev/null +++ b/ts/entrypoints/extractor.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/extract/knowledge-extract.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/graph-embeddings-query.mjs b/ts/entrypoints/graph-embeddings-query.mjs new file mode 100644 index 00000000..449b0002 --- /dev/null +++ b/ts/entrypoints/graph-embeddings-query.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/query/embeddings/qdrant-graph-service.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/graph-embeddings-store.mjs b/ts/entrypoints/graph-embeddings-store.mjs new file mode 100644 index 00000000..e883b49a --- /dev/null +++ b/ts/entrypoints/graph-embeddings-store.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/storage/embeddings/graph-embeddings-service.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/graph-rag.mjs b/ts/entrypoints/graph-rag.mjs new file mode 100644 index 00000000..62c1b755 --- /dev/null +++ b/ts/entrypoints/graph-rag.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/retrieval/graph-rag-service.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/pdf-decoder.mjs b/ts/entrypoints/pdf-decoder.mjs new file mode 100644 index 00000000..a00ed8de --- /dev/null +++ b/ts/entrypoints/pdf-decoder.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/decoding/pdf-decoder.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/text-completion-azure-openai.mjs b/ts/entrypoints/text-completion-azure-openai.mjs new file mode 100644 index 00000000..1ad91468 --- /dev/null +++ b/ts/entrypoints/text-completion-azure-openai.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/model/text-completion/azure-openai.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/text-completion-mistral.mjs b/ts/entrypoints/text-completion-mistral.mjs new file mode 100644 index 00000000..5073e926 --- /dev/null +++ b/ts/entrypoints/text-completion-mistral.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/model/text-completion/mistral.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/text-completion-ollama.mjs b/ts/entrypoints/text-completion-ollama.mjs new file mode 100644 index 00000000..31f4384a --- /dev/null +++ b/ts/entrypoints/text-completion-ollama.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/model/text-completion/ollama.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/text-completion-openai-compatible.mjs b/ts/entrypoints/text-completion-openai-compatible.mjs new file mode 100644 index 00000000..d02e0d85 --- /dev/null +++ b/ts/entrypoints/text-completion-openai-compatible.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/model/text-completion/openai-compatible.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/triples-query.mjs b/ts/entrypoints/triples-query.mjs new file mode 100644 index 00000000..ac3f7c0b --- /dev/null +++ b/ts/entrypoints/triples-query.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/query/triples/falkordb-service.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/entrypoints/triples-store.mjs b/ts/entrypoints/triples-store.mjs new file mode 100644 index 00000000..a0f76cad --- /dev/null +++ b/ts/entrypoints/triples-store.mjs @@ -0,0 +1,6 @@ +import("../packages/flow/dist/storage/triples/falkordb-service.js") + .then((m) => m.run()) + .catch((err) => { + console.error(err); + process.exit(1); + }); diff --git a/ts/package.json b/ts/package.json index b842aabf..c6c70a55 100644 --- a/ts/package.json +++ b/ts/package.json @@ -30,9 +30,13 @@ "doc-embeddings-query": "tsx scripts/run-doc-embeddings-query.ts", "graph-rag": "tsx scripts/run-graph-rag.ts", "document-rag": "tsx scripts/run-document-rag.ts", - "create-test-pdf": "tsx scripts/create-test-pdf.ts" + "create-test-pdf": "tsx scripts/create-test-pdf.ts", + "llm:azure-openai": "tsx scripts/run-llm-azure-openai.ts", + "llm:openai-compat": "tsx scripts/run-llm-openai-compatible.ts", + "llm:mistral": "tsx scripts/run-llm-mistral.ts" }, "devDependencies": { + "falkordb": "^5.0.0", "nats": "^2.29.0", "pdf-lib": "^1.17.1", "tsx": "^4.21.0", diff --git a/ts/packages/base/src/messaging/consumer.ts b/ts/packages/base/src/messaging/consumer.ts index f8bbe9d4..960c7780 100644 --- a/ts/packages/base/src/messaging/consumer.ts +++ b/ts/packages/base/src/messaging/consumer.ts @@ -73,8 +73,9 @@ export class Consumer { private async consumeLoop(flow: FlowContext): Promise { while (this.running) { + let msg: Message | null = null; try { - const msg = await this.backend!.receive(2000); + msg = await this.backend!.receive(2000); if (!msg) continue; await this.handleWithRetry(msg, flow); @@ -82,7 +83,13 @@ export class Consumer { } catch (err) { if (!this.running) break; console.error("[Consumer] Error in consume loop:", err); - // Brief pause before retry + if (msg) { + try { + await this.backend!.negativeAcknowledge(msg); + } catch (nakErr) { + console.error("[Consumer] Failed to nak message:", nakErr); + } + } await sleep(1000); } } diff --git a/ts/packages/flow/package.json b/ts/packages/flow/package.json index dd11bd94..88244405 100644 --- a/ts/packages/flow/package.json +++ b/ts/packages/flow/package.json @@ -18,6 +18,7 @@ "falkordb": "^5.0.0", "fastify": "^5.2.0", "ollama": "^0.6.3", + "@mistralai/mistralai": "^1.0.0", "openai": "^4.85.0", "pdfjs-dist": "^5.6.205" }, diff --git a/ts/packages/flow/src/extract/knowledge-extract.ts b/ts/packages/flow/src/extract/knowledge-extract.ts index 65c1ecf2..35bf46f3 100644 --- a/ts/packages/flow/src/extract/knowledge-extract.ts +++ b/ts/packages/flow/src/extract/knowledge-extract.ts @@ -93,64 +93,71 @@ export class KnowledgeExtractService extends FlowProcessor { // --- Extract relationships --- try { - const relPrompt = await promptClient.request({ - name: "extract-relationships", - variables: { text }, - }); + const relPrompt = await promptClient.request( + { name: "extract-relationships", variables: { text } }, + { timeoutMs: 10_000 }, + ); if (!relPrompt.error) { - const relCompletion = await llmClient.request({ - system: relPrompt.system, - prompt: relPrompt.prompt, - }); + let relationships: ExtractedRelationship[] | null = null; + for (let attempt = 0; attempt < 3; attempt++) { + const relCompletion = await llmClient.request( + { system: relPrompt.system, prompt: relPrompt.prompt }, + { timeoutMs: 120_000 }, + ); - if (!relCompletion.error && relCompletion.response) { - const relationships = parseJsonResponse(relCompletion.response); - - if (relationships) { - for (const rel of relationships) { - if (!rel.subject || !rel.predicate || !rel.object) continue; - - const subjectIri = toEntityIri(rel.subject); - const predicateIri = toEntityIri(rel.predicate); - const objectIri = toEntityIri(rel.object); - - // Main relationship triple - allTriples.push({ s: subjectIri, p: predicateIri, o: objectIri }); - - // rdfs:label triples for each entity - allTriples.push({ - s: subjectIri, - p: iriTerm(RDFS_LABEL), - o: literalTerm(rel.subject), - }); - allTriples.push({ - s: predicateIri, - p: iriTerm(RDFS_LABEL), - o: literalTerm(rel.predicate), - }); - allTriples.push({ - s: objectIri, - p: iriTerm(RDFS_LABEL), - o: literalTerm(rel.object), - }); - - // Entity contexts for subject and object - allEntityContexts.push({ - entity: subjectIri, - context: text, - chunkId: msg.documentId, - }); - allEntityContexts.push({ - entity: objectIri, - context: text, - chunkId: msg.documentId, - }); - } - - console.log(`[KnowledgeExtract] Extracted ${relationships.length} relationships`); + if (!relCompletion.error && relCompletion.response) { + relationships = parseJsonResponse(relCompletion.response); + if (relationships) break; + console.warn(`[KnowledgeExtract] Relationship parse failed, attempt ${attempt + 1}/3`); + } else { + break; // LLM error, don't retry } } + + if (relationships) { + for (const rel of relationships) { + if (!rel.subject || !rel.predicate || !rel.object) continue; + + const subjectIri = toEntityIri(rel.subject); + const predicateIri = toEntityIri(rel.predicate); + const objectIri = toEntityIri(rel.object); + + // Main relationship triple + allTriples.push({ s: subjectIri, p: predicateIri, o: objectIri }); + + // rdfs:label triples for each entity + allTriples.push({ + s: subjectIri, + p: iriTerm(RDFS_LABEL), + o: literalTerm(rel.subject), + }); + allTriples.push({ + s: predicateIri, + p: iriTerm(RDFS_LABEL), + o: literalTerm(rel.predicate), + }); + allTriples.push({ + s: objectIri, + p: iriTerm(RDFS_LABEL), + o: literalTerm(rel.object), + }); + + // Entity contexts for subject and object + allEntityContexts.push({ + entity: subjectIri, + context: text, + chunkId: msg.documentId, + }); + allEntityContexts.push({ + entity: objectIri, + context: text, + chunkId: msg.documentId, + }); + } + + console.log(`[KnowledgeExtract] Extracted ${relationships.length} relationships`); + } } } catch (err) { console.error("[KnowledgeExtract] Relationship extraction failed:", err); @@ -158,51 +165,58 @@ export class KnowledgeExtractService extends FlowProcessor { // --- Extract definitions --- try { - const defPrompt = await promptClient.request({ - name: "extract-definitions", - variables: { text }, - }); + const defPrompt = await promptClient.request( + { name: "extract-definitions", variables: { text } }, + { timeoutMs: 10_000 }, + ); if (!defPrompt.error) { - const defCompletion = await llmClient.request({ - system: defPrompt.system, - prompt: defPrompt.prompt, - }); + let definitions: ExtractedDefinition[] | null = null; + for (let attempt = 0; attempt < 3; attempt++) { + const defCompletion = await llmClient.request( + { system: defPrompt.system, prompt: defPrompt.prompt }, + { timeoutMs: 120_000 }, + ); - if (!defCompletion.error && defCompletion.response) { - const definitions = parseJsonResponse(defCompletion.response); - - if (definitions) { - for (const def of definitions) { - if (!def.entity || !def.definition) continue; - - const entityIri = toEntityIri(def.entity); - - // Definition triple - allTriples.push({ - s: entityIri, - p: iriTerm(SKOS_DEFINITION), - o: literalTerm(def.definition), - }); - - // Label triple - allTriples.push({ - s: entityIri, - p: iriTerm(RDFS_LABEL), - o: literalTerm(def.entity), - }); - - // Entity context - allEntityContexts.push({ - entity: entityIri, - context: text, - chunkId: msg.documentId, - }); - } - - console.log(`[KnowledgeExtract] Extracted ${definitions.length} definitions`); + if (!defCompletion.error && defCompletion.response) { + definitions = parseJsonResponse(defCompletion.response); + if (definitions) break; + console.warn(`[KnowledgeExtract] Definition parse failed, attempt ${attempt + 1}/3`); + } else { + break; // LLM error, don't retry } } + + if (definitions) { + for (const def of definitions) { + if (!def.entity || !def.definition) continue; + + const entityIri = toEntityIri(def.entity); + + // Definition triple + allTriples.push({ + s: entityIri, + p: iriTerm(SKOS_DEFINITION), + o: literalTerm(def.definition), + }); + + // Label triple + allTriples.push({ + s: entityIri, + p: iriTerm(RDFS_LABEL), + o: literalTerm(def.entity), + }); + + // Entity context + allEntityContexts.push({ + entity: entityIri, + context: text, + chunkId: msg.documentId, + }); + } + + console.log(`[KnowledgeExtract] Extracted ${definitions.length} definitions`); + } } } catch (err) { console.error("[KnowledgeExtract] Definition extraction failed:", err); @@ -245,23 +259,49 @@ function literalTerm(value: string): Term { /** * Parse JSON from LLM output, handling markdown code fences and malformed output. + * Uses progressive fallback: direct parse, array extraction, truncated array repair, single object wrap. */ function parseJsonResponse(raw: string): T | null { - try { - // Strip markdown code fences - let cleaned = raw.trim(); - - // Remove ```json ... ``` or ``` ... ``` - const fenceMatch = cleaned.match(/^```(?:json)?\s*\n?([\s\S]*?)\n?```$/); - if (fenceMatch) { - cleaned = fenceMatch[1].trim(); - } - - return JSON.parse(cleaned) as T; - } catch { - console.warn("[KnowledgeExtract] Failed to parse JSON from LLM response:", raw.slice(0, 200)); - return null; + // Attempt 1: direct parse after stripping fences + let cleaned = raw.trim(); + const fenceMatch = cleaned.match(/^```(?:json)?\s*\n?([\s\S]*?)\n?```$/); + if (fenceMatch) { + cleaned = fenceMatch[1].trim(); } + + try { + return JSON.parse(cleaned) as T; + } catch { /* fall through */ } + + // Attempt 2: extract first JSON array from the text + const arrayMatch = cleaned.match(/\[[\s\S]*\]/); + if (arrayMatch) { + try { + return JSON.parse(arrayMatch[0]) as T; + } catch { /* fall through */ } + + // Attempt 3: try to fix truncated array by closing it after the last complete object + const partial = arrayMatch[0]; + const lastBrace = partial.lastIndexOf('}'); + if (lastBrace > 0) { + const truncated = partial.slice(0, lastBrace + 1) + ']'; + try { + return JSON.parse(truncated) as T; + } catch { /* fall through */ } + } + } + + // Attempt 4: extract first JSON object, wrap in array + const objMatch = cleaned.match(/\{[\s\S]*?\}/); + if (objMatch) { + try { + const obj = JSON.parse(objMatch[0]); + return [obj] as unknown as T; + } catch { /* fall through */ } + } + + console.warn("[KnowledgeExtract] Failed to parse JSON from LLM response:", raw.slice(0, 300)); + return null; } export async function run(): Promise { diff --git a/ts/packages/flow/src/index.ts b/ts/packages/flow/src/index.ts index 51e255cd..d23f68a5 100644 --- a/ts/packages/flow/src/index.ts +++ b/ts/packages/flow/src/index.ts @@ -79,3 +79,12 @@ export { DocumentRagService } from "./retrieval/document-rag-service.js"; // Flow manager service export { FlowManagerService } from "./flow-manager/service.js"; + +// Azure OpenAI text completion +export { AzureOpenAIProcessor } from "./model/text-completion/azure-openai.js"; + +// OpenAI-compatible text completion +export { OpenAICompatibleProcessor } from "./model/text-completion/openai-compatible.js"; + +// Mistral text completion +export { MistralProcessor } from "./model/text-completion/mistral.js"; diff --git a/ts/packages/flow/src/model/text-completion/azure-openai.ts b/ts/packages/flow/src/model/text-completion/azure-openai.ts new file mode 100644 index 00000000..770c037d --- /dev/null +++ b/ts/packages/flow/src/model/text-completion/azure-openai.ts @@ -0,0 +1,156 @@ +/** + * Azure OpenAI text completion service. + * + * Env: + * AZURE_TOKEN (required – Azure OpenAI API key) + * AZURE_ENDPOINT (required – e.g. https://my-resource.openai.azure.com) + * AZURE_MODEL (default: gpt-4o) + * AZURE_API_VERSION (default: 2024-12-01-preview) + */ + +import { AzureOpenAI } from "openai"; +import { + LlmService, + type ProcessorConfig, + type LlmResult, + type LlmChunk, + TooManyRequestsError, +} from "@trustgraph/base"; + +export class AzureOpenAIProcessor extends LlmService { + private client: AzureOpenAI; + private readonly defaultModel: string; + private readonly defaultTemperature: number; + private readonly maxOutput: number; + + constructor( + config: ProcessorConfig & { + model?: string; + apiKey?: string; + endpoint?: string; + apiVersion?: string; + temperature?: number; + maxOutput?: number; + }, + ) { + super(config); + + this.defaultModel = config.model ?? process.env.AZURE_MODEL ?? "gpt-4o"; + this.defaultTemperature = config.temperature ?? 0.0; + this.maxOutput = config.maxOutput ?? 4096; + + const apiKey = config.apiKey ?? process.env.AZURE_TOKEN; + if (!apiKey) throw new Error("Azure OpenAI API key not specified"); + + const endpoint = config.endpoint ?? process.env.AZURE_ENDPOINT; + if (!endpoint) throw new Error("Azure OpenAI endpoint not specified"); + + const apiVersion = + config.apiVersion ?? + process.env.AZURE_API_VERSION ?? + "2024-12-01-preview"; + + this.client = new AzureOpenAI({ apiKey, apiVersion, endpoint }); + + console.log("[AzureOpenAI] LLM service initialized"); + } + + async generateContent( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): Promise { + const modelName = model ?? this.defaultModel; + const temp = temperature ?? this.defaultTemperature; + + try { + const resp = await this.client.chat.completions.create({ + model: modelName, + messages: [ + { role: "system", content: system }, + { role: "user", content: prompt }, + ], + temperature: temp, + max_completion_tokens: this.maxOutput, + }); + + return { + text: resp.choices[0].message.content ?? "", + inToken: resp.usage?.prompt_tokens ?? 0, + outToken: resp.usage?.completion_tokens ?? 0, + model: modelName, + }; + } catch (err) { + if ((err as any)?.status === 429) { + throw new TooManyRequestsError(); + } + throw err; + } + } + + override supportsStreaming(): boolean { + return true; + } + + async *generateContentStream( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): AsyncGenerator { + const modelName = model ?? this.defaultModel; + const temp = temperature ?? this.defaultTemperature; + + try { + const stream = await this.client.chat.completions.create({ + model: modelName, + messages: [ + { role: "system", content: system }, + { role: "user", content: prompt }, + ], + temperature: temp, + max_completion_tokens: this.maxOutput, + stream: true, + stream_options: { include_usage: true }, + }); + + let totalInputTokens = 0; + let totalOutputTokens = 0; + + for await (const chunk of stream) { + if (chunk.choices?.[0]?.delta?.content) { + yield { + text: chunk.choices[0].delta.content, + inToken: null, + outToken: null, + model: modelName, + isFinal: false, + }; + } + + if (chunk.usage) { + totalInputTokens = chunk.usage.prompt_tokens; + totalOutputTokens = chunk.usage.completion_tokens; + } + } + + yield { + text: "", + inToken: totalInputTokens, + outToken: totalOutputTokens, + model: modelName, + isFinal: true, + }; + } catch (err) { + if ((err as any)?.status === 429) { + throw new TooManyRequestsError(); + } + throw err; + } + } +} + +export async function run(): Promise { + await AzureOpenAIProcessor.launch("text-completion"); +} diff --git a/ts/packages/flow/src/model/text-completion/mistral.ts b/ts/packages/flow/src/model/text-completion/mistral.ts new file mode 100644 index 00000000..cef090b2 --- /dev/null +++ b/ts/packages/flow/src/model/text-completion/mistral.ts @@ -0,0 +1,144 @@ +/** + * Mistral text completion service. + * + * Env: + * MISTRAL_TOKEN (required – Mistral API key) + * MISTRAL_MODEL (default: ministral-8b-latest) + */ + +import { Mistral } from "@mistralai/mistralai"; +import { + LlmService, + type ProcessorConfig, + type LlmResult, + type LlmChunk, + TooManyRequestsError, +} from "@trustgraph/base"; + +export class MistralProcessor extends LlmService { + private client: Mistral; + private readonly defaultModel: string; + private readonly defaultTemperature: number; + private readonly maxOutput: number; + + constructor( + config: ProcessorConfig & { + model?: string; + apiKey?: string; + temperature?: number; + maxOutput?: number; + }, + ) { + super(config); + + this.defaultModel = + config.model ?? process.env.MISTRAL_MODEL ?? "ministral-8b-latest"; + this.defaultTemperature = config.temperature ?? 0.0; + this.maxOutput = config.maxOutput ?? 4096; + + const apiKey = config.apiKey ?? process.env.MISTRAL_TOKEN; + if (!apiKey) throw new Error("Mistral API key not specified"); + + this.client = new Mistral({ apiKey }); + + console.log("[Mistral] LLM service initialized"); + } + + async generateContent( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): Promise { + const modelName = model ?? this.defaultModel; + const temp = temperature ?? this.defaultTemperature; + + try { + const resp = await this.client.chat.complete({ + model: modelName, + messages: [ + { role: "system", content: system }, + { role: "user", content: prompt }, + ], + temperature: temp, + maxTokens: this.maxOutput, + }); + + return { + text: (resp.choices?.[0]?.message?.content as string) ?? "", + inToken: resp.usage?.promptTokens ?? 0, + outToken: resp.usage?.completionTokens ?? 0, + model: modelName, + }; + } catch (err) { + if ((err as any)?.statusCode === 429 || (err as any)?.status === 429) { + throw new TooManyRequestsError(); + } + throw err; + } + } + + override supportsStreaming(): boolean { + return true; + } + + async *generateContentStream( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): AsyncGenerator { + const modelName = model ?? this.defaultModel; + const temp = temperature ?? this.defaultTemperature; + + try { + const stream = await this.client.chat.stream({ + model: modelName, + messages: [ + { role: "system", content: system }, + { role: "user", content: prompt }, + ], + temperature: temp, + maxTokens: this.maxOutput, + }); + + let totalInputTokens = 0; + let totalOutputTokens = 0; + + for await (const chunk of stream) { + const delta = chunk.data?.choices?.[0]?.delta; + if (delta?.content) { + yield { + text: delta.content as string, + inToken: null, + outToken: null, + model: modelName, + isFinal: false, + }; + } + + if (chunk.data?.usage) { + totalInputTokens = chunk.data.usage.promptTokens ?? 0; + totalOutputTokens = chunk.data.usage.completionTokens ?? 0; + } + } + + yield { + text: "", + inToken: totalInputTokens, + outToken: totalOutputTokens, + model: modelName, + isFinal: true, + }; + } catch (err) { + if ((err as any)?.statusCode === 429 || (err as any)?.status === 429) { + throw new TooManyRequestsError(); + } + throw err; + } + } +} + +export async function run(): Promise { + await MistralProcessor.launch("text-completion"); +} diff --git a/ts/packages/flow/src/model/text-completion/openai-compatible.ts b/ts/packages/flow/src/model/text-completion/openai-compatible.ts new file mode 100644 index 00000000..73ed47cf --- /dev/null +++ b/ts/packages/flow/src/model/text-completion/openai-compatible.ts @@ -0,0 +1,139 @@ +/** + * OpenAI-compatible text completion service (generic local server). + * + * Works with LM Studio, llama.cpp, vLLM, Ollama OpenAI-compat endpoint, etc. + * + * Env: + * OPENAI_COMPAT_URL (required – e.g. http://localhost:1234/v1) + * OPENAI_COMPAT_KEY (default: sk-no-key-required) + * OPENAI_COMPAT_MODEL (default: default) + */ + +import OpenAI from "openai"; +import { + LlmService, + type ProcessorConfig, + type LlmResult, + type LlmChunk, +} from "@trustgraph/base"; + +export class OpenAICompatibleProcessor extends LlmService { + private client: OpenAI; + private readonly defaultModel: string; + private readonly defaultTemperature: number; + private readonly maxOutput: number; + + constructor( + config: ProcessorConfig & { + model?: string; + apiKey?: string; + baseUrl?: string; + temperature?: number; + maxOutput?: number; + }, + ) { + super(config); + + this.defaultModel = + config.model ?? process.env.OPENAI_COMPAT_MODEL ?? "default"; + this.defaultTemperature = config.temperature ?? 0.0; + this.maxOutput = config.maxOutput ?? 4096; + + const baseURL = config.baseUrl ?? process.env.OPENAI_COMPAT_URL; + if (!baseURL) + throw new Error( + "OpenAI-compatible server URL not specified (set OPENAI_COMPAT_URL)", + ); + + const apiKey = + config.apiKey ?? process.env.OPENAI_COMPAT_KEY ?? "sk-no-key-required"; + + this.client = new OpenAI({ baseURL, apiKey }); + + console.log("[OpenAI-Compatible] LLM service initialized"); + } + + async generateContent( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): Promise { + const modelName = model ?? this.defaultModel; + const temp = temperature ?? this.defaultTemperature; + + const resp = await this.client.chat.completions.create({ + model: modelName, + messages: [ + { role: "system", content: system }, + { role: "user", content: prompt }, + ], + temperature: temp, + max_tokens: this.maxOutput, + }); + + return { + text: resp.choices[0].message.content ?? "", + inToken: resp.usage?.prompt_tokens ?? 0, + outToken: resp.usage?.completion_tokens ?? 0, + model: modelName, + }; + } + + override supportsStreaming(): boolean { + return true; + } + + async *generateContentStream( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): AsyncGenerator { + const modelName = model ?? this.defaultModel; + const temp = temperature ?? this.defaultTemperature; + + const stream = await this.client.chat.completions.create({ + model: modelName, + messages: [ + { role: "system", content: system }, + { role: "user", content: prompt }, + ], + temperature: temp, + max_tokens: this.maxOutput, + stream: true, + }); + + let totalInputTokens = 0; + let totalOutputTokens = 0; + + for await (const chunk of stream) { + if (chunk.choices?.[0]?.delta?.content) { + yield { + text: chunk.choices[0].delta.content, + inToken: null, + outToken: null, + model: modelName, + isFinal: false, + }; + } + + if (chunk.usage) { + totalInputTokens = chunk.usage.prompt_tokens; + totalOutputTokens = chunk.usage.completion_tokens; + } + } + + yield { + text: "", + inToken: totalInputTokens, + outToken: totalOutputTokens, + model: modelName, + isFinal: true, + }; + } +} + +export async function run(): Promise { + await OpenAICompatibleProcessor.launch("text-completion"); +} diff --git a/ts/packages/workbench/src/App.tsx b/ts/packages/workbench/src/App.tsx index ada5af84..5012c5fe 100644 --- a/ts/packages/workbench/src/App.tsx +++ b/ts/packages/workbench/src/App.tsx @@ -3,6 +3,9 @@ import { RootLayout } from "@/components/layout/root-layout"; import ChatPage from "@/pages/chat"; import LibraryPage from "@/pages/library"; import GraphPage from "@/pages/graph"; +import PromptsPage from "@/pages/prompts"; +import TokenCostPage from "@/pages/token-cost"; +import KnowledgeCoresPage from "@/pages/knowledge-cores"; import FlowsPage from "@/pages/flows"; import SettingsPage from "@/pages/settings"; import { NotificationToasts } from "@/components/notification-toasts"; @@ -16,6 +19,9 @@ export default function App() { } /> } /> } /> + } /> + } /> + } /> } /> } /> diff --git a/ts/packages/workbench/src/components/layout/sidebar.tsx b/ts/packages/workbench/src/components/layout/sidebar.tsx index 8907831f..309f269f 100644 --- a/ts/packages/workbench/src/components/layout/sidebar.tsx +++ b/ts/packages/workbench/src/components/layout/sidebar.tsx @@ -3,6 +3,9 @@ import { MessageSquareText, LibraryBig, Rotate3d, + MessageCircleCode, + Coins, + BrainCircuit, Workflow, Settings, TestTube2, @@ -155,6 +158,9 @@ export function Sidebar() { + + + diff --git a/ts/packages/workbench/src/hooks/use-prompts.ts b/ts/packages/workbench/src/hooks/use-prompts.ts new file mode 100644 index 00000000..0ac50a8e --- /dev/null +++ b/ts/packages/workbench/src/hooks/use-prompts.ts @@ -0,0 +1,50 @@ +import { useCallback, useEffect, useState } from "react"; +import { useSocket } from "@/providers/socket-provider"; +import { useConnectionState } from "@/providers/socket-provider"; + +export function usePrompts() { + const socket = useSocket(); + const connectionState = useConnectionState(); + const [prompts, setPrompts] = useState>([]); + const [systemPrompt, setSystemPrompt] = useState(""); + const [loading, setLoading] = useState(false); + + const loadPrompts = useCallback(async () => { + try { + setLoading(true); + const list = await socket.config().getPrompts(); + setPrompts(Array.isArray(list) ? list : []); + } catch (err) { + console.error("Failed to load prompts:", err); + } finally { + setLoading(false); + } + }, [socket]); + + const loadSystemPrompt = useCallback(async () => { + try { + const sp = await socket.config().getSystemPrompt(); + setSystemPrompt(typeof sp === "string" ? sp : JSON.stringify(sp, null, 2)); + } catch (err) { + console.error("Failed to load system prompt:", err); + } + }, [socket]); + + const getPrompt = useCallback(async (id: string) => { + return socket.config().getPrompt(id); + }, [socket]); + + // Auto-load when connected + useEffect(() => { + const connected = + connectionState.status === "connected" || + connectionState.status === "authenticated" || + connectionState.status === "unauthenticated"; + if (connected) { + loadPrompts(); + loadSystemPrompt(); + } + }, [connectionState.status, loadPrompts, loadSystemPrompt]); + + return { prompts, systemPrompt, loading, loadPrompts, loadSystemPrompt, getPrompt }; +} diff --git a/ts/packages/workbench/src/pages/knowledge-cores.tsx b/ts/packages/workbench/src/pages/knowledge-cores.tsx new file mode 100644 index 00000000..685c7942 --- /dev/null +++ b/ts/packages/workbench/src/pages/knowledge-cores.tsx @@ -0,0 +1,244 @@ +import { useCallback, useEffect, useState } from "react"; +import { + BrainCircuit, + Loader2, + RefreshCw, + Download, + Trash2, + AlertTriangle, +} from "lucide-react"; +import { cn } from "@/lib/utils"; +import { useSocket } from "@/providers/socket-provider"; +import { useConnectionState } from "@/providers/socket-provider"; +import { useNotification } from "@/providers/notification-provider"; +import { useSessionStore } from "@/hooks/use-session-store"; +import { Dialog } from "@/components/ui/dialog"; + +// --------------------------------------------------------------------------- +// Delete confirmation dialog +// --------------------------------------------------------------------------- + +function DeleteCoreDialog({ + open, + coreId, + onClose, + onConfirm, +}: { + open: boolean; + coreId: string; + onClose: () => void; + onConfirm: () => void; +}) { + return ( + + + + + } + > +
+ +

+ Are you sure you want to delete knowledge core{" "} + {coreId}? + This action cannot be undone. +

+
+ + ); +} + +// --------------------------------------------------------------------------- +// Knowledge Cores page +// --------------------------------------------------------------------------- + +export default function KnowledgeCoresPage() { + const socket = useSocket(); + const connectionState = useConnectionState(); + const notify = useNotification(); + const flowId = useSessionStore((s) => s.flowId); + + const [cores, setCores] = useState([]); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + const [deleteTarget, setDeleteTarget] = useState(null); + const [actionInProgress, setActionInProgress] = useState(null); + + const loadCores = useCallback(async () => { + try { + setLoading(true); + setError(null); + const ids = await socket.knowledge().getKnowledgeCores(); + setCores(Array.isArray(ids) ? ids : []); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + setError(msg); + console.error("Failed to load knowledge cores:", err); + } finally { + setLoading(false); + } + }, [socket]); + + // Auto-load when connected + useEffect(() => { + const connected = + connectionState.status === "connected" || + connectionState.status === "authenticated" || + connectionState.status === "unauthenticated"; + if (connected) { + loadCores(); + } + }, [connectionState.status, loadCores]); + + const handleLoad = useCallback( + async (id: string) => { + setActionInProgress(id); + try { + await socket.knowledge().loadKgCore(id, flowId); + notify.success("Core loaded", `Knowledge core "${id}" has been loaded.`); + } catch (err) { + notify.error( + "Failed to load core", + err instanceof Error ? err.message : String(err), + ); + } finally { + setActionInProgress(null); + } + }, + [socket, flowId, notify], + ); + + const handleDelete = useCallback(async () => { + if (!deleteTarget) return; + setActionInProgress(deleteTarget); + try { + await socket.knowledge().deleteKgCore(deleteTarget); + notify.success("Core deleted", `Knowledge core "${deleteTarget}" has been deleted.`); + await loadCores(); + } catch (err) { + notify.error( + "Failed to delete core", + err instanceof Error ? err.message : String(err), + ); + } finally { + setActionInProgress(null); + setDeleteTarget(null); + } + }, [socket, deleteTarget, notify, loadCores]); + + return ( +
+ {/* Header */} +
+
+ +

Knowledge Cores

+ + {cores.length} core{cores.length !== 1 ? "s" : ""} + +
+ + +
+ + {/* Content */} + {loading && cores.length === 0 && ( +
+ + Loading knowledge cores... +
+ )} + + {error && ( +

+ {error} +

+ )} + + {!loading && !error && cores.length === 0 && ( +
+ +

No knowledge cores available.

+
+ )} + + {cores.length > 0 && ( +
+ + + + + + + + + {cores.map((id) => ( + + + + + ))} + +
Core IDActions
+ {id} + +
+ + +
+
+
+ )} + + {/* Delete confirmation dialog */} + setDeleteTarget(null)} + onConfirm={handleDelete} + /> +
+ ); +} diff --git a/ts/packages/workbench/src/pages/prompts.tsx b/ts/packages/workbench/src/pages/prompts.tsx new file mode 100644 index 00000000..1ac7cf96 --- /dev/null +++ b/ts/packages/workbench/src/pages/prompts.tsx @@ -0,0 +1,215 @@ +import { useCallback, useState } from "react"; +import { + MessageCircleCode, + Loader2, + RefreshCw, + ChevronRight, + X, + FileText, + Terminal, +} from "lucide-react"; +import { cn } from "@/lib/utils"; +import { usePrompts } from "@/hooks/use-prompts"; + +// --------------------------------------------------------------------------- +// Prompts page +// --------------------------------------------------------------------------- + +type Tab = "templates" | "system"; + +export default function PromptsPage() { + const { prompts, systemPrompt, loading, loadPrompts, loadSystemPrompt, getPrompt } = usePrompts(); + + const [activeTab, setActiveTab] = useState("templates"); + const [selectedPromptId, setSelectedPromptId] = useState(null); + const [promptDetail, setPromptDetail] = useState(""); + const [loadingDetail, setLoadingDetail] = useState(false); + + const handleSelectPrompt = useCallback( + async (id: string) => { + setSelectedPromptId(id); + setLoadingDetail(true); + try { + const detail = await getPrompt(id); + setPromptDetail( + typeof detail === "string" ? detail : JSON.stringify(detail, null, 2), + ); + } catch (err) { + console.error("Failed to load prompt detail:", err); + setPromptDetail("Error loading prompt."); + } finally { + setLoadingDetail(false); + } + }, + [getPrompt], + ); + + const handleRefresh = useCallback(() => { + loadPrompts(); + loadSystemPrompt(); + }, [loadPrompts, loadSystemPrompt]); + + return ( +
+ {/* Header */} +
+
+ +

Prompts

+
+ + +
+ + {/* Tabs */} +
+ + +
+ + {/* Templates tab */} + {activeTab === "templates" && ( +
+ {loading && prompts.length === 0 && ( +
+ + Loading prompts... +
+ )} + + {!loading && prompts.length === 0 && ( +
+ +

No prompt templates found.

+
+ )} + + {prompts.length > 0 && ( +
+ {/* Prompt list */} +
+
+

+ Templates ({prompts.length}) +

+
+
+ {prompts.map((p) => { + const id = p.id ?? (p as Record).name ?? String(p); + return ( + + ); + })} +
+
+ + {/* Prompt detail */} +
+ {selectedPromptId ? ( + <> +
+

+ {selectedPromptId} +

+ +
+
+ {loadingDetail ? ( +
+ + Loading... +
+ ) : ( + 
+                          {promptDetail}
+
+ )} +
+ + ) : ( +
+ Select a template to view its contents. +
+ )} +
+
+ )} +
+ )} + + {/* System Prompt tab */} + {activeTab === "system" && ( +
+
+

+ System Prompt +

+
+
+ {loading ? ( +
+ + Loading... +
+ ) : systemPrompt ? ( + 
+                {systemPrompt}
+
+ ) : ( +

No system prompt configured.

+ )} +
+
+ )} +
+ ); +} diff --git a/ts/packages/workbench/src/pages/settings.tsx b/ts/packages/workbench/src/pages/settings.tsx index 056d03bc..f0f413c3 100644 --- a/ts/packages/workbench/src/pages/settings.tsx +++ b/ts/packages/workbench/src/pages/settings.tsx @@ -51,7 +51,7 @@ function Section({ // --------------------------------------------------------------------------- export default function SettingsPage() { - const { settings, updateSetting } = useSettings(); + const { settings, updateSetting, updateFeatureSwitches } = useSettings(); const connectionState = useConnectionState(); const socket = useSocket(); const { flows } = useFlows(); @@ -318,6 +318,32 @@ export default function SettingsPage() { + {/* Feature Switches */} +
} + > + {Object.entries(settings.featureSwitches).map(([key, enabled]) => ( +
+
+

{key.replace(/([A-Z])/g, " $1").trim()}

+
+ +
+ ))} +
+ {/* About */}
([]); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + + const loadCosts = useCallback(async () => { + try { + setLoading(true); + setError(null); + const data = await socket.config().getTokenCosts(); + setCosts( + Array.isArray(data) + ? data.map((d: Record) => ({ + model: String(d.model ?? ""), + input_price: Number(d.input_price ?? 0), + output_price: Number(d.output_price ?? 0), + })) + : [], + ); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + setError(msg); + console.error("Failed to load token costs:", err); + } finally { + setLoading(false); + } + }, [socket]); + + // Auto-load when connected + useEffect(() => { + const connected = + connectionState.status === "connected" || + connectionState.status === "authenticated" || + connectionState.status === "unauthenticated"; + if (connected) { + loadCosts(); + } + }, [connectionState.status, loadCosts]); + + const formatPrice = (price: number) => { + if (price == null) return "--"; + return `$${price.toFixed(2)}`; + }; + + return ( +
+ {/* Header */} +
+
+ +

Token Cost

+ + {costs.length} model{costs.length !== 1 ? "s" : ""} + +
+ + +
+ + {/* Content */} + {loading && costs.length === 0 && ( +
+ + Loading token costs... +
+ )} + + {error && ( +

+ {error} +

+ )} + + {!loading && !error && costs.length === 0 && ( +
+ +

No token cost data available.

+
+ )} + + {costs.length > 0 && ( +
+ + + + + + + + + + {costs.map((cost) => ( + + + + + + ))} + +
ModelInput Price ($/1M tokens)Output Price ($/1M tokens)
+ {cost.model} + + {formatPrice(cost.input_price)} + + {formatPrice(cost.output_price)} +
+
+ )} +
+ ); +} diff --git a/ts/pnpm-lock.yaml b/ts/pnpm-lock.yaml index 487806da..b6cccf66 100644 --- a/ts/pnpm-lock.yaml +++ b/ts/pnpm-lock.yaml @@ -8,6 +8,9 @@ importers: .: devDependencies: + falkordb: + specifier: ^5.0.0 + version: 5.0.1 nats: specifier: ^2.29.0 version: 2.29.3 @@ -98,6 +101,9 @@ importers: '@fastify/websocket': specifier: ^11.0.0 version: 11.2.0 + '@mistralai/mistralai': + specifier: ^1.0.0 + version: 1.15.1 '@qdrant/js-client-rest': specifier: ^1.13.0 version: 1.17.0(typescript@5.9.3) @@ -667,6 +673,9 @@ packages: '@jridgewell/trace-mapping@0.3.31': resolution: {integrity: sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==} + '@mistralai/mistralai@1.15.1': + resolution: {integrity: sha512-fb995eiz3r0KsBGtRjFV+/iLbX+UpfalxpF+YitT3R6ukrPD4PN+FGwwmYcRFhNAzVzDUtTVxQYnjQWEnwV5nw==} + '@modelcontextprotocol/sdk@1.29.0': resolution: {integrity: sha512-zo37mZA9hJWpULgkRpowewez1y6ML5GsXJPY8FI0tBBCd77HEvza4jDqRKOXgHNn867PVGCyTdzqpz0izu5ZjQ==} engines: {node: '>=18'} @@ -2962,6 +2971,15 @@ snapshots: '@jridgewell/resolve-uri': 3.1.2 '@jridgewell/sourcemap-codec': 1.5.5 + '@mistralai/mistralai@1.15.1': + dependencies: + ws: 8.20.0 + zod: 3.25.76 + zod-to-json-schema: 3.25.2(zod@3.25.76) + transitivePeerDependencies: + - bufferutil + - utf-8-validate + '@modelcontextprotocol/sdk@1.29.0(zod@3.25.76)': dependencies: '@hono/node-server': 1.19.12(hono@4.12.10) diff --git a/ts/scripts/run-llm-azure-openai.ts b/ts/scripts/run-llm-azure-openai.ts new file mode 100644 index 00000000..2ba96ad3 --- /dev/null +++ b/ts/scripts/run-llm-azure-openai.ts @@ -0,0 +1,18 @@ +/** + * Start the Azure OpenAI text-completion service. + * + * Usage: AZURE_TOKEN=... AZURE_ENDPOINT=... pnpm tsx scripts/run-llm-azure-openai.ts + * + * Env: + * NATS_URL (default: nats://localhost:4222) + * AZURE_TOKEN (required) + * AZURE_ENDPOINT (required) + * AZURE_MODEL (default: gpt-4o) + * AZURE_API_VERSION (default: 2024-12-01-preview) + */ +import { run } from "../packages/flow/src/model/text-completion/azure-openai.js"; + +run().catch((err) => { + console.error("Azure OpenAI LLM service failed:", err); + process.exit(1); +}); diff --git a/ts/scripts/run-llm-mistral.ts b/ts/scripts/run-llm-mistral.ts new file mode 100644 index 00000000..88b5d077 --- /dev/null +++ b/ts/scripts/run-llm-mistral.ts @@ -0,0 +1,16 @@ +/** + * Start the Mistral text-completion service. + * + * Usage: MISTRAL_TOKEN=... pnpm tsx scripts/run-llm-mistral.ts + * + * Env: + * NATS_URL (default: nats://localhost:4222) + * MISTRAL_TOKEN (required) + * MISTRAL_MODEL (default: ministral-8b-latest) + */ +import { run } from "../packages/flow/src/model/text-completion/mistral.js"; + +run().catch((err) => { + console.error("Mistral LLM service failed:", err); + process.exit(1); +}); diff --git a/ts/scripts/run-llm-openai-compatible.ts b/ts/scripts/run-llm-openai-compatible.ts new file mode 100644 index 00000000..866aaafd --- /dev/null +++ b/ts/scripts/run-llm-openai-compatible.ts @@ -0,0 +1,17 @@ +/** + * Start the OpenAI-compatible text-completion service. + * + * Usage: OPENAI_COMPAT_URL=http://localhost:1234/v1 pnpm tsx scripts/run-llm-openai-compatible.ts + * + * Env: + * NATS_URL (default: nats://localhost:4222) + * OPENAI_COMPAT_URL (required) + * OPENAI_COMPAT_KEY (default: sk-no-key-required) + * OPENAI_COMPAT_MODEL (default: default) + */ +import { run } from "../packages/flow/src/model/text-completion/openai-compatible.js"; + +run().catch((err) => { + console.error("OpenAI-compatible LLM service failed:", err); + process.exit(1); +}); diff --git a/ts/scripts/test-pipeline.ts b/ts/scripts/test-pipeline.ts index 090066ae..99e6b121 100644 --- a/ts/scripts/test-pipeline.ts +++ b/ts/scripts/test-pipeline.ts @@ -551,7 +551,12 @@ async function testFullPipeline(): Promise { console.log(` FalkorDB: no nodes found (count=${count})`); } } catch (err) { - console.log(` FalkorDB check failed: ${err}`); + const errStr = String(err); + if (errStr.includes("Cannot find package") || errStr.includes("MODULE_NOT_FOUND")) { + console.log(" FalkorDB check skipped: falkordb package not available at workspace root"); + } else { + console.log(` FalkorDB check failed: ${err}`); + } } // 6. Verify embeddings in Qdrant