udpated to include multi modal usage

package-lock.json

@@ -11,7 +11,8 @@
       "dependencies": {
         "@grpc/grpc-js": "^1.14.2",
         "@grpc/proto-loader": "^0.8.0",
-        "@langchain/core": "^1.1.12",
+        "@langchain/core": "^1.1.13",
+        "@langchain/google-genai": "^2.1.8",
         "@langchain/langgraph": "^1.0.15",
         "@langchain/openai": "^1.2.1",
         "@nestjs/common": "^11.1.3",
@@ -985,6 +986,15 @@
         "node": "^18.18.0 || ^20.9.0 || >=21.1.0"
       }
     },
+    "node_modules/@google/generative-ai": {
+      "version": "0.24.1",
+      "resolved": "https://registry.npmjs.org/@google/generative-ai/-/generative-ai-0.24.1.tgz",
+      "integrity": "sha512-MqO+MLfM6kjxcKoy0p1wRzG3b4ZZXtPI+z2IE26UogS2Cm/XHO+7gGRBh6gcJsOiIVoH93UwKvW4HdgiOZCy9Q==",
+      "license": "Apache-2.0",
+      "engines": {
+        "node": ">=18.0.0"
+      }
+    },
     "node_modules/@grpc/grpc-js": {
       "version": "1.14.2",
       "resolved": "https://registry.npmjs.org/@grpc/grpc-js/-/grpc-js-1.14.2.tgz",
@@ -2149,9 +2159,9 @@
       }
     },
     "node_modules/@langchain/core": {
-      "version": "1.1.12",
-      "resolved": "https://registry.npmjs.org/@langchain/core/-/core-1.1.12.tgz",
-      "integrity": "sha512-sHWLvhyLi3fntlg3MEPB89kCjxEX7/+imlIYJcp6uFGCAZfGxVWklqp22HwjT1szorUBYrkO8u0YA554ReKxGQ==",
+      "version": "1.1.13",
+      "resolved": "https://registry.npmjs.org/@langchain/core/-/core-1.1.13.tgz",
+      "integrity": "sha512-CmTES4DNfNs7PisGm/is4RxOf1NAWCkhi+RrBBHb/gB5nZVFd+dfmXSomKoiBQ1DOdCUz1k9RX4DzSUbwg1swg==",
       "license": "MIT",
       "dependencies": {
         "@cfworker/json-schema": "^4.0.2",
@@ -2227,6 +2237,35 @@
         }
       }
     },
+    "node_modules/@langchain/google-genai": {
+      "version": "2.1.8",
+      "resolved": "https://registry.npmjs.org/@langchain/google-genai/-/google-genai-2.1.8.tgz",
+      "integrity": "sha512-bDN8A9sRH71c/A94S/ioFgTdBRhUouqf0y5+FHMB5zZvgGuJ7765o9lYe84T53j/Kze/i8xeyJ7d7O/34azJuA==",
+      "license": "MIT",
+      "dependencies": {
+        "@google/generative-ai": "^0.24.0",
+        "uuid": "^11.1.0"
+      },
+      "engines": {
+        "node": ">=20"
+      },
+      "peerDependencies": {
+        "@langchain/core": "1.1.13"
+      }
+    },
+    "node_modules/@langchain/google-genai/node_modules/uuid": {
+      "version": "11.1.0",
+      "resolved": "https://registry.npmjs.org/uuid/-/uuid-11.1.0.tgz",
+      "integrity": "sha512-0/A9rDy9P7cJ+8w1c9WD9V//9Wj15Ce2MPz8Ri6032usz+NfePxx5AcN3bN+r6ZL6jEo066/yNYB3tn4pQEx+A==",
+      "funding": [
+        "https://github.com/sponsors/broofa",
+        "https://github.com/sponsors/ctavan"
+      ],
+      "license": "MIT",
+      "bin": {
+        "uuid": "dist/esm/bin/uuid"
+      }
+    },
     "node_modules/@langchain/langgraph": {
       "version": "1.0.15",
       "resolved": "https://registry.npmjs.org/@langchain/langgraph/-/langgraph-1.0.15.tgz",

package.json

@@ -22,7 +22,8 @@
   "dependencies": {
     "@grpc/grpc-js": "^1.14.2",
     "@grpc/proto-loader": "^0.8.0",
-    "@langchain/core": "^1.1.12",
+    "@langchain/core": "^1.1.13",
+    "@langchain/google-genai": "^2.1.8",
     "@langchain/langgraph": "^1.0.15",
     "@langchain/openai": "^1.2.1",
     "@nestjs/common": "^11.1.3",

src/FFB/ffb-production.controller.ts

@@ -17,6 +17,8 @@ export class FFBProductionController {
   }
 
 
+
+
   /** Vector search endpoint */
   @Get('search')
   async search(@Query('q') q: string, @Query('k') k?: string) {

src/FFB/ffb.gateway.ts

@@ -53,6 +53,22 @@ export class FFBGateway
         return { event: 'sent', data: response }; // Acknowledgement if needed
     }
 
+    @SubscribeMessage('switch_model')
+    handleSwitchModel(
+        @MessageBody() data: { provider: 'openai' | 'gemini' },
+        @ConnectedSocket() client: Socket,
+    ) {
+        this.logger.log(`Switching model for ${client.id} to ${data.provider}`);
+        this.langchainService.switchModel(client.id, data.provider);
+    }
+
+    @SubscribeMessage('get_model')
+    handleGetModel(@ConnectedSocket() client: Socket) {
+        const model = this.langchainService.getCurrentModel(client.id);
+        console.log(model)
+        return { event: 'current_model', data: model };
+    }
+
     emitThought(socketId: string, data: any) {
         this.server.to(socketId).emit('agent_thought', data);
     }

src/FFB/services/config/agent-state.ts

@@ -0,0 +1,32 @@
+import { Annotation } from "@langchain/langgraph";
+import { BaseMessage } from "@langchain/core/messages";
+
+export const AgentState = Annotation.Root({
+    messages: Annotation<BaseMessage[]>({
+        reducer: (x, y) => x.concat(y),
+        default: () => [],
+    }),
+    activeIntent: Annotation<string>({
+        reducer: (x, y) => y ?? x ?? "General",
+        default: () => "General",
+    }),
+    entryCategory: Annotation<string>({
+        reducer: (x, y) => y ?? x,
+        default: () => "OutOfScope",
+    }),
+    entityStore: Annotation<Record<string, any>>({
+        reducer: (x, y) => ({ ...x, ...y }),
+        default: () => ({}),
+    }),
+    actionPayload: Annotation<any>({
+        reducer: (x, y) => y ?? x,
+        default: () => null,
+    }),
+    finalResponse: Annotation<string>({
+        reducer: (x, y) => y ?? x,
+    }),
+    socketId: Annotation<string>({
+        reducer: (x, y) => y ?? x,
+        default: () => "default",
+    })
+});

src/FFB/services/config/langchain-config.ts

@@ -0,0 +1,95 @@
+import { z } from "zod";
+
+export const SCHEMAS = {
+    ENTRY: z.object({
+        category: z.enum(['InScope-Actionable', 'InScope-NeedsGuidance', 'InScope-Meta', 'OutOfScope']),
+        reasoning: z.string()
+    }),
+    ROUTER: z.object({
+        intent: z.enum(['Semantic', 'Aggregate']),
+        entities: z.object({
+            site: z.string().nullable().describe("The site name mentioned, or null"),
+            date: z.string().nullable().describe("The date mentioned, or null"),
+        }),
+        reasoning: z.string()
+    }),
+    AGGREGATION: z.object({
+        matchStage: z.object({
+            site: z.string().nullable(),
+            startDate: z.string().nullable(),
+            endDate: z.string().nullable(),
+        }),
+        aggregationType: z.enum(["sum", "avg", "count"]),
+        fieldToAggregate: z.enum(["quantity", "weight"])
+    })
+};
+
+export const PROMPTS = {
+    ENTRY: (lastMessage: string) => `
+You are the Entry Node for a strictly defined FFB Production Data Agent.
+Your job is to CLASSIFY the user's message into one of four categories.
+
+IDENTITY:
+- You retrieve, summarize, and aggregate records from a vector database.
+- You DO NOT answer general knowledge, creative, or unrelated questions.
+- You NEVER behave like a general-purpose chatbot.
+
+CATEGORIES:
+1. InScope-Actionable: A valid, clear request for querying or aggregating data.
+2. InScope-NeedsGuidance: On-topic but vague, incomplete, or malformed request (e.g. mentions a site but no question).
+3. InScope-Meta: User asks about capabilities, "what can you do", "what did I ask before", or how to use the agent.
+4. OutOfScope: Any request unrelated to FFB data querying or aggregation (e.g. "write a poem", "who is the president", "general chat").
+
+LOGIC:
+- IF questions about capabilities OR memory ("what did I ask") -> InScope-Meta
+- IF unrelated to data/production -> OutOfScope
+- IF mentions domain entities (site, date) BUT lacks clear intent -> InScope-NeedsGuidance
+- ELSE -> InScope-Actionable
+
+User Message: "${lastMessage}"
+`,
+    ROUTER: (lastMessage: string) => `
+You are an Application Router for a production database.
+Analyze the user input and route to: [Semantic, Aggregate].
+
+INTENTS:
+- Aggregate: Asking for numbers, totals, averages, counts (e.g., "How much...", "Total weight").
+- Semantic: Asking for specific records, qualitative descriptions, issues, "what happened", "find info about" (e.g., "Show me records for Site A").
+
+User Input: "${lastMessage}"
+`,
+    META: (lastMessage: string, context: string) => `
+You are the FFB Production Agent.
+User Input: "${lastMessage}"
+
+Your Capabilities:
+1. Querying specific production logs.
+2. Aggregating data (totals, averages).
+3. Summarizing production events.
+
+INSTRUCTIONS:
+- If the user asks about capabilities, list them.
+- If the user asks "what did I ask?", summarize the RELEVANT previous user messages from the context.
+- Do NOT help with off-topic or general questions.
+- Be professional and concise.
+
+Context (Previous Valid Messages):
+${context}
+`,
+    GUIDANCE: `I noticed you're asking about production data, but I need a bit more detail.\n\n` +
+        `I can help you with:\n` +
+        `- **Aggregations**: "What is the total production weight for Site A?"\n` +
+        `- **Specifics**: "What happened at Site B on Jan 12?"\n` +
+        `\nCould you clarify what you're looking for?`,
+
+    REFUSAL: `I'm sorry, but I can only assist with FFB production data querying and aggregation.\n\n` +
+        `Please ask me about production weights, site activities, or specific dates.`,
+
+    SYNTHESIS: (lastMessage: string, payload: any) => `
+User Question: "${lastMessage}"
+Data Context: ${JSON.stringify(payload)}
+
+Synthesize a natural language answer based STRICTLY on the Data Context.
+Cite the source (e.g., "Based on aggregation results...").
+`
+};

src/FFB/services/ffb-langchain.service.ts

@@ -1,84 +1,110 @@
-import { Injectable } from '@nestjs/common';
+import { Injectable, Inject, forwardRef } from '@nestjs/common';
 import { ChatOpenAI } from '@langchain/openai';
+import { ChatGoogleGenerativeAI } from '@langchain/google-genai';
+import { BaseChatModel } from "@langchain/core/language_models/chat_models";
+import { StateGraph, START, END } from "@langchain/langgraph";
+import { BaseMessage, HumanMessage } from "@langchain/core/messages";
 import { FFBVectorService } from './ffb-vector.service';
-import { z } from "zod";
-import { StateGraph, START, END, Annotation } from "@langchain/langgraph";
-import { BaseMessage, HumanMessage, AIMessage } from "@langchain/core/messages";
-import { forwardRef, Inject } from '@nestjs/common';
 import { FFBGateway } from '../ffb.gateway';
-import { ThoughtPayload } from '../ffb-production.schema';
-// State Definition using Annotation
-const AgentState = Annotation.Root({
-    messages: Annotation<BaseMessage[]>({
-        reducer: (x, y) => x.concat(y),
-        default: () => [],
-    }),
-    activeIntent: Annotation<string>({
-        reducer: (x, y) => y ?? x ?? "General",
-        default: () => "General",
-    }),
-    entityStore: Annotation<Record<string, any>>({
-        reducer: (x, y) => ({ ...x, ...y }),
-        default: () => ({}),
-    }),
-    actionPayload: Annotation<any>({
-        reducer: (x, y) => y ?? x,
-        default: () => null,
-    }),
-    finalResponse: Annotation<string>({
-        reducer: (x, y) => y ?? x,
-    }),
-    socketId: Annotation<string>({
-        reducer: (x, y) => y ?? x,
-        default: () => "default",
-    })
-});
+
+// Config & Utils
+import { AgentState } from './config/agent-state';
+import { SessionManager } from './utils/session-manager';
+
+// Nodes
+import { entryNode } from './nodes/entry.node';
+import { routerNode } from './nodes/router.node';
+import { guidanceNode } from './nodes/guidance.node';
+import { metaNode } from './nodes/meta.node';
+import { refusalNode } from './nodes/refusal.node';
+import { vectorSearchNode } from './nodes/vector-search.node';
+import { aggregationNode } from './nodes/aggregation.node';
+import { synthesisNode } from './nodes/synthesis.node';
 
 @Injectable()
 export class FFBLangChainService {
-    private model: ChatOpenAI;
+    private openaiModel: BaseChatModel;
+    private geminiModel: BaseChatModel;
     private graph: any;
-    private sessions: Map<string, BaseMessage[]> = new Map();
+    private sessionManager: SessionManager;
 
     constructor(
         private readonly vectorService: FFBVectorService,
         @Inject(forwardRef(() => FFBGateway))
         private readonly gateway: FFBGateway
     ) {
-        this.model = new ChatOpenAI({
-            modelName: 'gpt-4o',
+        this.openaiModel = new ChatOpenAI({
+            modelName: 'gpt-4o-mini',
             apiKey: process.env.OPENAI_API_KEY,
             temperature: 0
         });
 
+        this.geminiModel = new ChatGoogleGenerativeAI({
+            model: 'gemini-2.5-flash',
+            apiKey: process.env.GOOGLE_API_KEY,
+            temperature: 0
+        });
+
+        this.sessionManager = new SessionManager();
         this.initGraph();
     }
 
+    private getModel(socketId: string): BaseChatModel {
+        const provider = this.sessionManager.getModelProvider(socketId);
+        return provider === 'gemini' ? this.geminiModel : this.openaiModel;
+    }
+
+    switchModel(socketId: string, provider: 'openai' | 'gemini') {
+        this.sessionManager.setModelProvider(socketId, provider);
+    }
+
+    getCurrentModel(socketId: string) {
+        const provider = this.sessionManager.getModelProvider(socketId);
+        return {
+            provider: provider,
+            modelName: provider === 'gemini' ? 'gemini-2.5-flash' : 'gpt-4o-mini'
+        };
+    }
+
+
+
     private initGraph() {
         const graph = new StateGraph(AgentState)
-            .addNode("router_node", this.routerNode.bind(this))
-            .addNode("clarifier_node", this.clarifierNode.bind(this))
-            .addNode("general_node", this.generalNode.bind(this))
-            .addNode("vector_search_node", this.vectorSearchNode.bind(this))
-            .addNode("aggregation_node", this.aggregationNode.bind(this))
-            .addNode("synthesis_node", this.synthesisNode.bind(this));
+            .addNode("entry_node", (state) => entryNode(state, this.getModel(state.socketId), this.gateway))
+            .addNode("guidance_node", (state) => guidanceNode(state))
+            .addNode("meta_node", (state) => metaNode(state, this.getModel(state.socketId)))
+            .addNode("refusal_node", (state) => refusalNode(state))
+            .addNode("router_node", (state) => routerNode(state, this.getModel(state.socketId), this.gateway))
+            .addNode("vector_search_node", (state) => vectorSearchNode(state, this.vectorService, this.gateway))
+            .addNode("aggregation_node", (state) => aggregationNode(state, this.getModel(state.socketId), this.vectorService, this.gateway))
+            .addNode("synthesis_node", (state) => synthesisNode(state, this.getModel(state.socketId), this.gateway));
 
         // Add Edges
-        graph.addEdge(START, "router_node");
+        graph.addEdge(START, "entry_node");
+
+        graph.addConditionalEdges(
+            "entry_node",
+            (state) => state.entryCategory,
+            {
+                "InScope-Actionable": "router_node",
+                "InScope-NeedsGuidance": "guidance_node",
+                "InScope-Meta": "meta_node",
+                "OutOfScope": "refusal_node"
+            }
+        );
 
         graph.addConditionalEdges(
             "router_node",
             (state) => state.activeIntent,
             {
-                Clarify: "clarifier_node",
-                General: "general_node",
                 Semantic: "vector_search_node",
                 Aggregate: "aggregation_node"
             }
         );
 
-        graph.addEdge("clarifier_node", END);
-        graph.addEdge("general_node", END);
+        graph.addEdge("guidance_node", END);
+        graph.addEdge("meta_node", END);
+        graph.addEdge("refusal_node", END);
         graph.addEdge("vector_search_node", "synthesis_node");
         graph.addEdge("aggregation_node", "synthesis_node");
         graph.addEdge("synthesis_node", END);
@@ -86,244 +112,49 @@ export class FFBLangChainService {
         this.graph = graph.compile();
     }
 
-    // --- NODE IMPLEMENTATIONS ---
-
-    private async routerNode(state: typeof AgentState.State): Promise<Partial<typeof AgentState.State>> {
-        const lastMessage = state.messages[state.messages.length - 1].content as string;
-
-        // Change this in your routerNode:
-        const routerSchema = z.object({
-            intent: z.enum(['General', 'Clarify', 'Semantic', 'Aggregate']),
-            entities: z.object({
-                // Use .nullable() instead of .optional() for OpenAI Strict mode
-                // Or ensure they are always provided by the LLM
-                site: z.string().nullable().describe("The site name mentioned, or null"),
-                date: z.string().nullable().describe("The date mentioned, or null"),
-            }), // Remove .optional() here; the object itself must be returned
-            reasoning: z.string()
-        });
-
-        let payload: ThoughtPayload = {
-            node: 'router_node',
-            status: 'processing',
-            message: 'Analyzing user intent...',
-            input: lastMessage
-        }
-
-        this.gateway.emitThought(state.socketId, payload);
-
-        const routerPrompt = `
-You are an Application Router for a production database.
-Analyze the user input and route to: [General, Clarify, Semantic, Aggregate].
-
-INTENT DEFINITIONS:
-- Aggregate: Use if the user asks for numbers, totals, averages, or counts (e.g., "How much...", "Total weight").
-- Semantic: Use if the user asks for specific records, qualitative descriptions, issues, "what happened", or "find info about" (e.g., "Show me records for Site A", "What were the notes on block X?").
-- Clarify: Use ONLY if the user names an entity (like a Site) but provides NO verb or question.
-- General: Use for greetings or off-topic chat.
-
-STRICT RULES:
-1. If "Site" is mentioned alone (e.g., "Site A"), route to 'Clarify'.
-2. If the user asks for data or "what happened" regarding a site, route to 'Semantic'.
-3. Do NOT route to 'Clarify' if there is a clear question.
-
-User Input: "${lastMessage}"
-`;
-
-        const structuredLlm = this.model.withStructuredOutput(routerSchema);
-        const result = await structuredLlm.invoke(routerPrompt);
-
-        // Merge extracted entities with existing store
-        this.gateway.emitThought(state.socketId, {
-            node: 'router_node',
-            status: 'completed',
-            result: result
-        });
-
-        return {
-            activeIntent: result.intent as any,
-            entityStore: result.entities || {},
-            socketId: state.socketId
-        };
-    }
-
-    private async clarifierNode(state: typeof AgentState.State): Promise<Partial<typeof AgentState.State>> {
-        const prompt = `User mentioned ${JSON.stringify(state.entityStore)}. Ask them to clarify what they want to know (e.g., total production, specific issues, etc.).`;
-
-        let payload: ThoughtPayload = {
-            node: 'clarifier_node',
-            status: 'processing',
-            message: 'Asking for clarification',
-            context: state.entityStore
-        }
-
-        this.gateway.emitThought(state.socketId, payload);
-
-        const response = await this.model.invoke(prompt);
-        return {
-            messages: [response]
-        };
-    }
-
-    private async generalNode(state: typeof AgentState.State): Promise<Partial<typeof AgentState.State>> {
-        const lastMessage = state.messages[state.messages.length - 1];
-        const response = await this.model.invoke([
-            new HumanMessage("You are a helpful assistant. Reply to: " + lastMessage.content)
-        ]);
-        return {
-            messages: [response]
-        };
-    }
-
-    private async vectorSearchNode(state: typeof AgentState.State): Promise<Partial<typeof AgentState.State>> {
-        const lastMessage = state.messages[state.messages.length - 1].content as string;
-        const filter: Record<string, any> = {};
-
-        if (state.entityStore && state.entityStore.site) {
-            filter.site = state.entityStore.site;
-        }
-
-        const results = await this.vectorService.vectorSearch(lastMessage, 5, filter);
-
-        let payload: ThoughtPayload = {
-            node: 'vector_search_node',
-            status: 'completed',
-            query: lastMessage,
-            filter: filter,
-            resultsCount: results.length
-        }
-        this.gateway.emitThought(state.socketId, payload);
-
-        return {
-            actionPayload: { type: 'search', query: lastMessage, results }
-        };
-    }
-
-    private async aggregationNode(state: typeof AgentState.State): Promise<Partial<typeof AgentState.State>> {
-        const lastMessage = state.messages[state.messages.length - 1].content as string;
-
-        const pipelineSchema = z.object({
-            matchStage: z.object({
-                site: z.string().nullable(),
-                startDate: z.string().nullable(),
-                endDate: z.string().nullable(),
-            }),
-            aggregationType: z.enum(["sum", "avg", "count"]),
-            fieldToAggregate: z.enum(["quantity", "weight"])
-        });
-
-        const structuredLlm = this.model.withStructuredOutput(pipelineSchema);
-        const params = await structuredLlm.invoke(`Extract aggregation parameters for: "${lastMessage}". Context: ${JSON.stringify(state.entityStore)}`);
-
-        const pipeline: any[] = [];
-        const match: any = {};
-
-        // Check for null instead of undefined
-        if (params.matchStage.site !== null) {
-            match.site = params.matchStage.site;
-        }
-
-        if (params.matchStage.startDate !== null || params.matchStage.endDate !== null) {
-            match.productionDate = {};
-            if (params.matchStage.startDate !== null) {
-                match.productionDate.$gte = new Date(params.matchStage.startDate);
-            }
-            if (params.matchStage.endDate !== null) {
-                match.productionDate.$lte = new Date(params.matchStage.endDate);
-            }
-        }
-
-        if (Object.keys(match).length > 0) {
-            pipeline.push({ $match: match });
-        }
-
-        const group: any = { _id: null };
-        const operator = `$${params.aggregationType}`;
-        group.totalValue = { [operator]: `$${params.fieldToAggregate}` };
-
-        pipeline.push({ $group: group });
-
-        const results = await this.vectorService.aggregate(pipeline);
-
-        let payload: ThoughtPayload = {
-            node: `aggregation_node`,
-            status: 'completed',
-            pipeline: pipeline,
-            results: results
-        }
-        this.gateway.emitThought(state.socketId, {
-            node: 'aggregation_node',
-            status: 'completed',
-            pipeline: pipeline,
-            results: results
-        });
-
-        return {
-            actionPayload: { type: 'aggregate', pipeline, results }
-        };
-    }
-
-    private async synthesisNode(state: typeof AgentState.State): Promise<Partial<typeof AgentState.State>> {
-        const lastMessage = state.messages[state.messages.length - 1].content as string;
-        const payload = state.actionPayload;
-
-        const prompt = `
-        User Question: "${lastMessage}"
-        Data Context: ${JSON.stringify(payload)}
-        
-        Synthesize a natural language answer based STRICTLY on the Data Context.
-        Cite the source (e.g., "Based on aggregation results...").
-        `;
-
-        let thoughtPayload: ThoughtPayload = {
-            node: 'synthesis_node',
-            status: 'processing',
-            message: 'Synthesizing final response',
-            dataContextLength: JSON.stringify(payload).length
-        }
-        this.gateway.emitThought(state.socketId, thoughtPayload);
-
-        const response = await this.model.invoke(prompt);
-        return {
-            messages: [response]
-        };
-    }
-
     // --- MAIN ENTRY POINT ---
 
     createSession(socketId: string) {
-        this.sessions.set(socketId, []);
-        console.log(`Session created for ${socketId}`);
+        this.sessionManager.createSession(socketId);
     }
 
     deleteSession(socketId: string) {
-        this.sessions.delete(socketId);
-        console.log(`Session deleted for ${socketId}`);
+        this.sessionManager.deleteSession(socketId);
     }
 
     async chat(socketId: string, message: string): Promise<string> {
         try {
-            // Get history or init empty
-            const history = this.sessions.get(socketId) || [];
+            // Get session & filter valid history
+            const session = this.sessionManager.getSession(socketId);
+            const validHistory = this.sessionManager.getValidHistory(socketId);
+            const userMsg = new HumanMessage(message);
 
             const inputs = {
-                messages: [...history, new HumanMessage(message)],
-                entityStore: {},
+                messages: [...validHistory, userMsg],
+                entityStore: session.entityStore,
                 socketId: socketId
             };
 
             const result = await this.graph.invoke(inputs);
 
             const allMessages = result.messages as BaseMessage[];
+            const updatedEntityStore = result.entityStore as Record<string, any>;
+            const classification = result.entryCategory as string;
 
-            // Update history (keep all messages for context window? Or truncate?)
-            // For now, keep all. Memory optimization might be needed later.
-            this.sessions.set(socketId, allMessages);
-
+            // Get the AI response (last message)
             const agentMessages = allMessages.filter((m: BaseMessage) => m._getType() === 'ai');
             const lastResponse = agentMessages[agentMessages.length - 1];
 
-            return lastResponse.content as string;
+            // Update Session Storage
+            this.sessionManager.updateSession(
+                socketId,
+                userMsg,
+                lastResponse,
+                classification,
+                updatedEntityStore
+            );
+
+            return lastResponse?.content as string || "I'm sorry, I encountered an error.";
 
         } catch (error) {
             console.error('Error calling LangGraph:', error);

src/FFB/services/nodes/aggregation.node.ts

@@ -0,0 +1,60 @@
+import { AgentState } from "../config/agent-state";
+import { PROMPTS, SCHEMAS } from "../config/langchain-config";
+import { BaseChatModel } from "@langchain/core/language_models/chat_models";
+import { FFBGateway } from "../../ffb.gateway";
+import { ThoughtPayload } from "../../ffb-production.schema";
+import { FFBVectorService } from "../ffb-vector.service";
+
+export const aggregationNode = async (
+    state: typeof AgentState.State,
+    model: BaseChatModel,
+    vectorService: FFBVectorService,
+    gateway: FFBGateway
+): Promise<Partial<typeof AgentState.State>> => {
+    const lastMessage = state.messages[state.messages.length - 1].content as string;
+
+    const structuredLlm = model.withStructuredOutput(SCHEMAS.AGGREGATION);
+    const params = await structuredLlm.invoke(`Extract aggregation parameters for: "${lastMessage}". Context: ${JSON.stringify(state.entityStore)}`);
+
+    const pipeline: any[] = [];
+    const match: any = {};
+
+    // Check for null instead of undefined
+    if (params.matchStage.site !== null) {
+        match.site = params.matchStage.site;
+    }
+
+    if (params.matchStage.startDate !== null || params.matchStage.endDate !== null) {
+        match.productionDate = {};
+        if (params.matchStage.startDate !== null) {
+            match.productionDate.$gte = new Date(params.matchStage.startDate);
+        }
+        if (params.matchStage.endDate !== null) {
+            match.productionDate.$lte = new Date(params.matchStage.endDate);
+        }
+    }
+
+    if (Object.keys(match).length > 0) {
+        pipeline.push({ $match: match });
+    }
+
+    const group: any = { _id: null };
+    const operator = `$${params.aggregationType}`;
+    group.totalValue = { [operator]: `$${params.fieldToAggregate}` };
+
+    pipeline.push({ $group: group });
+
+    const results = await vectorService.aggregate(pipeline);
+
+    let payload: ThoughtPayload = {
+        node: `aggregation_node`,
+        status: 'completed',
+        pipeline: pipeline,
+        results: results
+    }
+    gateway.emitThought(state.socketId, payload);
+
+    return {
+        actionPayload: { type: 'aggregate', pipeline, results }
+    };
+};

src/FFB/services/nodes/entry.node.ts

@@ -0,0 +1,33 @@
+import { AgentState } from "../config/agent-state";
+import { PROMPTS, SCHEMAS } from "../config/langchain-config";
+import { BaseChatModel } from "@langchain/core/language_models/chat_models";
+import { FFBGateway } from "../../ffb.gateway";
+
+export const entryNode = async (
+    state: typeof AgentState.State,
+    model: BaseChatModel,
+    gateway: FFBGateway
+): Promise<Partial<typeof AgentState.State>> => {
+    const lastMessage = state.messages[state.messages.length - 1].content as string;
+
+    gateway.emitThought(state.socketId, {
+        node: 'entry_node',
+        status: 'processing',
+        message: 'Classifying user intent...',
+        input: lastMessage
+    });
+
+    const structuredLlm = model.withStructuredOutput(SCHEMAS.ENTRY);
+    const result = await structuredLlm.invoke(PROMPTS.ENTRY(lastMessage));
+
+    gateway.emitThought(state.socketId, {
+        node: 'entry_node',
+        status: 'completed',
+        result: result
+    });
+
+    return {
+        entryCategory: result.category,
+        socketId: state.socketId
+    };
+};

src/FFB/services/nodes/guidance.node.ts

@@ -0,0 +1,9 @@
+import { AgentState } from "../config/agent-state";
+import { PROMPTS } from "../config/langchain-config";
+import { AIMessage } from "@langchain/core/messages";
+
+export const guidanceNode = async (
+    state: typeof AgentState.State
+): Promise<Partial<typeof AgentState.State>> => {
+    return { messages: [new AIMessage(PROMPTS.GUIDANCE)] };
+};

src/FFB/services/nodes/meta.node.ts

@@ -0,0 +1,16 @@
+import { AgentState } from "../config/agent-state";
+import { BaseChatModel } from "@langchain/core/language_models/chat_models";
+import { PROMPTS } from "../config/langchain-config";
+
+export const metaNode = async (
+    state: typeof AgentState.State,
+    model: BaseChatModel
+): Promise<Partial<typeof AgentState.State>> => {
+    const lastMessage = state.messages[state.messages.length - 1].content as string;
+
+    // Construct context from messages
+    const context = state.messages.map(m => `${m._getType()}: ${m.content}`).join('\n');
+
+    const response = await model.invoke(PROMPTS.META(lastMessage, context));
+    return { messages: [response] };
+};

src/FFB/services/nodes/refusal.node.ts

@@ -0,0 +1,9 @@
+import { AgentState } from "../config/agent-state";
+import { PROMPTS } from "../config/langchain-config";
+import { AIMessage } from "@langchain/core/messages";
+
+export const refusalNode = async (
+    state: typeof AgentState.State
+): Promise<Partial<typeof AgentState.State>> => {
+    return { messages: [new AIMessage(PROMPTS.REFUSAL)] };
+};

src/FFB/services/nodes/router.node.ts

@@ -0,0 +1,37 @@
+import { AgentState } from "../config/agent-state";
+import { PROMPTS, SCHEMAS } from "../config/langchain-config";
+import { BaseChatModel } from "@langchain/core/language_models/chat_models";
+import { FFBGateway } from "../../ffb.gateway";
+import { ThoughtPayload } from "../../ffb-production.schema";
+
+export const routerNode = async (
+    state: typeof AgentState.State,
+    model: BaseChatModel,
+    gateway: FFBGateway
+): Promise<Partial<typeof AgentState.State>> => {
+    const lastMessage = state.messages[state.messages.length - 1].content as string;
+
+    let payload: ThoughtPayload = {
+        node: 'router_node',
+        status: 'processing',
+        message: 'Routing actionable request...',
+        input: lastMessage
+    }
+
+    gateway.emitThought(state.socketId, payload);
+
+    const structuredLlm = model.withStructuredOutput(SCHEMAS.ROUTER);
+    const result = await structuredLlm.invoke(PROMPTS.ROUTER(lastMessage));
+
+    gateway.emitThought(state.socketId, {
+        node: 'router_node',
+        status: 'completed',
+        result: result
+    });
+
+    return {
+        activeIntent: result.intent as any,
+        entityStore: result.entities || {},
+        socketId: state.socketId
+    };
+};

src/FFB/services/nodes/synthesis.node.ts

@@ -0,0 +1,27 @@
+import { AgentState } from "../config/agent-state";
+import { PROMPTS } from "../config/langchain-config";
+import { BaseChatModel } from "@langchain/core/language_models/chat_models";
+import { FFBGateway } from "../../ffb.gateway";
+import { ThoughtPayload } from "../../ffb-production.schema";
+
+export const synthesisNode = async (
+    state: typeof AgentState.State,
+    model: BaseChatModel,
+    gateway: FFBGateway
+): Promise<Partial<typeof AgentState.State>> => {
+    const lastMessage = state.messages[state.messages.length - 1].content as string;
+    const payload = state.actionPayload;
+
+    let thoughtPayload: ThoughtPayload = {
+        node: 'synthesis_node',
+        status: 'processing',
+        message: 'Synthesizing final response',
+        dataContextLength: JSON.stringify(payload).length
+    }
+    gateway.emitThought(state.socketId, thoughtPayload);
+
+    const response = await model.invoke(PROMPTS.SYNTHESIS(lastMessage, payload));
+    return {
+        messages: [response]
+    };
+};

src/FFB/services/nodes/vector-search.node.ts

@@ -0,0 +1,32 @@
+import { AgentState } from "../config/agent-state";
+import { FFBGateway } from "../../ffb.gateway";
+import { ThoughtPayload } from "../../ffb-production.schema";
+import { FFBVectorService } from "../ffb-vector.service";
+
+export const vectorSearchNode = async (
+    state: typeof AgentState.State,
+    vectorService: FFBVectorService,
+    gateway: FFBGateway
+): Promise<Partial<typeof AgentState.State>> => {
+    const lastMessage = state.messages[state.messages.length - 1].content as string;
+    const filter: Record<string, any> = {};
+
+    if (state.entityStore && state.entityStore.site) {
+        filter.site = state.entityStore.site;
+    }
+
+    const results = await vectorService.vectorSearch(lastMessage, 5, filter);
+
+    let payload: ThoughtPayload = {
+        node: 'vector_search_node',
+        status: 'completed',
+        query: lastMessage,
+        filter: filter,
+        resultsCount: results.length
+    }
+    gateway.emitThought(state.socketId, payload);
+
+    return {
+        actionPayload: { type: 'search', query: lastMessage, results }
+    };
+};

src/FFB/services/utils/session-manager.ts

@@ -0,0 +1,74 @@
+import { BaseMessage } from "@langchain/core/messages";
+
+export interface SessionItem {
+    message: BaseMessage;
+    category: string;
+}
+
+export interface SessionData {
+    items: SessionItem[];
+    entityStore: Record<string, any>;
+    modelProvider: 'openai' | 'gemini';
+}
+
+export class SessionManager {
+    private sessions: Map<string, SessionData> = new Map();
+
+    createSession(socketId: string) {
+        this.sessions.set(socketId, {
+            items: [],
+            entityStore: {},
+            modelProvider: 'openai' // Default
+        });
+        console.log(`Session created for ${socketId}`);
+    }
+
+    deleteSession(socketId: string) {
+        this.sessions.delete(socketId);
+        console.log(`Session deleted for ${socketId}`);
+    }
+
+    getSession(socketId: string): SessionData {
+        let session = this.sessions.get(socketId);
+        if (!session) {
+            this.createSession(socketId);
+            return this.sessions.get(socketId)!;
+        }
+        return session;
+    }
+
+    getValidHistory(socketId: string): BaseMessage[] {
+        const session = this.getSession(socketId);
+        return session.items
+            .filter(item => item.category !== 'OutOfScope')
+            .map(item => item.message);
+    }
+
+    updateSession(socketId: string, userMsg: BaseMessage, aiMsg: BaseMessage | undefined, category: string, entityStore: Record<string, any>) {
+        const session = this.getSession(socketId);
+
+        // Add User Message
+        session.items.push({ message: userMsg, category });
+
+        // Add AI Message if exists
+        if (aiMsg) {
+            session.items.push({ message: aiMsg, category });
+        }
+
+        // Update Entity Store
+        session.entityStore = entityStore;
+
+        this.sessions.set(socketId, session);
+    }
+
+    setModelProvider(socketId: string, provider: 'openai' | 'gemini') {
+        const session = this.getSession(socketId);
+        session.modelProvider = provider;
+        this.sessions.set(socketId, session); // Ensure map is updated if needed (though object ref works)
+        console.log(`Updated model provider for ${socketId} to ${provider}`);
+    }
+
+    getModelProvider(socketId: string): 'openai' | 'gemini' {
+        return this.getSession(socketId).modelProvider;
+    }
+}