server.js

import express from "express";
import cors from "cors";
import Groq from "groq-sdk";
import { config } from "dotenv";
import path from "path";
import fs from "fs";
import { fileURLToPath } from "url";
import { Pinecone as PineconeClient } from "@pinecone-database/pinecone";
import { PDFLoader } from "@langchain/community/document_loaders/fs/pdf";
import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters";
import multer from "multer";

config();

// Configure multer for file uploads
const upload = multer({ dest: "uploads/" });

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

const app = express();
app.use(cors());
app.use(express.json());
app.use(express.static(path.join(__dirname, "public")));

// Initialize Pinecone
const pc = new PineconeClient();
const INDEX_NAME = "rag-embedded-index";
const NAMESPACE = "default";
const index = pc.index(INDEX_NAME);

// Initialize Groq
const groq = new Groq({ apiKey: process.env.GROQ_API_KEY });

// RAG Search function
async function searchVectorStore(query, topK = 3) {
    const results = await index.namespace(NAMESPACE).searchRecords({
        query: {
            topK: topK,
            inputs: { text: query }
        },
        rerank: {
            model: "bge-reranker-v2-m3",
            topN: topK,
            rankFields: ["content"]
        }
    });

    return results.result.hits.map(hit => {
        const fields = hit.fields;
        return {
            id: hit._id,
            score: hit._score,
            content: fields.content,
            source: fields.source
        };
    });
}

// Store messages per session
const sessions = new Map();

const getOrCreateSession = (sessionId) => {
    if (!sessions.has(sessionId)) {
        sessions.set(sessionId, [
            {
                role: "system",
                content: `You are a helpful AI assistant with access to a knowledge base about AI Engineering skills.
When answering questions, use the rag_search tool to find relevant information from the knowledge base.
Always base your answers on the retrieved context when available.
If the context doesn't contain relevant information, say so honestly.
Keep responses concise and helpful.
Current date: ${new Date().toUTCString()}`
            }
        ]);
    }
    return sessions.get(sessionId);
};

app.post("/api/chat", async (req, res) => {
    try {
        const { message, sessionId = "default" } = req.body;

        if (!message) {
            return res.status(400).json({ error: "Message is required" });
        }

        const messages = getOrCreateSession(sessionId);
        messages.push({ role: "user", content: message });

        let response;
        let toolCalls;
        let toolsUsed = [];
        let context = "";

        do {
            response = await groq.chat.completions.create({
                model: "llama-3.3-70b-versatile",
                temperature: 0,
                tool_choice: "auto",
                tools: [
                    {
                        type: "function",
                        function: {
                            name: "rag_search",
                            description: "Search the knowledge base for information about AI engineering skills, machine learning, transformers, RAG, prompt engineering, and related topics.",
                            parameters: {
                                type: "object",
                                properties: {
                                    query: {
                                        type: "string",
                                        description: "The search query to find relevant information",
                                    }
                                },
                                required: ["query"]
                            }
                        }
                    }
                ],
                messages: messages
            });

            toolCalls = response.choices[0].message.tool_calls;
            messages.push(response.choices[0].message);

            if (toolCalls && toolCalls.length > 0) {
                for (const toolCall of toolCalls) {
                    if (toolCall.function.name === "rag_search") {
                        toolsUsed.push("rag_search");
                        const args = JSON.parse(toolCall.function.arguments);

                        // Search vector store
                        const results = await searchVectorStore(args.query, 3);
                        context = results.map((r, i) =>
                            `[Source ${i + 1}] (Score: ${r.score.toFixed(3)})\n${r.content}`
                        ).join("\n\n");

                        messages.push({
                            role: "tool",
                            tool_call_id: toolCall.id,
                            content: context || "No relevant information found in the knowledge base."
                        });
                    }
                }
            }
        } while (toolCalls && toolCalls.length > 0);

        const assistantMessage = response.choices[0].message.content;

        res.json({
            response: assistantMessage,
            toolsUsed: toolsUsed.length > 0 ? toolsUsed : undefined,
            context: context || undefined
        });

    } catch (error) {
        console.error("Error:", error);
        res.status(500).json({ error: "Internal server error" });
    }
});

app.post("/api/clear", (req, res) => {
    const { sessionId = "default" } = req.body;
    sessions.delete(sessionId);
    res.json({ success: true });
});

// Ingest document endpoint
app.post("/api/ingest", upload.single("file"), async (req, res) => {
    try {
        if (!req.file) {
            return res.status(400).json({ error: "No file uploaded" });
        }

        const filePath = req.file.path;
        const originalName = req.file.originalname;
        const mimeType = req.file.mimetype;

        console.log(`Ingesting file: ${originalName}`);

        let textContent = "";

        // Handle PDF files
        if (mimeType === "application/pdf" || originalName.endsWith(".pdf")) {
            const loader = new PDFLoader(filePath, { splitPages: false });
            const docs = await loader.load();
            textContent = docs.map(doc => doc.pageContent).join("\n\n");
        }
        // Handle text files
        else if (mimeType === "text/plain" || originalName.endsWith(".txt")) {
            textContent = fs.readFileSync(filePath, "utf-8");
        }
        // Handle other text-based files
        else {
            textContent = fs.readFileSync(filePath, "utf-8");
        }

        if (!textContent || textContent.trim().length === 0) {
            fs.unlinkSync(filePath); // Clean up
            return res.status(400).json({ error: "Could not extract text from file" });
        }

        // Chunk the document
        const splitter = new RecursiveCharacterTextSplitter({
            chunkSize: 500,
            chunkOverlap: 100
        });
        const chunks = await splitter.splitText(textContent);

        console.log(`Created ${chunks.length} chunks from ${originalName}`);

        // Get current record count to generate unique IDs
        const stats = await index.describeIndexStats();
        const startId = stats.totalRecordCount || 0;

        // Create records for Pinecone
        const records = chunks.map((chunk, i) => ({
            _id: `doc-${Date.now()}-${startId + i}`,
            content: chunk,
            chunkIndex: i,
            source: originalName,
            ingestedAt: new Date().toISOString()
        }));

        // Upsert to Pinecone
        console.log("Upserting records to Pinecone...");
        await index.namespace(NAMESPACE).upsertRecords(records);

        // Clean up uploaded file
        fs.unlinkSync(filePath);

        // Wait a bit for indexing
        await new Promise(resolve => setTimeout(resolve, 3000));

        // Get updated stats
        const newStats = await index.describeIndexStats();

        res.json({
            success: true,
            message: `Successfully ingested ${originalName}`,
            chunksCreated: chunks.length,
            totalRecords: newStats.totalRecordCount
        });

    } catch (error) {
        console.error("Ingest error:", error);
        // Clean up file if it exists
        if (req.file && fs.existsSync(req.file.path)) {
            fs.unlinkSync(req.file.path);
        }
        res.status(500).json({ error: "Failed to ingest document", details: error.message });
    }
});

// Health check
app.get("/api/health", async (req, res) => {
    try {
        const stats = await index.describeIndexStats();
        res.json({
            status: "ok",
            index: INDEX_NAME,
            records: stats.totalRecordCount
        });
    } catch (error) {
        res.status(500).json({ status: "error", error: error.message });
    }
});

const PORT = process.env.PORT || 3000;
app.listen(PORT, () => {
    console.log(`RAG Chat Server running on http://localhost:${PORT}`);
});