feat: add context to queries

CHANGELOG.md

@@ -18,4 +18,4 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added a settings manager to manage settings.
 - Added a conversation store to manage conversations.
 - Created a Foundry ApplicationV2 to interact with the AI.
-- 
+- Added context to queries.

src/ai/client.js

@@ -144,17 +144,93 @@ export class Client {
      */
     #formatChatInput(context, title, query) {
         return [
+            ...query,
             {
                 role: 'system',
-                content: `You are a helpful AI assistant named AIde, running within the FoundryVTT environment.
+                content: `
+# AIde System Prompt
 
-<synopsis>
+You are AIde, an AI discussion assistant integrated into Foundry VTT to help users craft and manage their tabletop 
+roleplaying adventures. You were created by nivthefox and are currently interfacing with a user through Foundry VTT's
+ module system.
+
+## Core Identity & Purpose
+- You are a helpful AI assistant focused on tabletop roleplaying game adventure creation and management
+- Your responses should be useful for game masters and players using Foundry VTT
+- You maintain a friendly, professional demeanor while remaining focused on TTRPG-related assistance
+
+## Context Awareness
+- You have access to Foundry VTT journal entries provided as context in your conversations
+- When referencing context documents, they are formatted as: <JournalEntry title="[Document Title]">[Content]</JournalEntry>
+- You should use this context to inform your responses while maintaining coherence
+
+## Conversation Capabilities
+- You engage in back-and-forth dialogue about adventure creation, game mechanics, and world-building
+- You can access previous messages in the conversation for continuity
+- You aim to provide specific, actionable advice based on the user's needs
+
+## Technical Understanding
+- You are aware of Foundry VTT's capabilities and limitations
+- You understand common TTRPG terms and concepts
+- You can reference and explain game mechanics when relevant
+
+## Response Guidelines
+1. Keep most responses short and focused - no more than 2-3 sentences
+2. Do not provide lists of options or questions unless specifically asked
+3. Ask at most ONE follow-up question, and make it specific rather than open-ended
+4. Let the user drive the depth and pace of the conversation
+5. Wait for the user to request more detail before providing it
+6. Format responses using markdown syntax as specified in the formatting section
+7. When referencing context documents, cite them specifically by title
+8. Consider the game system and world context when providing advice
+
+## Conversation Flow
+- Wait for users to explicitly ask for details before providing them
+- When the user provides information, acknowledge it and ask for ONE specific detail to build on
+- Never provide outlines or lists of topics unless specifically requested
+- Focus responses on the immediate topic at hand
+- If you notice yourself writing a list or outline, STOP and rephrase as a simple question
+
+Example good flows:
+\`\`\`
+User: "I want to create a rival kingdom"
+You: "That sounds interesting! Where are they located relative to Thes?"
+
+User: "They're jealous of the residuum mines"
+You: "Ah, a kingdom envious of Thes's resources. What methods do they use to try to get their hands on the residuum?"
+\`\`\`
+
+Example bad flows:
+\`\`\`
+User: "I want to create a rival kingdom"
+You: "Great! Here's everything about kingdoms: location, government, economy..."
+
+User: "They're jealous of the residuum mines"
+You: "Let's outline all possible aspects of rivalry and resource competition..."
+\`\`\`
+
+## Ethical Guidelines
+1. Focus on creative and constructive adventure creation
+2. Avoid generating harmful or inappropriate content
+3. Respect intellectual property and copyright
+4. Maintain user privacy and confidentiality
+5. Do not provide advice that could compromise game security or player safety
+
+## Limitations
+- You cannot directly modify Foundry VTT content
+- You cannot access external websites or resources
+- You cannot execute code or system commands
+- You are limited to the context provided in the current conversation
+
+Remember: Your primary goal is to help users create engaging and enjoyable tabletop roleplaying experiences within Foundry VTT.
+
+<metadata>
 The user is running the following game system: ${game.system.title}
 The user is running the following game world: ${game.world.title}
 The user's name is: ${game.user.name}
 The title of this conversation is: ${title}
 The current time is: ${new Date().toLocaleString()} 
-</synopsis>
+</metadata>
 
 <formatting>
 Use markdown to add emphasis and structure to your messages:
@@ -168,11 +244,9 @@ Use markdown to add emphasis and structure to your messages:
 </formatting>
 
 <context>
-Use this context to answer the user's question:
-${context.map(doc => `# ${doc.title}\n${doc.content}`).join('\n\n')}
+${context.map(doc => `<JournalEntry title="${doc.name}">${doc.text.content}</JournalEntry>`).join('\n')}
 </context>`
-            },
-            ...query
+            }
         ];
     }
 

src/ai/provider/deepinfra.js

@@ -151,29 +151,33 @@ export class DeepInfra {
     /**
      * @param {string} model
      * @param {string} id
-     * @param {Chunk[]} chunks
+     * @param {Chunk[]} inputs
      * @returns {Promise<EmbeddingDocument>}
      */
-    async embed(model, id, chunks) {
-        const response = await fetch(`${this.#baseUrl}/${model}`, {
-            method: 'POST',
-            headers: {
-                'Content-Type': 'application/json',
-                'Authorization': `Bearer ${this.#apiKey}`
-            },
-            body: JSON.stringify({
-                inputs: chunks
-            })
-        });
+    async embed(model, id, inputs) {
+        const data = {
+            id,
+            vectors: []
+        };
 
-        if (!response.ok) {
-            throw new Error(`DeepInfra API error: ${response.status}`);
+        for (const input of inputs) {
+            const response = await fetch(`${this.#baseUrl}/embeddings`, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                    'Authorization': `Bearer ${this.#apiKey}`
+                },
+                body: JSON.stringify({input, model})
+            });
+
+            if (!response.ok) {
+                throw new Error(`DeepInfra API error: ${response.status}`);
+            }
+
+            const output = await response.json();
+            data.vectors.push(output.data[0].embedding);
         }
 
-        const data = await response.json();
-        return {
-            id,
-            vectors: data.embeddings
-        };
+        return data;
     }
 }

src/app/app.js

@@ -26,7 +26,8 @@ export class App {
         ctx.Hooks.once('setup', () => this.setup(ctx, id));
         ctx.Hooks.once('ready', () => this.ready(ctx, id));
         ctx.Hooks.on('renderSidebarTab', (app, html) =>
-            renderChatWithAIButton(app, html, this.conversationStore, this.chatClient));
+            renderChatWithAIButton(app, html, this.conversationStore, this.chatClient, this.embeddingClient,
+                this.vectorStore, this.documentManager));
     }
 
     async ready() {
@@ -41,18 +42,22 @@ export class App {
         this.chatClient = Client.create(providerSettings.chat);
         this.embeddingClient = Client.create(providerSettings.embedding);
         this.conversationStore = new ConversationStore(ctx);
-        this.vectorStore = new VectorStore(ctx);
+
+        const lookups = game.settings.get('aide', 'VectorStoreLookups');
+        const maxWeight = game.settings.get('aide', 'VectorStoreMaxWeight');
+        const queryBoostFactor = game.settings.get('aide', 'VectorStoreQueryBoostFactor');
+        this.vectorStore = new VectorStore(lookups, maxWeight, queryBoostFactor);
 
         // Initialize Document Manager
         const managerSettings = this.settings.getDocumentManagerSettings();
         this.documentManager = new DocumentManager(ctx, managerSettings, this.embeddingClient, this.vectorStore);
 
-        // todo: only rebuild if necessary
-        await this.documentManager.rebuildVectorStore();
-
         // Initialize Conversation Store
         await this.conversationStore.initialize();
 
+        // todo: only rebuild if necessary
+        await this.documentManager.rebuildVectorStore();
+
         // Register model choices
         const chatModels = await this.chatClient.getChatModels();
         this.settings.setChoices('ChatModel', chatModels);

src/document/manager.js

@@ -68,7 +68,10 @@ export class DocumentManager {
      */
     async chunks(id) {
         const doc = await this.getDocument(id);
-        return this.#chunks(doc);
+        if (!this.#indexable(doc)) {
+            return [];
+        }
+        return this.calculateChunks(doc.text.content);
     }
 
     /**
@@ -164,7 +167,7 @@ export class DocumentManager {
             return;
         }
 
-        const chunks = await this.#chunks(changed);
+        const chunks = this.calculateChunks(changed.text.content);
         if (chunks.length === 0) {
             return;
         }
@@ -173,12 +176,7 @@ export class DocumentManager {
         await this.#store.add(vectors);
     }
 
-    #chunks(doc) {
-        if (!doc || !this.#indexable(doc)) {
-            return [];
-        }
-
-        const content = doc.text.content;
+    calculateChunks(content) {
         const tokens = this.#tokenize(content);
         const chunks = [];
 

src/document/vector_store.js

@@ -92,23 +92,31 @@ export class VectorStore {
 
         return Array.from(this.#cache.entries())
             .map(([id, documentVectors]) => {
+                console.log(`Processing document ${id} with ${documentVectors.length} vectors`);
+
                 // For each query vector, calculate similarities with all document chunks
                 const queryScores = queries.map(queryVector => {
                     const similarities = documentVectors.map(docVector => ({
                         similarity: this.#calculateSimilarity(queryVector, docVector)
                     }));
+                    console.log(`Similarities for doc ${id}:`, similarities);
 
                     const maxSim = similarities.reduce((max, curr) =>
                         curr.similarity > max.similarity ? curr : max
                     );
+
                     const avgSim = similarities.reduce((sum, curr) =>
                         sum + curr.similarity, 0) / similarities.length;
 
-                    return (maxSim.similarity * this.maxWeight) + (avgSim * avgWeight);
+                    const score = (maxSim.similarity * this.maxWeight) + (avgSim * avgWeight);
+                    console.log(`Scores for doc ${id}: max=${maxSim.similarity}, avg=${avgSim}, final=${score}`);
+
+                    return score;
                 });
 
                 // Take the maximum score across all query vectors
                 const score = queryScores.reduce((sum, score) => sum + score, 0) / queryScores.length;
+                console.log(`Final score for doc ${id}: ${score}`);
 
                 return { id, score };
             })

src/foundry/context.js

@@ -2,6 +2,7 @@ export const context = {
     FilePicker,
     Hooks,
     fetch: fetch.bind(window),
+    fromUuid: fromUuid.bind(window),
 
     get foundry() { return foundry; },
     get game() { return game; },

src/types.js

@@ -16,7 +16,7 @@
 
 /**
  * @typedef {Object} ContextDocument
- * @property {string} uuid
+ * @property {string} id
  * @property {string} [title]
  * @property {string} [content]
  */

src/ui/chat.js

@@ -42,10 +42,13 @@ export class Chat extends HandlebarsApplicationMixin(ApplicationV2) {
     };
 
     // Public Methods
-    constructor(conversationStore, chatClient, options = {}) {
+    constructor(conversationStore, chatClient, embeddingClient, vectorStore, documentManager, options = {}) {
         super(options);
         this.chatClient = chatClient;
         this.conversationStore = conversationStore;
+        this.documentManager = documentManager;
+        this.embeddingClient = embeddingClient;
+        this.vectorStore = vectorStore;
         this.#initializeMarkdownConverter();
     }
 
@@ -79,6 +82,27 @@ export class Chat extends HandlebarsApplicationMixin(ApplicationV2) {
         this.#setupRenameEvent(this.element);
     }
 
+    async #determineContext(content) {
+        let embeddableContent = this.#activeConversation.messages
+            .reduce((acc, message) => `${acc}${message.content}\n\n`, '');
+        embeddableContent += content;
+
+        const model = game.settings.get('aide', 'EmbeddingModel');
+        const chunks = this.documentManager.calculateChunks(embeddableContent);
+        const embeddings = await this.embeddingClient.embed(model, this.#activeConversation.id, chunks);
+
+        const results = this.vectorStore.findSimilar(embeddings.vectors);
+
+        const context = [];
+        for (const result of results) {
+            const document = await fromUuid(result.id);
+            if (document) {
+                context.push(document);
+            }
+        }
+        return context;
+    }
+
     #formatMessageContent(content) {
         if (!content || content === '<p><br class="ProseMirror-trailingBreak"></p>') {
             return null;
@@ -275,9 +299,12 @@ export class Chat extends HandlebarsApplicationMixin(ApplicationV2) {
         this.#waitingForResponse = true;
         await this.render(false);
 
+        const context = await this.#determineContext(content);
+
+
         // Get AI response
         const model = game.settings.get('aide', 'ChatModel');
-        const response = await this.chatClient.generate(model, [],
+        const response = await this.chatClient.generate(model, context,
             this.#activeConversation.title, conversation.messages, true);
 
         // Initialize AI message

src/ui/sidebar.js

@@ -1,6 +1,7 @@
 import { Chat } from './chat.js';
 
-export async function renderChatWithAIButton(app, html, conversationStore, chatClient) {
+export async function renderChatWithAIButton(app, html, conversationStore, chatClient,
+    embeddingClient, vectorStore, documentManager) {
     if (!(app instanceof JournalDirectory)) {
         return;
     }
@@ -15,6 +16,7 @@ export async function renderChatWithAIButton(app, html, conversationStore, chatC
 
     const button = targetElement.find('.aide.chat-with-ai');
     button.click(() => {
-        new Chat(conversationStore, chatClient).render(true);
+        new Chat(conversationStore, chatClient, embeddingClient, vectorStore, documentManager)
+            .render(true);
     });
 }

templates/applications/Chat.hbs

@@ -7,10 +7,13 @@
                     <button class="settings" data-action="settings">
                         <i class="fas fa-sliders"></i>
                     </button>
+                    <!--
+                    // todo: consider how to handle context
                     <button class="context" data-action="context">
                         <i class="fas fa-paperclip"></i>
                         {{#if context.documents.length}}<span>{{ context.documents.length }}</span>{{/if}}
                     </button>
+                    //-->
             </aside>
         </h2>
         {{#each messages as | message | }}