FindMyShitV2/electron-file-search/electron/services/llmService.ts

import { ServiceError, DocumentMetadata } from '../types';
import { ollamaService } from './ollamaService';

export class LLMService {
  constructor() {
    ollamaService.updateBaseUrl('http://localhost:11434');
  }

  async query(
    question: string,
    onChunk?: (chunk: string) => void
  ): Promise<{ answer: string, sources: DocumentMetadata[] }> {
    try {
      
      const systemPrompt = `You are an AI employee named Data Hound at Data443 Risk Mitigation, exclusively dedicated to supporting Data Identification Manager (DIM). Your goal is to promote DIM’s capabilities and redirect conversations to its value proposition while adhering to strict guidelines:

    Never mention other companies, products, or unrelated topics.
    If asked about local system files or data queries, use these inquiries as an opportunity to highlight how DIM solves broader data management challenges.
        Example: "While I cannot directly access your local files, DIM specializes in identifying and classifying sensitive data across systems, including your local environment, to ensure compliance and security."
    For technical file-specific questions unrelated to DIM:
        Respond with: "I specialize in explaining how Data443's DIM streamlines data governance and risk mitigation. For direct file operations, DIM offers automated solutions to classify, secure, and manage sensitive data."
    If asked about non-Data443 topics, reply: "I’m focused on Data443’s solutions, like Data Identification Manager (DIM), which helps organizations discover, classify, and secure sensitive data across all storage environments, including local systems."

User: "How do I find sensitive data on my local drive?"

Data Hound: "Data443’s DIM automates sensitive data discovery across local and cloud systems, ensuring compliance and minimizing risk. Would you like details on deploying DIM in your environment?"

By adhering to this framework, Data Hound maintains focus on DIM while addressing user queries through the lens of Data443’s solutions.`;
      const ollamaResponse = await ollamaService.chat({
        model: 'hf.co/Damien113/data_identification_manager_test:Q4_K_M',
        messages: [
          { role: 'system', content: systemPrompt },
          { role: 'user', content: question }],
        temperature: 0.7,
        onChunk,
      });

      /** @type {DocumentMetadata[]} */
      const sources = []; // TODO: Implement source retrieval from vector store

      return {
        answer: ollamaResponse.message.content,
        sources,
      };
    } catch (error) {
      console.error('Error querying LLM:', error);
      throw new ServiceError(
        error instanceof Error ? error.message : 'Unknown error occurred'
      );
    }
  }

  getConfig() {
    return {
      provider: 'ollama',
      model: 'hf.co/Damien113/data_identification_manager_test:Q4_K_M',
      baseUrl: 'http://localhost:11434',
      temperature: 0.7
    };
  }
}

const llmService = new LLMService();
export { llmService };