@deepagents/text2sql 0.14.0 → 0.15.1
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/index.d.ts +0 -1
- package/dist/index.d.ts.map +1 -1
- package/dist/index.js +2012 -3287
- package/dist/index.js.map +4 -4
- package/dist/lib/adapters/adapter.d.ts +2 -0
- package/dist/lib/adapters/adapter.d.ts.map +1 -1
- package/dist/lib/adapters/bigquery/bigquery.d.ts +41 -0
- package/dist/lib/adapters/bigquery/bigquery.d.ts.map +1 -0
- package/dist/lib/adapters/bigquery/constraint.bigquery.grounding.d.ts +11 -0
- package/dist/lib/adapters/bigquery/constraint.bigquery.grounding.d.ts.map +1 -0
- package/dist/lib/adapters/bigquery/index.d.ts +35 -0
- package/dist/lib/adapters/bigquery/index.d.ts.map +1 -0
- package/dist/lib/adapters/bigquery/index.js +1321 -0
- package/dist/lib/adapters/bigquery/index.js.map +7 -0
- package/dist/lib/adapters/bigquery/indexes.bigquery.grounding.d.ts +14 -0
- package/dist/lib/adapters/bigquery/indexes.bigquery.grounding.d.ts.map +1 -0
- package/dist/lib/adapters/bigquery/info.bigquery.grounding.d.ts +9 -0
- package/dist/lib/adapters/bigquery/info.bigquery.grounding.d.ts.map +1 -0
- package/dist/lib/adapters/bigquery/row-count.bigquery.grounding.d.ts +14 -0
- package/dist/lib/adapters/bigquery/row-count.bigquery.grounding.d.ts.map +1 -0
- package/dist/lib/adapters/bigquery/table.bigquery.grounding.d.ts +15 -0
- package/dist/lib/adapters/bigquery/table.bigquery.grounding.d.ts.map +1 -0
- package/dist/lib/adapters/bigquery/view.bigquery.grounding.d.ts +12 -0
- package/dist/lib/adapters/bigquery/view.bigquery.grounding.d.ts.map +1 -0
- package/dist/lib/adapters/groundings/index.js +11 -2058
- package/dist/lib/adapters/groundings/index.js.map +4 -4
- package/dist/lib/adapters/mysql/index.js +11 -2058
- package/dist/lib/adapters/mysql/index.js.map +4 -4
- package/dist/lib/adapters/postgres/column-stats.postgres.grounding.d.ts +2 -4
- package/dist/lib/adapters/postgres/column-stats.postgres.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/postgres/column-values.postgres.grounding.d.ts +0 -8
- package/dist/lib/adapters/postgres/column-values.postgres.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/postgres/index.js +139 -2068
- package/dist/lib/adapters/postgres/index.js.map +4 -4
- package/dist/lib/adapters/postgres/row-count.postgres.grounding.d.ts +0 -3
- package/dist/lib/adapters/postgres/row-count.postgres.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/spreadsheet/index.js +11 -56
- package/dist/lib/adapters/spreadsheet/index.js.map +4 -4
- package/dist/lib/adapters/sqlite/index.js +11 -2058
- package/dist/lib/adapters/sqlite/index.js.map +4 -4
- package/dist/lib/adapters/sqlserver/index.js +11 -2058
- package/dist/lib/adapters/sqlserver/index.js.map +4 -4
- package/dist/lib/agents/result-tools.d.ts +1 -3
- package/dist/lib/agents/result-tools.d.ts.map +1 -1
- package/dist/lib/fragments/schema.d.ts +2 -0
- package/dist/lib/fragments/schema.d.ts.map +1 -1
- package/dist/lib/fs/index.d.ts +2 -1
- package/dist/lib/fs/index.d.ts.map +1 -1
- package/dist/lib/fs/mssql/ddl.mssql-fs.d.ts +2 -0
- package/dist/lib/fs/mssql/ddl.mssql-fs.d.ts.map +1 -0
- package/dist/lib/fs/mssql/mssql-fs.d.ts +56 -0
- package/dist/lib/fs/mssql/mssql-fs.d.ts.map +1 -0
- package/dist/lib/fs/scoped-fs.d.ts +2 -0
- package/dist/lib/fs/scoped-fs.d.ts.map +1 -1
- package/dist/lib/fs/{sqlite-fs.d.ts → sqlite/sqlite-fs.d.ts} +2 -0
- package/dist/lib/fs/sqlite/sqlite-fs.d.ts.map +1 -0
- package/dist/lib/fs/tracked-fs.d.ts +2 -0
- package/dist/lib/fs/tracked-fs.d.ts.map +1 -1
- package/dist/lib/instructions.d.ts.map +1 -1
- package/dist/lib/sql.d.ts +2 -2
- package/dist/lib/sql.d.ts.map +1 -1
- package/dist/lib/synthesis/index.js +254 -2119
- package/dist/lib/synthesis/index.js.map +4 -4
- package/package.json +17 -10
- package/dist/lib/agents/text2sql.agent.d.ts +0 -3
- package/dist/lib/agents/text2sql.agent.d.ts.map +0 -1
- package/dist/lib/fs/sqlite-fs.d.ts.map +0 -1
|
@@ -552,2066 +552,19 @@ var ColumnValuesGrounding = class extends AbstractGrounding {
|
|
|
552
552
|
};
|
|
553
553
|
|
|
554
554
|
// packages/text2sql/src/lib/adapters/groundings/report.grounding.ts
|
|
555
|
-
import { groq
|
|
556
|
-
import { tool } from "ai";
|
|
557
|
-
import
|
|
558
|
-
import z from "zod";
|
|
559
|
-
import "@deepagents/agent";
|
|
560
|
-
|
|
561
|
-
|
|
562
|
-
|
|
563
|
-
|
|
564
|
-
|
|
565
|
-
|
|
566
|
-
|
|
567
|
-
|
|
568
|
-
import { defineCommand } from "just-bash";
|
|
569
|
-
import spawn from "nano-spawn";
|
|
570
|
-
import "bash-tool";
|
|
571
|
-
import spawn2 from "nano-spawn";
|
|
572
|
-
import {
|
|
573
|
-
createBashTool
|
|
574
|
-
} from "bash-tool";
|
|
575
|
-
import dedent from "dedent";
|
|
576
|
-
import YAML from "yaml";
|
|
577
|
-
import { DatabaseSync } from "node:sqlite";
|
|
578
|
-
import { groq } from "@ai-sdk/groq";
|
|
579
|
-
import {
|
|
580
|
-
NoSuchToolError,
|
|
581
|
-
Output,
|
|
582
|
-
convertToModelMessages,
|
|
583
|
-
createUIMessageStream,
|
|
584
|
-
generateId as generateId2,
|
|
585
|
-
generateText,
|
|
586
|
-
smoothStream,
|
|
587
|
-
stepCountIs,
|
|
588
|
-
streamText
|
|
589
|
-
} from "ai";
|
|
590
|
-
import chalk2 from "chalk";
|
|
591
|
-
import "zod";
|
|
592
|
-
import "@deepagents/agent";
|
|
593
|
-
var defaultTokenizer = {
|
|
594
|
-
encode(text) {
|
|
595
|
-
return encode(text);
|
|
596
|
-
},
|
|
597
|
-
count(text) {
|
|
598
|
-
return encode(text).length;
|
|
599
|
-
}
|
|
600
|
-
};
|
|
601
|
-
var ModelsRegistry = class {
|
|
602
|
-
#cache = /* @__PURE__ */ new Map();
|
|
603
|
-
#loaded = false;
|
|
604
|
-
#tokenizers = /* @__PURE__ */ new Map();
|
|
605
|
-
#defaultTokenizer = defaultTokenizer;
|
|
606
|
-
/**
|
|
607
|
-
* Load models data from models.dev API
|
|
608
|
-
*/
|
|
609
|
-
async load() {
|
|
610
|
-
if (this.#loaded) return;
|
|
611
|
-
const response = await fetch("https://models.dev/api.json");
|
|
612
|
-
if (!response.ok) {
|
|
613
|
-
throw new Error(`Failed to fetch models: ${response.statusText}`);
|
|
614
|
-
}
|
|
615
|
-
const data = await response.json();
|
|
616
|
-
for (const [providerId, provider] of Object.entries(data)) {
|
|
617
|
-
for (const [modelId, model] of Object.entries(provider.models)) {
|
|
618
|
-
const info2 = {
|
|
619
|
-
id: model.id,
|
|
620
|
-
name: model.name,
|
|
621
|
-
family: model.family,
|
|
622
|
-
cost: model.cost,
|
|
623
|
-
limit: model.limit,
|
|
624
|
-
provider: providerId
|
|
625
|
-
};
|
|
626
|
-
this.#cache.set(`${providerId}:${modelId}`, info2);
|
|
627
|
-
}
|
|
628
|
-
}
|
|
629
|
-
this.#loaded = true;
|
|
630
|
-
}
|
|
631
|
-
/**
|
|
632
|
-
* Get model info by ID
|
|
633
|
-
* @param modelId - Model ID (e.g., "openai:gpt-4o")
|
|
634
|
-
*/
|
|
635
|
-
get(modelId) {
|
|
636
|
-
return this.#cache.get(modelId);
|
|
637
|
-
}
|
|
638
|
-
/**
|
|
639
|
-
* Check if a model exists in the registry
|
|
640
|
-
*/
|
|
641
|
-
has(modelId) {
|
|
642
|
-
return this.#cache.has(modelId);
|
|
643
|
-
}
|
|
644
|
-
/**
|
|
645
|
-
* List all available model IDs
|
|
646
|
-
*/
|
|
647
|
-
list() {
|
|
648
|
-
return [...this.#cache.keys()];
|
|
649
|
-
}
|
|
650
|
-
/**
|
|
651
|
-
* Register a custom tokenizer for specific model families
|
|
652
|
-
* @param family - Model family name (e.g., "llama", "claude")
|
|
653
|
-
* @param tokenizer - Tokenizer implementation
|
|
654
|
-
*/
|
|
655
|
-
registerTokenizer(family, tokenizer) {
|
|
656
|
-
this.#tokenizers.set(family, tokenizer);
|
|
657
|
-
}
|
|
658
|
-
/**
|
|
659
|
-
* Set the default tokenizer used when no family-specific tokenizer is registered
|
|
660
|
-
*/
|
|
661
|
-
setDefaultTokenizer(tokenizer) {
|
|
662
|
-
this.#defaultTokenizer = tokenizer;
|
|
663
|
-
}
|
|
664
|
-
/**
|
|
665
|
-
* Get the appropriate tokenizer for a model
|
|
666
|
-
*/
|
|
667
|
-
getTokenizer(modelId) {
|
|
668
|
-
const model = this.get(modelId);
|
|
669
|
-
if (model) {
|
|
670
|
-
const familyTokenizer = this.#tokenizers.get(model.family);
|
|
671
|
-
if (familyTokenizer) {
|
|
672
|
-
return familyTokenizer;
|
|
673
|
-
}
|
|
674
|
-
}
|
|
675
|
-
return this.#defaultTokenizer;
|
|
676
|
-
}
|
|
677
|
-
/**
|
|
678
|
-
* Estimate token count and cost for given text and model
|
|
679
|
-
* @param modelId - Model ID to use for pricing (e.g., "openai:gpt-4o")
|
|
680
|
-
* @param input - Input text (prompt)
|
|
681
|
-
*/
|
|
682
|
-
estimate(modelId, input) {
|
|
683
|
-
const model = this.get(modelId);
|
|
684
|
-
if (!model) {
|
|
685
|
-
throw new Error(
|
|
686
|
-
`Model "${modelId}" not found. Call load() first or check model ID.`
|
|
687
|
-
);
|
|
688
|
-
}
|
|
689
|
-
const tokenizer = this.getTokenizer(modelId);
|
|
690
|
-
const tokens = tokenizer.count(input);
|
|
691
|
-
const cost = tokens / 1e6 * model.cost.input;
|
|
692
|
-
return {
|
|
693
|
-
model: model.id,
|
|
694
|
-
provider: model.provider,
|
|
695
|
-
tokens,
|
|
696
|
-
cost,
|
|
697
|
-
limits: {
|
|
698
|
-
context: model.limit.context,
|
|
699
|
-
output: model.limit.output,
|
|
700
|
-
exceedsContext: tokens > model.limit.context
|
|
701
|
-
},
|
|
702
|
-
fragments: []
|
|
703
|
-
};
|
|
704
|
-
}
|
|
705
|
-
};
|
|
706
|
-
var _registry = null;
|
|
707
|
-
function getModelsRegistry() {
|
|
708
|
-
if (!_registry) {
|
|
709
|
-
_registry = new ModelsRegistry();
|
|
710
|
-
}
|
|
711
|
-
return _registry;
|
|
712
|
-
}
|
|
713
|
-
function isFragment(data) {
|
|
714
|
-
return typeof data === "object" && data !== null && "name" in data && "data" in data && typeof data.name === "string";
|
|
715
|
-
}
|
|
716
|
-
function isFragmentObject(data) {
|
|
717
|
-
return typeof data === "object" && data !== null && !Array.isArray(data) && !isFragment(data);
|
|
718
|
-
}
|
|
719
|
-
function isMessageFragment(fragment2) {
|
|
720
|
-
return fragment2.type === "message";
|
|
721
|
-
}
|
|
722
|
-
function fragment(name, ...children) {
|
|
723
|
-
return {
|
|
724
|
-
name,
|
|
725
|
-
data: children
|
|
726
|
-
};
|
|
727
|
-
}
|
|
728
|
-
function user(content) {
|
|
729
|
-
const message2 = typeof content === "string" ? {
|
|
730
|
-
id: generateId(),
|
|
731
|
-
role: "user",
|
|
732
|
-
parts: [{ type: "text", text: content }]
|
|
733
|
-
} : content;
|
|
734
|
-
return {
|
|
735
|
-
id: message2.id,
|
|
736
|
-
name: "user",
|
|
737
|
-
data: "content",
|
|
738
|
-
type: "message",
|
|
739
|
-
persist: true,
|
|
740
|
-
codec: {
|
|
741
|
-
decode() {
|
|
742
|
-
return message2;
|
|
743
|
-
},
|
|
744
|
-
encode() {
|
|
745
|
-
return message2;
|
|
746
|
-
}
|
|
747
|
-
}
|
|
748
|
-
};
|
|
749
|
-
}
|
|
750
|
-
function assistant(message2) {
|
|
751
|
-
return {
|
|
752
|
-
id: message2.id,
|
|
753
|
-
name: "assistant",
|
|
754
|
-
data: "content",
|
|
755
|
-
type: "message",
|
|
756
|
-
persist: true,
|
|
757
|
-
codec: {
|
|
758
|
-
decode() {
|
|
759
|
-
return message2;
|
|
760
|
-
},
|
|
761
|
-
encode() {
|
|
762
|
-
return message2;
|
|
763
|
-
}
|
|
764
|
-
}
|
|
765
|
-
};
|
|
766
|
-
}
|
|
767
|
-
function message(content) {
|
|
768
|
-
const message2 = typeof content === "string" ? {
|
|
769
|
-
id: generateId(),
|
|
770
|
-
role: "user",
|
|
771
|
-
parts: [{ type: "text", text: content }]
|
|
772
|
-
} : content;
|
|
773
|
-
return {
|
|
774
|
-
id: message2.id,
|
|
775
|
-
name: message2.role,
|
|
776
|
-
data: "content",
|
|
777
|
-
type: "message",
|
|
778
|
-
persist: true,
|
|
779
|
-
codec: {
|
|
780
|
-
decode() {
|
|
781
|
-
return message2;
|
|
782
|
-
},
|
|
783
|
-
encode() {
|
|
784
|
-
return message2;
|
|
785
|
-
}
|
|
786
|
-
}
|
|
787
|
-
};
|
|
788
|
-
}
|
|
789
|
-
function assistantText(content, options) {
|
|
790
|
-
const id = options?.id ?? crypto.randomUUID();
|
|
791
|
-
return assistant({
|
|
792
|
-
id,
|
|
793
|
-
role: "assistant",
|
|
794
|
-
parts: [{ type: "text", text: content }]
|
|
795
|
-
});
|
|
796
|
-
}
|
|
797
|
-
var LAZY_ID = Symbol.for("@deepagents/context:lazy-id");
|
|
798
|
-
function isLazyFragment(fragment2) {
|
|
799
|
-
return LAZY_ID in fragment2;
|
|
800
|
-
}
|
|
801
|
-
function lastAssistantMessage(content) {
|
|
802
|
-
return {
|
|
803
|
-
name: "assistant",
|
|
804
|
-
type: "message",
|
|
805
|
-
persist: true,
|
|
806
|
-
data: "content",
|
|
807
|
-
[LAZY_ID]: {
|
|
808
|
-
type: "last-assistant",
|
|
809
|
-
content
|
|
810
|
-
}
|
|
811
|
-
};
|
|
812
|
-
}
|
|
813
|
-
var ContextRenderer = class {
|
|
814
|
-
options;
|
|
815
|
-
constructor(options = {}) {
|
|
816
|
-
this.options = options;
|
|
817
|
-
}
|
|
818
|
-
/**
|
|
819
|
-
* Check if data is a primitive (string, number, boolean).
|
|
820
|
-
*/
|
|
821
|
-
isPrimitive(data) {
|
|
822
|
-
return typeof data === "string" || typeof data === "number" || typeof data === "boolean";
|
|
823
|
-
}
|
|
824
|
-
/**
|
|
825
|
-
* Group fragments by name for groupFragments option.
|
|
826
|
-
*/
|
|
827
|
-
groupByName(fragments) {
|
|
828
|
-
const groups = /* @__PURE__ */ new Map();
|
|
829
|
-
for (const fragment2 of fragments) {
|
|
830
|
-
const existing = groups.get(fragment2.name) ?? [];
|
|
831
|
-
existing.push(fragment2);
|
|
832
|
-
groups.set(fragment2.name, existing);
|
|
833
|
-
}
|
|
834
|
-
return groups;
|
|
835
|
-
}
|
|
836
|
-
/**
|
|
837
|
-
* Remove null/undefined from fragments and fragment data recursively.
|
|
838
|
-
* This protects renderers from nullish values and ensures they are ignored
|
|
839
|
-
* consistently across all output formats.
|
|
840
|
-
*/
|
|
841
|
-
sanitizeFragments(fragments) {
|
|
842
|
-
const sanitized = [];
|
|
843
|
-
for (const fragment2 of fragments) {
|
|
844
|
-
const cleaned = this.sanitizeFragment(fragment2, /* @__PURE__ */ new WeakSet());
|
|
845
|
-
if (cleaned) {
|
|
846
|
-
sanitized.push(cleaned);
|
|
847
|
-
}
|
|
848
|
-
}
|
|
849
|
-
return sanitized;
|
|
850
|
-
}
|
|
851
|
-
sanitizeFragment(fragment2, seen) {
|
|
852
|
-
const data = this.sanitizeData(fragment2.data, seen);
|
|
853
|
-
if (data == null) {
|
|
854
|
-
return null;
|
|
855
|
-
}
|
|
856
|
-
return {
|
|
857
|
-
...fragment2,
|
|
858
|
-
data
|
|
859
|
-
};
|
|
860
|
-
}
|
|
861
|
-
sanitizeData(data, seen) {
|
|
862
|
-
if (data == null) {
|
|
863
|
-
return void 0;
|
|
864
|
-
}
|
|
865
|
-
if (isFragment(data)) {
|
|
866
|
-
return this.sanitizeFragment(data, seen) ?? void 0;
|
|
867
|
-
}
|
|
868
|
-
if (Array.isArray(data)) {
|
|
869
|
-
if (seen.has(data)) {
|
|
870
|
-
return void 0;
|
|
871
|
-
}
|
|
872
|
-
seen.add(data);
|
|
873
|
-
const cleaned = [];
|
|
874
|
-
for (const item of data) {
|
|
875
|
-
const sanitizedItem = this.sanitizeData(item, seen);
|
|
876
|
-
if (sanitizedItem != null) {
|
|
877
|
-
cleaned.push(sanitizedItem);
|
|
878
|
-
}
|
|
879
|
-
}
|
|
880
|
-
return cleaned;
|
|
881
|
-
}
|
|
882
|
-
if (isFragmentObject(data)) {
|
|
883
|
-
if (seen.has(data)) {
|
|
884
|
-
return void 0;
|
|
885
|
-
}
|
|
886
|
-
seen.add(data);
|
|
887
|
-
const cleaned = {};
|
|
888
|
-
for (const [key, value] of Object.entries(data)) {
|
|
889
|
-
const sanitizedValue = this.sanitizeData(value, seen);
|
|
890
|
-
if (sanitizedValue != null) {
|
|
891
|
-
cleaned[key] = sanitizedValue;
|
|
892
|
-
}
|
|
893
|
-
}
|
|
894
|
-
return cleaned;
|
|
895
|
-
}
|
|
896
|
-
return data;
|
|
897
|
-
}
|
|
898
|
-
/**
|
|
899
|
-
* Template method - dispatches value to appropriate handler.
|
|
900
|
-
*/
|
|
901
|
-
renderValue(key, value, ctx) {
|
|
902
|
-
if (value == null) {
|
|
903
|
-
return "";
|
|
904
|
-
}
|
|
905
|
-
if (isFragment(value)) {
|
|
906
|
-
return this.renderFragment(value, ctx);
|
|
907
|
-
}
|
|
908
|
-
if (Array.isArray(value)) {
|
|
909
|
-
return this.renderArray(key, value, ctx);
|
|
910
|
-
}
|
|
911
|
-
if (isFragmentObject(value)) {
|
|
912
|
-
return this.renderObject(key, value, ctx);
|
|
913
|
-
}
|
|
914
|
-
return this.renderPrimitive(key, String(value), ctx);
|
|
915
|
-
}
|
|
916
|
-
/**
|
|
917
|
-
* Render all entries of an object.
|
|
918
|
-
*/
|
|
919
|
-
renderEntries(data, ctx) {
|
|
920
|
-
return Object.entries(data).map(([key, value]) => this.renderValue(key, value, ctx)).filter(Boolean);
|
|
921
|
-
}
|
|
922
|
-
};
|
|
923
|
-
var XmlRenderer = class extends ContextRenderer {
|
|
924
|
-
render(fragments) {
|
|
925
|
-
const sanitized = this.sanitizeFragments(fragments);
|
|
926
|
-
return sanitized.map((f) => this.#renderTopLevel(f)).filter(Boolean).join("\n");
|
|
927
|
-
}
|
|
928
|
-
#renderTopLevel(fragment2) {
|
|
929
|
-
if (this.isPrimitive(fragment2.data)) {
|
|
930
|
-
return this.#leafRoot(fragment2.name, String(fragment2.data));
|
|
931
|
-
}
|
|
932
|
-
if (Array.isArray(fragment2.data)) {
|
|
933
|
-
return this.#renderArray(fragment2.name, fragment2.data, 0);
|
|
934
|
-
}
|
|
935
|
-
if (isFragment(fragment2.data)) {
|
|
936
|
-
const child = this.renderFragment(fragment2.data, { depth: 1, path: [] });
|
|
937
|
-
return this.#wrap(fragment2.name, [child]);
|
|
938
|
-
}
|
|
939
|
-
if (isFragmentObject(fragment2.data)) {
|
|
940
|
-
return this.#wrap(
|
|
941
|
-
fragment2.name,
|
|
942
|
-
this.renderEntries(fragment2.data, { depth: 1, path: [] })
|
|
943
|
-
);
|
|
944
|
-
}
|
|
945
|
-
return "";
|
|
946
|
-
}
|
|
947
|
-
#renderArray(name, items, depth) {
|
|
948
|
-
const fragmentItems = items.filter(isFragment);
|
|
949
|
-
const nonFragmentItems = items.filter((item) => !isFragment(item));
|
|
950
|
-
const children = [];
|
|
951
|
-
for (const item of nonFragmentItems) {
|
|
952
|
-
if (item != null) {
|
|
953
|
-
if (isFragmentObject(item)) {
|
|
954
|
-
children.push(
|
|
955
|
-
this.#wrapIndented(
|
|
956
|
-
pluralize.singular(name),
|
|
957
|
-
this.renderEntries(item, { depth: depth + 2, path: [] }),
|
|
958
|
-
depth + 1
|
|
959
|
-
)
|
|
960
|
-
);
|
|
961
|
-
} else {
|
|
962
|
-
children.push(
|
|
963
|
-
this.#leaf(pluralize.singular(name), String(item), depth + 1)
|
|
964
|
-
);
|
|
965
|
-
}
|
|
966
|
-
}
|
|
967
|
-
}
|
|
968
|
-
if (this.options.groupFragments && fragmentItems.length > 0) {
|
|
969
|
-
const groups = this.groupByName(fragmentItems);
|
|
970
|
-
for (const [groupName, groupFragments] of groups) {
|
|
971
|
-
const groupChildren = groupFragments.map(
|
|
972
|
-
(frag) => this.renderFragment(frag, { depth: depth + 2, path: [] })
|
|
973
|
-
);
|
|
974
|
-
const pluralName = pluralize.plural(groupName);
|
|
975
|
-
children.push(this.#wrapIndented(pluralName, groupChildren, depth + 1));
|
|
976
|
-
}
|
|
977
|
-
} else {
|
|
978
|
-
for (const frag of fragmentItems) {
|
|
979
|
-
children.push(
|
|
980
|
-
this.renderFragment(frag, { depth: depth + 1, path: [] })
|
|
981
|
-
);
|
|
982
|
-
}
|
|
983
|
-
}
|
|
984
|
-
return this.#wrap(name, children);
|
|
985
|
-
}
|
|
986
|
-
#leafRoot(tag, value) {
|
|
987
|
-
const safe = this.#escape(value);
|
|
988
|
-
if (safe.includes("\n")) {
|
|
989
|
-
return `<${tag}>
|
|
990
|
-
${this.#indent(safe, 2)}
|
|
991
|
-
</${tag}>`;
|
|
992
|
-
}
|
|
993
|
-
return `<${tag}>${safe}</${tag}>`;
|
|
994
|
-
}
|
|
995
|
-
renderFragment(fragment2, ctx) {
|
|
996
|
-
const { name, data } = fragment2;
|
|
997
|
-
if (this.isPrimitive(data)) {
|
|
998
|
-
return this.#leaf(name, String(data), ctx.depth);
|
|
999
|
-
}
|
|
1000
|
-
if (isFragment(data)) {
|
|
1001
|
-
const child = this.renderFragment(data, { ...ctx, depth: ctx.depth + 1 });
|
|
1002
|
-
return this.#wrapIndented(name, [child], ctx.depth);
|
|
1003
|
-
}
|
|
1004
|
-
if (Array.isArray(data)) {
|
|
1005
|
-
return this.#renderArrayIndented(name, data, ctx.depth);
|
|
1006
|
-
}
|
|
1007
|
-
if (isFragmentObject(data)) {
|
|
1008
|
-
const children = this.renderEntries(data, {
|
|
1009
|
-
...ctx,
|
|
1010
|
-
depth: ctx.depth + 1
|
|
1011
|
-
});
|
|
1012
|
-
return this.#wrapIndented(name, children, ctx.depth);
|
|
1013
|
-
}
|
|
1014
|
-
return "";
|
|
1015
|
-
}
|
|
1016
|
-
#renderArrayIndented(name, items, depth) {
|
|
1017
|
-
const fragmentItems = items.filter(isFragment);
|
|
1018
|
-
const nonFragmentItems = items.filter((item) => !isFragment(item));
|
|
1019
|
-
const children = [];
|
|
1020
|
-
for (const item of nonFragmentItems) {
|
|
1021
|
-
if (item != null) {
|
|
1022
|
-
if (isFragmentObject(item)) {
|
|
1023
|
-
children.push(
|
|
1024
|
-
this.#wrapIndented(
|
|
1025
|
-
pluralize.singular(name),
|
|
1026
|
-
this.renderEntries(item, { depth: depth + 2, path: [] }),
|
|
1027
|
-
depth + 1
|
|
1028
|
-
)
|
|
1029
|
-
);
|
|
1030
|
-
} else {
|
|
1031
|
-
children.push(
|
|
1032
|
-
this.#leaf(pluralize.singular(name), String(item), depth + 1)
|
|
1033
|
-
);
|
|
1034
|
-
}
|
|
1035
|
-
}
|
|
1036
|
-
}
|
|
1037
|
-
if (this.options.groupFragments && fragmentItems.length > 0) {
|
|
1038
|
-
const groups = this.groupByName(fragmentItems);
|
|
1039
|
-
for (const [groupName, groupFragments] of groups) {
|
|
1040
|
-
const groupChildren = groupFragments.map(
|
|
1041
|
-
(frag) => this.renderFragment(frag, { depth: depth + 2, path: [] })
|
|
1042
|
-
);
|
|
1043
|
-
const pluralName = pluralize.plural(groupName);
|
|
1044
|
-
children.push(this.#wrapIndented(pluralName, groupChildren, depth + 1));
|
|
1045
|
-
}
|
|
1046
|
-
} else {
|
|
1047
|
-
for (const frag of fragmentItems) {
|
|
1048
|
-
children.push(
|
|
1049
|
-
this.renderFragment(frag, { depth: depth + 1, path: [] })
|
|
1050
|
-
);
|
|
1051
|
-
}
|
|
1052
|
-
}
|
|
1053
|
-
return this.#wrapIndented(name, children, depth);
|
|
1054
|
-
}
|
|
1055
|
-
renderPrimitive(key, value, ctx) {
|
|
1056
|
-
return this.#leaf(key, value, ctx.depth);
|
|
1057
|
-
}
|
|
1058
|
-
renderArray(key, items, ctx) {
|
|
1059
|
-
if (!items.length) {
|
|
1060
|
-
return "";
|
|
1061
|
-
}
|
|
1062
|
-
const itemTag = pluralize.singular(key);
|
|
1063
|
-
const children = items.filter((item) => item != null).map((item) => {
|
|
1064
|
-
if (isFragment(item)) {
|
|
1065
|
-
return this.renderFragment(item, { ...ctx, depth: ctx.depth + 1 });
|
|
1066
|
-
}
|
|
1067
|
-
if (isFragmentObject(item)) {
|
|
1068
|
-
return this.#wrapIndented(
|
|
1069
|
-
itemTag,
|
|
1070
|
-
this.renderEntries(item, { ...ctx, depth: ctx.depth + 2 }),
|
|
1071
|
-
ctx.depth + 1
|
|
1072
|
-
);
|
|
1073
|
-
}
|
|
1074
|
-
return this.#leaf(itemTag, String(item), ctx.depth + 1);
|
|
1075
|
-
});
|
|
1076
|
-
return this.#wrapIndented(key, children, ctx.depth);
|
|
1077
|
-
}
|
|
1078
|
-
renderObject(key, obj, ctx) {
|
|
1079
|
-
const children = this.renderEntries(obj, { ...ctx, depth: ctx.depth + 1 });
|
|
1080
|
-
return this.#wrapIndented(key, children, ctx.depth);
|
|
1081
|
-
}
|
|
1082
|
-
#escape(value) {
|
|
1083
|
-
if (value == null) {
|
|
1084
|
-
return "";
|
|
1085
|
-
}
|
|
1086
|
-
return value.replaceAll(/&/g, "&").replaceAll(/</g, "<").replaceAll(/>/g, ">").replaceAll(/"/g, """).replaceAll(/'/g, "'");
|
|
1087
|
-
}
|
|
1088
|
-
#indent(text, spaces) {
|
|
1089
|
-
if (!text.trim()) {
|
|
1090
|
-
return "";
|
|
1091
|
-
}
|
|
1092
|
-
const padding = " ".repeat(spaces);
|
|
1093
|
-
return text.split("\n").map((line) => line.length ? padding + line : padding).join("\n");
|
|
1094
|
-
}
|
|
1095
|
-
#leaf(tag, value, depth) {
|
|
1096
|
-
const safe = this.#escape(value);
|
|
1097
|
-
const pad = " ".repeat(depth);
|
|
1098
|
-
if (safe.includes("\n")) {
|
|
1099
|
-
return `${pad}<${tag}>
|
|
1100
|
-
${this.#indent(safe, (depth + 1) * 2)}
|
|
1101
|
-
${pad}</${tag}>`;
|
|
1102
|
-
}
|
|
1103
|
-
return `${pad}<${tag}>${safe}</${tag}>`;
|
|
1104
|
-
}
|
|
1105
|
-
#wrap(tag, children) {
|
|
1106
|
-
const content = children.filter(Boolean).join("\n");
|
|
1107
|
-
if (!content) {
|
|
1108
|
-
return "";
|
|
1109
|
-
}
|
|
1110
|
-
return `<${tag}>
|
|
1111
|
-
${content}
|
|
1112
|
-
</${tag}>`;
|
|
1113
|
-
}
|
|
1114
|
-
#wrapIndented(tag, children, depth) {
|
|
1115
|
-
const content = children.filter(Boolean).join("\n");
|
|
1116
|
-
if (!content) {
|
|
1117
|
-
return "";
|
|
1118
|
-
}
|
|
1119
|
-
const pad = " ".repeat(depth);
|
|
1120
|
-
return `${pad}<${tag}>
|
|
1121
|
-
${content}
|
|
1122
|
-
${pad}</${tag}>`;
|
|
1123
|
-
}
|
|
1124
|
-
};
|
|
1125
|
-
var ContextStore = class {
|
|
1126
|
-
};
|
|
1127
|
-
var ContextEngine = class {
|
|
1128
|
-
/** Non-message fragments (role, hints, etc.) - not persisted in graph */
|
|
1129
|
-
#fragments = [];
|
|
1130
|
-
/** Pending message fragments to be added to graph */
|
|
1131
|
-
#pendingMessages = [];
|
|
1132
|
-
#store;
|
|
1133
|
-
#chatId;
|
|
1134
|
-
#userId;
|
|
1135
|
-
#branchName;
|
|
1136
|
-
#branch = null;
|
|
1137
|
-
#chatData = null;
|
|
1138
|
-
#initialized = false;
|
|
1139
|
-
/** Initial metadata to merge on first initialization */
|
|
1140
|
-
#initialMetadata;
|
|
1141
|
-
constructor(options) {
|
|
1142
|
-
if (!options.chatId) {
|
|
1143
|
-
throw new Error("chatId is required");
|
|
1144
|
-
}
|
|
1145
|
-
if (!options.userId) {
|
|
1146
|
-
throw new Error("userId is required");
|
|
1147
|
-
}
|
|
1148
|
-
this.#store = options.store;
|
|
1149
|
-
this.#chatId = options.chatId;
|
|
1150
|
-
this.#userId = options.userId;
|
|
1151
|
-
this.#branchName = "main";
|
|
1152
|
-
this.#initialMetadata = options.metadata;
|
|
1153
|
-
}
|
|
1154
|
-
/**
|
|
1155
|
-
* Initialize the chat and branch if they don't exist.
|
|
1156
|
-
*/
|
|
1157
|
-
async #ensureInitialized() {
|
|
1158
|
-
if (this.#initialized) {
|
|
1159
|
-
return;
|
|
1160
|
-
}
|
|
1161
|
-
this.#chatData = await this.#store.upsertChat({
|
|
1162
|
-
id: this.#chatId,
|
|
1163
|
-
userId: this.#userId
|
|
1164
|
-
});
|
|
1165
|
-
if (this.#initialMetadata) {
|
|
1166
|
-
this.#chatData = await this.#store.updateChat(this.#chatId, {
|
|
1167
|
-
metadata: {
|
|
1168
|
-
...this.#chatData.metadata,
|
|
1169
|
-
...this.#initialMetadata
|
|
1170
|
-
}
|
|
1171
|
-
});
|
|
1172
|
-
this.#initialMetadata = void 0;
|
|
1173
|
-
}
|
|
1174
|
-
this.#branch = await this.#store.getActiveBranch(this.#chatId);
|
|
1175
|
-
this.#initialized = true;
|
|
1176
|
-
}
|
|
1177
|
-
/**
|
|
1178
|
-
* Create a new branch from a specific message.
|
|
1179
|
-
* Shared logic between rewind() and btw().
|
|
1180
|
-
*/
|
|
1181
|
-
async #createBranchFrom(messageId, switchTo) {
|
|
1182
|
-
const branches = await this.#store.listBranches(this.#chatId);
|
|
1183
|
-
const samePrefix = branches.filter(
|
|
1184
|
-
(it) => it.name === this.#branchName || it.name.startsWith(`${this.#branchName}-v`)
|
|
1185
|
-
);
|
|
1186
|
-
const newBranchName = `${this.#branchName}-v${samePrefix.length + 1}`;
|
|
1187
|
-
const newBranch = {
|
|
1188
|
-
id: crypto.randomUUID(),
|
|
1189
|
-
chatId: this.#chatId,
|
|
1190
|
-
name: newBranchName,
|
|
1191
|
-
headMessageId: messageId,
|
|
1192
|
-
isActive: false,
|
|
1193
|
-
createdAt: Date.now()
|
|
1194
|
-
};
|
|
1195
|
-
await this.#store.createBranch(newBranch);
|
|
1196
|
-
if (switchTo) {
|
|
1197
|
-
await this.#store.setActiveBranch(this.#chatId, newBranch.id);
|
|
1198
|
-
this.#branch = { ...newBranch, isActive: true };
|
|
1199
|
-
this.#branchName = newBranchName;
|
|
1200
|
-
this.#pendingMessages = [];
|
|
1201
|
-
}
|
|
1202
|
-
const chain = await this.#store.getMessageChain(messageId);
|
|
1203
|
-
return {
|
|
1204
|
-
id: newBranch.id,
|
|
1205
|
-
name: newBranch.name,
|
|
1206
|
-
headMessageId: newBranch.headMessageId,
|
|
1207
|
-
isActive: switchTo,
|
|
1208
|
-
messageCount: chain.length,
|
|
1209
|
-
createdAt: newBranch.createdAt
|
|
1210
|
-
};
|
|
1211
|
-
}
|
|
1212
|
-
/**
|
|
1213
|
-
* Rewind to a message without clearing pending messages.
|
|
1214
|
-
* Used internally when saving an update to an existing message.
|
|
1215
|
-
*/
|
|
1216
|
-
async #rewindForUpdate(messageId) {
|
|
1217
|
-
const pendingBackup = [...this.#pendingMessages];
|
|
1218
|
-
await this.rewind(messageId);
|
|
1219
|
-
this.#pendingMessages = pendingBackup;
|
|
1220
|
-
}
|
|
1221
|
-
/**
|
|
1222
|
-
* Get the current chat ID.
|
|
1223
|
-
*/
|
|
1224
|
-
get chatId() {
|
|
1225
|
-
return this.#chatId;
|
|
1226
|
-
}
|
|
1227
|
-
/**
|
|
1228
|
-
* Get the current branch name.
|
|
1229
|
-
*/
|
|
1230
|
-
get branch() {
|
|
1231
|
-
return this.#branchName;
|
|
1232
|
-
}
|
|
1233
|
-
/**
|
|
1234
|
-
* Get metadata for the current chat.
|
|
1235
|
-
* Returns null if the chat hasn't been initialized yet.
|
|
1236
|
-
*/
|
|
1237
|
-
get chat() {
|
|
1238
|
-
if (!this.#chatData) {
|
|
1239
|
-
return null;
|
|
1240
|
-
}
|
|
1241
|
-
return {
|
|
1242
|
-
id: this.#chatData.id,
|
|
1243
|
-
userId: this.#chatData.userId,
|
|
1244
|
-
createdAt: this.#chatData.createdAt,
|
|
1245
|
-
updatedAt: this.#chatData.updatedAt,
|
|
1246
|
-
title: this.#chatData.title,
|
|
1247
|
-
metadata: this.#chatData.metadata
|
|
1248
|
-
};
|
|
1249
|
-
}
|
|
1250
|
-
/**
|
|
1251
|
-
* Add fragments to the context.
|
|
1252
|
-
*
|
|
1253
|
-
* - Message fragments (user/assistant) are queued for persistence
|
|
1254
|
-
* - Non-message fragments (role/hint) are kept in memory for system prompt
|
|
1255
|
-
*/
|
|
1256
|
-
set(...fragments) {
|
|
1257
|
-
for (const fragment2 of fragments) {
|
|
1258
|
-
if (isMessageFragment(fragment2)) {
|
|
1259
|
-
this.#pendingMessages.push(fragment2);
|
|
1260
|
-
} else {
|
|
1261
|
-
this.#fragments.push(fragment2);
|
|
1262
|
-
}
|
|
1263
|
-
}
|
|
1264
|
-
return this;
|
|
1265
|
-
}
|
|
1266
|
-
// Unset a fragment by ID (not implemented yet)
|
|
1267
|
-
unset(fragmentId) {
|
|
1268
|
-
}
|
|
1269
|
-
/**
|
|
1270
|
-
* Render all fragments using the provided renderer.
|
|
1271
|
-
* @internal Use resolve() instead for public API.
|
|
1272
|
-
*/
|
|
1273
|
-
render(renderer) {
|
|
1274
|
-
return renderer.render(this.#fragments);
|
|
1275
|
-
}
|
|
1276
|
-
/**
|
|
1277
|
-
* Resolve context into AI SDK-ready format.
|
|
1278
|
-
*
|
|
1279
|
-
* - Initializes chat and branch if needed
|
|
1280
|
-
* - Loads message history from the graph (walking parent chain)
|
|
1281
|
-
* - Separates context fragments for system prompt
|
|
1282
|
-
* - Combines with pending messages
|
|
1283
|
-
*
|
|
1284
|
-
* @example
|
|
1285
|
-
* ```ts
|
|
1286
|
-
* const context = new ContextEngine({ store, chatId: 'chat-1', userId: 'user-1' })
|
|
1287
|
-
* .set(role('You are helpful'), user('Hello'));
|
|
1288
|
-
*
|
|
1289
|
-
* const { systemPrompt, messages } = await context.resolve();
|
|
1290
|
-
* await generateText({ system: systemPrompt, messages });
|
|
1291
|
-
* ```
|
|
1292
|
-
*/
|
|
1293
|
-
async resolve(options) {
|
|
1294
|
-
await this.#ensureInitialized();
|
|
1295
|
-
const systemPrompt = options.renderer.render(this.#fragments);
|
|
1296
|
-
const messages = [];
|
|
1297
|
-
if (this.#branch?.headMessageId) {
|
|
1298
|
-
const chain = await this.#store.getMessageChain(
|
|
1299
|
-
this.#branch.headMessageId
|
|
1300
|
-
);
|
|
1301
|
-
for (const msg of chain) {
|
|
1302
|
-
messages.push(message(msg.data).codec?.decode());
|
|
1303
|
-
}
|
|
1304
|
-
}
|
|
1305
|
-
for (let i = 0; i < this.#pendingMessages.length; i++) {
|
|
1306
|
-
const fragment2 = this.#pendingMessages[i];
|
|
1307
|
-
if (isLazyFragment(fragment2)) {
|
|
1308
|
-
this.#pendingMessages[i] = await this.#resolveLazyFragment(fragment2);
|
|
1309
|
-
}
|
|
1310
|
-
}
|
|
1311
|
-
for (const fragment2 of this.#pendingMessages) {
|
|
1312
|
-
if (!fragment2.codec) {
|
|
1313
|
-
throw new Error(
|
|
1314
|
-
`Fragment "${fragment2.name}" is missing codec. Lazy fragments must be resolved before decode.`
|
|
1315
|
-
);
|
|
1316
|
-
}
|
|
1317
|
-
const decoded = fragment2.codec.decode();
|
|
1318
|
-
messages.push(decoded);
|
|
1319
|
-
}
|
|
1320
|
-
return { systemPrompt, messages };
|
|
1321
|
-
}
|
|
1322
|
-
/**
|
|
1323
|
-
* Save pending messages to the graph.
|
|
1324
|
-
*
|
|
1325
|
-
* Each message is added as a node with parentId pointing to the previous message.
|
|
1326
|
-
* The branch head is updated to point to the last message.
|
|
1327
|
-
*
|
|
1328
|
-
* @example
|
|
1329
|
-
* ```ts
|
|
1330
|
-
* context.set(user('Hello'));
|
|
1331
|
-
* // AI responds...
|
|
1332
|
-
* context.set(assistant('Hi there!'));
|
|
1333
|
-
* await context.save(); // Persist to graph
|
|
1334
|
-
* ```
|
|
1335
|
-
*/
|
|
1336
|
-
async save() {
|
|
1337
|
-
await this.#ensureInitialized();
|
|
1338
|
-
if (this.#pendingMessages.length === 0) {
|
|
1339
|
-
return;
|
|
1340
|
-
}
|
|
1341
|
-
for (let i = 0; i < this.#pendingMessages.length; i++) {
|
|
1342
|
-
const fragment2 = this.#pendingMessages[i];
|
|
1343
|
-
if (isLazyFragment(fragment2)) {
|
|
1344
|
-
this.#pendingMessages[i] = await this.#resolveLazyFragment(fragment2);
|
|
1345
|
-
}
|
|
1346
|
-
}
|
|
1347
|
-
for (const fragment2 of this.#pendingMessages) {
|
|
1348
|
-
if (fragment2.id) {
|
|
1349
|
-
const existing = await this.#store.getMessage(fragment2.id);
|
|
1350
|
-
if (existing && existing.parentId) {
|
|
1351
|
-
await this.#rewindForUpdate(existing.parentId);
|
|
1352
|
-
fragment2.id = crypto.randomUUID();
|
|
1353
|
-
break;
|
|
1354
|
-
}
|
|
1355
|
-
}
|
|
1356
|
-
}
|
|
1357
|
-
let parentId = this.#branch.headMessageId;
|
|
1358
|
-
const now = Date.now();
|
|
1359
|
-
for (const fragment2 of this.#pendingMessages) {
|
|
1360
|
-
if (!fragment2.codec) {
|
|
1361
|
-
throw new Error(
|
|
1362
|
-
`Fragment "${fragment2.name}" is missing codec. Lazy fragments must be resolved before encode.`
|
|
1363
|
-
);
|
|
1364
|
-
}
|
|
1365
|
-
const messageData = {
|
|
1366
|
-
id: fragment2.id ?? crypto.randomUUID(),
|
|
1367
|
-
chatId: this.#chatId,
|
|
1368
|
-
parentId,
|
|
1369
|
-
name: fragment2.name,
|
|
1370
|
-
type: fragment2.type,
|
|
1371
|
-
data: fragment2.codec.encode(),
|
|
1372
|
-
createdAt: now
|
|
1373
|
-
};
|
|
1374
|
-
await this.#store.addMessage(messageData);
|
|
1375
|
-
parentId = messageData.id;
|
|
1376
|
-
}
|
|
1377
|
-
await this.#store.updateBranchHead(this.#branch.id, parentId);
|
|
1378
|
-
this.#branch.headMessageId = parentId;
|
|
1379
|
-
this.#pendingMessages = [];
|
|
1380
|
-
}
|
|
1381
|
-
/**
|
|
1382
|
-
* Resolve a lazy fragment by finding the appropriate ID.
|
|
1383
|
-
*/
|
|
1384
|
-
async #resolveLazyFragment(fragment2) {
|
|
1385
|
-
const lazy = fragment2[LAZY_ID];
|
|
1386
|
-
if (lazy.type === "last-assistant") {
|
|
1387
|
-
const lastId = await this.#getLastAssistantId();
|
|
1388
|
-
return assistantText(lazy.content, { id: lastId ?? crypto.randomUUID() });
|
|
1389
|
-
}
|
|
1390
|
-
throw new Error(`Unknown lazy fragment type: ${lazy.type}`);
|
|
1391
|
-
}
|
|
1392
|
-
/**
|
|
1393
|
-
* Find the most recent assistant message ID (pending or persisted).
|
|
1394
|
-
*/
|
|
1395
|
-
async #getLastAssistantId() {
|
|
1396
|
-
for (let i = this.#pendingMessages.length - 1; i >= 0; i--) {
|
|
1397
|
-
const msg = this.#pendingMessages[i];
|
|
1398
|
-
if (msg.name === "assistant" && !isLazyFragment(msg)) {
|
|
1399
|
-
return msg.id;
|
|
1400
|
-
}
|
|
1401
|
-
}
|
|
1402
|
-
if (this.#branch?.headMessageId) {
|
|
1403
|
-
const chain = await this.#store.getMessageChain(
|
|
1404
|
-
this.#branch.headMessageId
|
|
1405
|
-
);
|
|
1406
|
-
for (let i = chain.length - 1; i >= 0; i--) {
|
|
1407
|
-
if (chain[i].name === "assistant") {
|
|
1408
|
-
return chain[i].id;
|
|
1409
|
-
}
|
|
1410
|
-
}
|
|
1411
|
-
}
|
|
1412
|
-
return void 0;
|
|
1413
|
-
}
|
|
1414
|
-
/**
|
|
1415
|
-
* Estimate token count and cost for the full context.
|
|
1416
|
-
*
|
|
1417
|
-
* Includes:
|
|
1418
|
-
* - System prompt fragments (role, hints, etc.)
|
|
1419
|
-
* - Persisted chat messages (from store)
|
|
1420
|
-
* - Pending messages (not yet saved)
|
|
1421
|
-
*
|
|
1422
|
-
* @param modelId - Model ID (e.g., "openai:gpt-4o", "anthropic:claude-3-5-sonnet")
|
|
1423
|
-
* @param options - Optional settings
|
|
1424
|
-
* @returns Estimate result with token counts, costs, and per-fragment breakdown
|
|
1425
|
-
*/
|
|
1426
|
-
async estimate(modelId, options = {}) {
|
|
1427
|
-
await this.#ensureInitialized();
|
|
1428
|
-
const renderer = options.renderer ?? new XmlRenderer();
|
|
1429
|
-
const registry = getModelsRegistry();
|
|
1430
|
-
await registry.load();
|
|
1431
|
-
const model = registry.get(modelId);
|
|
1432
|
-
if (!model) {
|
|
1433
|
-
throw new Error(
|
|
1434
|
-
`Model "${modelId}" not found. Call load() first or check model ID.`
|
|
1435
|
-
);
|
|
1436
|
-
}
|
|
1437
|
-
const tokenizer = registry.getTokenizer(modelId);
|
|
1438
|
-
const fragmentEstimates = [];
|
|
1439
|
-
for (const fragment2 of this.#fragments) {
|
|
1440
|
-
const rendered = renderer.render([fragment2]);
|
|
1441
|
-
const tokens = tokenizer.count(rendered);
|
|
1442
|
-
const cost = tokens / 1e6 * model.cost.input;
|
|
1443
|
-
fragmentEstimates.push({
|
|
1444
|
-
id: fragment2.id,
|
|
1445
|
-
name: fragment2.name,
|
|
1446
|
-
tokens,
|
|
1447
|
-
cost
|
|
1448
|
-
});
|
|
1449
|
-
}
|
|
1450
|
-
if (this.#branch?.headMessageId) {
|
|
1451
|
-
const chain = await this.#store.getMessageChain(
|
|
1452
|
-
this.#branch.headMessageId
|
|
1453
|
-
);
|
|
1454
|
-
for (const msg of chain) {
|
|
1455
|
-
const content = String(msg.data);
|
|
1456
|
-
const tokens = tokenizer.count(content);
|
|
1457
|
-
const cost = tokens / 1e6 * model.cost.input;
|
|
1458
|
-
fragmentEstimates.push({
|
|
1459
|
-
name: msg.name,
|
|
1460
|
-
id: msg.id,
|
|
1461
|
-
tokens,
|
|
1462
|
-
cost
|
|
1463
|
-
});
|
|
1464
|
-
}
|
|
1465
|
-
}
|
|
1466
|
-
for (const fragment2 of this.#pendingMessages) {
|
|
1467
|
-
const content = String(fragment2.data);
|
|
1468
|
-
const tokens = tokenizer.count(content);
|
|
1469
|
-
const cost = tokens / 1e6 * model.cost.input;
|
|
1470
|
-
fragmentEstimates.push({
|
|
1471
|
-
name: fragment2.name,
|
|
1472
|
-
id: fragment2.id,
|
|
1473
|
-
tokens,
|
|
1474
|
-
cost
|
|
1475
|
-
});
|
|
1476
|
-
}
|
|
1477
|
-
const totalTokens = fragmentEstimates.reduce((sum, f) => sum + f.tokens, 0);
|
|
1478
|
-
const totalCost = fragmentEstimates.reduce((sum, f) => sum + f.cost, 0);
|
|
1479
|
-
return {
|
|
1480
|
-
model: model.id,
|
|
1481
|
-
provider: model.provider,
|
|
1482
|
-
tokens: totalTokens,
|
|
1483
|
-
cost: totalCost,
|
|
1484
|
-
limits: {
|
|
1485
|
-
context: model.limit.context,
|
|
1486
|
-
output: model.limit.output,
|
|
1487
|
-
exceedsContext: totalTokens > model.limit.context
|
|
1488
|
-
},
|
|
1489
|
-
fragments: fragmentEstimates
|
|
1490
|
-
};
|
|
1491
|
-
}
|
|
1492
|
-
/**
|
|
1493
|
-
* Rewind to a specific message by ID.
|
|
1494
|
-
*
|
|
1495
|
-
* Creates a new branch from that message, preserving the original branch.
|
|
1496
|
-
* The new branch becomes active.
|
|
1497
|
-
*
|
|
1498
|
-
* @param messageId - The message ID to rewind to
|
|
1499
|
-
* @returns The new branch info
|
|
1500
|
-
*
|
|
1501
|
-
* @example
|
|
1502
|
-
* ```ts
|
|
1503
|
-
* context.set(user('What is 2 + 2?', { id: 'q1' }));
|
|
1504
|
-
* context.set(assistant('The answer is 5.', { id: 'wrong' })); // Oops!
|
|
1505
|
-
* await context.save();
|
|
1506
|
-
*
|
|
1507
|
-
* // Rewind to the question, creates new branch
|
|
1508
|
-
* const newBranch = await context.rewind('q1');
|
|
1509
|
-
*
|
|
1510
|
-
* // Now add correct answer on new branch
|
|
1511
|
-
* context.set(assistant('The answer is 4.'));
|
|
1512
|
-
* await context.save();
|
|
1513
|
-
* ```
|
|
1514
|
-
*/
|
|
1515
|
-
async rewind(messageId) {
|
|
1516
|
-
await this.#ensureInitialized();
|
|
1517
|
-
const message2 = await this.#store.getMessage(messageId);
|
|
1518
|
-
if (!message2) {
|
|
1519
|
-
throw new Error(`Message "${messageId}" not found`);
|
|
1520
|
-
}
|
|
1521
|
-
if (message2.chatId !== this.#chatId) {
|
|
1522
|
-
throw new Error(`Message "${messageId}" belongs to a different chat`);
|
|
1523
|
-
}
|
|
1524
|
-
return this.#createBranchFrom(messageId, true);
|
|
1525
|
-
}
|
|
1526
|
-
/**
|
|
1527
|
-
* Create a checkpoint at the current position.
|
|
1528
|
-
*
|
|
1529
|
-
* A checkpoint is a named pointer to the current branch head.
|
|
1530
|
-
* Use restore() to return to this point later.
|
|
1531
|
-
*
|
|
1532
|
-
* @param name - Name for the checkpoint
|
|
1533
|
-
* @returns The checkpoint info
|
|
1534
|
-
*
|
|
1535
|
-
* @example
|
|
1536
|
-
* ```ts
|
|
1537
|
-
* context.set(user('I want to learn a new skill.'));
|
|
1538
|
-
* context.set(assistant('Would you like coding or cooking?'));
|
|
1539
|
-
* await context.save();
|
|
1540
|
-
*
|
|
1541
|
-
* // Save checkpoint before user's choice
|
|
1542
|
-
* const cp = await context.checkpoint('before-choice');
|
|
1543
|
-
* ```
|
|
1544
|
-
*/
|
|
1545
|
-
async checkpoint(name) {
|
|
1546
|
-
await this.#ensureInitialized();
|
|
1547
|
-
if (!this.#branch?.headMessageId) {
|
|
1548
|
-
throw new Error("Cannot create checkpoint: no messages in conversation");
|
|
1549
|
-
}
|
|
1550
|
-
const checkpoint = {
|
|
1551
|
-
id: crypto.randomUUID(),
|
|
1552
|
-
chatId: this.#chatId,
|
|
1553
|
-
name,
|
|
1554
|
-
messageId: this.#branch.headMessageId,
|
|
1555
|
-
createdAt: Date.now()
|
|
1556
|
-
};
|
|
1557
|
-
await this.#store.createCheckpoint(checkpoint);
|
|
1558
|
-
return {
|
|
1559
|
-
id: checkpoint.id,
|
|
1560
|
-
name: checkpoint.name,
|
|
1561
|
-
messageId: checkpoint.messageId,
|
|
1562
|
-
createdAt: checkpoint.createdAt
|
|
1563
|
-
};
|
|
1564
|
-
}
|
|
1565
|
-
/**
|
|
1566
|
-
* Restore to a checkpoint by creating a new branch from that point.
|
|
1567
|
-
*
|
|
1568
|
-
* @param name - Name of the checkpoint to restore
|
|
1569
|
-
* @returns The new branch info
|
|
1570
|
-
*
|
|
1571
|
-
* @example
|
|
1572
|
-
* ```ts
|
|
1573
|
-
* // User chose cooking, but wants to try coding path
|
|
1574
|
-
* await context.restore('before-choice');
|
|
1575
|
-
*
|
|
1576
|
-
* context.set(user('I want to learn coding.'));
|
|
1577
|
-
* context.set(assistant('Python is a great starting language!'));
|
|
1578
|
-
* await context.save();
|
|
1579
|
-
* ```
|
|
1580
|
-
*/
|
|
1581
|
-
async restore(name) {
|
|
1582
|
-
await this.#ensureInitialized();
|
|
1583
|
-
const checkpoint = await this.#store.getCheckpoint(this.#chatId, name);
|
|
1584
|
-
if (!checkpoint) {
|
|
1585
|
-
throw new Error(
|
|
1586
|
-
`Checkpoint "${name}" not found in chat "${this.#chatId}"`
|
|
1587
|
-
);
|
|
1588
|
-
}
|
|
1589
|
-
return this.rewind(checkpoint.messageId);
|
|
1590
|
-
}
|
|
1591
|
-
/**
|
|
1592
|
-
* Switch to a different branch by name.
|
|
1593
|
-
*
|
|
1594
|
-
* @param name - Branch name to switch to
|
|
1595
|
-
*
|
|
1596
|
-
* @example
|
|
1597
|
-
* ```ts
|
|
1598
|
-
* // List branches (via store)
|
|
1599
|
-
* const branches = await store.listBranches(context.chatId);
|
|
1600
|
-
* console.log(branches); // [{name: 'main', ...}, {name: 'main-v2', ...}]
|
|
1601
|
-
*
|
|
1602
|
-
* // Switch to original branch
|
|
1603
|
-
* await context.switchBranch('main');
|
|
1604
|
-
* ```
|
|
1605
|
-
*/
|
|
1606
|
-
async switchBranch(name) {
|
|
1607
|
-
await this.#ensureInitialized();
|
|
1608
|
-
const branch = await this.#store.getBranch(this.#chatId, name);
|
|
1609
|
-
if (!branch) {
|
|
1610
|
-
throw new Error(`Branch "${name}" not found in chat "${this.#chatId}"`);
|
|
1611
|
-
}
|
|
1612
|
-
await this.#store.setActiveBranch(this.#chatId, branch.id);
|
|
1613
|
-
this.#branch = { ...branch, isActive: true };
|
|
1614
|
-
this.#branchName = name;
|
|
1615
|
-
this.#pendingMessages = [];
|
|
1616
|
-
}
|
|
1617
|
-
/**
|
|
1618
|
-
* Create a parallel branch from the current position ("by the way").
|
|
1619
|
-
*
|
|
1620
|
-
* Use this when you want to fork the conversation without leaving
|
|
1621
|
-
* the current branch. Common use case: user wants to ask another
|
|
1622
|
-
* question while waiting for the model to respond.
|
|
1623
|
-
*
|
|
1624
|
-
* Unlike rewind(), this method:
|
|
1625
|
-
* - Uses the current HEAD (no messageId needed)
|
|
1626
|
-
* - Does NOT switch to the new branch
|
|
1627
|
-
* - Keeps pending messages intact
|
|
1628
|
-
*
|
|
1629
|
-
* @returns The new branch info (does not switch to it)
|
|
1630
|
-
* @throws Error if no messages exist in the conversation
|
|
1631
|
-
*
|
|
1632
|
-
* @example
|
|
1633
|
-
* ```ts
|
|
1634
|
-
* // User asked a question, model is generating...
|
|
1635
|
-
* context.set(user('What is the weather?'));
|
|
1636
|
-
* await context.save();
|
|
1637
|
-
*
|
|
1638
|
-
* // User wants to ask something else without waiting
|
|
1639
|
-
* const newBranch = await context.btw();
|
|
1640
|
-
* // newBranch = { name: 'main-v2', ... }
|
|
1641
|
-
*
|
|
1642
|
-
* // Later, switch to the new branch and add the question
|
|
1643
|
-
* await context.switchBranch(newBranch.name);
|
|
1644
|
-
* context.set(user('Also, what time is it?'));
|
|
1645
|
-
* await context.save();
|
|
1646
|
-
* ```
|
|
1647
|
-
*/
|
|
1648
|
-
async btw() {
|
|
1649
|
-
await this.#ensureInitialized();
|
|
1650
|
-
if (!this.#branch?.headMessageId) {
|
|
1651
|
-
throw new Error("Cannot create btw branch: no messages in conversation");
|
|
1652
|
-
}
|
|
1653
|
-
return this.#createBranchFrom(this.#branch.headMessageId, false);
|
|
1654
|
-
}
|
|
1655
|
-
/**
|
|
1656
|
-
* Update metadata for the current chat.
|
|
1657
|
-
*
|
|
1658
|
-
* @param updates - Partial metadata to merge (title, metadata)
|
|
1659
|
-
*
|
|
1660
|
-
* @example
|
|
1661
|
-
* ```ts
|
|
1662
|
-
* await context.updateChat({
|
|
1663
|
-
* title: 'Coding Help Session',
|
|
1664
|
-
* metadata: { tags: ['python', 'debugging'] }
|
|
1665
|
-
* });
|
|
1666
|
-
* ```
|
|
1667
|
-
*/
|
|
1668
|
-
async updateChat(updates) {
|
|
1669
|
-
await this.#ensureInitialized();
|
|
1670
|
-
const storeUpdates = {};
|
|
1671
|
-
if (updates.title !== void 0) {
|
|
1672
|
-
storeUpdates.title = updates.title;
|
|
1673
|
-
}
|
|
1674
|
-
if (updates.metadata !== void 0) {
|
|
1675
|
-
storeUpdates.metadata = {
|
|
1676
|
-
...this.#chatData?.metadata,
|
|
1677
|
-
...updates.metadata
|
|
1678
|
-
};
|
|
1679
|
-
}
|
|
1680
|
-
this.#chatData = await this.#store.updateChat(this.#chatId, storeUpdates);
|
|
1681
|
-
}
|
|
1682
|
-
/**
|
|
1683
|
-
* Track token usage for the current chat.
|
|
1684
|
-
* Accumulates usage metrics in chat.metadata.usage.
|
|
1685
|
-
*
|
|
1686
|
-
* @param usage - Token usage from AI SDK (LanguageModelUsage)
|
|
1687
|
-
*
|
|
1688
|
-
* @example
|
|
1689
|
-
* ```ts
|
|
1690
|
-
* // In onFinish callback
|
|
1691
|
-
* const usage = await result.totalUsage;
|
|
1692
|
-
* await context.trackUsage(usage);
|
|
1693
|
-
* ```
|
|
1694
|
-
*/
|
|
1695
|
-
async trackUsage(usage) {
|
|
1696
|
-
await this.#ensureInitialized();
|
|
1697
|
-
const freshChatData = await this.#store.getChat(this.#chatId);
|
|
1698
|
-
const currentUsage = freshChatData?.metadata?.usage ?? {};
|
|
1699
|
-
const updatedUsage = mergeWith(
|
|
1700
|
-
{},
|
|
1701
|
-
currentUsage,
|
|
1702
|
-
usage,
|
|
1703
|
-
(a, b) => typeof a === "number" || typeof b === "number" ? (a ?? 0) + (b ?? 0) : void 0
|
|
1704
|
-
);
|
|
1705
|
-
this.#chatData = await this.#store.updateChat(this.#chatId, {
|
|
1706
|
-
metadata: {
|
|
1707
|
-
...freshChatData?.metadata,
|
|
1708
|
-
usage: updatedUsage
|
|
1709
|
-
}
|
|
1710
|
-
});
|
|
1711
|
-
}
|
|
1712
|
-
/**
|
|
1713
|
-
* Consolidate context fragments (no-op for now).
|
|
1714
|
-
*
|
|
1715
|
-
* This is a placeholder for future functionality that merges context fragments
|
|
1716
|
-
* using specific rules. Currently, it does nothing.
|
|
1717
|
-
*
|
|
1718
|
-
* @experimental
|
|
1719
|
-
*/
|
|
1720
|
-
consolidate() {
|
|
1721
|
-
return void 0;
|
|
1722
|
-
}
|
|
1723
|
-
/**
|
|
1724
|
-
* Extract skill mounts from available_skills fragments.
|
|
1725
|
-
* Returns unified mount array where entries with `name` are individual skills.
|
|
1726
|
-
*
|
|
1727
|
-
* @example
|
|
1728
|
-
* ```ts
|
|
1729
|
-
* const context = new ContextEngine({ store, chatId, userId })
|
|
1730
|
-
* .set(skills({ paths: [{ host: './skills', sandbox: '/skills' }] }));
|
|
1731
|
-
*
|
|
1732
|
-
* const { mounts } = context.getSkillMounts();
|
|
1733
|
-
* // mounts: [{ name: 'bi-dashboards', host: './skills/bi-dashboards/SKILL.md', sandbox: '/skills/bi-dashboards/SKILL.md' }]
|
|
1734
|
-
*
|
|
1735
|
-
* // Extract skills only (entries with name)
|
|
1736
|
-
* const skills = mounts.filter(m => m.name);
|
|
1737
|
-
* ```
|
|
1738
|
-
*/
|
|
1739
|
-
getSkillMounts() {
|
|
1740
|
-
for (const fragment2 of this.#fragments) {
|
|
1741
|
-
if (fragment2.name === "available_skills" && fragment2.metadata?.mounts) {
|
|
1742
|
-
return { mounts: fragment2.metadata.mounts };
|
|
1743
|
-
}
|
|
1744
|
-
}
|
|
1745
|
-
return { mounts: [] };
|
|
1746
|
-
}
|
|
1747
|
-
/**
|
|
1748
|
-
* Inspect the full context state for debugging.
|
|
1749
|
-
* Returns a JSON-serializable object with context information.
|
|
1750
|
-
*
|
|
1751
|
-
* @param options - Inspection options (modelId and renderer required)
|
|
1752
|
-
* @returns Complete inspection data including estimates, rendered output, fragments, and graph
|
|
1753
|
-
*
|
|
1754
|
-
* @example
|
|
1755
|
-
* ```ts
|
|
1756
|
-
* const inspection = await context.inspect({
|
|
1757
|
-
* modelId: 'openai:gpt-4o',
|
|
1758
|
-
* renderer: new XmlRenderer(),
|
|
1759
|
-
* });
|
|
1760
|
-
* console.log(JSON.stringify(inspection, null, 2));
|
|
1761
|
-
*
|
|
1762
|
-
* // Or write to file for analysis
|
|
1763
|
-
* await fs.writeFile('context-debug.json', JSON.stringify(inspection, null, 2));
|
|
1764
|
-
* ```
|
|
1765
|
-
*/
|
|
1766
|
-
async inspect(options) {
|
|
1767
|
-
await this.#ensureInitialized();
|
|
1768
|
-
const { renderer } = options;
|
|
1769
|
-
const estimateResult = await this.estimate(options.modelId, { renderer });
|
|
1770
|
-
const rendered = renderer.render(this.#fragments);
|
|
1771
|
-
const persistedMessages = [];
|
|
1772
|
-
if (this.#branch?.headMessageId) {
|
|
1773
|
-
const chain = await this.#store.getMessageChain(
|
|
1774
|
-
this.#branch.headMessageId
|
|
1775
|
-
);
|
|
1776
|
-
persistedMessages.push(...chain);
|
|
1777
|
-
}
|
|
1778
|
-
const graph = await this.#store.getGraph(this.#chatId);
|
|
1779
|
-
return {
|
|
1780
|
-
estimate: estimateResult,
|
|
1781
|
-
rendered,
|
|
1782
|
-
fragments: {
|
|
1783
|
-
context: [...this.#fragments],
|
|
1784
|
-
pending: [...this.#pendingMessages],
|
|
1785
|
-
persisted: persistedMessages
|
|
1786
|
-
},
|
|
1787
|
-
graph,
|
|
1788
|
-
meta: {
|
|
1789
|
-
chatId: this.#chatId,
|
|
1790
|
-
branch: this.#branchName,
|
|
1791
|
-
timestamp: Date.now()
|
|
1792
|
-
}
|
|
1793
|
-
};
|
|
1794
|
-
}
|
|
1795
|
-
};
|
|
1796
|
-
function persona(input) {
|
|
1797
|
-
return {
|
|
1798
|
-
name: "persona",
|
|
1799
|
-
data: {
|
|
1800
|
-
name: input.name,
|
|
1801
|
-
...input.role && { role: input.role },
|
|
1802
|
-
...input.objective && { objective: input.objective },
|
|
1803
|
-
...input.tone && { tone: input.tone }
|
|
1804
|
-
}
|
|
1805
|
-
};
|
|
1806
|
-
}
|
|
1807
|
-
function pass(part) {
|
|
1808
|
-
return { type: "pass", part };
|
|
1809
|
-
}
|
|
1810
|
-
function runGuardrailChain(part, guardrails, context) {
|
|
1811
|
-
let currentPart = part;
|
|
1812
|
-
for (const guardrail2 of guardrails) {
|
|
1813
|
-
const result = guardrail2.handle(currentPart, context);
|
|
1814
|
-
if (result.type === "fail" || result.type === "stop") {
|
|
1815
|
-
return result;
|
|
1816
|
-
}
|
|
1817
|
-
currentPart = result.part;
|
|
1818
|
-
}
|
|
1819
|
-
return pass(currentPart);
|
|
1820
|
-
}
|
|
1821
|
-
var SKILLS_INSTRUCTIONS = dedent`A skill is a set of local instructions to follow that is stored in a \`SKILL.md\` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill.
|
|
1822
|
-
|
|
1823
|
-
### How to use skills
|
|
1824
|
-
- Discovery: The list below shows the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths.
|
|
1825
|
-
- Trigger rules: If the user names a skill (with \`$SkillName\` or plain text) OR the task clearly matches a skill's description shown below, you must use that skill for that turn before doing anything else. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned.
|
|
1826
|
-
- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback.
|
|
1827
|
-
- How to use a skill (progressive disclosure):
|
|
1828
|
-
1) After deciding to use a skill, open its \`SKILL.md\`. Read only enough to follow the workflow.
|
|
1829
|
-
2) If \`SKILL.md\` points to extra folders such as \`references/\`, load only the specific files needed for the request; don't bulk-load everything.
|
|
1830
|
-
3) If \`scripts/\` exist, prefer running or patching them instead of retyping large code blocks.
|
|
1831
|
-
4) If \`assets/\` or templates exist, reuse them instead of recreating from scratch.
|
|
1832
|
-
- Coordination and sequencing:
|
|
1833
|
-
- If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them.
|
|
1834
|
-
- Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why.
|
|
1835
|
-
- Context hygiene:
|
|
1836
|
-
- Keep context small: summarize long sections instead of pasting them; only load extra files when needed.
|
|
1837
|
-
- Avoid deep reference-chasing: prefer opening only files directly linked from \`SKILL.md\` unless you're blocked.
|
|
1838
|
-
- When variants exist (frameworks, providers, domains), pick only the relevant reference file(s) and note that choice.
|
|
1839
|
-
- Safety and fallback: If a skill can't be applied cleanly (missing files, unclear instructions), state the issue, pick the next-best approach, and continue.`;
|
|
1840
|
-
var ddl_sqlite_default = "-- Context Store DDL for SQLite\n-- This schema implements a DAG-based message history with branching and checkpoints.\n\n-- Performance PRAGMAs (session-level, run on each connection)\nPRAGMA journal_mode = WAL;\nPRAGMA synchronous = NORMAL;\nPRAGMA cache_size = -64000;\nPRAGMA temp_store = MEMORY;\nPRAGMA mmap_size = 268435456;\n\n-- Integrity\nPRAGMA foreign_keys = ON;\n\n-- Chats table\n-- createdAt/updatedAt: DEFAULT for insert, inline SET for updates\nCREATE TABLE IF NOT EXISTS chats (\n id TEXT PRIMARY KEY,\n userId TEXT NOT NULL,\n title TEXT,\n metadata TEXT,\n createdAt INTEGER NOT NULL DEFAULT (strftime('%s', 'now') * 1000),\n updatedAt INTEGER NOT NULL DEFAULT (strftime('%s', 'now') * 1000)\n);\n\nCREATE INDEX IF NOT EXISTS idx_chats_updatedAt ON chats(updatedAt);\nCREATE INDEX IF NOT EXISTS idx_chats_userId ON chats(userId);\n-- Composite index for listChats(): WHERE userId = ? ORDER BY updatedAt DESC\nCREATE INDEX IF NOT EXISTS idx_chats_userId_updatedAt ON chats(userId, updatedAt DESC);\n\n-- Messages table (nodes in the DAG)\nCREATE TABLE IF NOT EXISTS messages (\n id TEXT PRIMARY KEY,\n chatId TEXT NOT NULL,\n parentId TEXT,\n name TEXT NOT NULL,\n type TEXT,\n data TEXT NOT NULL,\n createdAt INTEGER NOT NULL,\n FOREIGN KEY (chatId) REFERENCES chats(id) ON DELETE CASCADE,\n FOREIGN KEY (parentId) REFERENCES messages(id)\n);\n\nCREATE INDEX IF NOT EXISTS idx_messages_chatId ON messages(chatId);\nCREATE INDEX IF NOT EXISTS idx_messages_parentId ON messages(parentId);\n-- Composite index for recursive CTE parent traversal in getMessageChain()\nCREATE INDEX IF NOT EXISTS idx_messages_chatId_parentId ON messages(chatId, parentId);\n\n-- Branches table (pointers to head messages)\nCREATE TABLE IF NOT EXISTS branches (\n id TEXT PRIMARY KEY,\n chatId TEXT NOT NULL,\n name TEXT NOT NULL,\n headMessageId TEXT,\n isActive INTEGER NOT NULL DEFAULT 0,\n createdAt INTEGER NOT NULL,\n FOREIGN KEY (chatId) REFERENCES chats(id) ON DELETE CASCADE,\n FOREIGN KEY (headMessageId) REFERENCES messages(id),\n UNIQUE(chatId, name)\n);\n\nCREATE INDEX IF NOT EXISTS idx_branches_chatId ON branches(chatId);\n-- Composite index for getActiveBranch(): WHERE chatId = ? AND isActive = 1\nCREATE INDEX IF NOT EXISTS idx_branches_chatId_isActive ON branches(chatId, isActive);\n\n-- Checkpoints table (pointers to message nodes)\nCREATE TABLE IF NOT EXISTS checkpoints (\n id TEXT PRIMARY KEY,\n chatId TEXT NOT NULL,\n name TEXT NOT NULL,\n messageId TEXT NOT NULL,\n createdAt INTEGER NOT NULL,\n FOREIGN KEY (chatId) REFERENCES chats(id) ON DELETE CASCADE,\n FOREIGN KEY (messageId) REFERENCES messages(id),\n UNIQUE(chatId, name)\n);\n\nCREATE INDEX IF NOT EXISTS idx_checkpoints_chatId ON checkpoints(chatId);\n\n-- FTS5 virtual table for full-text search\n-- messageId/chatId/name are UNINDEXED (stored but not searchable, used for filtering/joining)\n-- Only 'content' is indexed for full-text search\nCREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(\n messageId UNINDEXED,\n chatId UNINDEXED,\n name UNINDEXED,\n content,\n tokenize='porter unicode61'\n);\n";
|
|
1841
|
-
var SqliteContextStore = class extends ContextStore {
|
|
1842
|
-
#db;
|
|
1843
|
-
#statements = /* @__PURE__ */ new Map();
|
|
1844
|
-
/**
|
|
1845
|
-
* Get or create a prepared statement.
|
|
1846
|
-
* Statements are cached for the lifetime of the store to avoid
|
|
1847
|
-
* repeated SQL parsing and compilation overhead.
|
|
1848
|
-
*/
|
|
1849
|
-
#stmt(sql) {
|
|
1850
|
-
let stmt = this.#statements.get(sql);
|
|
1851
|
-
if (!stmt) {
|
|
1852
|
-
stmt = this.#db.prepare(sql);
|
|
1853
|
-
this.#statements.set(sql, stmt);
|
|
1854
|
-
}
|
|
1855
|
-
return stmt;
|
|
1856
|
-
}
|
|
1857
|
-
constructor(path3) {
|
|
1858
|
-
super();
|
|
1859
|
-
this.#db = new DatabaseSync(path3);
|
|
1860
|
-
this.#db.exec(ddl_sqlite_default);
|
|
1861
|
-
}
|
|
1862
|
-
/**
|
|
1863
|
-
* Execute a function within a transaction.
|
|
1864
|
-
* Automatically commits on success or rolls back on error.
|
|
1865
|
-
*/
|
|
1866
|
-
#useTransaction(fn) {
|
|
1867
|
-
this.#db.exec("BEGIN TRANSACTION");
|
|
1868
|
-
try {
|
|
1869
|
-
const result = fn();
|
|
1870
|
-
this.#db.exec("COMMIT");
|
|
1871
|
-
return result;
|
|
1872
|
-
} catch (error) {
|
|
1873
|
-
this.#db.exec("ROLLBACK");
|
|
1874
|
-
throw error;
|
|
1875
|
-
}
|
|
1876
|
-
}
|
|
1877
|
-
// ==========================================================================
|
|
1878
|
-
// Chat Operations
|
|
1879
|
-
// ==========================================================================
|
|
1880
|
-
async createChat(chat) {
|
|
1881
|
-
return this.#useTransaction(() => {
|
|
1882
|
-
const row = this.#db.prepare(
|
|
1883
|
-
`INSERT INTO chats (id, userId, title, metadata)
|
|
1884
|
-
VALUES (?, ?, ?, ?)
|
|
1885
|
-
RETURNING *`
|
|
1886
|
-
).get(
|
|
1887
|
-
chat.id,
|
|
1888
|
-
chat.userId,
|
|
1889
|
-
chat.title ?? null,
|
|
1890
|
-
chat.metadata ? JSON.stringify(chat.metadata) : null
|
|
1891
|
-
);
|
|
1892
|
-
this.#db.prepare(
|
|
1893
|
-
`INSERT INTO branches (id, chatId, name, headMessageId, isActive, createdAt)
|
|
1894
|
-
VALUES (?, ?, 'main', NULL, 1, ?)`
|
|
1895
|
-
).run(crypto.randomUUID(), chat.id, Date.now());
|
|
1896
|
-
return {
|
|
1897
|
-
id: row.id,
|
|
1898
|
-
userId: row.userId,
|
|
1899
|
-
title: row.title ?? void 0,
|
|
1900
|
-
metadata: row.metadata ? JSON.parse(row.metadata) : void 0,
|
|
1901
|
-
createdAt: row.createdAt,
|
|
1902
|
-
updatedAt: row.updatedAt
|
|
1903
|
-
};
|
|
1904
|
-
});
|
|
1905
|
-
}
|
|
1906
|
-
async upsertChat(chat) {
|
|
1907
|
-
return this.#useTransaction(() => {
|
|
1908
|
-
const row = this.#db.prepare(
|
|
1909
|
-
`INSERT INTO chats (id, userId, title, metadata)
|
|
1910
|
-
VALUES (?, ?, ?, ?)
|
|
1911
|
-
ON CONFLICT(id) DO UPDATE SET id = excluded.id
|
|
1912
|
-
RETURNING *`
|
|
1913
|
-
).get(
|
|
1914
|
-
chat.id,
|
|
1915
|
-
chat.userId,
|
|
1916
|
-
chat.title ?? null,
|
|
1917
|
-
chat.metadata ? JSON.stringify(chat.metadata) : null
|
|
1918
|
-
);
|
|
1919
|
-
this.#db.prepare(
|
|
1920
|
-
`INSERT OR IGNORE INTO branches (id, chatId, name, headMessageId, isActive, createdAt)
|
|
1921
|
-
VALUES (?, ?, 'main', NULL, 1, ?)`
|
|
1922
|
-
).run(crypto.randomUUID(), chat.id, Date.now());
|
|
1923
|
-
return {
|
|
1924
|
-
id: row.id,
|
|
1925
|
-
userId: row.userId,
|
|
1926
|
-
title: row.title ?? void 0,
|
|
1927
|
-
metadata: row.metadata ? JSON.parse(row.metadata) : void 0,
|
|
1928
|
-
createdAt: row.createdAt,
|
|
1929
|
-
updatedAt: row.updatedAt
|
|
1930
|
-
};
|
|
1931
|
-
});
|
|
1932
|
-
}
|
|
1933
|
-
async getChat(chatId) {
|
|
1934
|
-
const row = this.#db.prepare("SELECT * FROM chats WHERE id = ?").get(chatId);
|
|
1935
|
-
if (!row) {
|
|
1936
|
-
return void 0;
|
|
1937
|
-
}
|
|
1938
|
-
return {
|
|
1939
|
-
id: row.id,
|
|
1940
|
-
userId: row.userId,
|
|
1941
|
-
title: row.title ?? void 0,
|
|
1942
|
-
metadata: row.metadata ? JSON.parse(row.metadata) : void 0,
|
|
1943
|
-
createdAt: row.createdAt,
|
|
1944
|
-
updatedAt: row.updatedAt
|
|
1945
|
-
};
|
|
1946
|
-
}
|
|
1947
|
-
async updateChat(chatId, updates) {
|
|
1948
|
-
const setClauses = ["updatedAt = strftime('%s', 'now') * 1000"];
|
|
1949
|
-
const params = [];
|
|
1950
|
-
if (updates.title !== void 0) {
|
|
1951
|
-
setClauses.push("title = ?");
|
|
1952
|
-
params.push(updates.title ?? null);
|
|
1953
|
-
}
|
|
1954
|
-
if (updates.metadata !== void 0) {
|
|
1955
|
-
setClauses.push("metadata = ?");
|
|
1956
|
-
params.push(JSON.stringify(updates.metadata));
|
|
1957
|
-
}
|
|
1958
|
-
params.push(chatId);
|
|
1959
|
-
const row = this.#db.prepare(
|
|
1960
|
-
`UPDATE chats SET ${setClauses.join(", ")} WHERE id = ? RETURNING *`
|
|
1961
|
-
).get(...params);
|
|
1962
|
-
return {
|
|
1963
|
-
id: row.id,
|
|
1964
|
-
userId: row.userId,
|
|
1965
|
-
title: row.title ?? void 0,
|
|
1966
|
-
metadata: row.metadata ? JSON.parse(row.metadata) : void 0,
|
|
1967
|
-
createdAt: row.createdAt,
|
|
1968
|
-
updatedAt: row.updatedAt
|
|
1969
|
-
};
|
|
1970
|
-
}
|
|
1971
|
-
async listChats(options) {
|
|
1972
|
-
const params = [];
|
|
1973
|
-
const whereClauses = [];
|
|
1974
|
-
let limitClause = "";
|
|
1975
|
-
if (options?.userId) {
|
|
1976
|
-
whereClauses.push("c.userId = ?");
|
|
1977
|
-
params.push(options.userId);
|
|
1978
|
-
}
|
|
1979
|
-
if (options?.metadata) {
|
|
1980
|
-
whereClauses.push(`json_extract(c.metadata, '$.' || ?) = ?`);
|
|
1981
|
-
params.push(options.metadata.key);
|
|
1982
|
-
params.push(
|
|
1983
|
-
typeof options.metadata.value === "boolean" ? options.metadata.value ? 1 : 0 : options.metadata.value
|
|
1984
|
-
);
|
|
1985
|
-
}
|
|
1986
|
-
const whereClause = whereClauses.length > 0 ? `WHERE ${whereClauses.join(" AND ")}` : "";
|
|
1987
|
-
if (options?.limit !== void 0) {
|
|
1988
|
-
limitClause = " LIMIT ?";
|
|
1989
|
-
params.push(options.limit);
|
|
1990
|
-
if (options.offset !== void 0) {
|
|
1991
|
-
limitClause += " OFFSET ?";
|
|
1992
|
-
params.push(options.offset);
|
|
1993
|
-
}
|
|
1994
|
-
}
|
|
1995
|
-
const rows = this.#db.prepare(
|
|
1996
|
-
`SELECT
|
|
1997
|
-
c.id,
|
|
1998
|
-
c.userId,
|
|
1999
|
-
c.title,
|
|
2000
|
-
c.metadata,
|
|
2001
|
-
c.createdAt,
|
|
2002
|
-
c.updatedAt,
|
|
2003
|
-
COUNT(DISTINCT m.id) as messageCount,
|
|
2004
|
-
COUNT(DISTINCT b.id) as branchCount
|
|
2005
|
-
FROM chats c
|
|
2006
|
-
LEFT JOIN messages m ON m.chatId = c.id
|
|
2007
|
-
LEFT JOIN branches b ON b.chatId = c.id
|
|
2008
|
-
${whereClause}
|
|
2009
|
-
GROUP BY c.id
|
|
2010
|
-
ORDER BY c.updatedAt DESC${limitClause}`
|
|
2011
|
-
).all(...params);
|
|
2012
|
-
return rows.map((row) => ({
|
|
2013
|
-
id: row.id,
|
|
2014
|
-
userId: row.userId,
|
|
2015
|
-
title: row.title ?? void 0,
|
|
2016
|
-
metadata: row.metadata ? JSON.parse(row.metadata) : void 0,
|
|
2017
|
-
messageCount: row.messageCount,
|
|
2018
|
-
branchCount: row.branchCount,
|
|
2019
|
-
createdAt: row.createdAt,
|
|
2020
|
-
updatedAt: row.updatedAt
|
|
2021
|
-
}));
|
|
2022
|
-
}
|
|
2023
|
-
async deleteChat(chatId, options) {
|
|
2024
|
-
return this.#useTransaction(() => {
|
|
2025
|
-
const messageIds = this.#db.prepare("SELECT id FROM messages WHERE chatId = ?").all(chatId);
|
|
2026
|
-
let sql = "DELETE FROM chats WHERE id = ?";
|
|
2027
|
-
const params = [chatId];
|
|
2028
|
-
if (options?.userId !== void 0) {
|
|
2029
|
-
sql += " AND userId = ?";
|
|
2030
|
-
params.push(options.userId);
|
|
2031
|
-
}
|
|
2032
|
-
const result = this.#db.prepare(sql).run(...params);
|
|
2033
|
-
if (result.changes > 0 && messageIds.length > 0) {
|
|
2034
|
-
const placeholders = messageIds.map(() => "?").join(", ");
|
|
2035
|
-
this.#db.prepare(
|
|
2036
|
-
`DELETE FROM messages_fts WHERE messageId IN (${placeholders})`
|
|
2037
|
-
).run(...messageIds.map((m) => m.id));
|
|
2038
|
-
}
|
|
2039
|
-
return result.changes > 0;
|
|
2040
|
-
});
|
|
2041
|
-
}
|
|
2042
|
-
// ==========================================================================
|
|
2043
|
-
// Message Operations (Graph Nodes)
|
|
2044
|
-
// ==========================================================================
|
|
2045
|
-
async addMessage(message2) {
|
|
2046
|
-
if (message2.parentId === message2.id) {
|
|
2047
|
-
throw new Error(`Message ${message2.id} cannot be its own parent`);
|
|
2048
|
-
}
|
|
2049
|
-
this.#stmt(
|
|
2050
|
-
`INSERT INTO messages (id, chatId, parentId, name, type, data, createdAt)
|
|
2051
|
-
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
2052
|
-
ON CONFLICT(id) DO UPDATE SET
|
|
2053
|
-
name = excluded.name,
|
|
2054
|
-
type = excluded.type,
|
|
2055
|
-
data = excluded.data`
|
|
2056
|
-
).run(
|
|
2057
|
-
message2.id,
|
|
2058
|
-
message2.chatId,
|
|
2059
|
-
message2.parentId,
|
|
2060
|
-
message2.name,
|
|
2061
|
-
message2.type ?? null,
|
|
2062
|
-
JSON.stringify(message2.data),
|
|
2063
|
-
message2.createdAt
|
|
2064
|
-
);
|
|
2065
|
-
const content = typeof message2.data === "string" ? message2.data : JSON.stringify(message2.data);
|
|
2066
|
-
this.#stmt(`DELETE FROM messages_fts WHERE messageId = ?`).run(message2.id);
|
|
2067
|
-
this.#stmt(
|
|
2068
|
-
`INSERT INTO messages_fts(messageId, chatId, name, content)
|
|
2069
|
-
VALUES (?, ?, ?, ?)`
|
|
2070
|
-
).run(message2.id, message2.chatId, message2.name, content);
|
|
2071
|
-
}
|
|
2072
|
-
async getMessage(messageId) {
|
|
2073
|
-
const row = this.#stmt("SELECT * FROM messages WHERE id = ?").get(
|
|
2074
|
-
messageId
|
|
2075
|
-
);
|
|
2076
|
-
if (!row) {
|
|
2077
|
-
return void 0;
|
|
2078
|
-
}
|
|
2079
|
-
return {
|
|
2080
|
-
id: row.id,
|
|
2081
|
-
chatId: row.chatId,
|
|
2082
|
-
parentId: row.parentId,
|
|
2083
|
-
name: row.name,
|
|
2084
|
-
type: row.type ?? void 0,
|
|
2085
|
-
data: JSON.parse(row.data),
|
|
2086
|
-
createdAt: row.createdAt
|
|
2087
|
-
};
|
|
2088
|
-
}
|
|
2089
|
-
async getMessageChain(headId) {
|
|
2090
|
-
const rows = this.#stmt(
|
|
2091
|
-
`WITH RECURSIVE chain AS (
|
|
2092
|
-
SELECT *, 0 as depth FROM messages WHERE id = ?
|
|
2093
|
-
UNION ALL
|
|
2094
|
-
SELECT m.*, c.depth + 1 FROM messages m
|
|
2095
|
-
INNER JOIN chain c ON m.id = c.parentId
|
|
2096
|
-
WHERE c.depth < 100000
|
|
2097
|
-
)
|
|
2098
|
-
SELECT * FROM chain
|
|
2099
|
-
ORDER BY depth DESC`
|
|
2100
|
-
).all(headId);
|
|
2101
|
-
return rows.map((row) => ({
|
|
2102
|
-
id: row.id,
|
|
2103
|
-
chatId: row.chatId,
|
|
2104
|
-
parentId: row.parentId,
|
|
2105
|
-
name: row.name,
|
|
2106
|
-
type: row.type ?? void 0,
|
|
2107
|
-
data: JSON.parse(row.data),
|
|
2108
|
-
createdAt: row.createdAt
|
|
2109
|
-
}));
|
|
2110
|
-
}
|
|
2111
|
-
async hasChildren(messageId) {
|
|
2112
|
-
const row = this.#stmt(
|
|
2113
|
-
"SELECT EXISTS(SELECT 1 FROM messages WHERE parentId = ?) as hasChildren"
|
|
2114
|
-
).get(messageId);
|
|
2115
|
-
return row.hasChildren === 1;
|
|
2116
|
-
}
|
|
2117
|
-
async getMessages(chatId) {
|
|
2118
|
-
const chat = await this.getChat(chatId);
|
|
2119
|
-
if (!chat) {
|
|
2120
|
-
throw new Error(`Chat "${chatId}" not found`);
|
|
2121
|
-
}
|
|
2122
|
-
const activeBranch = await this.getActiveBranch(chatId);
|
|
2123
|
-
if (!activeBranch?.headMessageId) {
|
|
2124
|
-
return [];
|
|
2125
|
-
}
|
|
2126
|
-
return this.getMessageChain(activeBranch.headMessageId);
|
|
2127
|
-
}
|
|
2128
|
-
// ==========================================================================
|
|
2129
|
-
// Branch Operations
|
|
2130
|
-
// ==========================================================================
|
|
2131
|
-
async createBranch(branch) {
|
|
2132
|
-
this.#db.prepare(
|
|
2133
|
-
`INSERT INTO branches (id, chatId, name, headMessageId, isActive, createdAt)
|
|
2134
|
-
VALUES (?, ?, ?, ?, ?, ?)`
|
|
2135
|
-
).run(
|
|
2136
|
-
branch.id,
|
|
2137
|
-
branch.chatId,
|
|
2138
|
-
branch.name,
|
|
2139
|
-
branch.headMessageId,
|
|
2140
|
-
branch.isActive ? 1 : 0,
|
|
2141
|
-
branch.createdAt
|
|
2142
|
-
);
|
|
2143
|
-
}
|
|
2144
|
-
async getBranch(chatId, name) {
|
|
2145
|
-
const row = this.#db.prepare("SELECT * FROM branches WHERE chatId = ? AND name = ?").get(chatId, name);
|
|
2146
|
-
if (!row) {
|
|
2147
|
-
return void 0;
|
|
2148
|
-
}
|
|
2149
|
-
return {
|
|
2150
|
-
id: row.id,
|
|
2151
|
-
chatId: row.chatId,
|
|
2152
|
-
name: row.name,
|
|
2153
|
-
headMessageId: row.headMessageId,
|
|
2154
|
-
isActive: row.isActive === 1,
|
|
2155
|
-
createdAt: row.createdAt
|
|
2156
|
-
};
|
|
2157
|
-
}
|
|
2158
|
-
async getActiveBranch(chatId) {
|
|
2159
|
-
const row = this.#stmt(
|
|
2160
|
-
"SELECT * FROM branches WHERE chatId = ? AND isActive = 1"
|
|
2161
|
-
).get(chatId);
|
|
2162
|
-
if (!row) {
|
|
2163
|
-
return void 0;
|
|
2164
|
-
}
|
|
2165
|
-
return {
|
|
2166
|
-
id: row.id,
|
|
2167
|
-
chatId: row.chatId,
|
|
2168
|
-
name: row.name,
|
|
2169
|
-
headMessageId: row.headMessageId,
|
|
2170
|
-
isActive: true,
|
|
2171
|
-
createdAt: row.createdAt
|
|
2172
|
-
};
|
|
2173
|
-
}
|
|
2174
|
-
async setActiveBranch(chatId, branchId) {
|
|
2175
|
-
this.#db.prepare("UPDATE branches SET isActive = 0 WHERE chatId = ?").run(chatId);
|
|
2176
|
-
this.#db.prepare("UPDATE branches SET isActive = 1 WHERE id = ?").run(branchId);
|
|
2177
|
-
}
|
|
2178
|
-
async updateBranchHead(branchId, messageId) {
|
|
2179
|
-
this.#stmt("UPDATE branches SET headMessageId = ? WHERE id = ?").run(
|
|
2180
|
-
messageId,
|
|
2181
|
-
branchId
|
|
2182
|
-
);
|
|
2183
|
-
}
|
|
2184
|
-
async listBranches(chatId) {
|
|
2185
|
-
const rows = this.#db.prepare(
|
|
2186
|
-
`SELECT
|
|
2187
|
-
b.id,
|
|
2188
|
-
b.name,
|
|
2189
|
-
b.headMessageId,
|
|
2190
|
-
b.isActive,
|
|
2191
|
-
b.createdAt,
|
|
2192
|
-
COALESCE(
|
|
2193
|
-
(
|
|
2194
|
-
WITH RECURSIVE chain AS (
|
|
2195
|
-
SELECT id, parentId FROM messages WHERE id = b.headMessageId
|
|
2196
|
-
UNION ALL
|
|
2197
|
-
SELECT m.id, m.parentId FROM messages m
|
|
2198
|
-
INNER JOIN chain c ON m.id = c.parentId
|
|
2199
|
-
)
|
|
2200
|
-
SELECT COUNT(*) FROM chain
|
|
2201
|
-
),
|
|
2202
|
-
0
|
|
2203
|
-
) as messageCount
|
|
2204
|
-
FROM branches b
|
|
2205
|
-
WHERE b.chatId = ?
|
|
2206
|
-
ORDER BY b.createdAt ASC`
|
|
2207
|
-
).all(chatId);
|
|
2208
|
-
return rows.map((row) => ({
|
|
2209
|
-
id: row.id,
|
|
2210
|
-
name: row.name,
|
|
2211
|
-
headMessageId: row.headMessageId,
|
|
2212
|
-
isActive: row.isActive === 1,
|
|
2213
|
-
messageCount: row.messageCount,
|
|
2214
|
-
createdAt: row.createdAt
|
|
2215
|
-
}));
|
|
2216
|
-
}
|
|
2217
|
-
// ==========================================================================
|
|
2218
|
-
// Checkpoint Operations
|
|
2219
|
-
// ==========================================================================
|
|
2220
|
-
async createCheckpoint(checkpoint) {
|
|
2221
|
-
this.#db.prepare(
|
|
2222
|
-
`INSERT INTO checkpoints (id, chatId, name, messageId, createdAt)
|
|
2223
|
-
VALUES (?, ?, ?, ?, ?)
|
|
2224
|
-
ON CONFLICT(chatId, name) DO UPDATE SET
|
|
2225
|
-
messageId = excluded.messageId,
|
|
2226
|
-
createdAt = excluded.createdAt`
|
|
2227
|
-
).run(
|
|
2228
|
-
checkpoint.id,
|
|
2229
|
-
checkpoint.chatId,
|
|
2230
|
-
checkpoint.name,
|
|
2231
|
-
checkpoint.messageId,
|
|
2232
|
-
checkpoint.createdAt
|
|
2233
|
-
);
|
|
2234
|
-
}
|
|
2235
|
-
async getCheckpoint(chatId, name) {
|
|
2236
|
-
const row = this.#db.prepare("SELECT * FROM checkpoints WHERE chatId = ? AND name = ?").get(chatId, name);
|
|
2237
|
-
if (!row) {
|
|
2238
|
-
return void 0;
|
|
2239
|
-
}
|
|
2240
|
-
return {
|
|
2241
|
-
id: row.id,
|
|
2242
|
-
chatId: row.chatId,
|
|
2243
|
-
name: row.name,
|
|
2244
|
-
messageId: row.messageId,
|
|
2245
|
-
createdAt: row.createdAt
|
|
2246
|
-
};
|
|
2247
|
-
}
|
|
2248
|
-
async listCheckpoints(chatId) {
|
|
2249
|
-
const rows = this.#db.prepare(
|
|
2250
|
-
`SELECT id, name, messageId, createdAt
|
|
2251
|
-
FROM checkpoints
|
|
2252
|
-
WHERE chatId = ?
|
|
2253
|
-
ORDER BY createdAt DESC`
|
|
2254
|
-
).all(chatId);
|
|
2255
|
-
return rows.map((row) => ({
|
|
2256
|
-
id: row.id,
|
|
2257
|
-
name: row.name,
|
|
2258
|
-
messageId: row.messageId,
|
|
2259
|
-
createdAt: row.createdAt
|
|
2260
|
-
}));
|
|
2261
|
-
}
|
|
2262
|
-
async deleteCheckpoint(chatId, name) {
|
|
2263
|
-
this.#db.prepare("DELETE FROM checkpoints WHERE chatId = ? AND name = ?").run(chatId, name);
|
|
2264
|
-
}
|
|
2265
|
-
// ==========================================================================
|
|
2266
|
-
// Search Operations
|
|
2267
|
-
// ==========================================================================
|
|
2268
|
-
async searchMessages(chatId, query, options) {
|
|
2269
|
-
const limit = options?.limit ?? 20;
|
|
2270
|
-
const roles = options?.roles;
|
|
2271
|
-
let sql = `
|
|
2272
|
-
SELECT
|
|
2273
|
-
m.id,
|
|
2274
|
-
m.chatId,
|
|
2275
|
-
m.parentId,
|
|
2276
|
-
m.name,
|
|
2277
|
-
m.type,
|
|
2278
|
-
m.data,
|
|
2279
|
-
m.createdAt,
|
|
2280
|
-
fts.rank,
|
|
2281
|
-
snippet(messages_fts, 3, '<mark>', '</mark>', '...', 32) as snippet
|
|
2282
|
-
FROM messages_fts fts
|
|
2283
|
-
JOIN messages m ON m.id = fts.messageId
|
|
2284
|
-
WHERE messages_fts MATCH ?
|
|
2285
|
-
AND fts.chatId = ?
|
|
2286
|
-
`;
|
|
2287
|
-
const params = [query, chatId];
|
|
2288
|
-
if (roles && roles.length > 0) {
|
|
2289
|
-
const placeholders = roles.map(() => "?").join(", ");
|
|
2290
|
-
sql += ` AND fts.name IN (${placeholders})`;
|
|
2291
|
-
params.push(...roles);
|
|
2292
|
-
}
|
|
2293
|
-
sql += " ORDER BY fts.rank LIMIT ?";
|
|
2294
|
-
params.push(limit);
|
|
2295
|
-
const rows = this.#db.prepare(sql).all(...params);
|
|
2296
|
-
return rows.map((row) => ({
|
|
2297
|
-
message: {
|
|
2298
|
-
id: row.id,
|
|
2299
|
-
chatId: row.chatId,
|
|
2300
|
-
parentId: row.parentId,
|
|
2301
|
-
name: row.name,
|
|
2302
|
-
type: row.type ?? void 0,
|
|
2303
|
-
data: JSON.parse(row.data),
|
|
2304
|
-
createdAt: row.createdAt
|
|
2305
|
-
},
|
|
2306
|
-
rank: row.rank,
|
|
2307
|
-
snippet: row.snippet
|
|
2308
|
-
}));
|
|
2309
|
-
}
|
|
2310
|
-
// ==========================================================================
|
|
2311
|
-
// Visualization Operations
|
|
2312
|
-
// ==========================================================================
|
|
2313
|
-
async getGraph(chatId) {
|
|
2314
|
-
const messageRows = this.#db.prepare(
|
|
2315
|
-
`SELECT id, parentId, name, data, createdAt
|
|
2316
|
-
FROM messages
|
|
2317
|
-
WHERE chatId = ?
|
|
2318
|
-
ORDER BY createdAt ASC`
|
|
2319
|
-
).all(chatId);
|
|
2320
|
-
const nodes = messageRows.map((row) => {
|
|
2321
|
-
const data = JSON.parse(row.data);
|
|
2322
|
-
const content = typeof data === "string" ? data : JSON.stringify(data);
|
|
2323
|
-
return {
|
|
2324
|
-
id: row.id,
|
|
2325
|
-
parentId: row.parentId,
|
|
2326
|
-
role: row.name,
|
|
2327
|
-
content: content.length > 50 ? content.slice(0, 50) + "..." : content,
|
|
2328
|
-
createdAt: row.createdAt
|
|
2329
|
-
};
|
|
2330
|
-
});
|
|
2331
|
-
const branchRows = this.#db.prepare(
|
|
2332
|
-
`SELECT name, headMessageId, isActive
|
|
2333
|
-
FROM branches
|
|
2334
|
-
WHERE chatId = ?
|
|
2335
|
-
ORDER BY createdAt ASC`
|
|
2336
|
-
).all(chatId);
|
|
2337
|
-
const branches = branchRows.map((row) => ({
|
|
2338
|
-
name: row.name,
|
|
2339
|
-
headMessageId: row.headMessageId,
|
|
2340
|
-
isActive: row.isActive === 1
|
|
2341
|
-
}));
|
|
2342
|
-
const checkpointRows = this.#db.prepare(
|
|
2343
|
-
`SELECT name, messageId
|
|
2344
|
-
FROM checkpoints
|
|
2345
|
-
WHERE chatId = ?
|
|
2346
|
-
ORDER BY createdAt ASC`
|
|
2347
|
-
).all(chatId);
|
|
2348
|
-
const checkpoints = checkpointRows.map((row) => ({
|
|
2349
|
-
name: row.name,
|
|
2350
|
-
messageId: row.messageId
|
|
2351
|
-
}));
|
|
2352
|
-
return {
|
|
2353
|
-
chatId,
|
|
2354
|
-
nodes,
|
|
2355
|
-
branches,
|
|
2356
|
-
checkpoints
|
|
2357
|
-
};
|
|
2358
|
-
}
|
|
2359
|
-
};
|
|
2360
|
-
var InMemoryContextStore = class extends SqliteContextStore {
|
|
2361
|
-
constructor() {
|
|
2362
|
-
super(":memory:");
|
|
2363
|
-
}
|
|
2364
|
-
};
|
|
2365
|
-
var Agent = class _Agent {
|
|
2366
|
-
#options;
|
|
2367
|
-
#guardrails = [];
|
|
2368
|
-
tools;
|
|
2369
|
-
constructor(options) {
|
|
2370
|
-
this.#options = options;
|
|
2371
|
-
this.tools = options.tools || {};
|
|
2372
|
-
this.#guardrails = options.guardrails || [];
|
|
2373
|
-
}
|
|
2374
|
-
async generate(contextVariables, config) {
|
|
2375
|
-
if (!this.#options.context) {
|
|
2376
|
-
throw new Error(`Agent ${this.#options.name} is missing a context.`);
|
|
2377
|
-
}
|
|
2378
|
-
if (!this.#options.model) {
|
|
2379
|
-
throw new Error(`Agent ${this.#options.name} is missing a model.`);
|
|
2380
|
-
}
|
|
2381
|
-
const { messages, systemPrompt } = await this.#options.context.resolve({
|
|
2382
|
-
renderer: new XmlRenderer()
|
|
2383
|
-
});
|
|
2384
|
-
return generateText({
|
|
2385
|
-
abortSignal: config?.abortSignal,
|
|
2386
|
-
providerOptions: this.#options.providerOptions,
|
|
2387
|
-
model: this.#options.model,
|
|
2388
|
-
system: systemPrompt,
|
|
2389
|
-
messages: await convertToModelMessages(messages),
|
|
2390
|
-
stopWhen: stepCountIs(25),
|
|
2391
|
-
tools: this.#options.tools,
|
|
2392
|
-
experimental_context: contextVariables,
|
|
2393
|
-
experimental_repairToolCall: repairToolCall,
|
|
2394
|
-
toolChoice: this.#options.toolChoice,
|
|
2395
|
-
onStepFinish: (step) => {
|
|
2396
|
-
const toolCall = step.toolCalls.at(-1);
|
|
2397
|
-
if (toolCall) {
|
|
2398
|
-
console.log(
|
|
2399
|
-
`Debug: ${chalk2.yellow("ToolCalled")}: ${toolCall.toolName}(${JSON.stringify(toolCall.input)})`
|
|
2400
|
-
);
|
|
2401
|
-
}
|
|
2402
|
-
}
|
|
2403
|
-
});
|
|
2404
|
-
}
|
|
2405
|
-
/**
|
|
2406
|
-
* Stream a response from the agent.
|
|
2407
|
-
*
|
|
2408
|
-
* When guardrails are configured, `toUIMessageStream()` is wrapped to provide
|
|
2409
|
-
* self-correction behavior. Direct access to fullStream/textStream bypasses guardrails.
|
|
2410
|
-
*
|
|
2411
|
-
* @example
|
|
2412
|
-
* ```typescript
|
|
2413
|
-
* const stream = await agent.stream({});
|
|
2414
|
-
*
|
|
2415
|
-
* // With guardrails - use toUIMessageStream for protection
|
|
2416
|
-
* await printer.readableStream(stream.toUIMessageStream());
|
|
2417
|
-
*
|
|
2418
|
-
* // Or use printer.stdout which uses toUIMessageStream internally
|
|
2419
|
-
* await printer.stdout(stream);
|
|
2420
|
-
* ```
|
|
2421
|
-
*/
|
|
2422
|
-
async stream(contextVariables, config) {
|
|
2423
|
-
if (!this.#options.context) {
|
|
2424
|
-
throw new Error(`Agent ${this.#options.name} is missing a context.`);
|
|
2425
|
-
}
|
|
2426
|
-
if (!this.#options.model) {
|
|
2427
|
-
throw new Error(`Agent ${this.#options.name} is missing a model.`);
|
|
2428
|
-
}
|
|
2429
|
-
const result = await this.#createRawStream(contextVariables, config);
|
|
2430
|
-
if (this.#guardrails.length === 0) {
|
|
2431
|
-
return result;
|
|
2432
|
-
}
|
|
2433
|
-
return this.#wrapWithGuardrails(result, contextVariables, config);
|
|
2434
|
-
}
|
|
2435
|
-
/**
|
|
2436
|
-
* Create a raw stream without guardrail processing.
|
|
2437
|
-
*/
|
|
2438
|
-
async #createRawStream(contextVariables, config) {
|
|
2439
|
-
const { messages, systemPrompt } = await this.#options.context.resolve({
|
|
2440
|
-
renderer: new XmlRenderer()
|
|
2441
|
-
});
|
|
2442
|
-
const runId = generateId2();
|
|
2443
|
-
return streamText({
|
|
2444
|
-
abortSignal: config?.abortSignal,
|
|
2445
|
-
providerOptions: this.#options.providerOptions,
|
|
2446
|
-
model: this.#options.model,
|
|
2447
|
-
system: systemPrompt,
|
|
2448
|
-
messages: await convertToModelMessages(messages),
|
|
2449
|
-
experimental_repairToolCall: repairToolCall,
|
|
2450
|
-
stopWhen: stepCountIs(50),
|
|
2451
|
-
experimental_transform: config?.transform ?? smoothStream(),
|
|
2452
|
-
tools: this.#options.tools,
|
|
2453
|
-
experimental_context: contextVariables,
|
|
2454
|
-
toolChoice: this.#options.toolChoice,
|
|
2455
|
-
onStepFinish: (step) => {
|
|
2456
|
-
const toolCall = step.toolCalls.at(-1);
|
|
2457
|
-
if (toolCall) {
|
|
2458
|
-
console.log(
|
|
2459
|
-
`Debug: (${runId}) ${chalk2.bold.yellow("ToolCalled")}: ${toolCall.toolName}(${JSON.stringify(toolCall.input)})`
|
|
2460
|
-
);
|
|
2461
|
-
}
|
|
2462
|
-
}
|
|
2463
|
-
});
|
|
2464
|
-
}
|
|
2465
|
-
/**
|
|
2466
|
-
* Wrap a StreamTextResult with guardrail protection on toUIMessageStream().
|
|
2467
|
-
*
|
|
2468
|
-
* When a guardrail fails:
|
|
2469
|
-
* 1. Accumulated text + feedback is appended as the assistant's self-correction
|
|
2470
|
-
* 2. The feedback is written to the output stream (user sees the correction)
|
|
2471
|
-
* 3. A new stream is started and the model continues from the correction
|
|
2472
|
-
*/
|
|
2473
|
-
#wrapWithGuardrails(result, contextVariables, config) {
|
|
2474
|
-
const maxRetries = config?.maxRetries ?? this.#options.maxGuardrailRetries ?? 3;
|
|
2475
|
-
const context = this.#options.context;
|
|
2476
|
-
const originalToUIMessageStream = result.toUIMessageStream.bind(result);
|
|
2477
|
-
result.toUIMessageStream = (options) => {
|
|
2478
|
-
return createUIMessageStream({
|
|
2479
|
-
generateId: generateId2,
|
|
2480
|
-
execute: async ({ writer }) => {
|
|
2481
|
-
let currentResult = result;
|
|
2482
|
-
let attempt = 0;
|
|
2483
|
-
const { mounts } = context.getSkillMounts();
|
|
2484
|
-
const guardrailContext = {
|
|
2485
|
-
availableTools: Object.keys(this.tools),
|
|
2486
|
-
availableSkills: mounts
|
|
2487
|
-
};
|
|
2488
|
-
while (attempt < maxRetries) {
|
|
2489
|
-
if (config?.abortSignal?.aborted) {
|
|
2490
|
-
writer.write({ type: "finish" });
|
|
2491
|
-
return;
|
|
2492
|
-
}
|
|
2493
|
-
attempt++;
|
|
2494
|
-
let accumulatedText = "";
|
|
2495
|
-
let guardrailFailed = false;
|
|
2496
|
-
let failureFeedback = "";
|
|
2497
|
-
const uiStream = currentResult === result ? originalToUIMessageStream(options) : currentResult.toUIMessageStream(options);
|
|
2498
|
-
for await (const part of uiStream) {
|
|
2499
|
-
const checkResult = runGuardrailChain(
|
|
2500
|
-
part,
|
|
2501
|
-
this.#guardrails,
|
|
2502
|
-
guardrailContext
|
|
2503
|
-
);
|
|
2504
|
-
if (checkResult.type === "fail") {
|
|
2505
|
-
guardrailFailed = true;
|
|
2506
|
-
failureFeedback = checkResult.feedback;
|
|
2507
|
-
console.log(
|
|
2508
|
-
chalk2.yellow(
|
|
2509
|
-
`[${this.#options.name}] Guardrail triggered (attempt ${attempt}/${maxRetries}): ${failureFeedback.slice(0, 50)}...`
|
|
2510
|
-
)
|
|
2511
|
-
);
|
|
2512
|
-
break;
|
|
2513
|
-
}
|
|
2514
|
-
if (checkResult.type === "stop") {
|
|
2515
|
-
console.log(
|
|
2516
|
-
chalk2.red(
|
|
2517
|
-
`[${this.#options.name}] Guardrail stopped - unrecoverable error, no retry`
|
|
2518
|
-
)
|
|
2519
|
-
);
|
|
2520
|
-
writer.write(part);
|
|
2521
|
-
writer.write({ type: "finish" });
|
|
2522
|
-
return;
|
|
2523
|
-
}
|
|
2524
|
-
if (checkResult.part.type === "text-delta") {
|
|
2525
|
-
accumulatedText += checkResult.part.delta;
|
|
2526
|
-
}
|
|
2527
|
-
writer.write(part);
|
|
2528
|
-
}
|
|
2529
|
-
if (!guardrailFailed) {
|
|
2530
|
-
writer.write({ type: "finish" });
|
|
2531
|
-
return;
|
|
2532
|
-
}
|
|
2533
|
-
if (attempt >= maxRetries) {
|
|
2534
|
-
console.error(
|
|
2535
|
-
chalk2.red(
|
|
2536
|
-
`[${this.#options.name}] Guardrail retry limit (${maxRetries}) exceeded.`
|
|
2537
|
-
)
|
|
2538
|
-
);
|
|
2539
|
-
writer.write({ type: "finish" });
|
|
2540
|
-
return;
|
|
2541
|
-
}
|
|
2542
|
-
writeText(writer, failureFeedback);
|
|
2543
|
-
const selfCorrectionText = accumulatedText + " " + failureFeedback;
|
|
2544
|
-
context.set(lastAssistantMessage(selfCorrectionText));
|
|
2545
|
-
await context.save();
|
|
2546
|
-
currentResult = await this.#createRawStream(
|
|
2547
|
-
contextVariables,
|
|
2548
|
-
config
|
|
2549
|
-
);
|
|
2550
|
-
}
|
|
2551
|
-
},
|
|
2552
|
-
onError: (error) => {
|
|
2553
|
-
const message2 = error instanceof Error ? error.message : String(error);
|
|
2554
|
-
return `Stream failed: ${message2}`;
|
|
2555
|
-
}
|
|
2556
|
-
});
|
|
2557
|
-
};
|
|
2558
|
-
return result;
|
|
2559
|
-
}
|
|
2560
|
-
clone(overrides) {
|
|
2561
|
-
return new _Agent({
|
|
2562
|
-
...this.#options,
|
|
2563
|
-
...overrides
|
|
2564
|
-
});
|
|
2565
|
-
}
|
|
2566
|
-
};
|
|
2567
|
-
function agent(options) {
|
|
2568
|
-
return new Agent(options);
|
|
2569
|
-
}
|
|
2570
|
-
var repairToolCall = async ({
|
|
2571
|
-
toolCall,
|
|
2572
|
-
tools,
|
|
2573
|
-
inputSchema,
|
|
2574
|
-
error
|
|
2575
|
-
}) => {
|
|
2576
|
-
console.log(
|
|
2577
|
-
`Debug: ${chalk2.yellow("RepairingToolCall")}: ${toolCall.toolName}`,
|
|
2578
|
-
error.name
|
|
2579
|
-
);
|
|
2580
|
-
if (NoSuchToolError.isInstance(error)) {
|
|
2581
|
-
return null;
|
|
2582
|
-
}
|
|
2583
|
-
const tool2 = tools[toolCall.toolName];
|
|
2584
|
-
const { output } = await generateText({
|
|
2585
|
-
model: groq("openai/gpt-oss-20b"),
|
|
2586
|
-
output: Output.object({ schema: tool2.inputSchema }),
|
|
2587
|
-
prompt: [
|
|
2588
|
-
`The model tried to call the tool "${toolCall.toolName}" with the following inputs:`,
|
|
2589
|
-
JSON.stringify(toolCall.input),
|
|
2590
|
-
`The tool accepts the following schema:`,
|
|
2591
|
-
JSON.stringify(inputSchema(toolCall)),
|
|
2592
|
-
"Please fix the inputs."
|
|
2593
|
-
].join("\n")
|
|
2594
|
-
});
|
|
2595
|
-
return { ...toolCall, input: JSON.stringify(output) };
|
|
2596
|
-
};
|
|
2597
|
-
function writeText(writer, text) {
|
|
2598
|
-
const feedbackPartId = generateId2();
|
|
2599
|
-
writer.write({
|
|
2600
|
-
id: feedbackPartId,
|
|
2601
|
-
type: "text-start"
|
|
2602
|
-
});
|
|
2603
|
-
writer.write({
|
|
2604
|
-
id: feedbackPartId,
|
|
2605
|
-
type: "text-delta",
|
|
2606
|
-
delta: ` ${text}`
|
|
2607
|
-
});
|
|
2608
|
-
writer.write({
|
|
2609
|
-
id: feedbackPartId,
|
|
2610
|
-
type: "text-end"
|
|
2611
|
-
});
|
|
2612
|
-
}
|
|
2613
|
-
|
|
2614
|
-
// packages/text2sql/src/lib/adapters/groundings/report.grounding.ts
|
|
555
|
+
import { groq } from "@ai-sdk/groq";
|
|
556
|
+
import { tool } from "ai";
|
|
557
|
+
import dedent from "dedent";
|
|
558
|
+
import z from "zod";
|
|
559
|
+
import "@deepagents/agent";
|
|
560
|
+
import {
|
|
561
|
+
ContextEngine,
|
|
562
|
+
InMemoryContextStore,
|
|
563
|
+
agent,
|
|
564
|
+
fragment,
|
|
565
|
+
persona,
|
|
566
|
+
user
|
|
567
|
+
} from "@deepagents/context";
|
|
2615
568
|
var ReportGrounding = class extends AbstractGrounding {
|
|
2616
569
|
#adapter;
|
|
2617
570
|
#model;
|
|
@@ -2620,7 +573,7 @@ var ReportGrounding = class extends AbstractGrounding {
|
|
|
2620
573
|
constructor(adapter, config = {}) {
|
|
2621
574
|
super("business_context");
|
|
2622
575
|
this.#adapter = adapter;
|
|
2623
|
-
this.#model = config.model ??
|
|
576
|
+
this.#model = config.model ?? groq("openai/gpt-oss-20b");
|
|
2624
577
|
this.#cache = config.cache;
|
|
2625
578
|
this.#forceRefresh = config.forceRefresh ?? false;
|
|
2626
579
|
}
|
|
@@ -2652,7 +605,7 @@ var ReportGrounding = class extends AbstractGrounding {
|
|
|
2652
605
|
}),
|
|
2653
606
|
fragment(
|
|
2654
607
|
"instructions",
|
|
2655
|
-
|
|
608
|
+
dedent`
|
|
2656
609
|
Write a business context that helps another agent answer questions accurately.
|
|
2657
610
|
|
|
2658
611
|
For EACH table, do queries ONE AT A TIME:
|
|
@@ -2842,25 +795,87 @@ var TableGrounding = class extends AbstractGrounding {
|
|
|
2842
795
|
// packages/text2sql/src/lib/adapters/postgres/column-stats.postgres.grounding.ts
|
|
2843
796
|
var PostgresColumnStatsGrounding = class extends ColumnStatsGrounding {
|
|
2844
797
|
#adapter;
|
|
798
|
+
#pgStatsCache = /* @__PURE__ */ new Map();
|
|
2845
799
|
constructor(adapter, config = {}) {
|
|
2846
800
|
super(config);
|
|
2847
801
|
this.#adapter = adapter;
|
|
2848
802
|
}
|
|
803
|
+
async execute(ctx) {
|
|
804
|
+
await this.#fetchAllPgStats(ctx);
|
|
805
|
+
await super.execute(ctx);
|
|
806
|
+
}
|
|
807
|
+
async #fetchAllPgStats(ctx) {
|
|
808
|
+
const allContainers = [...ctx.tables, ...ctx.views];
|
|
809
|
+
if (allContainers.length === 0) return;
|
|
810
|
+
const conditions = allContainers.map((container) => {
|
|
811
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(container.name);
|
|
812
|
+
return `(schemaname = '${this.#adapter.escapeString(schema)}' AND tablename = '${this.#adapter.escapeString(table2)}')`;
|
|
813
|
+
});
|
|
814
|
+
const rows = await this.#adapter.runQuery(`
|
|
815
|
+
SELECT
|
|
816
|
+
schemaname,
|
|
817
|
+
tablename,
|
|
818
|
+
attname,
|
|
819
|
+
null_frac,
|
|
820
|
+
n_distinct,
|
|
821
|
+
histogram_bounds::text,
|
|
822
|
+
correlation
|
|
823
|
+
FROM pg_stats
|
|
824
|
+
WHERE ${conditions.join(" OR ")}
|
|
825
|
+
`);
|
|
826
|
+
for (const row of rows) {
|
|
827
|
+
const tableName = allContainers.find((c) => {
|
|
828
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(c.name);
|
|
829
|
+
return schema === row.schemaname && table2 === row.tablename;
|
|
830
|
+
})?.name;
|
|
831
|
+
if (!tableName) continue;
|
|
832
|
+
let map = this.#pgStatsCache.get(tableName);
|
|
833
|
+
if (!map) {
|
|
834
|
+
map = /* @__PURE__ */ new Map();
|
|
835
|
+
this.#pgStatsCache.set(tableName, map);
|
|
836
|
+
}
|
|
837
|
+
map.set(row.attname, row);
|
|
838
|
+
}
|
|
839
|
+
}
|
|
2849
840
|
async collectStats(tableName, column2) {
|
|
841
|
+
const cached = this.#pgStatsCache.get(tableName)?.get(column2.name);
|
|
842
|
+
if (cached) {
|
|
843
|
+
return this.#fromPgStats(cached);
|
|
844
|
+
}
|
|
845
|
+
return this.#collectStatsLive(tableName, column2);
|
|
846
|
+
}
|
|
847
|
+
#fromPgStats(row) {
|
|
848
|
+
const bounds = row.histogram_bounds ? this.#parsePgArray(row.histogram_bounds) : void 0;
|
|
849
|
+
const nullFraction = this.#adapter.toNumber(row.null_frac);
|
|
850
|
+
const nDistinct = this.#adapter.toNumber(row.n_distinct);
|
|
851
|
+
const correlation = this.#adapter.toNumber(row.correlation);
|
|
852
|
+
const result = {
|
|
853
|
+
...bounds?.length && { min: bounds[0] },
|
|
854
|
+
...bounds?.length && { max: bounds[bounds.length - 1] },
|
|
855
|
+
...nullFraction != null && {
|
|
856
|
+
nullFraction: Math.max(0, Math.min(1, nullFraction))
|
|
857
|
+
},
|
|
858
|
+
...nDistinct != null && { nDistinct },
|
|
859
|
+
...correlation != null && {
|
|
860
|
+
correlation: Math.max(-1, Math.min(1, correlation))
|
|
861
|
+
}
|
|
862
|
+
};
|
|
863
|
+
return Object.keys(result).length > 0 ? result : void 0;
|
|
864
|
+
}
|
|
865
|
+
async #collectStatsLive(tableName, column2) {
|
|
2850
866
|
if (!this.#shouldCollectStats(column2.type)) {
|
|
2851
867
|
return void 0;
|
|
2852
868
|
}
|
|
2853
869
|
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
2854
870
|
const tableIdentifier = `${this.#adapter.quoteIdentifier(schema)}.${this.#adapter.quoteIdentifier(table2)}`;
|
|
2855
871
|
const columnIdentifier = this.#adapter.quoteIdentifier(column2.name);
|
|
2856
|
-
const
|
|
872
|
+
const rows = await this.#adapter.runQuery(`
|
|
2857
873
|
SELECT
|
|
2858
874
|
MIN(${columnIdentifier})::text AS min_value,
|
|
2859
875
|
MAX(${columnIdentifier})::text AS max_value,
|
|
2860
876
|
AVG(CASE WHEN ${columnIdentifier} IS NULL THEN 1.0 ELSE 0.0 END) AS null_fraction
|
|
2861
877
|
FROM ${tableIdentifier}
|
|
2862
|
-
|
|
2863
|
-
const rows = await this.#adapter.runQuery(sql);
|
|
878
|
+
`);
|
|
2864
879
|
if (!rows.length) {
|
|
2865
880
|
return void 0;
|
|
2866
881
|
}
|
|
@@ -2885,6 +900,37 @@ var PostgresColumnStatsGrounding = class extends ColumnStatsGrounding {
|
|
|
2885
900
|
normalized
|
|
2886
901
|
);
|
|
2887
902
|
}
|
|
903
|
+
#parsePgArray(text) {
|
|
904
|
+
const inner = text.startsWith("{") && text.endsWith("}") ? text.slice(1, -1) : text;
|
|
905
|
+
if (!inner) {
|
|
906
|
+
return [];
|
|
907
|
+
}
|
|
908
|
+
const values = [];
|
|
909
|
+
let current = "";
|
|
910
|
+
let inQuote = false;
|
|
911
|
+
for (let i = 0; i < inner.length; i++) {
|
|
912
|
+
const ch = inner[i];
|
|
913
|
+
if (inQuote) {
|
|
914
|
+
if (ch === '"' && inner[i + 1] === '"') {
|
|
915
|
+
current += '"';
|
|
916
|
+
i++;
|
|
917
|
+
} else if (ch === '"') {
|
|
918
|
+
inQuote = false;
|
|
919
|
+
} else {
|
|
920
|
+
current += ch;
|
|
921
|
+
}
|
|
922
|
+
} else if (ch === '"') {
|
|
923
|
+
inQuote = true;
|
|
924
|
+
} else if (ch === ",") {
|
|
925
|
+
values.push(current);
|
|
926
|
+
current = "";
|
|
927
|
+
} else {
|
|
928
|
+
current += ch;
|
|
929
|
+
}
|
|
930
|
+
}
|
|
931
|
+
values.push(current);
|
|
932
|
+
return values;
|
|
933
|
+
}
|
|
2888
934
|
};
|
|
2889
935
|
|
|
2890
936
|
// packages/text2sql/src/lib/adapters/postgres/constraint.postgres.grounding.ts
|
|
@@ -3141,6 +1187,9 @@ var PostgresColumnValuesGrounding = class extends ColumnValuesGrounding {
|
|
|
3141
1187
|
return values?.length ? values : void 0;
|
|
3142
1188
|
}
|
|
3143
1189
|
async collectLowCardinality(tableName, column2) {
|
|
1190
|
+
if (this.#isHighCardinality(column2)) {
|
|
1191
|
+
return void 0;
|
|
1192
|
+
}
|
|
3144
1193
|
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
3145
1194
|
const tableIdentifier = `${this.#adapter.quoteIdentifier(schema)}.${this.#adapter.quoteIdentifier(table2)}`;
|
|
3146
1195
|
const columnIdentifier = this.#adapter.quoteIdentifier(column2.name);
|
|
@@ -3164,6 +1213,12 @@ var PostgresColumnValuesGrounding = class extends ColumnValuesGrounding {
|
|
|
3164
1213
|
}
|
|
3165
1214
|
return values.length ? values : void 0;
|
|
3166
1215
|
}
|
|
1216
|
+
#isHighCardinality(column2) {
|
|
1217
|
+
const nDistinct = column2.stats?.nDistinct;
|
|
1218
|
+
if (nDistinct == null) return false;
|
|
1219
|
+
if (nDistinct > 0) return nDistinct > this.lowCardinalityLimit;
|
|
1220
|
+
return true;
|
|
1221
|
+
}
|
|
3167
1222
|
};
|
|
3168
1223
|
|
|
3169
1224
|
// packages/text2sql/src/lib/adapters/postgres/postgres.ts
|
|
@@ -3656,11 +1711,27 @@ var PostgresRowCountGrounding = class extends RowCountGrounding {
|
|
|
3656
1711
|
this.#adapter = adapter;
|
|
3657
1712
|
}
|
|
3658
1713
|
async getRowCount(tableName) {
|
|
1714
|
+
const estimate = await this.#getEstimatedRowCount(tableName);
|
|
1715
|
+
if (estimate != null && estimate > 0) {
|
|
1716
|
+
return Math.round(estimate);
|
|
1717
|
+
}
|
|
1718
|
+
return this.#getLiveRowCount(tableName);
|
|
1719
|
+
}
|
|
1720
|
+
async #getEstimatedRowCount(tableName) {
|
|
1721
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
1722
|
+
const rows = await this.#adapter.runQuery(`
|
|
1723
|
+
SELECT c.reltuples::bigint AS estimate
|
|
1724
|
+
FROM pg_class c
|
|
1725
|
+
JOIN pg_namespace n ON n.oid = c.relnamespace
|
|
1726
|
+
WHERE n.nspname = '${this.#adapter.escapeString(schema)}'
|
|
1727
|
+
AND c.relname = '${this.#adapter.escapeString(table2)}'
|
|
1728
|
+
`);
|
|
1729
|
+
return this.#adapter.toNumber(rows[0]?.estimate);
|
|
1730
|
+
}
|
|
1731
|
+
async #getLiveRowCount(tableName) {
|
|
3659
1732
|
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
3660
1733
|
const tableIdentifier = `${this.#adapter.quoteIdentifier(schema)}.${this.#adapter.quoteIdentifier(table2)}`;
|
|
3661
|
-
const rows = await this.#adapter.runQuery(
|
|
3662
|
-
`SELECT COUNT(*) as count FROM ${tableIdentifier}`
|
|
3663
|
-
);
|
|
1734
|
+
const rows = await this.#adapter.runQuery(`SELECT COUNT(*) as count FROM ${tableIdentifier}`);
|
|
3664
1735
|
return this.#adapter.toNumber(rows[0]?.count);
|
|
3665
1736
|
}
|
|
3666
1737
|
};
|